jax.tree_util.Partial

jax.tree_util.Partial#

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

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

使用它以与 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

未来部分调用的关键字参数字典
目录