fix main

Author: cgarciae

.github/workflows/flax_test.yml

@@ -127,20 +127,20 @@ jobs:
           uv pip install -U git+https://github.com/google-deepmind/dm-haiku.git
           # temporary: install jax nightly
           uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
-          uv run tests/run_all_tests.sh --only-doctest
+          uv run --no-sync tests/run_all_tests.sh --only-doctest --use-venv
         elif [[ "${{ matrix.test-type }}" == "pytest" ]]; then
           uv pip install -U tensorflow-datasets
           # temporary: install jax nightly
           uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
-          uv run tests/run_all_tests.sh --only-pytest
+          uv run --no-sync tests/run_all_tests.sh --only-pytest --use-venv
         elif [[ "${{ matrix.test-type }}" == "pytype" ]]; then
-        # temporary: install jax nightly
+          # temporary: install jax nightly
           uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
-          uv run tests/run_all_tests.sh --only-pytype
+          uv run --no-sync tests/run_all_tests.sh --only-pytype --use-venv
         elif [[ "${{ matrix.test-type }}" == "mypy" ]]; then
           # temporary: install jax nightly
           uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
-          uv run tests/run_all_tests.sh --only-mypy
+          uv run --no-sync tests/run_all_tests.sh --only-mypy --use-venv
         else
           echo "Unknown test type: ${{ matrix.test-type }}"
           exit 1

examples/nnx_toy_examples/hijax_demo.py

@@ -21,7 +21,7 @@
 from flax import nnx
 
 # ## Data
-# We create a simple dataset of points sampled from a parabola with some noise.
+# We create a simple dataset of points sampled from a parabola with some  noise.
 X = np.linspace(-jnp.pi, jnp.pi, 100)[:, None]
 Y = 0.8 * jnp.sin(X) + 0.1 + np.random.normal(0, 0.1, size=X.shape)
 

tests/nnx/spmd_test.py

@@ -197,7 +197,11 @@ def __call__(self, x: jax.Array):
   def test_eager_sharding_context(self, use_eager_sharding):
     rngs = nnx.Rngs(0)
     with nnx.use_eager_sharding(use_eager_sharding):
-      mesh = jax.make_mesh(((2, 2)), ("data", "model"))
+      mesh = jax.make_mesh(
+        (2, 2),
+        ('data', 'model'),
+        axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto),
+      )
       with jax.set_mesh(mesh):
         w = nnx.Param(
           rngs.lecun_normal()((4, 8)),

tests/run_all_tests.sh

@@ -65,7 +65,6 @@ echo "jax: $(python -c 'import jax; print(jax.__version__)')"
 echo "flax: $(python -c 'import flax; print(flax.__version__)')"
 echo "=========================="
 echo ""
-
 sh $(dirname "$0")/download_dataset_metadata.sh || exit
 
 # Instead of using set -e, we have a manual error trap that