bump flax version
# 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](https://github.com/google/jax/releases/tag/jax-v0.2.21), 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](https://github.com/google/flax/blob/v0.3.5/flax/core/scope.py#L753)).  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`.