code: def test_tied_weights_example(self):
      class LinearTranspose(nnx.Module):
          def __init__(
              self, dout: int, din: int, *, rngs: nnx.Rngs
          ) -> None:
              self.kernel = nnx.Param(
                  nnx.initializers.lecun_normal()(rngs(), (dout, din))
              )
  
          def __call__(self, x):
              return x @ self.kernel.T
  
      class Encoder(nnx.Module):
          def __init__(self, *, rngs: nnx.Rngs) -> None:
              self.embed = nnx.Embed(10, 2, rngs=rngs)
              ...
              self.linear_out = LinearTranspose(10, 2, rngs=rngs)
  
              # tie the weights
              self.linear_out.variables.kernel = (
                  self.embed.variables.embedding
              )
  
          def __call__(self, x):
              x = self.embed(x)
              ...
              return self.linear_out(x)
  
      model = Encoder(rngs=nnx.Rngs(0))
      state, static = model.split()
  
      assert len(state.flat_state()) == 1
  
      x = jax.random.randint(jax.random.key(0), (2,), 0, 10)
      y = model(x)
  
      assert y.shape == (2, 10)
  
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 12 for calling the TestGraphUtils obj's underlying member method code with example usage: obj.test_tied_weights_example() and returns: unknown, self