From a1b7b985936395ecde95cecd166c25c11591c814 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Tue, 23 Sep 2025 05:54:28 +0000 Subject: [PATCH 1/3] Revert "use Pytree as base type in nnx_wrappers" This reverts commit 8840296efac95908c722695929a19491138dc9f0. Note: This commit should be reverted when flax version is upgraded to anything >= 0.11.0. Right now, JAX 0.6.0 does not support any flax version >= 0.10.7, hence this commit is required as flax started exposing PyTree after 0.11.0. --- src/MaxText/layers/nnx_wrappers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/MaxText/layers/nnx_wrappers.py b/src/MaxText/layers/nnx_wrappers.py index 39baeaf985..a0e3fba9b0 100644 --- a/src/MaxText/layers/nnx_wrappers.py +++ b/src/MaxText/layers/nnx_wrappers.py @@ -29,7 +29,7 @@ from flax.nnx import variablelib from flax.nnx.bridge import module as bdg_module from flax.nnx.module import Module -from flax.nnx import Pytree +from flax.nnx import Object from flax.nnx.rnglib import Rngs import jax from jax import tree_util as jtu @@ -140,7 +140,7 @@ def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict: def _set_initializing(module: Module, initializing: bool): for _, value in graph.iter_graph(module): - if isinstance(value, Pytree): + if isinstance(value, Object): value._object__state._initializing = initializing # pylint: disable=protected-access From 70c715b14ddf4378339c42fbfd7203ff12ec7cdf Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 23 Sep 2025 01:03:13 -0500 Subject: [PATCH 2/3] Added a TODO to convert the import --- src/MaxText/layers/nnx_wrappers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/layers/nnx_wrappers.py b/src/MaxText/layers/nnx_wrappers.py index a0e3fba9b0..8ecc245aa3 100644 --- a/src/MaxText/layers/nnx_wrappers.py +++ b/src/MaxText/layers/nnx_wrappers.py @@ -29,7 +29,7 @@ from flax.nnx import variablelib from flax.nnx.bridge import module as bdg_module from flax.nnx.module import Module -from flax.nnx import Object +from flax.nnx import Object #TODO(gulsumgudukbay) convert this import to from flax.nnx import Pytree after ROCm/jax supports flax >= 0.11.0 from flax.nnx.rnglib import Rngs import jax from jax import tree_util as jtu From af61b464ce668866174e65c0e8ad354c332da063 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 23 Sep 2025 01:06:42 -0500 Subject: [PATCH 3/3] Update nnx_wrappers.py --- src/MaxText/layers/nnx_wrappers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/layers/nnx_wrappers.py b/src/MaxText/layers/nnx_wrappers.py index 8ecc245aa3..1b74b2c531 100644 --- a/src/MaxText/layers/nnx_wrappers.py +++ b/src/MaxText/layers/nnx_wrappers.py @@ -29,7 +29,7 @@ from flax.nnx import variablelib from flax.nnx.bridge import module as bdg_module from flax.nnx.module import Module -from flax.nnx import Object #TODO(gulsumgudukbay) convert this import to from flax.nnx import Pytree after ROCm/jax supports flax >= 0.11.0 +from flax.nnx import Object # TODO(gulsumgudukbay) convert this import to from flax.nnx import Pytree after ROCm/jax supports flax >= 0.11.0 from flax.nnx.rnglib import Rngs import jax from jax import tree_util as jtu