Make lists return true for is_data

Author: samanklesaria

examples/gemma/transformer.py

@@ -537,7 +537,7 @@ def __init__(
         dtype=config.dtype,
         rngs=rngs,
     )
-    self.layers = nnx.data([
+    self.layers = [
         modules.Block(
           config=config,
           attn_type=attn_type,
@@ -547,7 +547,7 @@ def __init__(
         for _, attn_type in zip(
             range(config.num_layers), config.attention_types
         )
-    ])
+    ]
     self.final_norm = layers.RMSNorm(
       config.embed_dim,
       scale_init=modules.maybe_with_partitioning(

examples/nnx_toy_examples/hijax_demo.py

@@ -141,7 +141,7 @@ def create_block(rngs, /):
       # self.blocks = nnx.stateful(create_block(rngs.fork(split=num_blocks)))
       self.blocks = create_block(rngs.fork(split=num_blocks))
     else:
-      self.blocks = nnx.data([Block(dhidden, dhidden, rngs=rngs) for i in range(num_blocks)])
+      self.blocks = [Block(dhidden, dhidden, rngs=rngs) for i in range(num_blocks)]
 
   def __call__(self, x: jax.Array, *, rngs: nnx.Rngs | None = None):
     self.count[...] += 1

flax/nnx/helpers.py

@@ -25,7 +25,6 @@
 from flax.nnx.module import Module
 from flax.nnx.proxy_caller import ApplyCaller
 from flax.nnx.rnglib import Rngs
-from flax.nnx.pytreelib import data
 from flax.nnx.statelib import State
 from flax.training.train_state import struct
 
@@ -120,7 +119,7 @@ def __init__(self, *fns: tp.Callable[..., tp.Any]):
     Args:
       *fns: A sequence of callables to apply.
     """
-    self.layers = data(list(fns))
+    self.layers = list(fns)
 
   def __call__(self, *args, rngs: tp.Optional[Rngs] = None, **kwargs) -> tp.Any:
     if len(self.layers) == 0:

flax/nnx/pytreelib.py

@@ -173,7 +173,7 @@ def is_data(value: tp.Any, /) -> bool:
     ... # ------ STATIC ------------
     >>> assert not nnx.is_data( 'hello' )                       # strings, arbitrary objects
     >>> assert not nnx.is_data( 42 )                            # int, float, bool, complex, etc.
-    >>> assert not nnx.is_data( [1, 2.0, 3j, jnp.array(1)] )    # list, dict, tuple, pytrees
+    >>> assert nnx.is_data( [1, 2.0, 3j, jnp.array(1)] )    # list, dict, tuple, pytrees
 
 
   Args:
@@ -183,7 +183,8 @@ def is_data(value: tp.Any, /) -> bool:
     A string representing the attribute status.
   """
   return (
-    graph.is_node_leaf(value)
+    isinstance(value, list)
+    or graph.is_node_leaf(value)
     or graph.is_graph_node(value)
     or type(value) in DATA_REGISTRY
   )

tests/nnx/graph_utils_test.py

@@ -24,8 +24,6 @@
 import jax.numpy as jnp
 
 
-
-
 class StatefulLinear(nnx.Module):
   def __init__(self, din, dout, rngs):
     self.w = nnx.Param(jax.random.uniform(rngs(), (din, dout)))
@@ -1007,17 +1005,6 @@ def f(m):
 
     self.assertEqual(m.a, 2)
 
-  def test_data_after_init(self):
-    test = self
-    class Foo(nnx.Module):
-      def __init__(self):
-        self.ls = []
-        self.ls.append(jnp.array(1))
-
-    with self.assertRaisesRegex(
-        ValueError, 'Found unexpected Arrays on value of type'
-    ):
-      m = Foo()
 
   def test_update_dict(self):
     node = {

tests/nnx/module_test.py

@@ -28,7 +28,6 @@
 
 A = TypeVar('A')
 
-
 class PytreeTest(absltest.TestCase):
   def test_pytree(self):
     class Foo(nnx.Pytree):
@@ -1001,13 +1000,10 @@ def __call__(self, x, *, rngs: nnx.Rngs):
 
     assert isinstance(y, jax.Array)
 
-  def test_modules_iterator(self):
-    class Foo(nnx.Module):
-      def __init__(self, *, rngs: nnx.Rngs):
-        self.submodules = nnx.data([
+        self.submodules = [
           {'a': nnx.Linear(1, 1, rngs=rngs)},
           {'b': nnx.Conv(1, 1, 1, rngs=rngs)},
-        ])
+        ]
         self.linear = nnx.Linear(1, 1, rngs=rngs)
         self.dropout = nnx.Dropout(0.5, rngs=rngs)
 
@@ -1030,10 +1026,10 @@ def __init__(self, *, rngs: nnx.Rngs):
   def test_children_modules_iterator(self):
     class Foo(nnx.Module):
       def __init__(self, *, rngs: nnx.Rngs):
-        self.submodules = nnx.data([
+        self.submodules = [
           {'a': nnx.Linear(1, 1, rngs=rngs)},
           {'b': nnx.Conv(1, 1, 1, rngs=rngs)},
-        ])
+        ]
         self.linear = nnx.Linear(1, 1, rngs=rngs)
         self.dropout = nnx.Dropout(0.5, rngs=rngs)
 

tests/nnx/nn/attention_test.py

@@ -51,10 +51,9 @@ class Model(nnx.Module):
       attention_kwargs: dict
 
       def __init__(self, attention_kwargs, rng):
-        self.attention_layers = nnx.data([
-          nnx.MultiHeadAttention(**attention_kwargs, rngs=rng) for i in range(3)
-        ])
-
+              self.attention_layers = [
+                nnx.MultiHeadAttention(**attention_kwargs, rngs=rng) for i in range(3)
+              ]
       def __call__(self, x, sow_weights=False):
         x = self.attention_layers[0](x, sow_weights=sow_weights)
         x = self.attention_layers[1](x)