multipers 2.3.3__cp312-cp312-macosx_11_0_arm64.whl

multipers/array_api/__init__.py

@@ -0,0 +1,62 @@
+import multipers.array_api.numpy as npapi
+
+
+def api_from_tensor(x, *, verbose: bool = False, strict=False):
+    if strict:
+        if npapi.is_tensor(x):
+            return npapi
+        import multipers.array_api.torch as torchapi
+
+        if torchapi.is_tensor(x):
+            return torchapi
+        raise ValueError(f"Unsupported (strict) type {type(x)=}")
+    if npapi.is_promotable(x):
+        if verbose:
+            print("using numpy backend")
+        return npapi
+    import multipers.array_api.torch as torchapi
+
+    if torchapi.is_promotable(x):
+        if verbose:
+            print("using torch backend")
+        return torchapi
+    raise ValueError(f"Unsupported type {type(x)=}")
+
+
+def api_from_tensors(*args):
+    assert len(args) > 0, "no tensor given"
+    import multipers.array_api.numpy as npapi
+
+    is_numpy = True
+    for x in args:
+        if not npapi.is_promotable(x):
+            is_numpy = False
+            break
+    if is_numpy:
+        return npapi
+
+    # only torch for now
+    import multipers.array_api.torch as torchapi
+
+    is_torch = True
+    for x in args:
+        if not torchapi.is_promotable(x):
+            is_torch = False
+            break
+    if is_torch:
+        return torchapi
+    raise ValueError(f"Incompatible types got {[type(x) for x in args]=}.")
+
+
+def to_numpy(x):
+    api = api_from_tensor(x)
+    return api.asnumpy(x)
+
+
+def check_keops():
+    import os
+
+    if os.name == "nt":
+        # see https://github.com/getkeops/keops/pull/421
+        return False
+    return npapi.check_keops()

multipers/array_api/numpy.py

@@ -0,0 +1,104 @@
+from contextlib import nullcontext
+
+import numpy as _np
+from scipy.spatial.distance import cdist
+
+backend = _np
+cat = _np.concatenate
+norm = _np.linalg.norm
+astensor = _np.asarray
+asnumpy = _np.asarray
+tensor = _np.array
+stack = _np.stack
+empty = _np.empty
+where = _np.where
+no_grad = nullcontext
+zeros = _np.zeros
+min = _np.min
+max = _np.max
+repeat_interleave = _np.repeat
+cdist = cdist  # type: ignore[no-redef]
+unique = _np.unique
+inf = _np.inf
+searchsorted = _np.searchsorted
+LazyTensor = None
+
+# Test keops
+_is_keops_available = None
+
+
+def check_keops():
+    global _is_keops_available, LazyTensor
+    if _is_keops_available is not None:
+        return _is_keops_available
+    import pykeops.numpy as pknp
+    from pykeops.numpy import LazyTensor as LT
+
+    formula = "SqNorm2(x - y)"
+    var = ["x = Vi(3)", "y = Vj(3)"]
+    expected_res = _np.array([63.0, 90.0])
+    x = _np.arange(1, 10).reshape(-1, 3).astype("float32")
+    y = _np.arange(3, 9).reshape(-1, 3).astype("float32")
+
+    my_conv = pknp.Genred(formula, var)
+    try:
+        _is_keops_available = _np.allclose(my_conv(x, y).flatten(), expected_res)
+        LazyTensor = LT
+    except:
+        from warnings import warn
+
+        warn("Could not initialize keops (numpy). using workarounds")
+        _is_keops_available = False
+
+    return _is_keops_available
+
+
+def from_numpy(x):
+    return _np.asarray(x)
+
+
+def ascontiguous(x):
+    return _np.ascontiguousarray(x)
+
+
+def sort(x, axis=-1):
+    return _np.sort(x, axis=axis)
+
+
+def device(x):  # type: ignore[no-unused-arg]
+    return None
+
+
+# type: ignore[no-unused-arg]
+def linspace(low, high, r, device=None, dtype=None):
+    return _np.linspace(low, high, r, dtype=dtype)
+
+
+def cartesian_product(*arrays, dtype=None):
+    mesh = _np.meshgrid(*arrays, indexing="ij")
+    coordinates = _np.stack(mesh, axis=-1).reshape(-1, len(arrays)).astype(dtype)
+    return coordinates
+
+
+def quantile_closest(x, q, axis=None):
+    return _np.quantile(x, q, axis=axis, method="closest_observation")
+
+
+def minvalues(x: _np.ndarray, **kwargs):
+    return _np.min(x, **kwargs)
+
+
+def maxvalues(x: _np.ndarray, **kwargs):
+    return _np.max(x, **kwargs)
+
+
+def is_tensor(x):
+    return isinstance(x, _np.ndarray)
+
+
+def is_promotable(x):
+    return isinstance(x, _np.ndarray | list | tuple)
+
+
+def has_grad(_):
+    return False

multipers/array_api/torch.py

@@ -0,0 +1,117 @@
+import numpy as _np
+import torch as _t
+
+backend = _t
+cat = _t.cat
+norm = _t.norm
+astensor = _t.as_tensor
+tensor = _t.tensor
+stack = _t.stack
+empty = _t.empty
+where = _t.where
+no_grad = _t.no_grad
+cdist = _t.cdist
+zeros = _t.zeros
+min = _t.min
+max = _t.max
+repeat_interleave = _t.repeat_interleave
+linspace = _t.linspace
+cartesian_product = _t.cartesian_prod
+inf = _t.inf
+searchsorted = _t.searchsorted
+LazyTensor = None
+
+
+_is_keops_available = None
+
+
+def check_keops():
+    global _is_keops_available, LazyTensor
+    if _is_keops_available is not None:
+        return _is_keops_available
+    try:
+        import pykeops.torch as pknp
+        from pykeops.torch import LazyTensor as LT
+
+        formula = "SqNorm2(x - y)"
+        var = ["x = Vi(3)", "y = Vj(3)"]
+        expected_res = _t.tensor([63.0, 90.0])
+        x = _t.arange(1, 10, dtype=_t.float32).view(-1, 3)
+        y = _t.arange(3, 9, dtype=_t.float32).view(-1, 3)
+
+        my_conv = pknp.Genred(formula, var)
+        _is_keops_available = _t.allclose(
+            my_conv(x, y).view(-1), expected_res.type(_t.float32)
+        )
+        LazyTensor = LT
+
+    except:
+        from warnings import warn
+
+        warn("Could not initialize keops (torch). using workarounds")
+
+        _is_keops_available = False
+
+    return _is_keops_available
+check_keops()
+
+
+def from_numpy(x):
+    return _t.from_numpy(x)
+
+
+def ascontiguous(x):
+    return _t.as_tensor(x).contiguous()
+
+
+def device(x):
+    return x.device
+
+
+def sort(x, axis=-1):
+    return _t.sort(x, dim=axis).values
+
+
+# in our context, this allows to get a correct gradient.
+def unique(x, assume_sorted=False, _mean=True):
+    if not x.requires_grad:
+        return x.unique(sorted=assume_sorted)
+    if x.ndim != 1:
+        raise ValueError(f"Got ndim!=1. {x=}")
+    if not assume_sorted:
+        x = x.sort().values
+    _, c = _t.unique(x, sorted=True, return_counts=True)
+    if _mean:
+        x = _t.segment_reduce(data=x, reduce="mean", lengths=c, unsafe=True, axis=0)
+    else:
+        c = _np.concatenate([[0], _np.cumsum(c[:-1])])
+        x = x[c]
+    return x
+
+
+def quantile_closest(x, q, axis=None):
+    return _t.quantile(x, q, dim=axis, interpolation="nearest")
+
+
+def minvalues(x: _t.Tensor, **kwargs):
+    return _t.min(x, **kwargs).values
+
+
+def maxvalues(x: _t.Tensor, **kwargs):
+    return _t.max(x, **kwargs).values
+
+
+def asnumpy(x):
+    return x.detach().numpy()
+
+
+def is_tensor(x):
+    return isinstance(x, _t.Tensor)
+
+
+def is_promotable(x):
+    return isinstance(x, _t.Tensor)
+
+
+def has_grad(x):
+    return x.requires_grad