使用Equinox模块结合jax.lax.scan时遇ArrayImpl不可哈希错误

我是Equinox和JAX的新手，想用它们模拟动态系统。但把系统模型作为Equinox模块传入jax.lax.scan时，出现了TypeError: unhashable type: 'ArrayImpl'错误。我知道JAX要求函数参数为纯函数，原本以为Equinox模块能满足这个特性。以下是复现错误的测试脚本：

import equinox as eqx
import jax
import jax.numpy as jnp


class EqxModel(eqx.Module):
    A: jax.Array
    B: jax.Array
    C: jax.Array
    D: jax.Array

    def __call__(self, states, inputs):
        x = states.reshape(-1, 1)
        u = inputs.reshape(-1, 1)
        x_next = self.A @ x + self.B @ u
        y = self.C @ x + self.D @ u
        return x_next.reshape(-1), y.reshape(-1)


def simulate(model, inputs, x0):
    xk = x0
    outputs = []
    for uk in inputs:
        xk, yk = model(xk, uk)
        outputs.append(yk)
    outputs = jnp.stack(outputs)
    return xk, outputs


A = jnp.array([[0.7, 1.0], [0.0, 1.0]])
B = jnp.array([[0.0], [1.0]])
C = jnp.array([[0.3, 0.0]])
D = jnp.array([[0.0]])
model = EqxModel(A, B, C, D)

# 测试仿真，这部分正常运行
inputs = jnp.array([[0.0], [1.0], [1.0], [1.0]])
x0 = jnp.zeros(2)
xk, outputs = simulate(model, inputs, x0)
assert jnp.allclose(xk, jnp.array([2.7, 3.0]))
assert jnp.allclose(outputs, jnp.array([[0.0], [0.0], [0.0], [0.3]]))

# 这一行会抛出TypeError
xk, outputs = jax.lax.scan(model, x0, inputs)

问题答疑 #

1. unhashable type: 'ArrayImpl'指的是什么？

是的，这个错误里的ArrayImpl就是你定义的A、B、C、D这些JAX数组。JAX数组属于不可哈希类型，而jax.lax.scan要求传入的函数必须是纯函数，且函数依赖的闭包变量必须可哈希。直接把Equinox模块传给scan时，模块内的数组会被当作闭包变量处理，因为数组不可哈希，所以触发了这个错误。

2. 解决方法

因为你的模型参数（A/B/C/D）在仿真时是静态的，推荐用以下两种方法解决：

方法一：使用eqx.filter_scan（Equinox专门适配的scan）

Equinox提供了eqx.filter_scan，可以直接处理带参数的模块，不需要额外包装：

# 替换报错的scan调用为以下代码
xk, outputs = eqx.filter_scan(model, x0, inputs)

方法二：用lambda或partial包装函数，把模块作为参数传入

如果坚持用jax.lax.scan，可以把模块作为额外参数传入，避免闭包捕获不可哈希的数组：

from functools import partial

# 定义包装函数
def scan_step(model, states, inputs):
    return model(states, inputs)

# 使用partial绑定模型，再传入scan
xk, outputs = jax.lax.scan(partial(scan_step, model), x0, inputs)

或者用lambda简化：

xk, outputs = jax.lax.scan(lambda s, u: model(s, u), x0, inputs)

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