code: def test_shared_node(getkey):
      class MyModule(eqx.Module):
          shared: eqx.nn.Shared
  
          def __init__(self):
              attention = eqx.nn.MultiheadAttention(
                  num_heads=3, query_size=12, key=getkey()
              )
              my_proj = eqx.nn.Linear(
                  12, 12, use_bias=False, key=getkey()
              )
              where = lambda pair: pair[1].key_proj
              get = lambda pair: pair[0]
              self.shared = eqx.nn.Shared(
                  (my_proj, attention), where, get
              )
  
          def __call__(self, x: Float[Array, "seq 12"]):
              nonlocal called
              called = True
              my_proj, attention = self.shared()
              eq = eqx.tree_equal(my_proj, attention.key_proj)
              x = attention(x, x, x)
              out = jax.vmap(my_proj)(x)
              return out, eq
  
      called = False
      module = MyModule()
      x = jr.normal(getkey(), (5, 12))
  
      @eqx.filter_jit
      @eqx.filter_grad(has_aux=True)
      def f(module, x):
          out, eq = module(x)
          return jnp.sum(out), eq
  
      d_module, eq = f(module, x)
      assert called
      assert eq
      module = eqx.apply_updates(module, d_module)
      d_module, eq = f(module, x)
      assert eq
      module = eqx.apply_updates(module, d_module)
  
file_path: equinox/tests/test_shared.py
imports: import equinox
  import jax
  import jax.numpy
  import jax.random
  import pytest
  from jaxtyping import Array, Float, Int
label_desc: Use image node_id 2 for calling a global function with example usage: test_shared_node(getkey) and returns: jnp,  eq, out,  eq