ndlike 0.2.2__tar.gz → 0.2.4__tar.gz

{ndlike-0.2.2/src/ndlike.egg-info → ndlike-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ndlike
-Version: 0.2.2
+Version: 0.2.4
 Summary: A toolkit for defining custom array types on top of NumPy, JAX, and PyTorch.
 Author-email: Clemens Damke <clemens@cortys.de>
 License-Expression: MIT

{ndlike-0.2.2 → ndlike-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ndlike"
-version = "0.2.2"
+version = "0.2.4"
 description = "A toolkit for defining custom array types on top of NumPy, JAX, and PyTorch."
 readme = "README.md"
 license = "MIT"

{ndlike-0.2.2 → ndlike-0.2.4}/src/ndlike/jax_functions.py

@@ -318,6 +318,34 @@ def jax_divmod(
     return jnp.divmod(x1, x2)  # ty: ignore[invalid-argument-type]
 
 
+@jax_function_dispatch("x")
+def jax_frexp(x: _JaxFamily | JaxArrayLike, /) -> tuple[jax.Array | _JaxFamily, jax.Array | _JaxFamily]:
+    """Decompose values into mantissas and exponents, mirroring ``jax.numpy.frexp``.
+
+    Args:
+        x: Input values.
+
+    Returns:
+        A tuple of floating-point mantissas and integer exponents for native
+        inputs, or the result supplied by a custom override.
+    """
+    return jnp.frexp(x)  # ty: ignore[invalid-argument-type]
+
+
+@jax_function_dispatch("x")
+def jax_modf(x: _JaxFamily | JaxArrayLike, /) -> tuple[jax.Array | _JaxFamily, jax.Array | _JaxFamily]:
+    """Split values into fractional and integral parts, mirroring ``jax.numpy.modf``.
+
+    Args:
+        x: Input values.
+
+    Returns:
+        A tuple of fractional and integral parts for native inputs, or the result
+        supplied by a custom override.
+    """
+    return jnp.modf(x)  # ty: ignore[invalid-argument-type]
+
+
 @jax_function_dispatch("x1", "x2")
 def jax_power(x1: _JaxFamily | JaxArrayLike, x2: _JaxFamily | JaxArrayLike, /) -> jax.Array | _JaxFamily:
     """Raise operands to powers elementwise, mirroring ``jax.numpy.power``.
@@ -966,6 +994,7 @@ __all__ = [
     "jax_equal",
     "jax_expand_dims",
     "jax_floor_divide",
+    "jax_frexp",
     "jax_function_dispatch",
     "jax_greater",
     "jax_greater_equal",
@@ -976,6 +1005,7 @@ __all__ = [
     "jax_matmul",
     "jax_matrix_transpose",
     "jax_mean",
+    "jax_modf",
     "jax_moveaxis",
     "jax_multiply",
     "jax_negative",

{ndlike-0.2.2 → ndlike-0.2.4}/src/ndlike/jax_operators.py

@@ -70,8 +70,8 @@ class JaxOperatorsMixin:  # noqa: PLW1641
     Subclasses can define ``__array_ufunc__`` to override or disable that behavior.
     """
 
-    # Prefer reflected operators in mixed NumPy infix expressions.
-    __array_priority__ = 1000
+    # Let NumPy handle its own infix operations without forced deferral.
+    __array_priority__ = 0
 
     __add__ = _binary(jf.jax_add)
     __radd__ = _binary(jf.jax_add, reflected=True)

{ndlike-0.2.2 → ndlike-0.2.4}/src/ndlike/protected_axis/jax.py

@@ -442,8 +442,8 @@ class JaxAxisProtected[J: JaxLike | jax.Array | np.ndarray](JaxLikeImplementatio
             operands: Original operands in expression order.
 
         Returns:
-            The field values to validate and reconstruct. For ``divmod``, this
-            hook is called separately for the quotient and remainder.
+            The field values to validate and reconstruct. For ``divmod``,
+            ``frexp``, and ``modf``, this hook is called separately for each output.
         """
         del func, operands
         return values

{ndlike-0.2.2 → ndlike-0.2.4}/src/ndlike/protected_axis/jax_functions.py

@@ -678,6 +678,8 @@ def _batch_comparison(values: dict[str, Any], template: JaxAxisProtectedInternal
         jf.jax_floor_divide,
         jf.jax_remainder,
         jf.jax_divmod,
+        jf.jax_frexp,
+        jf.jax_modf,
         jf.jax_power,
         jf.jax_matmul,
         jf.jax_bitwise_and,
@@ -721,9 +723,10 @@ def protected_operator_function(
                 operation = jax.vmap(operation, in_axes=in_axes, out_axes=-1)
         results[name] = operation(*field_args)
 
+    multiple_outputs = func in (jf.jax_divmod, jf.jax_frexp, jf.jax_modf)
     result_sets = (
         tuple({name: value[index] for name, value in results.items()} for index in range(2))
-        if func is jf.jax_divmod
+        if multiple_outputs
         else (results,)
     )
     reconstructed = []
@@ -734,4 +737,4 @@ def protected_operator_function(
         processed = owner._postprocess_elementwise_result(values, func=func, operands=args)  # noqa: SLF001
         _validate_operator_result(processed, template)
         reconstructed.append(template.create(processed))
-    return tuple(reconstructed) if func is jf.jax_divmod else reconstructed[0]
+    return tuple(reconstructed) if multiple_outputs else reconstructed[0]

{ndlike-0.2.2 → ndlike-0.2.4}/src/ndlike/protected_axis/numpy.py

@@ -377,7 +377,9 @@ class NumpyAxisProtected[T: NumpyLike | np.ndarray](NumpyLikeImplementation[T],
             else:
                 out_values.append(None)
 
-        if out is not None and len(protected_axes) != 1 and any(values is None for values in out_values):
+        if len(protected_axes) != 1 and any(
+            item is not None and values is None for item, values in zip(out_items, out_values, strict=True)
+        ):
             msg = "non-protected out is only supported for single-field protected objects."
             raise TypeError(msg)
 
@@ -398,7 +400,8 @@ class NumpyAxisProtected[T: NumpyLike | np.ndarray](NumpyLikeImplementation[T],
                 raise TypeError(msg)
             return normalize_axis(axis, batch_ndim)
 
-        results: dict[str, NumpyProtectedValue] = {}
+        output_count = ufunc.nout if method == "__call__" else 1
+        results: list[dict[str, NumpyProtectedValue]] = [{} for _ in range(output_count)]
         for name, axes_count in protected_axes.items():
             field_kwargs = dict(kwargs)
 
@@ -445,43 +448,49 @@ class NumpyAxisProtected[T: NumpyLike | np.ndarray](NumpyLikeImplementation[T],
 
             result = getattr(ufunc, method)(*field_input_items, **field_kwargs)
 
-            if out is not None:
-                continue
+            result_items = result if output_count > 1 else (result,)
+            for index, result_item in enumerate(result_items):
+                if out is not None and out_items[index] is not None:
+                    continue
 
-            if axes_count == 0 and not hasattr(result, "ndim"):
-                result = np.asarray(result)
+                item = result_item
+                if axes_count == 0 and not hasattr(item, "ndim"):
+                    item = np.asarray(item)
 
-            result_ndim = value_ndim(result)
-            result_shape = value_shape(result)
-            original_shape = value_shape(self_values[name])
-            if result_ndim < axes_count:
-                msg = f"Ufunc operation removed protected trailing axes for field {name!r}."
-                raise ValueError(msg)
-            if protected_shape(result_shape, axes_count) != protected_shape(original_shape, axes_count):
-                msg = f"Ufunc operation modified protected trailing axes for field {name!r}."
-                raise ValueError(msg)
+                result_ndim = value_ndim(item)
+                result_shape = value_shape(item)
+                original_shape = value_shape(self_values[name])
+                if result_ndim < axes_count:
+                    msg = f"Ufunc operation removed protected trailing axes for field {name!r}."
+                    raise ValueError(msg)
+                if protected_shape(result_shape, axes_count) != protected_shape(original_shape, axes_count):
+                    msg = f"Ufunc operation modified protected trailing axes for field {name!r}."
+                    raise ValueError(msg)
+
+                if isinstance(item, np.ndarray):
+                    item = self._postprocess_ufunc_result(item, ufunc=ufunc, method=method)
 
-            if isinstance(result, np.ndarray):
-                result = self._postprocess_ufunc_result(result, ufunc=ufunc, method=method)
+                results[index][name] = cast("NumpyProtectedValue", item)
 
-            results[name] = cast("NumpyProtectedValue", result)
+        outputs: list[object] = []
+        for index, values in enumerate(results):
+            if out is not None and out_items[index] is not None:
+                outputs.append(out_items[index])
+                continue
 
-        if out is not None:
-            if method == "__call__":
-                return out
-            return out_items[0] if isinstance(out, tuple) else out
+            expected_batch_shape: tuple[int, ...] | None = None
+            for name, value in values.items():
+                axes_count = protected_axes[name]
+                current_batch_shape = batch_shape(value_shape(value), axes_count)
+                if expected_batch_shape is None:
+                    expected_batch_shape = current_batch_shape
+                elif current_batch_shape != expected_batch_shape:
+                    msg = "Ufunc operation produced inconsistent batch-shapes across protected fields."
+                    raise ValueError(msg)
 
-        expected_batch_shape: tuple[int, ...] | None = None
-        for name, value in results.items():
-            axes_count = protected_axes[name]
-            current_batch_shape = batch_shape(value_shape(value), axes_count)
-            if expected_batch_shape is None:
-                expected_batch_shape = current_batch_shape
-            elif current_batch_shape != expected_batch_shape:
-                msg = "Ufunc operation produced inconsistent batch-shapes across protected fields."
-                raise ValueError(msg)
+            outputs.append(self.with_protected_values(values))
 
-        return self.with_protected_values(results)
+        return tuple(outputs) if output_count > 1 else outputs[0]
 
     def __array_function__(
         self,

{ndlike-0.2.2 → ndlike-0.2.4}/src/ndlike/torch_operators.py

@@ -66,8 +66,8 @@ class TorchOperatorsMixin:  # noqa: PLW1641
     Subclasses can define ``__array_ufunc__`` to override or disable that behavior.
     """
 
-    # Prefer reflected operators in mixed NumPy infix expressions.
-    __array_priority__ = 1000
+    # Let NumPy handle its own infix operations without forced deferral.
+    __array_priority__ = 0
 
     __add__ = _binary(torch.add)
     __radd__ = _binary(torch.add, reflected=True)

{ndlike-0.2.2 → ndlike-0.2.4/src/ndlike.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ndlike
-Version: 0.2.2
+Version: 0.2.4
 Summary: A toolkit for defining custom array types on top of NumPy, JAX, and PyTorch.
 Author-email: Clemens Damke <clemens@cortys.de>
 License-Expression: MIT