Added support for sharing layers

[patrick-kidger]

Added support for sharing layers between different parts of a model.

This has the API eqx.nn.Shared(pytree, where, get). At init time, then everything in where(pytree) is deleted -- these as the duplicate nodes we want to remove. At call time, everything in get(pytree) is inserted in their place -- these as the nodes we're keeping, and in doing so we restore a layer that we can evaluate.

Example usage like so:

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(...)
        where = lambda pair: pair[1].weight
        get = lambda pair: pair[0].weight
        self.shared = eqx.nn.Shared((embedding, linear), where, get)

    def __call__(self, tokens: Int[Array, "sequence"]):
        embedding, linear = self.shared()
        values = jax.vmap(embedding)(tokens)
        ...  # the rest of your language model goes here!
        return jax.vmap(linear)(values)

If we need to perform any kidn of transform, e.g. transposition, then that can be done just by placing it in get, e.g. get = lambda pair: jnp.transpose(pair[0].weight).