[Python] JAX lax.scan: How to iterate over layers and memory slices simultaneously without...

I'm trying to implement a framework that manages RNN networks with an arbitrary number of layers (it's all part of a library I'm building based on jax/equinox), the problem is that I can't find an implementation that allows me to manage multiple layers in a modular way, let me explain:

def __call__(self, x):

def forward(carry, x_t):

def layer_apply(h_t, layer):
return layer(h_t, x_t)

t_memory = lax.scan(layer_apply, carry, self.layers)

return t_memory[0], t_memory[1]

return lax.scan(forward, self.h, x)


this code is a function of a class called 'Recurrent_block' (an eqx.module), that accept as parameters a tuple of rnn layers, and creates a self.h which is a tuple of jnp.array all zero with the correct shape for the memory of every layer, so:

self.h = tuple of jnp.zeros, one for layer: (h(0), h(1)... h(n))

self.layers = tuple of eqx.Module classes, each with a call method that implements the RNN forward pass

x = jnp.array with shape (seq_len, batch_size, features)

(As you can see, everything is known at compile time)

Now: There are two huge problems in the code:

1. lax.scan doesn't accept xs having non-homogeneous shapes (w_h, w_h, and the bias have different shapes by nature), so I can't pass the layer tuple as xs to scan.
2. In the code, I'm nesting the entire temporal tuple in the layer call, not the current slice, and I don't know how to give it just the slice he needs(because the layer needs only his own temporal memory h_t)

I would prefer the solution to use lax.scan as it is very optimized for this type of operation from what I know, even multiple nested scans are fine, the only important thing is that it can be compiled in XLA without any strange fallbacks, and I absolutely don't want an hard-coded solution like this:

x = layer[0](h[0], x),

x = layer[1](h[1], x)

...

because that would obviously be very inconvenient. The point is that since I have everything known at compile time, it should be possible to jit the function, but I can't find a solution.

I have tried a few approaches:

1. using a index for layers and h

def __call__(self, x, i):


def forward(carry, x_t):

def layer_apply(h_t, i):
return self.layers(h_t, x_t)

t_memory = lax.scan(layer_apply, carry, i)

return t_memory[0], t_memory[1]

return lax.scan(forward, self.h, x)


where i is a index array made with arange, the error of this code is:

**layer_apply
return self.layers(h_t, x_t)
~~~~~~^^^
jax.errors.TracerIntegerConversionError: The __index__() method was called on traced array with shape int32[]
The error occurred while tracing the function layer_apply at test.py:37 for scan. This concrete value was not available in Python because it depends on the value of the argument i.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerIntegerConversionError**

1. using an unrolled for loop

def __call__(self, x):


def forward(carry, x_t):
for layer, h_t in zip(self.layers, carry):
x_t = layer(h_t, x_t)
return carry, x_t

return lax.scan(forward, self.h, x)[1]


this is probably my favorite one, but as you can see the problem is that carry is not updated, so it does not raise an error but it doesn't make sense.

if you want to test, here is a dummy code which works like mine:

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

class Rnn(eqx.Module): #layer class
dim : int = eqx.static_field()
w_x : jnp.array
w_h : jnp.array
b : jnp.array


def __init__(self, dim):
self.dim = dim
self.w_x = jnp.ones((dim,dim), dtype=jnp.float32)
self.w_h = jnp.ones((dim,dim), dtype=jnp.float32)
self.b = jnp.zeros((dim,), dtype=jnp.float32)

def __call__(self, h, x): #rnn forward
return x@self.w_x + h@self.w_h + self.b

dim = 2
batch_size = 3
seq_len = 3

x = jnp.ones((seq_len, 3, dim)) #input

layers = (Rnn(dim), Rnn(dim)) #tuple of layers

h = (jnp.zeros((batch_size, dim), dtype=jnp.float32), jnp.zeros((batch_size, dim), dtype=jnp.float32)) #temporal memory

def scan_fn():
return #your solution !

print(scan_fn()) #a general result for the iteration in every rnn layer

Continue reading...