jax.linear_transpose#

jax.linear_transpose(fun, *primals, reduce_axes=())[源代码]#

转置一个被保证为线性的函数。

对于线性函数，此变换等同于 vjp()，但避免了计算前向传递的开销。

转置函数的输出将始终具有与 primals 完全相同的 dtypes，即使某些值被截断（例如，从复数到浮点数，或从 float64 到 float32）。为避免截断，请在 primals 中使用与转置函数所需输出的完整范围相匹配的 dtypes。不支持整数 dtypes。

参数:

  • fun (Callable) – 要转置的线性函数。

  • *primals – 用于评估 fun(*primals) 的形状/dtypes 的数组、标量或（嵌套）标准 Python 容器（元组、列表、字典、命名元组，即 pytrees）的位置参数元组。这些参数可以是实数标量/ ndarrays，但这并非必需：仅访问 shape 和 dtype 属性。有关示例，请参见下文。（请注意，鸭子类型对象不能是命名元组，因为那些会被视为标准 Python 容器。）

返回:

一个计算 fun 的转置的可调用对象。此函数的有效输入必须具有与 fun(*primals) 的结果相同的形状/dtypes/结构。输出将是一个元组，其形状/dtypes/结构与 primals 相同。

返回类型:

可调用对象

>>> import jax
>>>
>>> f = lambda x, y: 0.5 * x - 0.5 * y
>>> scalar = jax.ShapeDtypeStruct(shape=(), dtype=np.dtype(np.float32))
>>> f_transpose = jax.linear_transpose(f, scalar, scalar)
>>> f_transpose(1.0)
(Array(0.5, dtype=float32), Array(-0.5, dtype=float32))