[JAX] Replace all jax.tree_* calls with jax.tree_util.tree_*
# What does this PR do?

Newer versions of JAX moved all tree utility methods to `jax.tree_util` and emit warnings for using the old locations under `jax`. This PR refactors all `jax.tree_*` calls to `jax.tree_util.tree_*`, thus disabling the warnings with no changes to the code functionality.

```python
import jax

inputs = jax.numpy.ones(10)

# Define identity function with tree_map -> FutureWarning 
jax.tree_map(lambda x: x, inputs)

# Define identity function with tree_util.tree_map -> no warning!
jax.tree_util.tree_map(lambda x: x, inputs)
```

```
FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
```

All `jax.tree_*` calls in Flax were updated in this PR: https://github.com/google/flax/pull/2325. However, the `FutureWarning` message will still be raised if using newer versions of JAX and older versions of Flax. 

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
_The documentation is not available anymore as the PR was closed or merged._<|||||>Sorry for being so late here! Let's merge it 