Marcus Chiam

Closing after #3434 merged

sure let me take a look!

> > Can specify the system version, Python version, Flax repo commit hash, and output of pip freeze from one of the machines? > > Following up here: python version:...

I wonder if this [`if` condition](https://github.com/google/flax/blob/main/examples/imagenet/train.py#L403) is not evaluating to True. What value are you getting for `steps_per_epoch`?

> On v5e we noticed that training and evaluation seem to work but eventually the code freezes. Here is the output: > > Could this be related to checkpointing? >...

Apologies, seems like PR #3386 (which fixes the `config.num_train_steps` and `config.steps_per_eval` to 1) linked to this issue and automatically closed it once it was submitted.

Ya I think removing the `mesh_sharding` function and just adding the `mesh` to `NamedSharding(mesh, PartitionSpec(...))` would be more clear. > add mesh only on the top-level with `with mesh`. I...

@cgarciae @marcvanzee

Would this work: ``` import flax.linen as nn from jax import numpy as jnp, Array from typing import Callable, Tuple, Any from flax.linen import initializers PRNGKey = Any Shape =...

hi @dasklol, could you clarify what you mean by "never change release tags"? Every time a release is published, the tag gets updated to reflect the new version number.