``glass_box_umap.jacobian``
===========================

.. py:module:: glass_box_umap.jacobian


Overview
--------

.. list-table:: Classes
   :header-rows: 0
   :widths: auto
   :class: summarytable

   * - :py:obj:`JacobianVerification <glass_box_umap.jacobian.JacobianVerification>`
     - Result of verifying that ``f(x) ≈ J(x) @ x``.


.. list-table:: Function
   :header-rows: 0
   :widths: auto
   :class: summarytable

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



Classes
-------
.. autoclass:: JacobianVerification







Functions
---------
.. py:function:: compute_jacobian(model: torch.nn.Module, x: torch.Tensor, batch_size: int = 1024) -> torch.Tensor

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

   :param model: Encoder network (will be deep-copied and set to eval mode).
   :param x: Input tensor of shape ``(n, in_dim)``.
   :param batch_size: Number of samples per Jacobian batch.

   :returns: Jacobian tensor of shape ``(n, out_dim, in_dim)``.


.. py:function:: project_jacobian(jacobian: torch.Tensor, proj_tensor: torch.Tensor) -> torch.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.

   :param jacobian: Jacobian tensor of shape ``(n, out_dim, in_dim_reduced)``.
   :param proj_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)``.


.. py:function:: reduce_contributions(contributions: numpy.typing.NDArray[numpy.floating], method: Literal['l2'] = 'l2') -> numpy.typing.NDArray[numpy.floating]

   Reduce per-feature contributions across embedding dimensions.

   :param contributions: Feature contributions with shape (n_samples, n_components, n_features).
   :param method: Reduction method. ``"l2"`` takes the L2 norm across components.

   :returns: Reduced contributions with shape (n_samples, n_features).


.. py:function:: verify_jacobian(Z: numpy.typing.NDArray[numpy.floating], J: numpy.typing.NDArray[numpy.floating], X: numpy.typing.NDArray[numpy.floating]) -> JacobianVerification

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

   :param Z: Embedding output, shape ``(n, out_dim)``.
   :param J: Jacobian, shape ``(n, out_dim, in_dim)``.
   :param X: Input data, shape ``(n, in_dim)``.

   :returns: A ``JacobianVerification`` with error diagnostics.