jax-orthogonal-net

name	jax-orthogonal-net
description	Use this skill when constructing, applying, training, saving, loading, or documenting a network-predicted orthogonal transformation on vectors in JAX plus Equinox, where a small residual MLP predicts the parameters of an orthogonal matrix from an input, using the matrix exponential, the Cayley transform, or a QR factorization.

JAX Orthogonal Net

This skill builds an input-dependent orthogonal transformation. A small residual MLP predicts the parameters of an orthogonal matrix Q(x) from an input x, and the matrix is then applied to a vector v. Because Q depends on x, the transformation is input-dependent and more expressive than a fixed orthogonal matrix. The skill is self-contained and depends only on jax, equinox, jaxtyping, optax, and pytest.

The three parameterizations are ported from local_disentanglement/networks/orthogonal_net.py in the local_coordinates repository. The standalone learnable matrix without a network is a separate skill, jax-orthogonal-matrix. This skill covers the network-predicted version.

• construct.py builds all three transforms and checks orthogonality.
• transform.py applies each transform forward and inverse and checks the roundtrip.
• train.py fits the network to a known input-dependent rotation field.
• save_load.py saves, lists, and reloads a model with a parity assertion.
• tests/ holds orthogonality, determinant, invertibility, and gradient tests.

Each skill ships its own isolated environment. Sync it once, then run the examples or tests.

uv sync --directory ~/.claude/skills/jax-orthogonal-net
uv run --directory ~/.claude/skills/jax-orthogonal-net python construct.py
uv run --directory ~/.claude/skills/jax-orthogonal-net python transform.py
uv run --directory ~/.claude/skills/jax-orthogonal-net python train.py
uv run --directory ~/.claude/skills/jax-orthogonal-net python save_load.py
uv run --directory ~/.claude/skills/jax-orthogonal-net pytest

Module Map

The package is jax_orthogonal_net.

• jax_orthogonal_net.mlp holds sinusoidal_embedding, ResidualBlock, and ResidualMLP. The MLP predicts the parameters of the orthogonal matrix.
• jax_orthogonal_net.orthogonal_net holds skew_from_params and the three classes VectorOrthogonalMatrixExp, VectorOrthogonalCayley, VectorOrthogonalQR.
• The package root re-exports the three classes, skew_from_params, ResidualMLP, and sinusoidal_embedding.

The Predicting Network

ResidualMLP is an input projection, a stack of n_blocks residual blocks, and an output projection to the parameter count. Each residual block is Linear, gelu, Linear with a skip connection. The block width is hidden_size. With use_time=True the network also takes a scalar sigma and concatenates a sinusoidal_embedding of it onto the input before the first projection.

The Three Parameterizations

Every class has the same interface. The constructor takes key and dim, plus the keyword arguments hidden_size, n_blocks, use_time, and embedding_size. The exp and Cayley classes also take output_scale. The QR class takes use_sign_correction.

The interface keeps a two-input split. get_matrix(x, sigma) builds Q(x), the call returns Q(x) @ v, and inverse returns Q(x).T @ Uv. The vector v is separate from the conditioning input x. The split supports pushing a tangent vector v through a transform defined at the point x. The call and inverse accept a key argument that is ignored, kept for interface parity with the source code. The config property returns the constructor keyword arguments for save_load.

Class	Group reached	Continuity in the parameters	Cost of one matrix
VectorOrthogonalMatrixExp	SO(n)	smooth everywhere	a matrix exponential
VectorOrthogonalCayley	SO(n) without rotations that have a -1 eigenvalue	smooth everywhere	one linear solve
VectorOrthogonalQR	O(n) with sign correction, one fixed component without it	discontinuous where a diagonal entry of R crosses zero	a QR factorization

VectorOrthogonalMatrixExp predicts the upper triangle of a skew-symmetric matrix A and returns expm(A). VectorOrthogonalCayley predicts the same skew upper triangle and returns solve(I + A, I - A). For both, the prediction is multiplied by output_scale, which defaults to 0.1, so Q(x) is close to the identity at initialization. VectorOrthogonalQR predicts a dense matrix and returns its QR factor. With use_sign_correction=True, the default, Q(x) reaches the full O(n). The sign correction is discontinuous where a diagonal entry of R crosses zero. Use the exp or Cayley class when a continuous parameterization matters.

Construct

Pass a key and a matrix size dim. The constructors expose hidden_size, n_blocks, use_time, and embedding_size.

import jax
from jax_orthogonal_net import VectorOrthogonalMatrixExp

key = jax.random.PRNGKey(0)
model = VectorOrthogonalMatrixExp(key, dim=6, hidden_size=64, n_blocks=3)

x = jax.random.normal(key, (6,))
Q = model.get_matrix(x)         # (6, 6), orthogonal, depends on x

See construct.py for a runnable version that checks Q(x) is orthogonal across a batch of inputs, that it genuinely depends on x, and that the time-conditioned variant responds to sigma.

Transform

The call applies Q(x) to a vector v. inverse applies Q(x).T, which recovers v because Q(x) is orthogonal.

v = jax.random.normal(key, (6,))
u = model(x, v)                 # Q(x) @ v
v_back = model.inverse(x, u)    # Q(x).T @ u, equals v

When use_time=True, pass sigma to every call.

timed = VectorOrthogonalMatrixExp(key, dim=6, use_time=True)
u = timed(x, v, sigma=jnp.array(0.5))
v_back = timed.inverse(x, u, sigma=jnp.array(0.5))

Each module acts on a single pair (x, v). Batch with jax.vmap(model) over stacks of inputs and vectors. See transform.py for the roundtrip check, the time-conditioned path, and the vmap example.

Train

train.py fits the network to a known input-dependent rotation field. A fixed linear map W sends each input x to skew parameters, and the target field is Q_true(x) = expm(skew(W @ x)). For each input x and probe vector v the target is Q_true(x) @ v. The network learns to reproduce the field by mean squared error. The training step is a plain eqx.filter_value_and_grad over the fit loss with an optax.adam optimizer. JIT lives on train_step, the outermost call.

@eqx.filter_jit
def train_step(model, opt_state, X, V, T, optim):
    loss, grads = eqx.filter_value_and_grad(fit_loss)(model, X, V, T)
    updates, opt_state = optim.update(grads, opt_state)
    model = eqx.apply_updates(model, updates)
    return model, opt_state, loss

The field lies in SO(n), so all three methods fit it. Select the method with --method.

uv run --directory <skill-dir> python train.py --method exp
uv run --directory <skill-dir> python train.py --method cayley
uv run --directory <skill-dir> python train.py --method qr

Save And Load

save_load.py saves one model per directory under MODEL_ROOT, which defaults to saved_models. Each directory holds weights.eqx and metadata.json.

eqx.tree_serialise_leaves writes only the array leaves. The static fields, dim and the network hyperparameters, are not in weights.eqx. The metadata.json file is the sole source of truth for rebuilding the skeleton. Each class exposes a config property that returns its constructor keyword arguments, and the metadata records the class name alongside that config.

metadata = {
    "name": name,
    "description": description,
    "saved_at": datetime.now(timezone.utc).isoformat(),
    "class": type(model).__name__,
    "config": model.config,
    "num_parameters": count_array_parameters(model),
}

To load, look the class up in REGISTRY, rebuild the skeleton from the saved config, then deserialize the leaves. The load path asserts that the rebuilt config matches the metadata.

cls = REGISTRY[metadata["class"]]
skeleton = cls(key, **metadata["config"])
model = eqx.tree_deserialise_leaves(model_dir / "weights.eqx", skeleton)

No pretrained checkpoint ships with this skill. The example saves a freshly built model and asserts an exact output roundtrip after reload.

Tests

Four files under tests/ exercise the transforms across all three classes, a range of dim, and both the plain and time-conditioned paths.

• test_orthogonality.py checks that Q(x) is orthogonal for a batch of inputs, with and without time conditioning.
• test_determinant.py checks det Q(x) is +1 for the exp and Cayley classes and has unit magnitude for the QR class.
• test_invertibility.py checks that inverse undoes the call and that the call preserves the vector norm.
• test_gradients.py checks that gradients of a scalar loss through the network and the parameterization are finite and not identically zero.

Run the suite with uv run --directory <skill-dir> pytest.

Notes For Agents

• Keep JAX code in this style. Equinox modules with parameters as pytree leaves, no eqx.filter_jit on inner methods. Only train_step is JIT'd, so XLA compiles the whole graph at once.
• dim, use_time, hidden_size, n_blocks, embedding_size, output_scale, and use_sign_correction are eqx.field(static=True). Static fields are not serialized. When reconstructing a saved model, rebuild the skeleton from metadata.json, never from the .eqx file alone.
• The two-input split is load-bearing. x conditions the matrix and v is transformed. They are separate arguments even when a caller passes the same vector for both.
• The sigma argument is required at every call when use_time=True and is unused when use_time=False. The key argument is always ignored and exists only for interface parity with the source code.
• The exp path uses jax.scipy.linalg.expm. A truncated Taylor series is not a safe substitute, because the network can predict large skew parameters and the series loses orthogonality far from the identity.
• The standalone learnable matrix without a network lives in the separate jax-orthogonal-matrix skill.