glass_box_umap.jacobian

Overview

Classes

JacobianVerification

Result of verifying that f(x) ≈ J(x) @ x.

Function

compute_jacobian(model, x, batch_size)	Compute the Jacobian of a model using vmap + jacrev with functional_call.
project_jacobian(jacobian, proj_tensor)	Map a Jacobian’s input axis through a linear projection.
reduce_contributions(contributions, method)	Reduce per-feature contributions across embedding dimensions.
verify_jacobian(Z, J, X)	Verify that f(x) ≈ J(x) @ x.

Classes

class JacobianVerification(z_range: tuple[float, float], reconstruction_range: tuple[float, float], max_error: float, mean_error: float, relative_error: float)

Result of verifying that f(x) ≈ J(x) @ x.

Attributes:

z_range : tuple[float, float]

(min, max) of the embedding output.

reconstruction_range : tuple[float, float]

(min, max) of the Jacobian reconstruction.

max_error : float

Maximum absolute error between embedding and reconstruction.

mean_error : float

Mean absolute error between embedding and reconstruction.

relative_error : float

Max error relative to the embedding’s magnitude.

Functions

compute_jacobian(model: Module, x: Tensor, batch_size: int = 1024) → Tensor

Compute the Jacobian of a model using vmap + jacrev with functional_call.

Parameters:

• model : Module

  Encoder network (will be deep-copied and set to eval mode).

• x : Tensor

  Input tensor of shape (n, in_dim).

• batch_size : int

  Number of samples per Jacobian batch.

Returns:

Jacobian tensor of shape (n, out_dim, in_dim).

Return type:

Tensor

project_jacobian(jacobian: Tensor, proj_tensor: Tensor) → Tensor

Map a Jacobian’s input axis through a linear projection.

Used to express a Jacobian computed in a reduced input space (e.g. PCA components) in terms of the original features, by right-multiplying with the projection matrix that maps reduced-space inputs back to original features.

Parameters:

• jacobian : Tensor

  Jacobian tensor of shape (n, out_dim, in_dim_reduced).

• proj_tensor : Tensor

  Projection matrix of shape (in_dim_reduced, in_dim_original), e.g. pca.components_.

Returns:

Jacobian of shape (n, out_dim, in_dim_original).

Return type:

Tensor

reduce_contributions(contributions: NDArray[floating], method: Literal['l2'] = 'l2') → NDArray[floating]

Reduce per-feature contributions across embedding dimensions.

Parameters:

• contributions : NDArray[floating]

  Feature contributions with shape (n_samples, n_components, n_features).

• method : Literal['l2']

  Reduction method. "l2" takes the L2 norm across components.

Returns:

Reduced contributions with shape (n_samples, n_features).

Return type:

NDArray[floating]

verify_jacobian(Z: NDArray[floating], J: NDArray[floating], X: NDArray[floating]) → JacobianVerification

Verify that f(x) ≈ J(x) @ x.

Parameters:

• Z : NDArray[floating]

  Embedding output, shape (n, out_dim).

• J : NDArray[floating]

  Jacobian, shape (n, out_dim, in_dim).

• X : NDArray[floating]

  Input data, shape (n, in_dim).

Returns:

A JacobianVerification with error diagnostics.

Return type:

JacobianVerification