Understanding "Interoperate with purely functional JAX code"

[sh0416]

Hi, I'm new in Flax.
I followed up the tutorial in JAX documentation and I made myself not to use stateful function such as class. link

However, when I came into the Flax, the linen module starts with class.
I think something went wrong and search out how to use Flax with pure functional.
The link provides a direction for someone who wants to stick with purely functional code, but I couldn't understand the phrase "Using Linen modules inside otherwise purely functional code".
If someone has an example or clearly understands what it means, could you provide a more detailed explanation for this?

Thanks,

[marcvanzee]

I followed up the tutorial in JAX documentation and I made myself not to use stateful function such as class. link

The solution in the tutorial you linked to actually does use a class, but that class doesn't maintain a state. The reason is that stateful classes often don't work well under JAX transformations (such as jit in the tutorial). The tutorial gives a nice example as well.

In Linen, exactly the same principle is applied: Modules are indeed classes, but they are stateless. You explicitly pass the state (params, rngs) when you initialize or apply a Module, and the new state is returned. I think it is useful to take a look at the Linen Introduction Colab, which explains these ideas in more detail.

[sh0416]

@marcvanzee Thanks for the reply, but I still confused about the definition of state.

I didn't clarify the difference between the code in jax and the one in flax, but I feel like that the module in flax has a state. Therefore, I just wonder the criterion of state in functional programming. Let's think about this line.

model = nn.Dense(features=3)

I suspect that the model has state. The content of state might store the information that the features is 3. However, I also think that the flax developer intend to the features is immutable so that the model pretends as stateless function. Is it right???

If my interpretation is wrong, please correct it.
Thanks in advance.

[marcvanzee]

Sorry my explanation was a bit confusing. Indeed, when you instantiate a module, this instance has a state. But this state is immutable: it only contains state variables that are set during initialization of the class (e.g., the number of features) and they don't change when calling init or apply. Most importantly, the stateful part of your ML model (the parameters, the RNGs) are not stored in the Module, but they are passed in explicitly when calling init or apply, and the new state is returned.

[sh0416]

@marcvanzee Yeap, I clearly understand your words. Thanks :)