如何处理JAX custom_vjp中非JAX对象（如scipy.sparse.csr_matrix）引发的错误？

嘿，我之前踩过一模一样的坑！JAX的自动微分系统对输入类型有严格要求，scipy的稀疏矩阵不属于JAX的可追踪数值类型（没有JAX要求的dtype等属性），直接放到custom_vjp装饰的函数里肯定会报错。下面给你两个可行的解决思路，附带修改后的代码示例：

方案1：转换为JAX原生稀疏矩阵（推荐，保留稀疏性） #

JAX提供了实验性的稀疏矩阵支持（jax.experimental.sparse），可以把scipy的CSR矩阵转换成JAX原生的稀疏矩阵类型，这样既能保留稀疏结构节省内存，又能被JAX的微分系统正常处理。

修改后的代码： #

import jax
import jax.numpy as jnp
import scipy.sparse as sp
from jax import custom_vjp
import jax.experimental.sparse as jsparse

@custom_vjp
def simple_function_vjp(sparse_matrix, vector):
    # 这里的sparse_matrix是JAX原生CSR类型，支持JAX运算和微分
    return sparse_matrix @ vector

def simple_function_fwd(sparse_matrix, vector):
    return simple_function_vjp(sparse_matrix, vector), (sparse_matrix, vector)

def simple_function_bwd(residuals, grad_output):
    sparse_matrix, vector = residuals
    # JAX稀疏矩阵的转置和乘法和scipy用法一致
    grad_vector = sparse_matrix.T @ grad_output
    # 如果需要返回稠密形式的梯度矩阵，直接用jnp.outer即可
    grad_sparse_matrix = jnp.outer(grad_output, vector)
    
    return (grad_sparse_matrix, grad_vector)

# 使用时先把scipy稀疏矩阵转换成JAX原生稀疏矩阵
scipy_csr = sp.csr_matrix([[1, 0], [0, 2]])
jax_csr = jsparse.CSR.from_scipy_sparse(scipy_csr)
test_vector = jnp.array([1.0, 2.0])

# 测试正向计算和反向梯度
forward_output = simple_function_vjp(jax_csr, test_vector)
gradients = jax.grad(lambda sm, v: simple_function_vjp(sm, v).sum())(jax_csr, test_vector)

print("正向输出:", forward_output)
print("梯度结果:", gradients)

方案2：转换为稠密JAX数组（适合小矩阵） #

如果你的稀疏矩阵规模不大，完全可以直接转换成JAX的稠密数组，这样操作起来更简单，不需要额外处理稀疏结构：

修改后的代码片段： #

# 在调用函数前转换scipy稀疏矩阵为JAX稠密数组
dense_jax_matrix = jnp.array(scipy_csr.todense())

# 后续函数可以直接用dense_jax_matrix作为输入，原custom_vjp代码不需要大改
forward_output = simple_function_vjp(dense_jax_matrix, test_vector)

小提示 #

• 尽量不要在custom_vjp装饰的函数内部做scipy到JAX的类型转换，因为JAX的追踪机制可能无法识别这类外部操作，导致微分出错。最好在函数调用前完成转换。
• 如果需要对稀疏矩阵本身求导，要注意JAX稀疏矩阵的梯度返回形式——默认可能是稠密数组，如果你需要稀疏形式的梯度，可能要额外处理。

备注：内容来源于stack exchange，提问作者James Yong