Add initial to max and min

pytensor/link/jax/dispatch/elemwise.py

@@ -23,6 +23,7 @@ def elemwise_fn(*inputs):
 @jax_funcify.register(CAReduce)
 def jax_funcify_CAReduce(op, **kwargs):
     axis = op.axis
+    initial = getattr(op, "initial", False)
     op_nfunc_spec = getattr(op, "nfunc_spec", None)
     scalar_nfunc_spec = getattr(op.scalar_op, "nfunc_spec", None)
     scalar_op_name = getattr(op.scalar_op, "name", None)
@@ -46,7 +47,9 @@ def careduce(x):
 
         if op_nfunc_spec:
             jax_op = getattr(jnp, op_nfunc_spec[0])
-            return jax_op(x, axis=axis).astype(acc_dtype)
+            # jnp.all/jnp.any don't take `initial`, so only pass it when set
+            kwargs = {"initial": op.initial_value(x.dtype)} if initial else {}
+            return jax_op(x, axis=axis, **kwargs).astype(acc_dtype)
 
         # The PyTensor `Op` didn't tell us which NumPy equivalent to use (or
         # there isn't one), so we use this fallback approach

pytensor/link/jax/dispatch/math.py

@@ -18,9 +18,7 @@ def jax_funcify_Max(op, **kwargs):
     axis = op.axis
 
     def max(x):
-        max_res = jnp.max(x, axis)
-
-        return max_res
+        return jnp.max(x, axis, initial=op.initial_value(x.dtype))
 
     return max
 

pytensor/link/mlx/dispatch/math.py

@@ -1,7 +1,7 @@
 import mlx.core as mx
 
 from pytensor.link.mlx.dispatch import mlx_funcify
-from pytensor.tensor.math import Argmax, Dot, Max
+from pytensor.tensor.math import Argmax, Dot, Max, Min
 
 
 @mlx_funcify.register(Dot)
@@ -13,8 +13,11 @@ def dot(x, y):
 
 
 @mlx_funcify.register(Max)
-def mlx_funcify_Max(op, node=None, **kwargs):
-    def max_fn(x):
+@mlx_funcify.register(Min)
+def mlx_funcify_Max_Min(op, node=None, **kwargs):
+    reduce_fn = mx.max if isinstance(op, Max) else mx.min
+
+    def max_min_fn(x):
         axes = op.axis
         if axes is None:
             reduce_axes = None
@@ -23,9 +26,16 @@ def max_fn(x):
 
         keepdims = getattr(op, "keepdims", False)
 
-        return mx.max(x, axis=reduce_axes, keepdims=keepdims)
+        if op.initial and x.size == 0:
+            # MLX shapes are concrete, so an empty reduction can be filled directly
+            reduced = range(x.ndim) if reduce_axes is None else reduce_axes
+            out_shape = [s for i, s in enumerate(x.shape) if i not in reduced]
+            fill = op.initial_value(node.outputs[0].type.dtype)
+            return mx.full(out_shape, fill, dtype=x.dtype)
+
+        return reduce_fn(x, axis=reduce_axes, keepdims=keepdims)
 
-    return max_fn
+    return max_min_fn
 
 
 @mlx_funcify.register(Argmax)

pytensor/link/numba/dispatch/elemwise.py

@@ -420,8 +420,10 @@ def careduce_add(x):
     acc_dtype_str = f"np.{acc_dtype.name}"
     careduce_fn_name = f"careduce_{scalar_op}"
 
-    if acc_dtype.kind in "ui" and not np.isfinite(identity):
-        if np.isposinf(identity):
+    if acc_dtype.kind in "uib" and not np.isfinite(identity):
+        if acc_dtype.kind == "b":
+            identity = bool(np.isposinf(identity))
+        elif np.isposinf(identity):
             identity = np.iinfo(acc_dtype).max
         else:
             identity = np.iinfo(acc_dtype).min
@@ -884,7 +886,9 @@ def _reduce_identity(identity, acc_dtype):
     integer accumulators, mirroring ``create_multiaxis_reducer``.
     """
     acc_dtype = np.dtype(acc_dtype)
-    if acc_dtype.kind in "ui" and not np.isfinite(identity):
+    if acc_dtype.kind == "b" and not np.isfinite(identity):
+        identity = bool(np.isposinf(identity))
+    elif acc_dtype.kind in "ui" and not np.isfinite(identity):
         identity = (
             np.iinfo(acc_dtype).max
             if np.isposinf(identity)
@@ -1147,7 +1151,7 @@ def impl(*outer_inputs):
         def ov_fused_elemwise_fn(*outer_inputs):
             return impl_fn
 
-    cache_version = 9
+    cache_version = 10
     if scalar_cache_key is None:
         key = None
     else:
@@ -1198,7 +1202,7 @@ def numba_funcify_CAReduce(op, node, **kwargs):
     )
     careduce_fn = numba_basic.numba_njit(careduce_py_fn, boundscheck=False)
 
-    cache_version = 6
+    cache_version = 7
     careduce_key = sha256(
         str(
             (

pytensor/tensor/math.py

@@ -365,10 +365,55 @@ def __str__(self):
 
 
 class NonZeroDimsCAReduce(FixedOpCAReduce):
+    """A reduction that raises on a zero-size reduced axis, like NumPy's ``max``.
+
+    With ``initial=True`` an empty reduction returns the identity of the scalar
+    ``Op`` (``-inf`` for ``maximum``, ``inf`` for ``minimum``) instead of raising,
+    like NumPy's ``initial``. For integer and boolean inputs the lowest/highest
+    value of the dtype is used.
+    """
+
+    __props__ = (*FixedOpCAReduce.__props__, "initial")
+
+    def __init__(self, scalar_op, axis=None, initial=False, **kwargs):
+        self.initial = bool(initial)
+        super().__init__(scalar_op, axis=axis, **kwargs)
+
+    def __str__(self):
+        if not self.initial:
+            return super().__str__()
+        return f"{type(self).__name__}{{{self._axis_str()}, initial}}"
+
+    def initial_value(self, dtype):
+        """The identity as a value of `dtype` (its bounds for ints and bools), or None."""
+        if not self.initial:
+            return None
+        identity = self.scalar_op.identity
+        dtype = np.dtype(dtype)
+        if dtype.kind in "fc":
+            return identity
+        if dtype.kind == "b":
+            return identity > 0
+        info = np.iinfo(dtype)
+        return info.max if identity > 0 else info.min
+
+    def perform(self, node, inp, out):
+        if not self.initial:
+            return super().perform(node, inp, out)
+        [x] = inp
+        dtype = node.outputs[0].type.dtype
+        out[0][0] = np.asarray(
+            self.ufunc.reduce(x, axis=self.axis, initial=self.initial_value(dtype)),
+            dtype=dtype,
+        )
+
     def _c_all(self, node, name, input_names, output_names, sub):
         setup, alloc, loop, cast = super()._c_all(
             node, name, input_names, output_names, sub
         )
+        if self.initial:
+            # Empty reductions return the scalar_op identity
+            return setup, alloc, loop, cast
 
         # We add an additional check for zero-sized dimensions (This seems like
         # something that could enabled in `elemwise_cgen.make_checks`.)
@@ -415,7 +460,8 @@ class MaxAndMinCAReduce(NonZeroDimsCAReduce):
 
     def clone(self, **kwargs):
         axis = kwargs.get("axis", self.axis)
-        return type(self)(axis=axis)
+        initial = kwargs.get("initial", self.initial)
+        return type(self)(axis=axis, initial=initial)
 
     def pullback(self, inputs, outputs, output_grads):
         # The strict-sense mathematical gradient of a maximum/minimum reduction
@@ -468,18 +514,18 @@ def pushforward(self, inputs, outputs, tangents):
 class Max(MaxAndMinCAReduce):
     nfunc_spec = ("max", 1, 1)
 
-    def __init__(self, axis):
-        super().__init__(ps.maximum, axis)
+    def __init__(self, axis, initial=False):
+        super().__init__(ps.maximum, axis, initial=initial)
 
 
 class Min(MaxAndMinCAReduce):
     nfunc_spec = ("min", 1, 1)
 
-    def __init__(self, axis):
-        super().__init__(ps.minimum, axis)
+    def __init__(self, axis, initial=False):
+        super().__init__(ps.minimum, axis, initial=initial)
 
 
-def max(x, axis=None, keepdims=False):
+def max(x, axis=None, keepdims=False, initial=None):
     """
     Returns maximum elements obtained by iterating over given axis.
 
@@ -492,20 +538,27 @@ def max(x, axis=None, keepdims=False):
         If this is set to True, the axes which are reduced are left in
         the result as dimensions with size one. With this option, the result
         will broadcast correctly against the original tensor.
+    initial: None or -inf
+        If ``-inf``, reducing an empty axis returns ``-inf`` (the lowest value
+        of the dtype for integer inputs) instead of raising. Other values are
+        not supported yet.
 
     Notes
     -----
-    We return an error as numpy when we reduce a dim with a shape of 0.
+    Without ``initial``, we return an error as numpy when we reduce a dim with
+    a shape of 0.
 
     """
-    out = Max(axis=axis)(x)
+    if initial is not None and float(initial) != -np.inf:
+        raise ValueError(f"initial must be None or -inf for max, got {initial}")
+    out = Max(axis=axis, initial=initial is not None)(x)
 
     if keepdims:
         out = makeKeepDims(x, out, axis)
     return out
 
 
-def min(x, axis=None, keepdims=False):
+def min(x, axis=None, keepdims=False, initial=None):
     """
     Returns minimum elements obtained by iterating over given axis.
 
@@ -518,18 +571,31 @@ def min(x, axis=None, keepdims=False):
         If this is set to True, the axes which are reduced are left in
         the result as dimensions with size one. With this option, the result
         will broadcast correctly against the original tensor.
+    initial: None or inf
+        If ``inf``, reducing an empty axis returns ``inf`` (the highest value
+        of the dtype for integer inputs) instead of raising. Other values are
+        not supported yet.
 
     """
     x = as_tensor_variable(x)
+    if initial is not None and float(initial) != np.inf:
+        raise ValueError(f"initial must be None or inf for min, got {initial}")
+    max_initial = None if initial is None else -np.inf
     str_x_type = str(x.dtype)
+    if str_x_type in int_dtypes and initial is not None:
+        # -max(-x) would overflow on the empty result
+        out = Min(axis=axis, initial=True)(x)
+        return makeKeepDims(x, out, axis) if keepdims else out
     if str_x_type.startswith("float") or str_x_type in int_dtypes:
-        return -max(-x, axis=axis, keepdims=keepdims)
+        return -max(-x, axis=axis, keepdims=keepdims, initial=max_initial)
     elif str_x_type in uint_dtypes:
         itype = np.iinfo(x.dtype)
         max_val = np.array(itype.max, dtype=itype.dtype)
-        return max_val - max(max_val - x, axis=axis, keepdims=keepdims)
+        return max_val - max(
+            max_val - x, axis=axis, keepdims=keepdims, initial=max_initial
+        )
     elif str_x_type == "bool":
-        return ~max(~x, axis=axis, keepdims=keepdims)
+        return ~max(~x, axis=axis, keepdims=keepdims, initial=max_initial)
     else:
         # Be careful about unsigned integers, complex
         raise NotImplementedError()

pytensor/tensor/rewriting/math.py

@@ -2167,13 +2167,20 @@ def local_reduce_chain(fgraph, node) -> list[TensorVariable] | None:
     if outer_op.scalar_op != inner_op.scalar_op:
         return None
 
+    # Keeping `initial` from either side only turns an empty-axis error into it
+    clone_kwargs = {}
+    if getattr(inner_op, "initial", False):
+        if not hasattr(outer_op, "initial"):
+            return None
+        clone_kwargs["initial"] = True
+
     outer_axis = outer_op.axis
     inner_axis = inner_op.axis
     [x] = inner_reduce.owner.inputs
     # check to see either the inner or outer prod is doing a
     # product over all axis, in which case we can remove it
     if outer_axis is None or inner_axis is None:
-        return [outer_op.clone(axis=None)(x)]
+        return [outer_op.clone(axis=None, **clone_kwargs)(x)]
 
     # Merge axis
     newaxis = list(inner_axis)
@@ -2186,7 +2193,7 @@ def local_reduce_chain(fgraph, node) -> list[TensorVariable] | None:
         newaxis.append(new_i)
 
     assert len(newaxis) == len(inner_axis) + len(outer_axis)
-    return [outer_op.clone(axis=sorted(newaxis))(x)]
+    return [outer_op.clone(axis=sorted(newaxis), **clone_kwargs)(x)]
 
 
 @register_canonicalize
@@ -2282,17 +2289,8 @@ def local_reduce_broadcastable(fgraph, node):
                     ii += 1
             new_reduced = reduced.dimshuffle(*pattern)
             if new_axis:
-                if type(node.op) is CAReduce:
-                    # This case handles `CAReduce` instances
-                    # (e.g. generated by `scalar_elemwise`), and not the
-                    # scalar `Op`-specific subclasses
-                    # TODO FIXME: This highlights a major design flaw in
-                    # `CAReduce` (or at least our use of it), and it needs
-                    # to be fixed
-                    new_op = node.op.__class__(node.op.scalar_op, axis=new_axis)
-                else:
-                    new_op = node.op.__class__(axis=new_axis)
-                return [new_op(new_reduced)]
+                # `clone` keeps the other props of the reduction, like `initial`
+                return [node.op.clone(axis=new_axis)(new_reduced)]
             else:
                 # -- in this case we can remove the reduction completely
                 return [new_reduced.astype(odtype)]
@@ -2352,21 +2350,29 @@ def local_careduce_of_alloc(fgraph, node):
     # them into a multiplication by their size, ``Prod``/``ProdWithoutZeros``
     # into a power; idempotent reductions (max/min/all/any) are unaffected, as
     # repeating a value doesn't change its maximum, minimum, or truth.
-    if isinstance(node.op, Sum | Prod | ProdWithoutZeros):
-        size_shapes = [shapes[a] for a in axis if a < offset or value_bcast[a - offset]]
-        if size_shapes:
-            size = variadic_mul(*size_shapes)
-            if value.dtype in ("float16", "float32"):
-                # Avoid a float64 upcast from the int64 shapes (or a float16
-                # downcast); either would be amplified by the mul/pow below.
-                size = size.astype("float32")
-            if isinstance(node.op, Sum):
-                value = value * size
-            elif isinstance(node.op, Prod):
-                value = value**size
-            else:
-                # An empty ProdWithoutZeros reduction is 0, not 1
-                value = switch(eq(size, 0), 0, value**size)
+    size_shapes = [shapes[a] for a in axis if a < offset or value_bcast[a - offset]]
+    if size_shapes and isinstance(node.op, Sum | Prod | ProdWithoutZeros):
+        size = variadic_mul(*size_shapes)
+        if value.dtype in ("float16", "float32"):
+            # Avoid a float64 upcast from the int64 shapes (or a float16
+            # downcast); either would be amplified by the mul/pow below.
+            size = size.astype("float32")
+        if isinstance(node.op, Sum):
+            value = value * size
+        elif isinstance(node.op, Prod):
+            value = value**size
+        else:
+            # An empty ProdWithoutZeros reduction is 0, not 1
+            value = switch(eq(size, 0), 0, value**size)
+
+    # With `initial`, a broadcast axis that is empty at runtime makes the result `initial`
+    if size_shapes and getattr(node.op, "initial", False):
+        dtype = node.outputs[0].dtype
+        value = switch(
+            eq(variadic_mul(*size_shapes), 0),
+            np.asarray(node.op.initial_value(dtype), dtype=dtype),
+            value,
+        )
 
     # The reduction may change the dtype; a single elemwise has no accumulation
     # error, so ignore acc_dtype and just cast to the reduction's output dtype.

pytensor/tensor/rewriting/uncanonicalize.py

@@ -64,7 +64,9 @@ def local_max_to_min(fgraph, node):
         ):
             neg_node = max.owner.inputs[0]
             if neg_node.owner and neg_node.owner.op == neg:
-                new = Min(max.owner.op.axis)(neg_node.owner.inputs[0])
+                new = Min(
+                    max.owner.op.axis, initial=getattr(max.owner.op, "initial", False)
+                )(neg_node.owner.inputs[0])
                 return [copy_stack_trace(node.outputs[0], new)]
 
     return False

pytensor/tensor/special.py

@@ -1,5 +1,7 @@
 from functools import reduce
 
+import numpy as np
+
 from pytensor.gradient import DisconnectedType, disconnected_type
 from pytensor.graph.replace import _vectorize_node
 from pytensor.tensor import as_tensor_variable
@@ -118,7 +120,7 @@ def __init__(self, *, axis, **kwargs):
         super().__init__(**kwargs)
 
     def build_inner_graph(self, x):
-        x_max = x.max(axis=self.axis, keepdims=True)
+        x_max = x.max(axis=self.axis, keepdims=True, initial=-np.inf)
         # Do not offset when x_max = -inf, to avoid nan in the output
         x_max = switch(isinf(x_max), 0.0, x_max)
         out = log(exp(x - x_max).sum(axis=self.axis, keepdims=True)) + x_max