Get the PyTree structure of flax model without initializing weights

[patil-suraj]

the jax.eval_shape function can be used to get the PyTree structure of flax model params and optimizer state (any JAX function for that matter) without having to actually initialize them.

class MLP(nn.Module):
    def setup(self):
        self.dense1 = nn.Dense(256, use_bias=False)
        self.dense2 = nn.Dense(10, use_bias=False)
    
    def __call__(self, x):
        return self.dense2(nn.relu(self.dense1(x)))

key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (1, 8))
m = MLP()
params_tree = jax.eval_shape(m.init, key, x)
params_tree

this should give

FrozenDict({
    params: {
        dense1: {
            kernel: ShapeDtypeStruct(shape=(8, 256), dtype=float32),
        },
        dense2: {
            kernel: ShapeDtypeStruct(shape=(256, 10), dtype=float32),
        },
    },
})

[marcvanzee]

Thanks for this great tip @patil-suraj jax.eval_shape indeed only computes the shapes, so nothing it actually initialized (0 FLOPS). It also helps you to catch shape errors quickly. In our previous API we had Module.init_by_shape, which did the same, but we didn't port it since it's typically more explicit to just use jax.jit or jax.eval_shape for abstract evaluation. No need to use abstractions if the original is just as short! 😄