如何为softmax编写JAX自定义向量Jacobian乘积(vjp)

Question

为了理解JAX的反向模式自动差异，我尝试为softmax编写一个custom_vjp，如下所示：

import jax
import jax.numpy as jnp
import numpy as np
@jax.custom_vjp
def stablesoftmax(x):
print(f"input: {x} shape: {x.shape}")
expc = jnp.exp(x - jnp.amax(x))
return expc / jnp.sum(expc)
def ssm_fwd(x):
s = stablesoftmax(x)
return s, s
def ssm_bwd(acts, d_dacts):
dacts_dinput = jnp.diag(acts) - jnp.outer(acts, acts)  # Jacobian
d_dinput = jnp.dot(d_dacts, dacts_dinput)  # Vector-Jacobian product
print(f"Saved activations:n{acts} shape: {acts.shape}")
print(f"d/d_acts:n{d_dacts} shape: {d_dacts.shape}")
print(f"d_acts/d_input (Jacobian of softmax):n{dacts_dinput} shape: {dacts_dinput.shape}")
print(f"d/d_input:n{d_dinput} shape: {d_dinput.shape}")
return d_dinput
stablesoftmax.defvjp(ssm_fwd, ssm_bwd)
print(f"JAX version: {jax.__version__}")
y = np.array([1., 2., 3.])
a = stablesoftmax(y)
softmax_jac_fun = jax.jacrev(stablesoftmax)
dsoftmax_dy = softmax_jac_fun(y)
print(f"Softmax Jacobian: {dsoftmax_dy}")

但是当我调用jacrev时，我得到了关于VJP结果的结构与softmax:的输入结构不匹配的错误

JAX version: 0.2.13
input: [1. 2. 3.] shape: (3,)
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
input: [1. 2. 3.] shape: (3,)
Saved activations:
[0.09003057 0.24472848 0.66524094] shape: (3,)
d/d_acts:
Traced<ShapedArray(float32[3])>with<BatchTrace(level=1/0)>
with val = array([[1., 0., 0.],
[0., 1., 0.],
[0., 0., 1.]], dtype=float32)
batch_dim = 0 shape: (3,)
d_acts/d_input (Jacobian of softmax):
[[ 0.08192507 -0.02203305 -0.05989202]
[-0.02203305  0.18483645 -0.1628034 ]
[-0.05989202 -0.1628034   0.22269544]] shape: (3, 3)
d/d_input:
Traced<ShapedArray(float32[3])>with<BatchTrace(level=1/0)>
with val = DeviceArray([[ 0.08192507, -0.02203305, -0.05989202],
[-0.02203305,  0.18483645, -0.1628034 ],
[-0.05989202, -0.1628034 ,  0.22269544]], dtype=float32)
batch_dim = 0 shape: (3,)
Traceback (most recent call last):
File "analysis/vjp_test.py", line 30, in <module>
dsoftmax_dy = softmax_jac_fun(y)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: TypeError: Custom VJP rule must produce an output with the same container (pytree) structure as the args tuple of the primal function, and in particular must produce a tuple of length equal to the number of arguments to the primal function, but got VJP output structure PyTreeDef(*) for primal input structure PyTreeDef((*,)).

然而，您可以看到，当我打印形状时，它们都有形状(3(，但JAX似乎不同意？(事实上，输入和输出是3乘3的矩阵，但这是因为JAX试图在jacrev中对JVM进行vmap，因此一次性撤回R(3(的整个基础(即3x3单位矩阵(。

注意：如果我直接使用jax.grad或jax.vjp，我会得到同样的错误。

Answer 1

根据custom_vjp文档：

bwd的输出必须是一个长度等于原始函数的自变量数量的元组

所以向后传递中的return语句应该如下所示：

def ssm_bwd(acts, d_dacts):
...
return (d_dinput,)