ohagi 1.0.0__py3-none-any.whl

ohagi/__init__.py

@@ -0,0 +1,65 @@
+"""OHAGI: composable constrained learning in JAX.
+
+Precision is controlled by the application. Legacy frontends retain their
+historical precision settings; importing the common package is lightweight.
+"""
+
+from importlib import import_module
+from ohagi.compat import install_aliases
+from ohagi.compat import plugins as _plugins
+from ohagi._projection_aliases import ALIASES as _PROJECTION_ALIASES
+
+install_aliases(_PROJECTION_ALIASES)
+__version__ = "1.0.0"
+_EXPORTS = {
+    "Problem": ("problems", "Problem"),
+    "Projection": ("problems", "Projection"),
+    "FunctionalProblem": ("problems", "FunctionalProblem"),
+    "SymbolicProblem": ("problems", "SymbolicProblem"),
+    "MatrixProblem": ("problems", "MatrixProblem"),
+    "ParametricQP": ("problems", "ParametricQP"),
+    "ParametricLP": ("problems", "ParametricLP"),
+    "MLP": ("nn", "MLP"),
+    "MLPBackbone": ("nn", "MLPBackbone"),
+    "Loss": ("losses", "Loss"),
+    "prepare_projection": ("solvers", "prepare_projection"),
+    "ChambollePockConfig": ("solvers", "ChambollePockConfig"),
+    "ADMMConfig": ("solvers", "ADMMConfig"),
+    "FischerBurmeisterConfig": ("solvers", "FischerBurmeisterConfig"),
+    "ProjectionConfig": ("solvers", "ProjectionConfig"),
+    "LinearEqualityConfig": ("solvers", "LinearEqualityConfig"),
+    "LinearSolverConfig": ("solvers", "LinearSolverConfig"),
+    "ImplicitGradientConfig": ("solvers", "ImplicitGradientConfig"),
+    "SolveStatus": ("solvers", "SolveStatus"),
+    "Distance": ("projections", "Distance"),
+    "WeightedDistance": ("projections", "WeightedDistance"),
+    "LocalQuadratic": ("projections", "LocalQuadratic"),
+    "Sequential": ("projections", "Sequential"),
+    "ProjectedModel": ("models", "ProjectedModel"),
+    "Trainer": ("training.engine", "Trainer"),
+    "Adam": ("optim", "Adam"),
+    "save_predictor": ("export", "save_predictor"),
+    "load_predictor": ("export", "load_predictor"),
+    "load_inference": ("inference", "load"),
+    "export_inference": ("inference", "export"),
+}
+__all__ = list(dict.fromkeys([*_EXPORTS, *_plugins.public_symbols(__name__)]))
+
+
+def __getattr__(name):
+    if name in _EXPORTS:
+        module, symbol = _EXPORTS[name]
+        value = getattr(import_module("ohagi." + module), symbol)
+    else:
+        try:
+            return _plugins.public_symbol(__name__, name)
+        except AttributeError:
+            pass
+        try:
+            value = import_module("ohagi." + name)
+        except ModuleNotFoundError as error:
+            if error.name != "ohagi." + name:
+                raise
+            raise AttributeError(name) from error
+    globals()[name] = value
+    return value

ohagi/_projection_aliases.py

@@ -0,0 +1,198 @@
+"""Historical core projection imports point to PURIN."""
+
+ALIASES = {
+    "ohagi._utils.validation": ("purin._utils.validation", False),
+    "ohagi.autodiff": ("purin.autodiff", True),
+    "ohagi.autodiff.complementarity": ("purin.autodiff.complementarity", False),
+    "ohagi.autodiff.fixed_point": ("purin.autodiff.fixed_point", False),
+    "ohagi.compilation": ("purin.compilation", True),
+    "ohagi.differentiation": ("purin.differentiation", True),
+    "ohagi.differentiation.policies": ("purin.differentiation.policies", False),
+    "ohagi.execution.runtime": ("purin.execution.runtime", False),
+    "ohagi.linalg": ("purin.linalg", True),
+    "ohagi.linalg.equilibration": ("purin.linalg.equilibration", False),
+    "ohagi.linalg.operators": ("purin.linalg.operators", False),
+    "ohagi.linalg.sparse": ("purin.linalg.sparse", True),
+    "ohagi.linalg.sparse.assembly": ("purin.linalg.sparse.assembly", False),
+    "ohagi.linalg.sparse.conversion": ("purin.linalg.sparse.conversion", False),
+    "ohagi.linalg.sparse.minres": ("purin.linalg.sparse.minres", False),
+    "ohagi.linalg.sparse.operators": ("purin.linalg.sparse.operators", False),
+    "ohagi.linalg.sparse.types": ("purin.linalg.sparse.types", False),
+    "ohagi.problems": ("purin.problems", True),
+    "ohagi.problems.canonical": ("purin.problems.canonical", False),
+    "ohagi.problems.compile": ("purin.problems.compile", False),
+    "ohagi.problems.expressions": ("purin.problems.expressions", False),
+    "ohagi.problems.functional": ("purin.problems.functional", False),
+    "ohagi.problems.problem": ("purin.problems.problem", False),
+    "ohagi.problems.projection": ("purin.problems.projection", False),
+    "ohagi.problems.qp": ("purin.problems.qp", False),
+    "ohagi.problems.sparsity": ("purin.problems.sparsity", False),
+    "ohagi.problems.symbolic": ("purin.problems.symbolic", False),
+    "ohagi.problems.types": ("purin.problems.types", False),
+    "ohagi.projections": ("purin.projections", True),
+    "ohagi.projections.adapters": ("purin.projections.adapters", False),
+    "ohagi.projections.diagnostics": ("purin.projections.diagnostics", False),
+    "ohagi.projections.distance": ("purin.projections.distance", False),
+    "ohagi.projections.local_quadratic": ("purin.projections.local_quadratic", False),
+    "ohagi.projections.sequential": ("purin.projections.sequential", False),
+    "ohagi.solvers": ("purin.solvers", True),
+    "ohagi.solvers.admm": ("purin.solvers.admm", True),
+    "ohagi.solvers.admm.config": ("purin.solvers.admm.config", False),
+    "ohagi.solvers.admm.factor": ("purin.solvers.admm.factor", False),
+    "ohagi.solvers.admm.kernels": ("purin.solvers.admm.kernels", False),
+    "ohagi.solvers.admm.rows": ("purin.solvers.admm.rows", False),
+    "ohagi.solvers.admm.solver": ("purin.solvers.admm.solver", False),
+    "ohagi.solvers.api": ("purin.solvers.api", False),
+    "ohagi.solvers.capabilities": ("purin.solvers.capabilities", False),
+    "ohagi.solvers.chambolle_pock": ("purin.solvers.chambolle_pock", True),
+    "ohagi.solvers.chambolle_pock.bounds": (
+        "purin.solvers.chambolle_pock.bounds",
+        False,
+    ),
+    "ohagi.solvers.chambolle_pock.config": (
+        "purin.solvers.chambolle_pock.config",
+        False,
+    ),
+    "ohagi.solvers.chambolle_pock.dense": ("purin.solvers.chambolle_pock.dense", False),
+    "ohagi.solvers.chambolle_pock.dense_implicit": (
+        "purin.solvers.chambolle_pock.dense_implicit",
+        False,
+    ),
+    "ohagi.solvers.chambolle_pock.dense_layer": (
+        "purin.solvers.chambolle_pock.dense_layer",
+        False,
+    ),
+    "ohagi.solvers.chambolle_pock.kernels": (
+        "purin.solvers.chambolle_pock.kernels",
+        False,
+    ),
+    "ohagi.solvers.chambolle_pock.solver": (
+        "purin.solvers.chambolle_pock.solver",
+        False,
+    ),
+    "ohagi.solvers.config": ("purin.solvers.config", False),
+    "ohagi.solvers.diagonal_qp": ("purin.solvers.diagonal_qp", False),
+    "ohagi.solvers.fischer_burmeister": ("purin.solvers.fischer_burmeister", True),
+    "ohagi.solvers.fischer_burmeister.common": (
+        "purin.solvers.fischer_burmeister.common",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.complementarity": (
+        "purin.solvers.fischer_burmeister.complementarity",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.config": (
+        "purin.solvers.fischer_burmeister.config",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.dense_semismooth": (
+        "purin.solvers.fischer_burmeister.dense_semismooth",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.newton": (
+        "purin.solvers.fischer_burmeister.newton",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.operators": (
+        "purin.solvers.fischer_burmeister.operators",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.reduced": (
+        "purin.solvers.fischer_burmeister.reduced",
+        False,
+    ),
+    "ohagi.solvers.fischer_burmeister.solver": (
+        "purin.solvers.fischer_burmeister.solver",
+        False,
+    ),
+    "ohagi.solvers.implicit": ("purin.solvers.implicit", False),
+    "ohagi.solvers.linear_equality": ("purin.solvers.linear_equality", True),
+    "ohagi.solvers.linear_equality.config": (
+        "purin.solvers.linear_equality.config",
+        False,
+    ),
+    "ohagi.solvers.linear_equality.solver": (
+        "purin.solvers.linear_equality.solver",
+        False,
+    ),
+    "ohagi.solvers.metric": ("purin.solvers.metric", False),
+    "ohagi.solvers.primal_dual_interior_point": (
+        "purin.solvers.primal_dual_interior_point",
+        True,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.assembly": (
+        "purin.solvers.primal_dual_interior_point.assembly",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.config": (
+        "purin.solvers.primal_dual_interior_point.config",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.derivatives": (
+        "purin.solvers.primal_dual_interior_point.derivatives",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.newton": (
+        "purin.solvers.primal_dual_interior_point.newton",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.operators": (
+        "purin.solvers.primal_dual_interior_point.operators",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference": (
+        "purin.solvers.primal_dual_interior_point.reference",
+        True,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.bounds": (
+        "purin.solvers.primal_dual_interior_point.reference.bounds",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.filter": (
+        "purin.solvers.primal_dual_interior_point.reference.filter",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.iteration": (
+        "purin.solvers.primal_dual_interior_point.reference.iteration",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.line_search": (
+        "purin.solvers.primal_dual_interior_point.reference.line_search",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.newton": (
+        "purin.solvers.primal_dual_interior_point.reference.newton",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.residuals": (
+        "purin.solvers.primal_dual_interior_point.reference.residuals",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.state": (
+        "purin.solvers.primal_dual_interior_point.reference.state",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.reference.step": (
+        "purin.solvers.primal_dual_interior_point.reference.step",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.residuals": (
+        "purin.solvers.primal_dual_interior_point.residuals",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.scalar": (
+        "purin.solvers.primal_dual_interior_point.scalar",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.solver": (
+        "purin.solvers.primal_dual_interior_point.solver",
+        False,
+    ),
+    "ohagi.solvers.primal_dual_interior_point.types": (
+        "purin.solvers.primal_dual_interior_point.types",
+        False,
+    ),
+    "ohagi.solvers.registry": ("purin.solvers.registry", False),
+    "ohagi.solvers.semismooth": ("purin.solvers.semismooth", False),
+    "ohagi.solvers.types": ("purin.solvers.types", False),
+}

ohagi/_utils/__init__.py

@@ -0,0 +1,8 @@
+"""Internal scalar validation helpers."""
+
+from .validation import (
+    nonnegative_float as nonnegative_float,
+    open_unit_interval as open_unit_interval,
+    positive_float as positive_float,
+    positive_integer as positive_integer,
+)

ohagi/_utils/serialization.py

@@ -0,0 +1,68 @@
+"""Shared JSON conversion, checksums, and atomic artifact writes."""
+
+from __future__ import annotations
+import hashlib
+import json
+import os
+import tempfile
+from collections.abc import Mapping
+from pathlib import Path
+import jax
+import numpy as np
+
+
+def json_safe(value):
+    """Convert nested numerical values into JSON-compatible Python values."""
+    if isinstance(value, np.generic):
+        return value.item()
+    if isinstance(value, (np.ndarray, jax.Array)):
+        return np.asarray(jax.device_get(value)).tolist()
+    if isinstance(value, Mapping):
+        return {str(key): json_safe(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [json_safe(item) for item in value]
+    return value
+
+
+def sha256(path):
+    """Return the SHA-256 digest of a file without loading it all into memory."""
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def atomic_write(path, data):
+    """Replace a file atomically using a temporary file in the same directory."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = None
+    try:
+        with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as handle:
+            temporary = Path(handle.name)
+            handle.write(data)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary, path)
+    finally:
+        if temporary is not None:
+            temporary.unlink(missing_ok=True)
+
+
+def write_json(path, value):
+    """Atomically write a deterministic, human-readable JSON document."""
+    data = json.dumps(json_safe(value), indent=2, sort_keys=True) + "\n"
+    atomic_write(path, data.encode("utf-8"))
+
+
+def timestamp():
+    """Compact filesystem-friendly run timestamp."""
+    import time
+
+    return time.strftime("%Y%m%d_%H%M%S")
+
+
+def resolve_path(path):
+    """Expand and resolve a filesystem path for artifact access."""
+    return Path(path).expanduser().resolve()

ohagi/_utils/tree.py

@@ -0,0 +1,23 @@
+"""Small PyTree predicates and selections used by training steps."""
+
+from __future__ import annotations
+import jax
+import jax.numpy as jnp
+
+
+def tree_all_finite(tree):
+    """Return a scalar JAX predicate indicating whether every pytree leaf is finite."""
+    leaves = jax.tree_util.tree_leaves(tree)
+    result = jnp.asarray(True)
+    for leaf in leaves:
+        result = result & jnp.all(jnp.isfinite(leaf))
+    return result
+
+
+def tree_select(condition, true_tree, false_tree):
+    """Select corresponding pytree leaves with a scalar or broadcastable predicate."""
+    return jax.tree_util.tree_map(
+        lambda true_value, false_value: jnp.where(condition, true_value, false_value),
+        true_tree,
+        false_tree,
+    )

ohagi/compat/__init__.py

@@ -0,0 +1,5 @@
+"""Identity-preserving legacy module aliases for the three frontend packages."""
+
+from .aliases import install_aliases, reexport_module, registered_aliases
+
+__all__ = ["install_aliases", "reexport_module", "registered_aliases"]

ohagi/compat/aliases.py

@@ -0,0 +1,87 @@
+"""Lazy aliases point to canonical OHAGI or PURIN modules.
+
+Loading the same Python file under two names duplicates classes and registries.
+This loader instead returns the already imported canonical module object. Exact
+alias catalogs prevent unrelated packages from being intercepted.
+"""
+
+import sys
+from importlib import import_module
+from importlib.abc import Loader, MetaPathFinder
+from importlib.util import spec_from_loader
+from importlib.machinery import PathFinder
+from threading import RLock
+
+
+class _AliasLoader(Loader):
+    def __init__(self, target):
+        self.target = target
+
+    def create_module(self, spec):
+        """Create the compatibility namespace requested by the import machinery."""
+        return import_module(self.target)
+
+    def exec_module(self, module):
+        """Populate an imported compatibility namespace with its registered resources."""
+        pass
+
+
+class _AliasFinder(MetaPathFinder):
+    def __init__(self):
+        self.aliases = {}
+        self.lock = RLock()
+
+    def find_spec(self, fullname, path=None, target=None):
+        """Resolve a registered compatibility namespace during module import."""
+        entry = self.aliases.get(fullname)
+        if entry is None:
+            if (
+                any(name.startswith(fullname + ".") for name in self.aliases)
+                and PathFinder.find_spec(fullname, path) is None
+            ):
+                from .plugins import _Namespace
+
+                return spec_from_loader(fullname, _Namespace(), is_package=True)
+            return None
+        destination, is_package = entry
+        return spec_from_loader(
+            fullname, _AliasLoader(destination), is_package=is_package
+        )
+
+
+_FINDER = _AliasFinder()
+
+
+def install_aliases(aliases):
+    """Install exact aliases, rejecting mixed old/new modules in one process."""
+    with _FINDER.lock:
+        for name, entry in aliases.items():
+            target, is_package = entry
+            if name == target:
+                continue
+            previous = _FINDER.aliases.get(name)
+            if previous is not None and previous != (target, is_package):
+                raise ImportError(f"Conflicting alias for {name}")
+            module = sys.modules.get(name)
+            if module is not None and module.__name__ != target:
+                raise ImportError(
+                    f"{name} is already loaded from another implementation; use a fresh Python process"
+                )
+        _FINDER.aliases.update({name: tuple(entry) for name, entry in aliases.items()})
+        if _FINDER not in sys.meta_path:
+            sys.meta_path.insert(0, _FINDER)
+
+
+def reexport_module(namespace, target):
+    """Preserve public and historical private attributes on a root facade."""
+    implementation = import_module(target)
+    for name, value in vars(implementation).items():
+        if not name.startswith("__") or name in {"__all__", "__version__", "__doc__"}:
+            namespace[name] = value
+    namespace["__getattr__"] = lambda name: getattr(implementation, name)
+    namespace["__dir__"] = lambda: sorted(set(namespace) | set(dir(implementation)))
+
+
+def registered_aliases():
+    """Return the currently registered compatibility module aliases."""
+    return dict(_FINDER.aliases)

ohagi/compat/paths.py

@@ -0,0 +1,9 @@
+"""Resolve module names stored before the OHAGI/PURIN package rename."""
+
+
+def canonical_module_name(name: str) -> str:
+    """Translate historical factory namespaces without changing attributes."""
+    for previous, current in (("hax", "ohagi"), ("projx", "purin")):
+        if name == previous or name.startswith(previous + "."):
+            return current + name[len(previous) :]
+    return name

ohagi/compat/plugins.py

@@ -0,0 +1,94 @@
+"""Lazy plugin discovery; the core has no knowledge of frontend packages."""
+
+from importlib.metadata import entry_points
+from functools import lru_cache
+from importlib import import_module
+from importlib.abc import MetaPathFinder, Loader
+from importlib.machinery import PathFinder
+from importlib.util import spec_from_loader
+import sys
+
+
+@lru_cache(maxsize=None)
+def _entries(group):
+    return tuple(entry_points(group=group))
+
+
+_symbols = {}
+
+
+def public_symbols(module):
+    """Return the public symbols registered for a frontend namespace."""
+    prefix = module + "."
+    names = set(_symbols)
+    names.update(ep.name for ep in _entries("ohagi.symbols"))
+    return sorted(
+        key[len(prefix) :]
+        for key in names
+        if key.startswith(prefix) and "." not in key[len(prefix) :]
+    )
+
+
+def register_symbols(values):
+    """Register frontend symbols for compatibility lookup."""
+    _symbols.update(values)
+    for key in values:
+        parent, name = key.rsplit(".", 1)
+        module = sys.modules.get(parent)
+        exports = getattr(module, "__all__", None) if module is not None else None
+        if isinstance(exports, list) and name not in exports:
+            exports.append(name)
+
+
+def public_symbol(module, name):
+    """Resolve a named public symbol from the compatibility registry."""
+    key = module + "." + name
+    destination = _symbols.get(key)
+    if destination is None:
+        matches = tuple(ep for ep in _entries("ohagi.symbols") if ep.name == key)
+        if not matches:
+            raise AttributeError(name)
+        destination = matches[0].value
+    owner, symbol = destination.split(":", 1)
+    return getattr(import_module(owner), symbol)
+
+
+class _Namespace(Loader):
+    def create_module(self, spec):
+        """Create the compatibility namespace requested by the import machinery."""
+        return None
+
+    def exec_module(self, module):
+        """Populate an imported compatibility namespace with its registered resources."""
+        pass
+
+
+class _Plugins(MetaPathFinder):
+    busy = False
+
+    def find_spec(self, fullname, path=None, target=None):
+        """Resolve a registered compatibility namespace during module import."""
+        if self.busy or not fullname.startswith("ohagi."):
+            return None
+        if PathFinder.find_spec(fullname, path) is not None:
+            return None
+        matches = tuple(
+            ep
+            for ep in _entries("ohagi.compat")
+            if fullname == ep.name or fullname.startswith(ep.name + ".")
+        )
+        if matches:
+            self.busy = True
+            try:
+                from .aliases import install_aliases, _FINDER
+
+                install_aliases(matches[0].load())
+                return _FINDER.find_spec(fullname, path, target)
+            finally:
+                self.busy = False
+        if any(ep.name.startswith(fullname + ".") for ep in _entries("ohagi.compat")):
+            return spec_from_loader(fullname, _Namespace(), is_package=True)
+        return None
+
+
+sys.meta_path.append(_Plugins())

ohagi/data/__init__.py

@@ -0,0 +1,4 @@
+"""Deterministic dataset selection, splitting and fixed-shape batching."""
+
+from .batching import fixed_batches as fixed_batches
+from ohagi.transforms.scaling import Standardizer as Standardizer

ohagi/data/batching.py

@@ -0,0 +1,70 @@
+"""Fixed shapes and explicit masks for JIT-compiled steps."""
+
+import numpy as np
+
+
+def fixed_batches(data, batch_size, *, indices=None, drop_last=False):
+    """Yield fixed-shape NumPy batches with explicit validity masks.
+
+    Args:
+        data: Nonempty mapping of arrays sharing the same sample dimension.
+            The reserved key ``mask`` must not be supplied.
+        batch_size: Positive number of rows in each emitted batch.
+        indices: Optional one-dimensional integer selection/order of sample rows.
+        drop_last: Discard incomplete final batches instead of padding them.
+
+    Yields:
+        Mapping of data arrays with leading dimension ``batch_size`` plus a
+        boolean ``mask``. Padded rows repeat the final selected observation and
+        have mask value False, so reducers must honor the mask.
+
+    Raises:
+        ValueError: If sizes, sample dimensions or selected indices are invalid.
+    """
+    if not isinstance(batch_size, int) or batch_size < 1:
+        raise ValueError("batch_size must be positive")
+    if not data:
+        raise ValueError("data must contain arrays")
+    arrays = {key: np.asarray(value) for key, value in data.items()}
+    if any(value.ndim == 0 for value in arrays.values()):
+        raise ValueError("data arrays require a sample dimension")
+    size = len(next(iter(arrays.values())))
+    if size < 1 or any(
+        value.ndim == 0 or len(value) != size for value in arrays.values()
+    ):
+        raise ValueError("data arrays must share a nonempty sample dimension")
+    if "mask" in arrays:
+        raise ValueError("mask is generated by fixed_batches")
+    indices = np.arange(size) if indices is None else np.asarray(indices)
+    if (
+        indices.ndim != 1
+        or not np.issubdtype(indices.dtype, np.integer)
+        or np.any(indices < 0)
+        or np.any(indices >= size)
+    ):
+        raise ValueError("invalid batch indices")
+    if drop_last:
+        indices = indices[: (len(indices) // batch_size) * batch_size]
+    for start in range(0, len(indices), batch_size):
+        selected = indices[start : start + batch_size]
+        count = len(selected)
+        padded = np.pad(selected, (0, batch_size - count), mode="edge")
+        yield {
+            **{key: value[padded] for key, value in arrays.items()},
+            "mask": np.arange(batch_size) < count,
+        }
+
+
+def trimmed_batches(values, batch_size):
+    """Reshape a feature matrix into complete fixed-size batches.
+
+    Args:
+        values: Two-dimensional array of shape ``(samples, features)``.
+        batch_size: Number of observations in each complete batch.
+
+    Returns:
+        Array of shape ``(samples // batch_size, batch_size, features)``.
+        Incomplete trailing rows are discarded; no padding or mask is added.
+    """
+    trimmed = values[: (values.shape[0] // batch_size) * batch_size]
+    return trimmed.reshape((-1, batch_size, values.shape[1]))