code: def test_module_list(self):
      rngs = nnx.Rngs(0)
      ls = [
          nnx.Linear(2, 2, rngs=rngs),
          nnx.BatchNorm(2, rngs=rngs),
      ]
  
      state, static = nnx.graph_utils.graph_flatten(ls)
  
      assert state["0"]["kernel"].shape == (2, 2)
      assert state["0"]["bias"].shape == (2,)
      assert state["1"]["scale"].shape == (2,)
      assert state["1"]["bias"].shape == (2,)
      assert state["1"]["mean"].shape == (2,)
      assert state["1"]["var"].shape == (2,)
  
file_path: flax/flax/experimental/nnx/tests/test_graph_utils.py
imports: import jax
  import pytest
  from flax.experimental import nnx
label_desc: Use image node_id 9 for calling the TestGraphUtils obj's underlying member method code with example usage: obj.test_module_list() without return types