statgpu 0.1.0__py3-none-any.whl

statgpu/backends/_array_ops.py

@@ -0,0 +1,529 @@
+"""
+Backend utilities for GLM loss functions.
+
+Provides wrapper functions that dispatch to numpy/cupy/torch
+based on the input array type, so GLM loss functions can use
+a single code path for all backends.
+"""
+
+import numpy as np
+
+from statgpu.backends._base import _resolve_backend
+
+
+def _xp(arr):
+    """Get the array module (numpy/cupy/torch) from array type."""
+    mod = type(arr).__module__
+    if mod.startswith("cupy"):
+        import cupy
+        return cupy
+    if mod.startswith("torch"):
+        import torch
+        return torch
+    import numpy
+    return numpy
+
+
+def _clip(arr, lo, hi):
+    """Clip array values."""
+    xp = _xp(arr)
+    if xp.__name__ == "torch":
+        if lo is not None and hi is not None:
+            return xp.clamp(arr, min=lo, max=hi)
+        if lo is not None:
+            return xp.clamp(arr, min=lo)
+        if hi is not None:
+            return xp.clamp(arr, max=hi)
+        return arr
+    return xp.clip(arr, lo, hi)
+
+
+def _exp(arr):
+    """Element-wise exponential."""
+    xp = _xp(arr)
+    return xp.exp(arr)
+
+
+def _log(arr):
+    """Element-wise natural log."""
+    xp = _xp(arr)
+    return xp.log(arr)
+
+
+def _log1p(arr):
+    """Element-wise log(1+x)."""
+    xp = _xp(arr)
+    return xp.log1p(arr)
+
+
+def _sigmoid(arr):
+    """Numerically stable sigmoid: 1 / (1 + exp(-x))."""
+    xp = _xp(arr)
+    # float32 overflows exp() at ~89; float64 at ~709
+    dtype = getattr(arr, 'dtype', None)
+    max_val = 88.0 if dtype is not None and '32' in str(dtype) else 700.0
+    z = _clip(arr, -max_val, max_val)
+    if xp.__name__ == "torch":
+        return xp.sigmoid(z)
+    return 1.0 / (1.0 + xp.exp(-z))
+
+
+def _softplus(x):
+    """Numerically stable softplus: log(1 + exp(x))."""
+    xp = _xp(x)
+    if xp.__name__ == "torch":
+        import torch.nn.functional as F
+        return F.softplus(x)
+    return xp.log1p(xp.exp(-xp.abs(x))) + _clip(x, 0.0, None)
+
+
+def _sum(arr):
+    """Sum of all elements."""
+    xp = _xp(arr)
+    return xp.sum(arr)
+
+
+def _eigvalsh(arr):
+    """Eigenvalues of a symmetric matrix (sorted ascending)."""
+    xp = _xp(arr)
+    return xp.linalg.eigvalsh(arr)
+
+
+def _zeros_like(arr):
+    """Create zeros array with same shape/type as arr."""
+    xp = _xp(arr)
+    return xp.zeros_like(arr)
+
+
+def _zeros(n, backend, ref_tensor=None, dtype=None):
+    """Create a 1-D zeros vector on the requested backend."""
+    backend = _resolve_backend(backend, ref_tensor)
+    if backend == "numpy":
+        return np.zeros(n, dtype=dtype)
+    if backend == "cupy":
+        import cupy as cp
+        out_dtype = (
+            dtype if dtype is not None else getattr(ref_tensor, "dtype", cp.float64)
+        )
+        return cp.zeros(n, dtype=out_dtype)
+    import torch
+    device = getattr(ref_tensor, "device", "cpu") if ref_tensor is not None else "cpu"
+    out_dtype = dtype or (
+        getattr(ref_tensor, "dtype", torch.float64)
+        if ref_tensor is not None
+        else torch.float64
+    )
+    return torch.zeros(n, device=device, dtype=out_dtype)
+
+
+def _copy_arr(arr):
+    """Copy array: .clone() for torch, .copy() for numpy/cupy."""
+    if hasattr(arr, "clone"):
+        return arr.clone()
+    return arr.copy()
+
+
+def _diag(reg, backend="auto", ref_tensor=None, dtype=None):
+    """Create a diagonal matrix on the requested backend."""
+    backend = _resolve_backend(backend, ref_tensor, reg)
+    if backend == "cupy":
+        import cupy as cp
+        out_dtype = dtype if dtype is not None else getattr(reg, "dtype", cp.float64)
+        return cp.diag(cp.asarray(reg, dtype=out_dtype))
+    if backend == "torch":
+        import torch
+        device = (
+            ref_tensor.device
+            if ref_tensor is not None
+            else getattr(reg, "device", "cpu")
+        )
+        out_dtype = dtype or (
+            ref_tensor.dtype
+            if ref_tensor is not None
+            and getattr(ref_tensor, "is_floating_point", lambda: False)()
+            else reg.dtype
+            if hasattr(reg, "is_floating_point")
+            and reg.is_floating_point()
+            else torch.float64
+        )
+        return torch.diag(torch.as_tensor(reg, dtype=out_dtype, device=device))
+    arr = np.asarray(reg, dtype=dtype) if dtype is not None else reg
+    return np.diag(arr)
+
+
+def _to_backend(arr, backend="auto", ref_tensor=None, dtype=None):
+    """Convert an array to the requested backend, matching ref_tensor when needed."""
+    backend = _resolve_backend(backend, ref_tensor, arr)
+    if backend == "cupy":
+        import cupy as cp
+        out_dtype = dtype
+        if out_dtype is None:
+            ref_dtype = getattr(ref_tensor, "dtype", None)
+            if ref_dtype is not None and 'float' in str(ref_dtype):
+                out_dtype = ref_dtype
+            else:
+                out_dtype = cp.float64
+        return cp.asarray(arr, dtype=out_dtype)
+    if backend == "torch":
+        import torch
+        device = (
+            ref_tensor.device
+            if ref_tensor is not None
+            else getattr(arr, "device", "cpu")
+        )
+        out_dtype = dtype or (
+            ref_tensor.dtype
+            if ref_tensor is not None
+            and getattr(ref_tensor, "is_floating_point", lambda: False)()
+            else arr.dtype
+            if hasattr(arr, "is_floating_point")
+            and arr.is_floating_point()
+            else torch.float64
+        )
+        return torch.as_tensor(arr, dtype=out_dtype, device=device)
+    return np.asarray(arr, dtype=dtype or float)
+
+
+def _solve_linear_system(A, b, backend="auto"):
+    """Solve a linear system, falling back to least squares if singular."""
+    backend = _resolve_backend(backend, A)
+    try:
+        if backend == "torch":
+            import torch
+            b_col = b.unsqueeze(1) if b.ndim == 1 else b
+            sol = torch.linalg.solve(A, b_col)
+            return sol.squeeze(1) if b.ndim == 1 else sol
+        if backend == "cupy":
+            import cupy as cp
+            return cp.linalg.solve(A, b)
+        return np.linalg.solve(A, b)
+    except (np.linalg.LinAlgError, RuntimeError):
+        # LinAlgError for numpy/cupy singular matrices
+        # RuntimeError for torch singular matrices
+        if backend == "torch":
+            import torch
+            b_col = b.unsqueeze(1) if b.ndim == 1 else b
+            sol = torch.linalg.lstsq(A, b_col).solution
+            return sol.squeeze(1) if b.ndim == 1 else sol
+        if backend == "cupy":
+            import cupy as cp
+            return cp.linalg.lstsq(A, b)[0]
+        return np.linalg.lstsq(A, b, rcond=None)[0]
+
+
+def _eye_like(n, ref):
+    """Create an identity matrix on the same backend/device as ref."""
+    backend = _resolve_backend("auto", ref)
+    if backend == "cupy":
+        import cupy as cp
+        return cp.eye(n, dtype=ref.dtype)
+    if backend == "torch":
+        import torch
+        return torch.eye(n, dtype=ref.dtype, device=ref.device)
+    return np.eye(n, dtype=getattr(ref, "dtype", np.float64))
+
+
+def _sync_scalars(*dev_vals, backend):
+    """Batch device scalars into Python floats with one backend sync point."""
+    backend = _resolve_backend(backend, *dev_vals)
+    if backend == "numpy":
+        return tuple(float(v) for v in dev_vals)
+    if backend == "torch":
+        import torch
+        ref = next(
+            (
+                v
+                for v in dev_vals
+                if type(v).__module__.startswith("torch")
+            ),
+            None,
+        )
+        device = getattr(ref, "device", None)
+        dtype = getattr(ref, "dtype", torch.float64)
+        stacked = torch.stack(
+            [torch.as_tensor(v, device=device, dtype=dtype) for v in dev_vals]
+        )
+        return tuple(stacked[i].item() for i in range(len(dev_vals)))
+    import cupy as cp
+    stacked = cp.stack([cp.asarray(v) for v in dev_vals])
+    return tuple(float(stacked[i]) for i in range(len(dev_vals)))
+
+
+def _abs_sum(x):
+    """Sum of absolute values, returned as a Python scalar."""
+    xp = _xp(x)
+    if xp.__name__ == "torch":
+        return float(xp.sum(xp.abs(x)).item())
+    return float(xp.sum(xp.abs(x)))
+
+
+def _abs_max(x):
+    """Max absolute value, returned as a Python scalar."""
+    xp = _xp(x)
+    if xp.__name__ == "torch":
+        return float(xp.max(xp.abs(x)).item())
+    return float(xp.max(xp.abs(x)))
+
+
+def _norm2(x):
+    """L2 norm, returned as a Python scalar."""
+    xp = _xp(x)
+    if xp.__name__ == "torch":
+        return float(xp.linalg.norm(x).item())
+    return float(xp.linalg.norm(x))
+
+
+def _dot(a, b):
+    """Dot product, returned as a Python scalar."""
+    val = a.dot(b)
+    return float(val.item() if hasattr(val, "item") else val)
+
+
+def _dot_dev(a, b):
+    """Dot product staying on device for GPU backends."""
+    if isinstance(a, np.ndarray):
+        return float(a.dot(b))
+    return a.dot(b)
+
+
+def _sum_sq(x):
+    """Sum of squares, returned as a Python scalar."""
+    xp = _xp(x)
+    val = xp.sum(x ** 2)
+    return float(val.item() if hasattr(val, "item") else val)
+
+
+def _sum_sq_dev(x):
+    """Sum of squares staying on device for GPU backends."""
+    xp = _xp(x)
+    val = xp.sum(x ** 2)
+    if xp.__name__ == "numpy":
+        return float(val)
+    return val
+
+
+def _norm2_dev(x):
+    """L2 norm staying on device for GPU backends."""
+    xp = _xp(x)
+    val = xp.linalg.norm(x)
+    if xp.__name__ == "numpy":
+        return float(val)
+    return val
+
+
+def _abs_sum_dev(x):
+    """Sum of absolute values staying on device for GPU backends."""
+    xp = _xp(x)
+    val = xp.sum(xp.abs(x))
+    if xp.__name__ == "numpy":
+        return float(val)
+    return val
+
+
+def _device_leq(a, b):
+    """Device-side a <= b comparison, returned as a Python bool."""
+    backend = _resolve_backend("auto", a, b)
+    if backend == "torch":
+        return bool((a <= b).item())
+    if backend == "cupy":
+        return bool(a <= b)
+    return a <= b
+
+
+def _device_gt(a, b):
+    """Device-side a > b comparison, returned as a Python bool."""
+    backend = _resolve_backend("auto", a, b)
+    if backend == "torch":
+        return bool((a > b).item())
+    if backend == "cupy":
+        return bool(a > b)
+    return a > b
+
+
+def _clip_grad_on_device(grad, coef_old, backend):
+    """Clip gradient entirely on the selected backend."""
+    # Lazy import to avoid circular dependency (backends <-> solvers)
+    from statgpu.solvers._constants import (
+        _GRAD_CLIP_COEF_FACTOR, _GRAD_CLIP_ABS_FLOOR, _GRAD_CLIP_MAX,
+    )
+    if backend == "numpy":
+        gn = float(np.linalg.norm(grad))
+        ca = float(np.sum(np.abs(coef_old)))
+        gmax = max(ca * _GRAD_CLIP_COEF_FACTOR + _GRAD_CLIP_ABS_FLOOR, _GRAD_CLIP_MAX)
+        if gn > gmax:
+            return grad * (gmax / gn)
+        return grad
+    if backend == "torch":
+        import torch
+        gn_sq = torch.sum(grad ** 2)
+        coef_abs = torch.sum(torch.abs(coef_old))
+        gmax = coef_abs * _GRAD_CLIP_COEF_FACTOR + _GRAD_CLIP_ABS_FLOOR
+        gmax = torch.clamp(gmax, min=_GRAD_CLIP_MAX)
+        scale = torch.where(
+            gn_sq > gmax * gmax,
+            gmax / torch.sqrt(gn_sq + 1e-30),
+            torch.ones(1, device=grad.device, dtype=grad.dtype),
+        )
+        return grad * scale
+    import cupy as cp
+    gn_sq = cp.sum(grad ** 2)
+    coef_abs = cp.sum(cp.abs(coef_old))
+    gmax = cp.maximum(coef_abs * _GRAD_CLIP_COEF_FACTOR + _GRAD_CLIP_ABS_FLOOR, _GRAD_CLIP_MAX)
+    scale = cp.where(
+        gn_sq > gmax * gmax,
+        gmax / cp.sqrt(gn_sq + 1e-30),
+        cp.ones(1, dtype=grad.dtype),
+    )
+    return grad * scale
+
+
+def _max_eigval_power(mat, n_iter=20, tol=1e-8):
+    """Largest eigenvalue of a symmetric matrix via power iteration.
+
+    Much faster than full eigendecomposition, especially on GPU
+    where cuSOLVER eigvalsh has high kernel compilation overhead.
+    O(p^2) vs O(p^3). Accuracy within 1% for 20 iterations.
+
+    Parameters
+    ----------
+    mat : 2-d array (p, p), symmetric positive semi-definite.
+    n_iter : int
+        Max power iterations.
+    tol : float
+        Early stopping tolerance on eigenvalue change.
+
+    Returns
+    -------
+    float : max eigenvalue estimate.
+    """
+    xp = _xp(mat)
+    p = mat.shape[0]
+    dtype = getattr(mat, 'dtype', None)
+    # Build a deterministic but non-constant seed vector to avoid the
+    # pathological case where an all-ones vector is orthogonal to the top
+    # eigenspace (e.g., [[1,-1],[-1,1]]).
+    if xp.__name__ == "torch":
+        v = xp.arange(1, p + 1, dtype=dtype, device=mat.device)
+    elif dtype is not None:
+        v = xp.arange(1, p + 1, dtype=dtype)
+    else:
+        v = xp.arange(1, p + 1, dtype=xp.float64)
+
+    v_norm = xp.sqrt(xp.dot(v, v))
+    v_norm_val = float(v_norm)
+    if v_norm_val < 1e-15:
+        return 1.0
+    v = v / v_norm
+
+    if xp.__name__ == "numpy":
+        lambda_old = 0.0
+        lambda_new = 0.0
+        for _ in range(n_iter):
+            v_new = mat @ v
+            # Cache dot(v_new, v_new) to avoid recomputing mat @ v.
+            nv2 = xp.dot(v_new, v_new)
+            v_norm_sq = float(nv2)
+            if v_norm_sq < 1e-30:
+                return 1.0
+            v_norm = v_norm_sq ** 0.5
+            v = v_new / v_norm
+            # lambda = v^T A v = v^T v_new (v_new = A v, already computed)
+            lambda_new = float(xp.dot(v, v_new))
+            if lambda_old > 0 and abs(lambda_new - lambda_old) < tol * abs(lambda_new):
+                break
+            lambda_old = lambda_new
+        return lambda_new
+
+    lambda_old = 0.0
+    lambda_val = 0.0
+    for i in range(n_iter):
+        v_new = mat @ v
+        dot_vn_vn = xp.dot(v_new, v_new)
+        v_norm_sq = float(dot_vn_vn.item() if hasattr(dot_vn_vn, "item") else dot_vn_vn)
+        if v_norm_sq < 1e-30:
+            return 1.0  # Zero matrix — same fallback as numpy path
+        v_norm = v_norm_sq ** 0.5
+        v = v_new / v_norm
+        lambda_new = xp.dot(v, v_new)
+        lambda_val = float(lambda_new.item() if hasattr(lambda_new, "item") else lambda_new)
+        if i > 0 and abs(lambda_val - lambda_old) < tol * abs(lambda_val):
+            return lambda_val
+        lambda_old = lambda_val
+    return lambda_val
+
+
+def _soft_threshold(w, thresh):
+    """Soft-thresholding operator: sign(w) * max(|w| - thresh, 0).
+
+    Works across numpy/cupy/torch.  ``thresh`` may be a scalar or an
+    array with the same shape as ``w`` (adaptive weights).
+
+    Uses ``xp.where`` for fewer intermediate arrays (2 vs 4 with
+    sign*clip formulation).
+    """
+    xp = _xp(w)
+    abs_w = xp.abs(w)
+    # +0.0 eliminates negative zeros from sign(w)
+    return (xp.where(abs_w > thresh, abs_w - thresh, 0.0) * xp.sign(w)) + 0.0
+
+
+def _scalar_tensor(val, ref_arr):
+    """Create a scalar value compatible with *ref_arr*'s backend/device.
+
+    For torch, returns a 0-d tensor on the same device and dtype.
+    For cupy/numpy, returns a plain Python float (scalars work directly).
+    """
+    xp = _xp(ref_arr)
+    if xp.__name__ == "torch":
+        import torch
+        return torch.tensor(val, dtype=ref_arr.dtype, device=ref_arr.device)
+    return float(val)
+
+
+def _xp_copy(arr):
+    """Copy array on the same backend.  `.clone()` for torch, `.copy()` for others."""
+    xp = _xp(arr)
+    if xp.__name__ == "torch":
+        return arr.clone()
+    return arr.copy()
+
+
+def _xp_zeros(shape, dtype, ref_arr):
+    """Create zeros array on the same device/dtype as *ref_arr*."""
+    xp = _xp(ref_arr)
+    if xp.__name__ == "torch":
+        import torch
+        return torch.zeros(shape, dtype=dtype or ref_arr.dtype, device=ref_arr.device)
+    return xp.zeros(shape, dtype=dtype or getattr(ref_arr, 'dtype', None))
+
+
+def _xp_asarray(arr, dtype, ref_arr):
+    """Convert array to the same backend/device as *ref_arr*.
+
+    Handles numpy→cupy, numpy→torch, and same-backend dtype casts.
+    """
+    xp = _xp(ref_arr)
+    if xp.__name__ == "torch":
+        import torch
+        if isinstance(arr, torch.Tensor):
+            out = arr.to(dtype=dtype, device=ref_arr.device)
+        else:
+            out = torch.as_tensor(np.asarray(arr, dtype=np.float64),
+                                  dtype=dtype, device=ref_arr.device)
+        return out
+    if xp.__name__ == "cupy":
+        # Convert torch dtypes to numpy for cupy compatibility
+        if hasattr(dtype, '__module__') and 'torch' in str(getattr(dtype, '__module__', '')):
+            from statgpu.backends._utils import _torch_dtype_to_np
+            dtype = _torch_dtype_to_np(dtype)
+        return xp.asarray(arr, dtype=dtype)
+    return np.asarray(arr, dtype=dtype)
+
+
+def _xp_eye(n, dtype, ref_arr):
+    """Create identity matrix on the same device/dtype as *ref_arr*."""
+    xp = _xp(ref_arr)
+    if xp.__name__ == "torch":
+        import torch
+        return torch.eye(n, dtype=dtype or ref_arr.dtype, device=ref_arr.device)
+    return xp.eye(n, dtype=dtype or getattr(ref_arr, 'dtype', None))

statgpu/backends/_base.py

@@ -0,0 +1,184 @@
+"""
+Abstract base class for compute backends.
+
+A backend wraps an array library (NumPy, CuPy, or PyTorch) and exposes a
+uniform interface so that model implementations can stay array-library agnostic.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+import numpy as np
+
+
+# ---------------------------------------------------------------------------
+# Array-type detection helpers (deferred imports to avoid hard deps)
+# ---------------------------------------------------------------------------
+
+def _is_cupy_array(x: Any) -> bool:
+    """Return True if *x* is a CuPy ndarray."""
+    try:
+        import cupy as cp
+        return isinstance(x, cp.ndarray)
+    except Exception:
+        return False
+
+
+def _is_torch_array(x: Any) -> bool:
+    """Return True if *x* is a PyTorch Tensor."""
+    try:
+        import torch
+        return isinstance(x, torch.Tensor)
+    except Exception:
+        return False
+
+
+def _resolve_backend(backend: str, *arrays) -> str:
+    """Resolve the named *backend* string to one of ``'numpy'``, ``'cupy'``,
+    ``'torch'``.
+
+    Accepts legacy aliases ``'cpu'`` → ``'numpy'`` and ``'cuda'``/``'gpu'`` → ``'cupy'``.
+    When *backend* is ``'auto'``, inspect *arrays* and return the
+    matching backend name based on the first recognised array type.
+    Falls back to ``'numpy'`` when no array matches.
+    """
+    backend_name = str(backend).strip().lower()
+    backend_name = {"cpu": "numpy", "cuda": "cupy", "gpu": "cupy"}.get(
+        backend_name, backend_name
+    )
+    if backend_name not in ("auto", "numpy", "cupy", "torch"):
+        raise ValueError(
+            "backend must be one of: 'auto', 'numpy', 'cupy', 'torch', "
+            "or legacy aliases 'cpu', 'cuda', 'gpu'"
+        )
+    if backend_name != "auto":
+        return backend_name
+
+    for arr in arrays:
+        if arr is not None:
+            if _is_torch_array(arr):
+                return "torch"
+            if _is_cupy_array(arr):
+                return "cupy"
+    return "numpy"
+
+
+class BackendBase(ABC):
+    """
+    Abstract base for compute backends.
+
+    Subclasses wrap a specific array library and expose:
+
+    * ``xp``        – the underlying array module (numpy / cupy / torch).
+    * ``asarray``   – convert arbitrary inputs to the backend's native array.
+    * ``to_numpy``  – convert the backend's arrays back to ``numpy.ndarray``.
+    * ``is_available`` – runtime check for the library being usable.
+
+    The ``xp`` object follows the NumPy array API so that operations such as
+    ``xp.linalg.solve``, ``xp.sum``, ``xp.exp`` etc. work without
+    library-specific branches in the calling code.
+    """
+
+    #: Short name used in repr and config ('numpy', 'cupy', 'torch').
+    name: str = ""
+
+    # ------------------------------------------------------------------
+    # Abstract interface
+    # ------------------------------------------------------------------
+
+    @property
+    @abstractmethod
+    def xp(self) -> Any:
+        """Return the array module (numpy / cupy / torch)."""
+
+    @abstractmethod
+    def asarray(self, x, dtype=None) -> Any:
+        """
+        Convert *x* to this backend's native array type.
+
+        Parameters
+        ----------
+        x : array-like, numpy.ndarray, cupy.ndarray, or torch.Tensor
+            Input data.
+        dtype : dtype-like, optional
+            Desired data type.
+
+        Returns
+        -------
+        array
+            Native array on the backend's device.
+        """
+
+    @abstractmethod
+    def to_numpy(self, x) -> np.ndarray:
+        """
+        Convert *x* to a ``numpy.ndarray``.
+
+        Parameters
+        ----------
+        x : array-like
+            A native array produced by this backend (or any array-like).
+
+        Returns
+        -------
+        numpy.ndarray
+        """
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Return True if this backend can be used in the current environment."""
+
+    # ------------------------------------------------------------------
+    # Convenience helpers (non-abstract, built on top of xp)
+    # ------------------------------------------------------------------
+
+    def solve(self, A, b):
+        """Solve the linear system *Ax = b*."""
+        return self.xp.linalg.solve(A, b)
+
+    def lstsq(self, A, b, rcond=None):
+        """Return the least-squares solution to *Ax ≈ b*."""
+        return self.xp.linalg.lstsq(A, b, rcond=rcond)
+
+    def astype(self, arr, dtype):
+        """Cast *arr* to *dtype* (backend-agnostic .astype / .to)."""
+        return arr.astype(dtype)
+
+    def concatenate(self, arrays, axis=0):
+        """Concatenate *arrays* along *axis* (.concatenate / .cat)."""
+        return self.xp.concatenate(arrays, axis=axis)
+
+    def take_along_axis(self, arr, indices, axis):
+        """Gather elements along *axis* (.take_along_axis / .take_along_dim)."""
+        return self.xp.take_along_axis(arr, indices, axis=axis)
+
+    def cummin(self, arr, axis=0):
+        """Cumulative minimum along *axis*."""
+        return self.xp.minimum.accumulate(arr, axis=axis)
+
+    def cummax(self, arr, axis=0):
+        """Cumulative maximum along *axis*."""
+        return self.xp.maximum.accumulate(arr, axis=axis)
+
+    def flip(self, arr, axis=0):
+        """Reverse the order of elements along *axis*."""
+        return self.xp.flip(arr, axis=axis)
+
+    def copy(self, arr):
+        """Return a copy of *arr*."""
+        return arr.copy()
+
+    def reshape(self, arr, shape):
+        """Reshape *arr* to *shape*."""
+        return arr.reshape(shape)
+
+    def logsumexp(self, arr, axis=None):
+        """Log-sum-exp along *axis*."""
+        import numpy as np
+        xp = self.xp
+        m = xp.max(arr, axis=axis, keepdims=True)
+        return xp.squeeze(m, axis=axis) + xp.log(xp.sum(xp.exp(arr - m), axis=axis))
+
+    def __repr__(self) -> str:
+        available = "available" if self.is_available() else "unavailable"
+        return f"{self.__class__.__name__}(name={self.name!r}, {available})"