code: def test_tied_weights(self):
      class Foo(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs) -> None:
              self.bar = nnx.Linear(2, 2, rngs=rngs)
              self.baz = nnx.Linear(2, 2, rngs=rngs)
  
              # tie the weights
              self.baz.variables.kernel = self.bar.variables.kernel
  
      node = Foo(rngs=nnx.Rngs(0))
      state, static = nnx.graph_utils.graph_flatten(node)
  
      assert len(state.flat_state()) == 3  # 2 bias + 1 kernel
  
      node2 = static.merge(state)
  
      assert node2.bar.variables.kernel is node2.baz.variables.kernel
  
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 11 for calling the TestGraphUtils obj's underlying member method code with example usage: obj.test_tied_weights() without return types