Embeddings¶
Subspace relationships between manifolds.
An embedding defines how one manifold sits inside another: it provides maps for injecting points into the ambient space (embed), and for pulling back cotangent vectors — i.e. gradients — from the ambient space to the subspace (pullback). Linear embeddings additionally support projection and translation.
Class Hierarchy¶
Base Embedding Classes¶
- class Embedding[source]¶
-
Defines how a smaller model space sits inside a larger one.
Use an embedding when you want to convert parameters from a constrained model to a more general one (
embed), or transform gradients computed in the general space back to the constrained space (pullback). For example, embedding a diagonal-covariance normal into a full-covariance normal.Mathematically, an injection \(\phi: \mathcal{M} \to \mathcal{N}\) from a submanifold into an ambient space. Subclasses implement
embed; the defaultpullbackis computed via JAX’s VJP.- abstract property sub_man: Sub¶
The submanifold.
- abstract property amb_man: Ambient¶
The ambient manifold.
- abstractmethod embed(coords: Array) Array[source]¶
Embed a point from the submanifold into the ambient manifold.
- pullback(at_point: Array, cotangent_vector: Array) Array[source]¶
Transform a gradient from the ambient space to the subspace.
This is what makes constrained optimization work: compute a gradient in the full model space, then pull it back to get the gradient in the restricted space.
Mathematically, given the embedding \(\phi: \mathcal M \to \mathcal N\), at \(w = \phi(v)\) the pullback is the mapping
from the cotangent space on \(\mathcal N\) at \(w\) to the cotangent space on \(\mathcal M\) at \(v\). Where \(\omega \in T^*_w \mathcal N\) is a cotangent vector, the pullback is given by\[ \phi^*: T^*_w \mathcal N \to T^*_v \mathcal M \]where \(\frac{\partial \phi}{\partial v}\) is the Jacobian of \(\phi\) at \(v\).\[ \phi^*(\omega) = \omega \cdot \frac{\partial \phi}{\partial v}, \]- Parameters:
at_point – Point on the submanifold
cotangent_vector – Cotangent vector on the ambient manifold
- Returns:
Pulled back cotangent vector on the submanifold
- class LinearEmbedding[source]¶
Bases:
Embedding,ABC,GenericAn embedding that additionally supports projection and translation.
Use a linear embedding when you need to move coordinates in both directions:
embedto go from the smaller space to the larger, andprojectto extract the relevant components back out. Translation lets you update just the subspace components of an ambient point without touching the rest.Mathematically, a linear embedding \(\phi: \mathcal{M} \to \mathcal{N}\) comes with:
A projection \(\pi: \mathcal{N} \to \mathcal{M}\) satisfying \(\pi \circ \phi = \text{id}_{\mathcal{M}}\), and
a translation \(\tau(p, q) = p + \phi(q)\).
For linear embeddings the pullback is location-independent and reduces to projection.
- abstractmethod project(coords: Array) Array[source]¶
Project from the ambient space to the subspace.
Composed Embeddings¶
- class ComposedEmbedding(sub_emb: Embedding[Sub, Mid], mid_emb: Embedding[Mid, Ambient])[source]¶
Bases:
Embedding,ABC,GenericChain two embeddings when a model space is nested two levels deep.
For example, a diagonal normal embeds into a full normal, which in turn embeds into a harmonium — composing these lets you pull gradients all the way back to the diagonal parameters.
Mathematically, \(\phi = \phi_2 \circ \phi_1: \mathcal{M} \to \mathcal{L} \to \mathcal{N}\), with pullback in reverse order: \(\phi^*(\omega) = \phi_1^*(\phi_2^*(\omega))\).
- property amb_man: Ambient¶
The ambient manifold.
- property sub_man: Sub¶
The submanifold.
- pullback(at_point: Array, cotangent_vector: Array) Array[source]¶
Transform a gradient from the ambient space to the subspace.
This is what makes constrained optimization work: compute a gradient in the full model space, then pull it back to get the gradient in the restricted space.
Mathematically, given the embedding \(\phi: \mathcal M \to \mathcal N\), at \(w = \phi(v)\) the pullback is the mapping
from the cotangent space on \(\mathcal N\) at \(w\) to the cotangent space on \(\mathcal M\) at \(v\). Where \(\omega \in T^*_w \mathcal N\) is a cotangent vector, the pullback is given by\[ \phi^*: T^*_w \mathcal N \to T^*_v \mathcal M \]where \(\frac{\partial \phi}{\partial v}\) is the Jacobian of \(\phi\) at \(v\).\[ \phi^*(\omega) = \omega \cdot \frac{\partial \phi}{\partial v}, \]- Parameters:
at_point – Point on the submanifold
cotangent_vector – Cotangent vector on the ambient manifold
- Returns:
Pulled back cotangent vector on the submanifold
- property mid_man: Mid¶
The intermediate manifold.
- class LinearComposedEmbedding(sub_emb: LinearEmbedding[Sub, Mid], mid_emb: LinearEmbedding[Mid, Ambient])[source]¶
Bases:
ComposedEmbedding,LinearEmbedding,GenericLinear version of
ComposedEmbedding, preserving projection and translation through the chain.Mathematically, \(\pi = \pi_1 \circ \pi_2\) and \(\tau(p, q) = \tau_2(p, \tau_1(0_{\mathcal{L}}, q))\).
- sub_emb: LinearEmbedding¶
- mid_emb: LinearEmbedding¶
Algebraic Embeddings¶
- class IdentityEmbedding(man: M)[source]¶
Bases:
LinearEmbedding,GenericThe trivial case where sub and ambient are the same manifold.
Used as the default embedding when a linear map operates on the full space without restriction (see
AmbientMap).- man: M¶
- property amb_man: M¶
The ambient manifold.
- property sub_man: M¶
The submanifold.
- class TupleEmbedding[source]¶
Bases:
LinearEmbedding,ABC,GenericEmbeds a single component of a
Tuplemanifold, used to isolate one factor of a product space.For example, in a harmonium (which is a
Tripleof observable, interaction, and latent), aTupleEmbeddinglets you extract or update just the observable parameters. Projection extracts thetup_idx-th component; embedding sets that component and zeros the rest; translation adds to that component only.