code: def test_shared_array(getkey):
      class MyModule(eqx.Module):
          shared: eqx.nn.Shared
  
          def __init__(self):
              embedding = eqx.nn.Embedding(
                  num_embeddings=3, embedding_size=4, key=getkey()
              )
              head = eqx.nn.Linear(4, 3, key=getkey())
              where = lambda pair: pair[1].weight
              get = lambda pair: pair[0].weight
              self.shared = eqx.nn.Shared((embedding, head), where, get)
  
          def __call__(self, token: Int[Array, ""]):
              nonlocal called
              called = True
              embedding, head = self.shared()
              assert embedding.weight is head.weight
              return head(embedding(token))
  
      called = False
      module = MyModule()
      module(jnp.array(0))
      assert called
  
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 1 for calling a global function with example usage: test_shared_array(getkey) and returns: head