Shape-dependent initialization gives concretization error at apply.

[GJBoth]

I'm working on implementing Hadamard Transforms. I use the following function to build a hadamard matrix:

def hadamard(key, shape, dtype=jnp.float32):
    lg2 = jnp.log2(shape[0])
    H = jnp.ones((1, ), dtype=dtype)
    for i in jnp.arange(lg2):
        H = jnp.vstack([jnp.hstack([H, H]), jnp.hstack([H, -H])])
    H = 2**(-lg2 / 2) * H
    return H

I cannot lower this to a fori_loop or while, because the shape changes every iteration. I then put it in a simple dense layer:

class HadamardTransform(nn.Module):
    n_hadamard: int
        
    @nn.compact
    def __call__(self, X):
        kernel = self.param("kernel", hadamard, (self.n_hadamard, ))
        z = jnp.dot(X, kernel)
        return z

I can initialize this model and my parameters look fine (i.e. its a pytree with a device array):

# fake test data
key = random.PRNGKey(42)
X = random.normal(key, (1, 4096))

# Instantiating model
model = HadamardTransform(4096)
params = model.init(key, X)

However, when I want to do the forward pass with

model.apply(params, X)

I get a concretizationtype error on the for loop of the hadamard function: 'ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.'

I'm not sure why this is happening - I can init the model, so I would expect forward to work. Is jitting the apply function also jitting the hadamard function?

[GJBoth]

It's definitely an issue with initialization. For example, using setup, this doesnt work:

class HadamardTransform(nn.Module):
    n_hadamard: int
        
    def setup(self):
        self.kernel = self.param("kernel", hadamard, (self.n_hadamard, ))
    
    def __call__(self, X):
        z = jnp.dot(X, self.kernel)
        return z

but this workaround I came up with does:

class HadamardTransform(nn.Module):
    n_hadamard: int
        
    def setup(self):
        H = hadamard(_, (self.n_hadamard, ))
        self.kernel = self.param("kernel", lambda _: H)
    
    def __call__(self, X):
        z = jnp.dot(X, self.kernel)
        return z

[jheek]

The error is not super informative here but it comes from JAX so we cannot fix that.
The issue is that you try to use a jax value which is dynamic in a place where you can only use static (python or numpy values). The jnp.arange() expects a constant because otherwise it has a dynamic shape which XLA does not support (yet).

So in general shapes should be static so you should use normal python or numpy ops on shapes (note the jnp -> np):

import numpy as np
lg2 = np.log2(shape[0])
H = jnp.ones((1, ), dtype=dtype)
for i in np.arange(lg2):
  ...

now the range will have constant size and the rest should work.