spacelearn 0.1.0__py3-none-any.whl

spacelearn/__init__.py

@@ -0,0 +1,16 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from ._find_subspaces import subspaces, principal_directions
+from ._combined import combined_loss, solve_subspaces, solve_dims
+
+__all__ = [
+    "principal_directions",
+    "subspaces",
+    "combined_loss",
+    "solve_subspaces",
+    "solve_dims",
+]

spacelearn/_combined.py

@@ -0,0 +1,261 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+from dataclasses import dataclass, field
+from ._find_subspaces import principal_directions, subspaces
+from .loss import alignment_loss, disentanglement_loss, isotropy_loss, stability_loss
+from .optim import fit_decoders, optimal_bins, k_per_q
+from .data import pool
+
+@dataclass
+class Space:
+    W: dict[str, torch.Tensor]
+    A: dict[str, torch.Tensor]
+    V: dict[str, torch.Tensor]
+    _names: list[str] = field(init=False)
+
+    def __post_init__(self):
+        if not (
+            self.W.keys() == self.A.keys() == self.V.keys()
+        ):
+            raise ValueError(
+                "W, A, and V must contain identical keys"
+            )
+
+        self._names = list(self.W.keys())
+
+    def __getitem__(self, key: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.W[key], self.A[key], self.V[key]
+    
+    def keys(self):
+        return self._names
+    
+    def values(self):
+        return [(self.W[k], self.A[k], self.V[k]) for k in self._names]
+    
+    def __len__(self):
+        return len(self._names)
+    
+    def items(self):
+        return [(k, (self.W[k], self.A[k], self.V[k])) for k in self._names]
+
+    def __iter__(self):
+        return iter(self._names)
+
+    def __contains__(self, key):
+        return key in self._names
+
+def combined_loss(
+    Z: torch.Tensor,
+    W: dict[str, torch.Tensor],
+    A: dict[str, torch.Tensor],
+    Y: dict[str, torch.Tensor],
+    W_prev: dict[str, torch.Tensor],
+    *,
+    n_dir: int = 32
+) -> torch.Tensor:
+    """
+    Computes the complete SPACE training objective by combining alignment,
+    disentanglement, isotropy, and stability losses into a single scalar.
+
+    #### USAGE:
+        loss = combined_loss(Z, W, A, Y, W_prev)
+
+        # custom isotropy sampling
+        loss = combined_loss(
+            Z, W, A, Y, W_prev,
+            n_dir=64
+        )
+
+    #### HOW IT WORKS:
+        loss =
+            alignment_loss(Z, W, A, Y)
+            + disentanglement_loss(Z, W)
+            + isotropy_loss(Z, W, n_dir=n_dir)
+            + stability_loss(W, W_prev)
+
+    #### ARGS:
+        - Z: Latent representations with shape (B, D).
+        - W: Dictionary of subspace basis matrices, one per quantity,
+             each shaped (k, D).
+        - A: Dictionary of fitted decoder matrices mapping subspace
+             coordinates to physical targets.
+        - Y: Dictionary of pooled physical targets.
+        - W_prev: Dictionary of subspace bases from the previous
+                  optimization step, used for stability regularization.
+        - n_dir: Number of random directions used by isotropy_loss when
+                 estimating residual isotropy. Defaults to 32.
+
+    #### RETURNS:
+        - A scalar tensor containing the summed training objective.
+    """
+    return (
+        alignment_loss(Z, W, A, Y)
+        + disentanglement_loss(Z, W)
+        + isotropy_loss(Z, W, n_dir=n_dir)
+        + stability_loss(W, W_prev)
+    )
+
+def solve_subspaces(
+    Z: torch.Tensor,
+    Y: dict[torch.Tensor],
+    *,
+    k: int | dict[str, int] = 32,
+) -> dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """
+    Solves the complete subspace-discovery pipeline for a collection of
+    physical quantities.
+
+    For each quantity, this function finds its principal latent
+    direction, constructs a k-dimensional subspace basis, and fits an
+    optimal linear decoder back to physical space.
+
+    #### USAGE:
+        WVA = solve_subspaces(Z, Y, k=32)
+
+        W, V, A = WVA["scaled"]
+
+    #### HOW IT WORKS:
+        1. principal_directions() identifies the strongest latent
+           direction associated with each target quantity
+        2. subspaces() expands these directions into k-dimensional
+           orthonormal subspace bases
+        3. fit_decoders() solves least-squares decoders from subspace
+           coordinates back to physical targets
+        4. Results are combined into:
+               {name: (W_q, V_q, A_q)}
+
+    #### ARGS:
+        - Z: Latent representations with shape (B, D).
+        - Y: Dictionary of pooled physical targets.
+        - k: Number of basis directions per subspace. Defaults to 32.
+
+    #### RETURNS:
+        - A dict / Space_dataclass mapping each quantity name to:
+            (W_q,  A_q, V_q)
+
+          where:
+            W_q = subspace basis matrix (k, D)
+            V_q = principal latent direction (D,)
+            A_q = fitted decoder matrix
+    """
+    wv    = principal_directions(Z, Y)
+    W     = subspaces(Z, Y, k=k)
+    A     = fit_decoders(W, Z, Y)
+    SPACE = Space(W, A, V={k: v[1] for k, v in wv.items()})
+    return SPACE
+
+def solve_dims(
+    Y: dict[str, torch.Tensor],
+    max_bins: int,
+    bin_thresholds: float | dict[str, float],
+    k_thresholds: float | dict[str, float],
+    *,
+    bin_default_threshold: float | None = None,
+    k_default_threshold: float | None = None,
+    pooling_mode: str = "avg",
+    pool_per_quantity: bool = False,
+    k_per_quantity: bool = False
+) -> dict[str, tuple[int, int]]:
+    """
+    Computes recommended pooling resolutions and subspace dimensions for
+    a collection of physical targets.
+
+    This combines optimal_bins() and k_per_q() into a single sizing
+    pipeline.
+
+    #### USAGE:
+        dims = solve_dims(
+            Y,
+            max_bins=16,
+            bin_thresholds=0.95,
+            k_thresholds=0.90
+        )
+
+        # per-quantity sizing
+        dims = solve_dims(
+            Y,
+            max_bins=16,
+            bin_thresholds=0.95,
+            k_thresholds=0.90,
+            pool_per_quantity=True,
+            k_per_quantity=True
+        )
+
+    #### HOW IT WORKS:
+        1. optimal_bins() determines the minimum pooling resolution
+           needed to retain the desired fraction of spatial variance
+        2. Targets are pooled using the selected resolution(s)
+        3. k_per_q() estimates the minimum subspace dimensionality
+           needed to retain the desired fraction of variance
+        4. Pooling resolutions and subspace dimensions are combined into
+           per-quantity (n_bins, k) recommendations
+
+    #### ARGS:
+        - Y: Dictionary of physical target tensors.
+        - max_bins: Maximum pooling grid side length considered.
+        - bin_thresholds: Variance-retention threshold(s) used by
+                          optimal_bins().
+        - k_thresholds: Variance-retention threshold(s) used by
+                        k_per_q().
+        - bin_default_threshold: Fallback threshold for quantities
+                                 missing from bin_thresholds.
+        - k_default_threshold: Fallback threshold for quantities
+                               missing from k_thresholds.
+        - pooling_mode: Pooling mode passed to pool().
+                         Defaults to "avg".
+        - pool_per_quantity: Whether pooling resolution is selected
+                             independently for each quantity.
+        - k_per_quantity: Whether subspace dimensionality is selected
+                          independently for each quantity.
+
+    #### RETURNS:
+        - A dict:
+              {name: (n_bins, k)}
+
+          mapping each quantity to its recommended pooling resolution
+          and subspace dimensionality.
+    """
+    """
+    Computes optimal binning and subspace dimensions for a collection of targets.
+    
+    If 'pool_per_quantity' is True, it returns a dictionary mapping each quantity 
+    to its optimal (n, k) pair. Otherwise, it returns a single global (n, k) tuple.
+    """
+    # optimal bin resolution
+    n = optimal_bins(
+        Y, max_bins, bin_thresholds, 
+        default_threshold=bin_default_threshold, 
+        per_quantity=pool_per_quantity
+    )
+    
+    # pool based on optimal bin resolution
+    if pool_per_quantity:
+        # n is a dict {name: n_q}
+        pooled = {q: pool(y, n[q], pooling_mode) for q, y in Y.items()}
+    else:
+        # n is a single global int
+        pooled = {q: pool(y, n, pooling_mode) for q, y in Y.items()}
+        
+    # optimal subspace sizes
+    k = k_per_q(
+        pooled, k_thresholds, 
+        default_threshold=k_default_threshold, 
+        per_quantity=k_per_quantity
+    )
+    
+    # aggregate n and k
+    if pool_per_quantity:
+        if k_per_quantity:
+            return {q: (n[q], k[q]) for q in Y.keys()}
+        else:
+            return {q: (n[q], k) for q in Y.keys()}
+    else:
+        if k_per_quantity:
+            return {q: (n, k[q]) for q in Y.keys()}
+        else:
+            return {q: (n, k) for q in Y.keys()}

spacelearn/_find_subspaces.py

@@ -0,0 +1,101 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+from .util import full, left
+
+def principal_directions(
+    z: torch.Tensor,  # (B, D)
+    pooled_targets: dict[str, torch.Tensor], # {name: (B, n^2)}
+) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
+    """
+    Finds the rank-1 pair (w, v) for multiple quantities such that:
+        (z @ w).unsqueeze(1) @ v  ~  y
+    via the cross-covariance SVD. This runs 'batched' across all quantities 
+    provided in the dictionary.
+
+    w : (D,)  — direction in latent space
+    v : (n^2,) — direction in target space
+
+    This is the analytical solution, the first left/right singular
+    vectors of C = z.T @ y are exactly the directions that maximise
+    the explained covariance between z and y under a rank-1 constraint.
+
+    HOW THIS WORKS IN DETAIL:
+
+    1. We compute the cross-covariance matrix C = z.T @ y
+    2. We want to find the vectors w (the latent directions) and v (the target directions)
+       that solve y = (z @ W) @ z.T
+    3. since theres no guaranteed solution we instead search the optimal solution
+       min_{w, v} |y - (z @ w).unsqueeze(1) @ v|^2_F which is minimizing the forbenius norm of the error
+       Note that the forbenious norm ||A||^2_F is equal to the trace of A.T @ A (let the trace fn be Tr())
+       This gives: Tr((y-zWv.T).T @ (y-zWv.T)) = Tr((y.T - v(zW).T)(Y-zWv.T))
+       ...
+       We end up having to maximize 2v.T(z.T@y).Tw
+       -> max_{w,v} v.T(z.T@y).T w with z.T @ y = C
+    4. using SVD we get w.T(UEV.T)v = (U.Tw).T E (V.T v)
+       Since UV are both orthogonal matrices we get unit vectors and this simplifies to
+       Sum_{i=1}^n sigma_i ^w_i ^v_i (let ^ denote the result of V.T v and U.T w)
+    5. To maximize this sum we need to put the weight on the first sinular value (making the unit vecs one hot vectors along the first singular direction)
+       Naturally this leaves us with the first singular vectors found by svd, as they 
+       align the input space along the direction in output space and scale by the larges singular value
+
+    6. In short: C gives how well any dimension in z correlates with any other dimension in y.
+                   SVD finds the basis in which te maximum covariance lies along one direction.
+                   -> leaves us with a rank-1 solution since we discared all othersingular vectors
+    """
+    # center z once
+    z_c = z - z.mean(0, keepdim=True) # (B, D)
+    
+    results = {}
+    
+    for name, y in pooled_targets.items():
+        # if not flat, flatten
+        if y.dim() > 2:
+            y = y.flatten(1)
+        
+        # center target
+        y_c = y - y.mean(0, keepdim=True) # (B, n^2)
+
+        # cross-covariance matrix
+        # gives correlation between latent and target dimensions
+        C = z_c.T @ y_c # (D, n^2)
+
+        # rank-1 SVD (keep only first singular vector)
+        # using full SVD and taking first column is fine for small n^2
+        U, _, Vh = full(C)
+
+        w = U[:, 0] # (D,)  latent direction
+        v = Vh[0, :] # (n^2,) target direction
+
+        # normalise to unit length
+        w = w / (w.norm() + 1e-8)
+        v = v / (v.norm() + 1e-8)
+        
+        results[name] = (w.detach(), v.detach())
+
+    return results
+
+def subspaces(Z, pooled_global, k: int) -> dict[str, torch.Tensor]:
+    """
+    Find the optimal orthogonal basis in the latent space, 
+    such that the span of subspaces capture the most variance of the pooled physics targets
+    """
+    print("Estimating subspaces...")
+    Z_c = Z - Z.mean(0, keepdim=True) # centered for variance calc
+    W   = {}
+    for name, target in pooled_global.items():
+        # if not flat, flatten
+        if target.dim() > 2:
+            target = target.flatten(1)
+        y_c     = target - target.mean(0, keepdim=True)    # center
+        C       = Z_c.T @ y_c                              # cross-cov-matrix
+        U, _    = left(C)                                  # svd to find the left singular vectors (projection directions from latent into the subspace)
+        _k = k if isinstance(k, int) else k[name]          # check if k is global or per quantity
+        W_q     = U[:, :_k].T                               # only keep the top k (subspace size constraint)
+        W_q, _  = torch.linalg.qr(W_q.T, mode="reduced")   # orthonormalise
+        W[name] = W_q.T                                    # save
+    return W

spacelearn/data/__init__.py

@@ -0,0 +1,11 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from ._pipeline import input_to_quantity
+from ._pool import pool
+from ._utils import ensure_dim, to_device
+
+__all__ = ["input_to_quantity", "pool", "ensure_dim", "to_device"]

spacelearn/data/_compute_Q.py

@@ -0,0 +1,70 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import inspect
+from typing import Callable, Any
+
+def compute_Q(*arg_functions: Callable, **kwarg_functions: Callable) -> Callable:
+    """
+    Creates a dependency-injecting wrapper that executes multiple functions 
+    simultaneously, passing only the required arguments to each.
+
+    #### USAGE:
+        # Define solvers
+        def solver_a(A: torch.Tensor, q: int): ...
+        def solver_b(A: torch.Tensor): ...
+
+        # Create the combined executor
+        executor = compute_Q(solver_a, solver_b=solver_b)
+
+        # Run both, injecting only the arguments they need
+        results = executor(A=my_matrix, q=5) 
+        # solver_a receives {'A': my_matrix, 'q': 5}
+        # solver_b receives {'A': my_matrix}
+
+    #### VALIDATION:
+        The function will raise a TypeError if multiple registered functions 
+        define the same parameter name but specify different type annotations.
+
+    #### ARGS:
+        - *arg_functions: Unnamed callables (names derived from __name__).
+        - **kwarg_functions: Named callables for the registry.
+    """
+    registry = {}
+    for f in arg_functions:
+        registry[getattr(f, "__name__", str(id(f)))] = f
+    registry.update(kwarg_functions)
+
+    # collect all fns and check that dublicated params all use the same type
+    type_hints = {}
+    registry_meta = {}
+
+    for name, f in registry.items():
+        sig = inspect.signature(f)
+        params = sig.parameters
+        registry_meta[name] = list(params.keys())
+
+        # validate types for params with the same names accross all fns
+        for p_name, param in params.items():
+            if p_name in type_hints and param.annotation != inspect.Parameter.empty:
+                if type_hints[p_name] != param.annotation:
+                    raise TypeError(
+                        f"Type mismatch for argument '{p_name}': "
+                        f"Expected {type_hints[p_name]}, but {name} defined it as {param.annotation}"
+                    )
+            type_hints[p_name] = param.annotation
+
+    def wrapped(**kwargs) -> dict[str, Any]:
+        results = {}
+        for name, f in registry.items():
+            relevant_args = {
+                k: kwargs[k] for k in registry_meta[name] if k in kwargs
+            }
+            results[name] = f(**relevant_args)
+        return results
+
+    return wrapped
+

spacelearn/data/_pipeline.py

@@ -0,0 +1,55 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+from typing import Callable
+from ._compute_Q import compute_Q
+from ._pool import pool
+
+
+def input_to_quantity(
+    n_bins: int,
+    mode: str,
+    *arg_functions: Callable,
+    **kwarg_functions: Callable
+) -> Callable:
+    """
+    Builds a callable that computes one or more quantities from injected
+    arguments (via `compute_Q`) and pools each resulting tensor into a
+    fixed number of bins.
+
+    #### USAGE:
+        def solver_a(A: torch.Tensor, q: int): ...
+        def solver_b(A: torch.Tensor): ...
+
+        # Create the combined executor: run solver_a and solver_b, then
+        # pool each of their outputs into 10 bins using mean pooling
+        executor = input_to_quantity(10, "mean", solver_a, solver_b=solver_b)
+
+        results = executor(A=my_matrix, q=5)
+        # results = {
+        #     "solver_a": pool(solver_a(A=my_matrix, q=5), 10, "mean"),
+        #     "solver_b": pool(solver_b(A=my_matrix), 10, "mean"),
+        # }
+
+    #### ARGS:
+        - n_bins: Number of bins each function's output is pooled into.
+        - mode: Pooling strategy passed to `pool` (e.g. "mean", "sum", "max").
+        - *arg_functions: Unnamed callables to register (names derived from
+          `__name__`); see `compute_Q`.
+        - **kwarg_functions: Named callables to register; see `compute_Q`.
+
+    #### RETURNS:
+        A callable that accepts keyword arguments, injects the relevant
+        subset into each registered function (see `compute_Q`), and returns
+        a dict mapping each function's name to its pooled `torch.Tensor`
+        output.
+    """ 
+    Q = compute_Q(*arg_functions, **kwarg_functions)
+    def wrapped(**kwargs) -> dict[str, torch.Tensor]:
+        data = Q(**kwargs)
+        return {name: pool(chunks, n_bins, mode) for name, chunks in data.items()}
+    return wrapped

spacelearn/data/_pool.py

@@ -0,0 +1,35 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+import torch.nn.functional as F
+
+def spatial_pool_targets(chunks: torch.Tensor, n_bins: int, mode: str = "avg") -> torch.Tensor:
+    """chunks: (B, 1, H, W) -> (B, 1, n_bins, n_bins)"""
+    if mode == "avg":
+        return F.adaptive_avg_pool2d(chunks, (n_bins, n_bins))
+    elif mode == "max":
+        return F.adaptive_max_pool2d(chunks, (n_bins, n_bins))
+    else:
+        raise ValueError(f"Invalid mode: {mode}")
+    
+def pool_targets(chunks: torch.Tensor, n_bins: int, mode: str = "avg") -> torch.Tensor:
+    """chunks: (B, 1, H, W) -> (B, 1, n_bins, n_bins)"""
+    if mode == "avg":
+        return F.adaptive_avg_pool1d(chunks, n_bins)
+    elif mode == "max":
+        return F.adaptive_max_pool1d(chunks, n_bins)
+    else:
+        raise ValueError(f"Invalid mode: {mode}")
+
+def pool(chunks: torch.Tensor, n_bins: int, mode: str = "avg") -> torch.Tensor:
+    ndims = chunks.dim()
+    if ndims == 4: # B, C, H, W
+        return spatial_pool_targets(chunks, n_bins, mode)
+    elif ndims == 2: # B, S
+        return pool_targets(chunks, n_bins, mode)
+    else:
+        raise ValueError(f"Invalid number of dimensions: {ndims}. Expected 2 (B, S) or 4 (B, C, H, W)")

spacelearn/data/_utils.py

@@ -0,0 +1,115 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from typing import Callable
+import torch
+import functools
+
+def ensure_dim(func: Callable = None, *, dims: int, targets: list[str] | None = None) -> Callable:
+    """
+    Wraps a function to force input tensors to have a minimum number of dimensions 
+    by prepending unit dimensions (1s) to their shape.
+
+    #### USAGE:
+        # Example 1: Force all inputs to be at least 4D (B, C, H, W)
+        @ensure_dim(dims=4)
+        def process_spatial(grid: torch.Tensor): ...
+
+        # Example 2: Only force specific inputs to be 4D
+        @ensure_dim(dims=4, targets=['pressure', 'velocity'])
+        def solve_physics(pressure: torch.Tensor, density: torch.Tensor):
+            # 'pressure' will be forced to 4D
+            # 'density' will remain untouched
+            ...
+
+    #### ARGS:
+        - function: The callable to wrap.
+        - dims: The target number of dimensions (e.g., 4 for spatial grids).
+        - targets: An optional list of argument names to target for expansion. 
+                   If None or empty, all positional and keyword tensors are expanded.
+
+    #### RETURNS:
+        - A wrapped callable that expands tensors to the requested dimensionality 
+          using `view()` before execution.
+    """
+    if func is None:
+        return functools.partial(ensure_dim, dims=dims, targets=targets)
+    # helper to expand a tensor using view
+    def expand(t):
+        if isinstance(t, torch.Tensor) and t.dim() < dims:
+            # adds dimensions to the front (e.g., to make (H, W) -> (1, 1, H, W))
+            new_shape = (1,) * (dims - t.dim()) + t.shape
+            return t.view(new_shape)
+        return t
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        
+
+        # positional arguments
+        new_args = tuple(expand(a) for a in args)
+
+        # keyword arguments
+        new_kwargs = {}
+        for k, v in kwargs.items():
+            if targets is None or not targets or k in targets:
+                new_kwargs[k] = expand(v)
+            else:
+                new_kwargs[k] = v
+                
+        return func(*new_args, **new_kwargs)
+    
+    return wrapped
+
+def to_device(func: Callable = None, *, device: str = "cpu", targets: list[str] = None):
+    """
+    Wraps a function to move input tensors to a target device before execution.
+
+    #### USAGE:
+        # Example 1: Move all tensor inputs to a device
+        @to_device(device="cuda")
+        def process(grid: torch.Tensor): ...
+
+        # Example 2: Used bare, defaults to "cpu"
+        @to_device
+        def process(grid: torch.Tensor): ...
+
+        # Example 3: Only move specific inputs
+        @to_device(device="cuda", targets=["pressure"])
+        def solve_physics(pressure: torch.Tensor, density: torch.Tensor):
+            # 'pressure' will be moved to "cuda"
+            # 'density' will remain untouched
+            ...
+
+    #### ARGS:
+        - func: The callable to wrap.
+        - device: The target device to move tensors to (e.g., "cpu", "cuda").
+        - targets: An optional list of argument names to target for moving.
+                   If None, all positional and keyword tensors are moved.
+
+    #### RETURNS:
+        - A wrapped callable that moves tensors to the requested device
+          using `.to()` before execution.
+    """
+    # This allows the decorator to be used as @to_device OR @to_device(device='cuda')
+    if func is None:
+        return functools.partial(to_device, device=device, targets=targets)
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        def move(t):
+            if isinstance(t, torch.Tensor):
+                return t.to(device)
+            return t
+
+        new_args = tuple(move(a) for a in args)
+        new_kwargs = {
+            k: move(v) if (targets is None or k in targets) else v
+            for k, v in kwargs.items()
+        }
+        return func(*new_args, **new_kwargs)
+
+    return wrapped

spacelearn/loss/__init__.py

@@ -0,0 +1,17 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from ._alignment_loss import alignment_loss
+from ._disentanglement_loss import disentanglement_loss
+from ._isotropy_loss import isotropy_loss
+from ._stability_loss import stability_loss
+
+__all__ = [
+    "alignment_loss",
+    "disentanglement_loss",
+    "isotropy_loss",
+    "stability_loss",
+]