jax.tree_util.Partial

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

functools.partial 的一个版本，可在 pytrees 中工作。

它用于部分函数求值，且与 JAX 的变换兼容，例如 Partial(func, *args, **kwargs)。

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

例如，以下是 Partial 的基本用法，类似于 functools.partial。

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

Pytree 兼容性意味着可以将生成的部分函数作为参数传递到转换后的 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	未来部分调用的参数元组
func	在未来部分调用中使用的函数对象
keywords	未来部分调用的关键字参数字典