jax-orthogonal-matrix

name	jax-orthogonal-matrix
description	Use this skill when constructing, applying, training, saving, loading, or documenting a standalone learnable orthogonal matrix in JAX plus Equinox, parameterized by the matrix exponential of a skew-symmetric matrix, the Cayley transform, or a QR factorization.

JAX Orthogonal Matrix

This skill parameterizes a learnable orthogonal matrix three ways. Each class is a plain Equinox module whose learnable leaf maps to an orthogonal matrix Q through a fixed differentiable construction. 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. That file wraps each construction in a neural network that predicts the parameters from an input. The network-predicted version is a separate skill, jax-orthogonal-net. This skill covers the standalone matrix only.

• construct.py builds all three parameterizations and checks orthogonality.
• transform.py applies each matrix forward and inverse and checks the roundtrip.
• train.py recovers a rotation by orthogonal Procrustes alignment.
• 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-matrix
uv run --directory ~/.claude/skills/jax-orthogonal-matrix python construct.py
uv run --directory ~/.claude/skills/jax-orthogonal-matrix python transform.py
uv run --directory ~/.claude/skills/jax-orthogonal-matrix python train.py
uv run --directory ~/.claude/skills/jax-orthogonal-matrix python save_load.py
uv run --directory ~/.claude/skills/jax-orthogonal-matrix pytest

Module Map

The package is jax_orthogonal_matrix.

• jax_orthogonal_matrix.parameterizations holds skew_from_params and the three classes OrthogonalMatrixExp, OrthogonalCayley, OrthogonalQR.
• skew_from_params(params, dim) builds a skew-symmetric matrix from its strict upper triangle of dim * (dim - 1) // 2 entries.
• The package root re-exports the three classes and skew_from_params.

The Three Parameterizations

Every class has the same interface. The constructor takes key and dim. The matrix property returns the orthogonal matrix Q. Calling the module returns Q @ v. The inverse method returns Q.T @ u, which undoes the call. The config property returns the constructor keyword arguments for save_load.

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

OrthogonalMatrixExp stores the upper triangle of a skew-symmetric matrix A and returns expm(A). Since A is skew-symmetric, expm(A) is orthogonal with determinant +1. The exp path uses jax.scipy.linalg.expm, which is reverse-mode differentiable and stays orthogonal even for large skew matrices.

OrthogonalCayley stores the same skew upper triangle and returns solve(I + A, I - A). The matrix I + A is always invertible for skew A, so the solve never hits a singular system. The Cayley transform cannot represent a rotation by pi, because such a rotation has a -1 eigenvalue.

OrthogonalQR stores a dense matrix W and returns the Q factor of its QR factorization. With use_sign_correction=True, the default, the columns of Q are multiplied by the signs of the diagonal of R. The correction makes the factorization unique and lets Q reach both determinant components of O(n). Without the correction the LAPACK Householder convention pins Q to a single determinant component fixed by the parity of dim. The 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 to any of the three constructors. The exp and Cayley classes initialize near the identity with small random skew parameters. The QR class initializes from a Gaussian, which gives a Haar-distributed orthogonal matrix when sign correction is on.

import jax
from jax_orthogonal_matrix import OrthogonalMatrixExp, OrthogonalCayley, OrthogonalQR

key = jax.random.PRNGKey(0)
layer = OrthogonalMatrixExp(key, dim=6)
Q = layer.matrix            # (6, 6), orthogonal

See construct.py for a runnable version that checks Q.T @ Q is the identity for all three classes.

Transform

Calling a module applies Q to a vector. inverse applies Q.T, which recovers the input because Q is orthogonal.

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

Each module acts on a single vector. Batch with jax.vmap(layer) over a stack of vectors. See transform.py for the roundtrip check and the vmap example.

Train

train.py runs orthogonal Procrustes. It builds a source point cloud X and a ground-truth rotation R_true drawn from SO(n), forms the target Y = X @ R_true.T, then trains a learnable orthogonal matrix Q to minimize the mean squared alignment error. The training step is a plain eqx.filter_value_and_grad over the alignment loss with an optax.adam optimizer. JIT lives on train_step, the outermost call.

@eqx.filter_jit
def train_step(layer, opt_state, X, Y, optim):
    loss, grads = eqx.filter_value_and_grad(alignment_loss)(layer, X, Y)
    updates, opt_state = optim.update(grads, opt_state)
    layer = eqx.apply_updates(layer, updates)
    return layer, opt_state, loss

R_true is built from a matrix exponential, so it lies in SO(n). The exp and cayley methods recover it to fp32 precision. The qr method is run with --method qr. Its loss decreases but it can stall before Q reaches R_true, since the sign correction is discontinuous. train.py reports that stall instead of asserting exact recovery for the QR 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 use_sign_correction, 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 parameterizations across all three classes and a range of dim.

• test_orthogonality.py checks that Q.T @ Q and Q @ Q.T are the identity.
• test_determinant.py checks det Q is +1 for the exp and Cayley classes and has unit magnitude for the QR class.
• test_invertibility.py checks that inverse undoes __call__ and that the call preserves the vector norm.
• test_gradients.py checks that gradients of a scalar loss through the parameters 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 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 exp path uses jax.scipy.linalg.expm. A truncated Taylor series is not a safe substitute, because it loses orthogonality far from the identity.
• The QR sign correction is load-bearing. With it on, OrthogonalQR reaches the full O(n) and the Gaussian init is Haar-distributed. With it off, Q is confined to one determinant component fixed by the parity of dim.
• Batching is external. Each module acts on one vector and is batched with jax.vmap. There is no batch_size field.
• The network-predicted version of these parameterizations lives in the separate jax-orthogonal-net skill.