bump flax version

[patil-suraj]

What does this PR do?

jax v0.2.21 introduced a slight breaking change, jnp.ndarray now only returns True for jax arrays and not numpy arrays. so for standard numpy array x, isinstance(x, jnp.ndarray) will
now return False . (cf here, the last point under breaking changes)

And the jax.random.split method returns numpy array instead of jax. We use this method during random init of the model to create the rngs dict.

Flax has a check for valid rngs where it expects jnp.ndarray , but since split returns np array and jax now does not treat np arrays as jnp.ndarray the check returns False and raises an error.

This has been fixed in flax==0.3.5 (cf here). So all of our models fail with jax>=0.2.21 and flax<0.3.5.

To fix this, this PR bumps the flax version to 0.3.5.

[LysandreJik]

LGTM