JAX中右乘矩阵时Jacobian维度交换问题求解

当你使用jax.jacrev计算右乘操作的雅可比时，得到的(2,3)维度是符合数学定义的：雅可比矩阵的维度为**(输出维度, 输入维度)**，每个元素对应输出分量对输入分量的偏导数。你的输入是3维向量，输出是2维向量，因此默认输出(2,3)是正确的。

如果你需要和W维度一致的(3,2)矩阵，有两种快速解决方式：

方法1：直接转置结果 #

这是最简洁的方案，对计算得到的雅可比矩阵做转置操作即可：

import jax
import jax.numpy as jnp

X = jnp.arange(3*10, dtype=float).reshape(10, 3)
W = jnp.ones((3, 2))

def funz(x):
    return jnp.dot((0.5 * x ** 2), W)

funz_jac = jax.jacrev(funz)

x = X[0]
jac = funz_jac(x).T  # 转置操作
print(jac.shape)  # Output: (3, 2)

方法2：修改函数输出（可选） #

如果你希望求导直接得到目标维度，可以调整函数输出为转置形式，但这种方式不如直接转置直观：

def funz_transposed(x):
    return jnp.dot((0.5 * x ** 2), W).T

funz_jac_transposed = jax.jacrev(funz_transposed)
jac = funz_jac_transposed(x)
print(jac.shape)  # Output: (3, 2)

对于批量输入场景，转置操作依然有效——JAX会保留批量维度，转置后维度会从(batch_size, 2, 3)变为(batch_size, 3, 2)。

内容的提问来源于stack exchange，提问作者Ouden