jax.tree_util.Partial

class jax.tree_util.Partial(func, *args, **kw)

functools.partial 的一个版本，可在 pytrees 中使用。

使用它进行部分函数求值，使其与 JAX 的转换兼容，例如，Partial(func, *args, **kwargs)。

（您需要显式选择启用此行为，因为我们不希望 functools.partial 的语义与普通函数闭包不同。）

例如，以下是在类似于 functools.partial 的方式中使用 Partial 的基本用法

>>> import jax.numpy as jnp
>>> add_one = Partial(jnp.add, 1)
>>> add_one(2)
Array(3, dtype=int32, weak_type=True)

Pytree 兼容性意味着生成的 partial 函数可以作为参数在转换后的 JAX 函数中传递，这对于标准的 functools.partial 函数是不可能的

>>> from jax import jit
>>> @jit
... def call_func(f, *args):
...   return f(*args)
...
>>> call_func(add_one, 2)
Array(3, dtype=int32, weak_type=True)

将零个参数传递给 Partial 有效地包装了原始函数，使其成为 JAX 转换函数中的有效参数

>>> call_func(Partial(jnp.add), 1, 2)
Array(3, dtype=int32, weak_type=True)

如果我们直接将 jnp.add 传递给 call_func，则会导致 TypeError。

请注意，如果在值被跟踪的上下文中使用 Partial 的结果，则在传递给部分求值函数时，会导致所有绑定参数被跟踪

>>> print_zero = Partial(print, 0)
>>> print_zero()
0
>>> call_func(print_zero)  
Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace...>

__init__()

方法

__init__()

属性

args	未来 partial 调用的参数元组
func	在未来 partial 调用中使用的函数对象
keywords	未来 partial 调用的关键字参数字典