Why is Flax Linear not identical to matrix multiplication?

[thijs-vanweezel]

Due to the novelty of Flax, NNX, and JAX, there’s not a lot of resources available. I’m running into the following peculiarity:

x = jnp.random.normal((1,512), key=KEY)
layer = nnx.Linear(512, 512, rngs=nnx.Rngs(KEY))
y1 = layer(x)
y2 = [email protected]() + layer.bias
print(y1==y2) # returns all False

My understanding is that matrix multiplication should be identical to a linear / fully connected layer. The discrepancy demonstrated here hinders the inspection of certain behavior (and the implementation of invertible dense layers using nnx.tensorsolve).

Does anyone know what causes this discrepancy?

[thijs-vanweezel]

The answer has been provided on Stackoverflow.

The matmul should be transposed;

y = x.squeeze()@layer.kernel + layer.bias

So, to invert a nnx.Linear operation using nnx.tensorsolve:

solve_batched = jax.vmap(jnp.linalg.tensorsolve)
solve_batched(
    a=jnp.broadcast_to(
        layer.kernel.value.T, # Note the transposition
        (y.shape[0],*layer.kernel.value.shape)), 
    b=y - layer.bias)