class_docstr: Used to tie together multiple nodes across a PyTree.
  
  Note that Equinox modules are Py**Trees** -- so the same layer, appearing in two
  difference parts of the tree, will be treated as two copies of this layer. For
  example,
  ```python
  class SubModel(eqx.Module):
      linear: eqx.nn.Linear
  
  class Model(eqx.Module):
      linear: eqx.nn.Linear
      submodel: SubModel
  
      def __init__(self):
          linear = eqx.nn.Linear(...)
          self.linear = linear
          self.submodel = SubModel(linear)
  ```
  is used to declare `model.linear` and `model.submodel.linear` as two separate
  layers. They will start with the same initial parameter values, and then update
  independently during training.
  
  For when we really do want to share layers or weights across different parts of a
  model, then `eqx.nn.Shared` exists as a way to easily express this in the PyTree
  paradigm.
  
  !!! Example
  
      It is common in many language models to have an initial embedding matrix at the
      start, and then to reuse this as the weight of the final linear transformation.
  
      ```python
      import equinox as eqx
      import jax.numpy as jnp
      from jaxtyping import Array, Int
  
      class LanguageModel(eqx.Module):
          shared: eqx.nn.Shared
  
          def __init__(self):
              embedding = eqx.nn.Embedding(...)
              linear = eqx.nn.Linear(...)
              # These two weights will now be tied together.
              where = lambda embed_and_lin: embed_and_lin[1].weight
              get = lambda embed_and_lin: embed_and_lin[0].weight
              self.shared = eqx.nn.Shared((embedding, linear), where, get)
  
          def __call__(self, tokens: Int[Array, "sequence"]):
              # Expand back out so we can evaluate these layers.
              embedding, linear = self.shared()
              assert embedding.weight is linear.weight  # same parameter!
              # Now go ahead and evaluate your language model.
              values = jax.vmap(embedding)(tokens)
              ...  # other layers, probably
              return jax.vmap(linear)(values)
      ```
  
      _(Side note: you will sometimes see some authors referring to transposing
      the embedding matrix prior to the final linear layer. This is because some
      other libraries store the weight matrices of linear layers the other way
      around. If that had been necessary here then we could have done it with
      `get = lambda embed_and_lin: jnp.transpose(embed_and_lin[0].weight)`.)_
class_docstr_tok: ["Used","to","tie","together","multiple","nodes","across","a","PyTree",".","Note","that","Equinox","modules","are","Py","*","*","Trees","*","*","--","so","the","same","layer",",","appearing","in","two","difference","parts","of","the","tree",",","will","be","treated","as","two","copies","of","this","layer",".","For","example",",","``","`","python","class","SubModel","(","eqx.Module",")",":","linear",":","eqx.nn.Linear","class","Model","(","eqx.Module",")",":","linear",":","eqx.nn.Linear","submodel",":","SubModel","def","__init__","(","self",")",":","linear","=","eqx.nn.Linear","(","...",")","self.linear","=","linear","self.submodel","=","SubModel","(","linear",")","``","`","is","used","to","declare","`","model.linear","`","and","`","model.submodel.linear","`","as","two","separate","layers",".","They","will","start","with","the","same","initial","parameter","values",",","and","then","update","independently","during","training",".","For","when","we","really","do","want","to","share","layers","or","weights","across","different","parts","of","a","model",",","then","`","eqx.nn.Shared","`","exists","as","a","way","to","easily","express","this","in","the","PyTree","paradigm",".","!","!","!","Example","It","is","common","in","many","language","models","to","have","an","initial","embedding","matrix","at","the","start",",","and","then","to","reuse","this","as","the","weight","of","the","final","linear","transformation",".","``","`","python","import","equinox","as","eqx","import","jax.numpy","as","jnp","from","jaxtyping","import","Array",",","Int","class","LanguageModel","(","eqx.Module",")",":","shared",":","eqx.nn.Shared","def","__init__","(","self",")",":","embedding","=","eqx.nn.Embedding","(","...",")","linear","=","eqx.nn.Linear","(","...",")","#","These","two","weights","will","now","be","tied","together",".","where","=","lambda","embed_and_lin",":","embed_and_lin","[","1","]",".weight","get","=","lambda","embed_and_lin",":","embed_and_lin","[","0","]",".weight","self.shared","=","eqx.nn.Shared","(","(","embedding",",","linear",")",",","where",",","get",")","def","__call__","(","self",",","tokens",":","Int","[","Array",",","``","sequence","''","]",")",":","#","Expand","back","out","so","we","can","evaluate","these","layers",".","embedding",",","linear","=","self.shared","(",")","assert","embedding.weight","is","linear.weight","#","same","parameter","!","#","Now","go","ahead","and","evaluate","your","language","model",".","values","=","jax.vmap","(","embedding",")","(","tokens",")","...","#","other","layers",",","probably","return","jax.vmap","(","linear",")","(","values",")","``","`","_","(","Side","note",":","you","will","sometimes","see","some","authors","referring","to","transposing","the","embedding","matrix","prior","to","the","final","linear","layer",".","This","is","because","some","other","libraries","store","the","weight","matrices","of","linear","layers","the","other","way","around",".","If","that","had","been","necessary","here","then","we","could","have","done","it","with","`","get","=","lambda","embed_and_lin",":","jnp.transpose","(","embed_and_lin","[","0","]",".weight",")","`",".",")","_"]
docstr: **Arguments:**
  
  None.
  
  **Returns:**
  
  A PyTree of the same structure as the original `pytree`, with `get(pytree)` in
  the place of the nodes at `where(pytree)`.
docstr_tok: ["*","*","Arguments",":","*","*","None",".","*","*","Returns",":","*","*","A","PyTree","of","the","same","structure","as","the","original","`","pytree","`",",","with","`","get","(","pytree",")","`","in","the","place","of","the","nodes","at","`","where","(","pytree",")","`","."]
code: def __call__(self):
      """**Arguments:**
  
      None.
  
      **Returns:**
  
      A PyTree of the same structure as the original `pytree`, with `get(pytree)` in
      the place of the nodes at `where(pytree)`.
      """
      return tree_at(self.where, self.pytree, self.get(self.pytree))
  
file_path: equinox/equinox/nn/_shared.py
imports: from collections.abc import Callable
  from jaxtyping import PyTree
  from .._eval_shape import filter_eval_shape
  from .._module import Module
  from .._tree import tree_at, tree_equal
label_desc: Use image node_id 2 for calling the Shared obj's underlying member method code with example usage: obj.__call__() and returns: tree_at