Standard MLP (equinox.nn.MLP) does not work with `apply_updates` function about equinox

[patrick-kidger]

As such I'm going to close #109 -- but I appreciate the willingness to get stuck in.

[patrick-kidger]

This is actually intended behaviour.

What's going on here is that the PyTree structure of the model may include non-JAX-arrays in its leaves, and indeed may contain arbitrary Python objects. Their interactions with other operations -- in this case you're trying to add two PyTrees together -- are resolved through filtering.

With respect to applying updates, this is covered in the FAQ. (See also the docs and Section 4 of the paper.)

It's important that we have pretty much everything be part of the PyTree structure. For example, someone may wish to define and use

class LearntActivation(eqx.Module):
    param: jnp.ndarray

    def __call__(self, x):
        return jax.nn.relu(param * x)

eqx.nn.MLP(..., activation=LearntActivation(...))

[callumtilbury]

Thanks for the quick response, @patrick-kidger ! I was trying to use the filtering approach, but I was facing some problems. I had a deeper dive, and I realise now where I went wrong. I appreciate the help!