From e3e9b6cbe06ff080eb1186fe18e848e39868abdb Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Tue, 29 Sep 2026 22:54:49 +0530
Subject: [PATCH 01/11] Add initial to max and min

---
 pytensor/link/jax/dispatch/elemwise.py      | 15 ++++
 pytensor/link/jax/dispatch/math.py          |  9 +++
 pytensor/link/mlx/dispatch/math.py          |  3 +
 pytensor/link/numba/dispatch/elemwise.py    | 14 ++--
 pytensor/link/pytorch/dispatch/elemwise.py  |  6 ++
 pytensor/tensor/math.py                     | 84 ++++++++++++++++++---
 pytensor/tensor/rewriting/math.py           |  8 ++
 pytensor/tensor/rewriting/uncanonicalize.py |  5 +-
 pytensor/tensor/special.py                  |  4 +-
 pytensor/tensor/variable.py                 |  8 +-
 tests/link/jax/test_math.py                 |  8 ++
 tests/link/numba/test_elemwise.py           |  7 ++
 tests/tensor/rewriting/test_math.py         | 21 ++++++
 tests/tensor/test_math.py                   | 26 +++++++
 tests/tensor/test_special.py                |  7 ++
 15 files changed, 202 insertions(+), 23 deletions(-)

diff --git a/pytensor/link/jax/dispatch/elemwise.py b/pytensor/link/jax/dispatch/elemwise.py
index dfdcd6d54e..c48e9ae295 100644
--- a/pytensor/link/jax/dispatch/elemwise.py
+++ b/pytensor/link/jax/dispatch/elemwise.py
@@ -20,9 +20,20 @@ def elemwise_fn(*inputs):
     return elemwise_fn
 
 
+def careduce_initial(initial, dtype):
+    """Cast the `initial` of a `Max`/`Min` to a value representable in `dtype`."""
+    if initial is None or jnp.issubdtype(dtype, jnp.floating):
+        return initial
+    if dtype == jnp.bool_:
+        return initial > 0
+    info = jnp.iinfo(dtype)
+    return info.max if initial > 0 else info.min
+
+
 @jax_funcify.register(CAReduce)
 def jax_funcify_CAReduce(op, **kwargs):
     axis = op.axis
+    initial = getattr(op, "initial", None)
     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,6 +57,10 @@ def careduce(x):
 
         if op_nfunc_spec:
             jax_op = getattr(jnp, op_nfunc_spec[0])
+            if initial is not None:
+                return jax_op(
+                    x, axis=axis, initial=careduce_initial(initial, x.dtype)
+                ).astype(acc_dtype)
             return jax_op(x, axis=axis).astype(acc_dtype)
 
         # The PyTensor `Op` didn't tell us which NumPy equivalent to use (or
diff --git a/pytensor/link/jax/dispatch/math.py b/pytensor/link/jax/dispatch/math.py
index 9aa6076ee7..a741612229 100644
--- a/pytensor/link/jax/dispatch/math.py
+++ b/pytensor/link/jax/dispatch/math.py
@@ -2,6 +2,7 @@
 import numpy as np
 
 from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.link.jax.dispatch.elemwise import careduce_initial
 from pytensor.tensor.math import Argmax, Dot, Max
 
 
@@ -16,6 +17,14 @@ def dot(x, y):
 @jax_funcify.register(Max)
 def jax_funcify_Max(op, **kwargs):
     axis = op.axis
+    initial = op.initial
+
+    if initial is not None:
+
+        def max(x):
+            return jnp.max(x, axis, initial=careduce_initial(initial, x.dtype))
+
+        return max
 
     def max(x):
         max_res = jnp.max(x, axis)
diff --git a/pytensor/link/mlx/dispatch/math.py b/pytensor/link/mlx/dispatch/math.py
index 2f927634ad..800f4ee4de 100644
--- a/pytensor/link/mlx/dispatch/math.py
+++ b/pytensor/link/mlx/dispatch/math.py
@@ -14,6 +14,9 @@ def dot(x, y):
 
 @mlx_funcify.register(Max)
 def mlx_funcify_Max(op, node=None, **kwargs):
+    if op.initial is not None:
+        raise NotImplementedError("Max with initial is not supported in MLX")
+
     def max_fn(x):
         axes = op.axis
         if axes is None:
diff --git a/pytensor/link/numba/dispatch/elemwise.py b/pytensor/link/numba/dispatch/elemwise.py
index 029b7e42e0..9ba35d994d 100644
--- a/pytensor/link/numba/dispatch/elemwise.py
+++ b/pytensor/link/numba/dispatch/elemwise.py
@@ -403,8 +403,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
@@ -766,7 +768,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)
@@ -994,7 +998,7 @@ def impl(*outer_inputs):
         def ov_fused_elemwise_fn(*outer_inputs):
             return impl_fn
 
-    cache_version = 4
+    cache_version = 5
     if scalar_cache_key is None:
         key = None
     else:
@@ -1045,7 +1049,7 @@ def numba_funcify_CAReduce(op, node, **kwargs):
     )
     careduce_fn = numba_basic.numba_njit(careduce_py_fn, boundscheck=False)
 
-    cache_version = 4
+    cache_version = 5
     careduce_key = sha256(
         str(
             (
diff --git a/pytensor/link/pytorch/dispatch/elemwise.py b/pytensor/link/pytorch/dispatch/elemwise.py
index be7f091c40..fd5ccf0039 100644
--- a/pytensor/link/pytorch/dispatch/elemwise.py
+++ b/pytensor/link/pytorch/dispatch/elemwise.py
@@ -105,6 +105,9 @@ def torch_any(x):
 
 @pytorch_funcify.register(Max)
 def pytorch_funcify_max(op, **kwargs):
+    if op.initial is not None:
+        raise NotImplementedError("Max with initial is not supported in PyTorch")
+
     def torch_max(x):
         if isinstance(op.axis, tuple):
             for d in sorted(op.axis, reverse=True):
@@ -118,6 +121,9 @@ def torch_max(x):
 
 @pytorch_funcify.register(Min)
 def pytorch_funcify_min(op, **kwargs):
+    if op.initial is not None:
+        raise NotImplementedError("Min with initial is not supported in PyTorch")
+
     def torch_min(x):
         if isinstance(op.axis, tuple):
             for d in sorted(op.axis, reverse=True):
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index bf76ab97a1..0e99fe1aab 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -369,6 +369,9 @@ 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 getattr(self, "initial", None) is not None:
+            # Empty reductions return the scalar_op identity, which is `initial`
+            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`.)
@@ -411,11 +414,50 @@ class MaxAndMinCAReduce(NonZeroDimsCAReduce):
 
     Subclasses only need to bind the appropriate scalar ``Op`` (``maximum`` or
     ``minimum``) in their ``__init__`` and set ``nfunc_spec``.
+
+    ``initial`` can be ``None`` (reducing an empty axis raises) or the identity of
+    the scalar ``Op`` (``-inf`` for ``Max``, ``inf`` for ``Min``), in which case an
+    empty reduction returns it, like NumPy's ``initial``. For integer and boolean
+    inputs the lowest/highest value of the dtype is used.
     """
 
+    __props__ = (*NonZeroDimsCAReduce.__props__, "initial")
+
+    def __init__(self, scalar_op, axis, initial=None):
+        if initial is not None:
+            initial = float(initial)
+            if initial != scalar_op.identity:
+                raise ValueError(
+                    f"initial must be None or {scalar_op.identity} for {type(self).__name__}, got {initial}"
+                )
+        self.initial = initial
+        super().__init__(scalar_op, axis)
+
     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 __str__(self):
+        if self.initial is None:
+            return super().__str__()
+        return f"{type(self).__name__}{{{self._axis_str()}, initial={self.initial}}}"
+
+    def perform(self, node, inp, out):
+        if self.initial is None:
+            return super().perform(node, inp, out)
+        [x] = inp
+        dtype = np.dtype(node.outputs[0].type.dtype)
+        if dtype.kind == "f":
+            initial = self.initial
+        elif dtype.kind == "b":
+            initial = self.initial > 0
+        else:
+            info = np.iinfo(dtype)
+            initial = info.max if self.initial > 0 else info.min
+        out[0][0] = np.asarray(
+            self.ufunc.reduce(x, axis=self.axis, initial=initial), dtype=dtype
+        )
 
     def pullback(self, inputs, outputs, output_grads):
         # The strict-sense mathematical gradient of a maximum/minimum reduction
@@ -468,18 +510,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=None):
+        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=None):
+        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 +534,25 @@ 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)
+    out = Max(axis=axis, initial=initial)(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 +565,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=initial)(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()
diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index 6005412b1a..bd99c11a09 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -2116,6 +2116,10 @@ def local_reduce_chain(fgraph, node) -> list[TensorVariable] | None:
     if outer_op.scalar_op != inner_op.scalar_op:
         return None
 
+    # Merging would change what an empty inner or outer reduction returns
+    if getattr(outer_op, "initial", None) != getattr(inner_op, "initial", None):
+        return None
+
     outer_axis = outer_op.axis
     inner_axis = inner_op.axis
     [x] = inner_reduce.owner.inputs
@@ -2272,6 +2276,10 @@ def local_careduce_of_alloc(fgraph, node):
         case _:
             return None
 
+    # A broadcast axis may be empty at runtime, where the result is `initial`
+    if getattr(node.op, "initial", None) is not None:
+        return None
+
     ndim = len(shapes)
     axis = node.op.axis
     axis = tuple(range(ndim)) if axis is None else axis
diff --git a/pytensor/tensor/rewriting/uncanonicalize.py b/pytensor/tensor/rewriting/uncanonicalize.py
index f9f4f950ee..b0bbbfff98 100644
--- a/pytensor/tensor/rewriting/uncanonicalize.py
+++ b/pytensor/tensor/rewriting/uncanonicalize.py
@@ -64,7 +64,10 @@ 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])
+                initial = getattr(max.owner.op, "initial", None)
+                new = Min(
+                    max.owner.op.axis, initial=None if initial is None else -initial
+                )(neg_node.owner.inputs[0])
                 return [copy_stack_trace(node.outputs[0], new)]
 
     return False
diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py
index 94968e1cfe..4a073d5a91 100644
--- a/pytensor/tensor/special.py
+++ b/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
diff --git a/pytensor/tensor/variable.py b/pytensor/tensor/variable.py
index 2f232f4772..3ce38ae45a 100644
--- a/pytensor/tensor/variable.py
+++ b/pytensor/tensor/variable.py
@@ -675,13 +675,13 @@ def std(self, axis=None, ddof=0, keepdims=False, corrected=False):
             self, axis=axis, ddof=ddof, keepdims=keepdims, corrected=corrected
         )
 
-    def min(self, axis=None, keepdims=False):
+    def min(self, axis=None, keepdims=False, initial=None):
         """See :func:`pytensor.tensor.math.min`."""
-        return pt.math.min(self, axis, keepdims=keepdims)
+        return pt.math.min(self, axis, keepdims=keepdims, initial=initial)
 
-    def max(self, axis=None, keepdims=False):
+    def max(self, axis=None, keepdims=False, initial=None):
         """See :func:`pytensor.tensor.math.max`."""
-        return pt.math.max(self, axis, keepdims=keepdims)
+        return pt.math.max(self, axis, keepdims=keepdims, initial=initial)
 
     def argmin(self, axis=None, keepdims=False):
         """See :func:`pytensor.tensor.math.argmin`."""
diff --git a/tests/link/jax/test_math.py b/tests/link/jax/test_math.py
index 9f0172675a..d994e226ec 100644
--- a/tests/link/jax/test_math.py
+++ b/tests/link/jax/test_math.py
@@ -4,6 +4,7 @@
 from pytensor.configdefaults import config
 from pytensor.tensor.math import Argmax, Max, maximum
 from pytensor.tensor.math import max as pt_max
+from pytensor.tensor.math import min as pt_min
 from pytensor.tensor.type import dvector, matrix, scalar, vector
 from tests.link.jax.test_basic import compare_jax_and_py
 
@@ -21,6 +22,13 @@ def test_jax_max_and_argmax():
     compare_jax_and_py([x], [out], [np.r_[1, 2]])
 
 
+@pytest.mark.parametrize("dtype", ["float64", "int64"])
+def test_jax_max_min_initial(dtype):
+    x = matrix("x", dtype=dtype)
+    outs = [pt_max(x, axis=0, initial=-np.inf), pt_min(x, axis=0, initial=np.inf)]
+    compare_jax_and_py([x], outs, [np.zeros((0, 3), dtype=dtype)])
+
+
 def test_dot():
     y = vector("y")
     y_test_value = np.r_[1.0, 2.0].astype(config.floatX)
diff --git a/tests/link/numba/test_elemwise.py b/tests/link/numba/test_elemwise.py
index 95c28e9434..80db426ced 100644
--- a/tests/link/numba/test_elemwise.py
+++ b/tests/link/numba/test_elemwise.py
@@ -408,6 +408,13 @@ def test_CAReduce_discrete_infinity_identity(axis):
     )
 
 
+@pytest.mark.parametrize("dtype", ["int8", "bool"])
+def test_CAReduce_discrete_empty_initial(dtype):
+    x = pt.matrix("x", dtype=dtype)
+    outs = [x.max(axis=0, initial=-np.inf), x.min(axis=0, initial=np.inf)]
+    compare_numba_and_py([x], outs, [np.zeros((0, 3), dtype=dtype)])
+
+
 def test_scalar_Elemwise_Clip():
     a = pt.scalar("a")
     b = pt.scalar("b")
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index f6683b621d..af9015036d 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -2989,6 +2989,17 @@ def test_all(self):
         [new_out] = local_reduce_chain.transform(fg, out.owner)
         assert equal_computations([new_out], [x.all(axis=(0, 2))])
 
+    def test_max_initial(self):
+        x = tensor3()
+        out = x.max(axis=-1, initial=-np.inf).max(axis=0)
+        fg = FunctionGraph([x], [out], clone=False)
+        assert local_reduce_chain.transform(fg, out.owner) is None
+
+        out = x.max(axis=-1, initial=-np.inf).max(axis=0, initial=-np.inf)
+        fg = FunctionGraph([x], [out], clone=False)
+        [new_out] = local_reduce_chain.transform(fg, out.owner)
+        assert equal_computations([new_out], [x.max(axis=(0, 2), initial=-np.inf)])
+
 
 class TestLocalSumProd:
     """Test sum/prod rewrites."""
@@ -3457,6 +3468,16 @@ def test_local_careduce_of_alloc_mixed_axes(self):
         result.assert_graph(pt.alloc(v * np.array([6]), 4, 3))
         result.assert_eval(v_val)
 
+    def test_local_careduce_of_alloc_initial(self):
+        # The broadcast axis may be empty at runtime, where the result is `initial`
+        v = vector("v")
+        n = scalar("n", dtype="int64")
+        out = Max(axis=0, initial=-np.inf)(pt.alloc(v, n, 3))
+        f = function([v, n], out, mode=get_default_mode().including("specialize"))
+        np.testing.assert_array_equal(
+            f(np.ones(3, dtype=v.dtype), 0), np.full(3, -np.inf)
+        )
+
     @pytest.mark.parametrize("reduce_op", [Max, Min, All, Any])
     def test_local_careduce_of_alloc_idempotent(self, reduce_op):
         # Idempotent reductions (max/min/all/any) over the alloc's broadcast axes
diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py
index f3f7257b3a..e0f0aa87cf 100644
--- a/tests/tensor/test_math.py
+++ b/tests/tensor/test_math.py
@@ -1266,6 +1266,32 @@ class TestMinMax:
     def setup_method(self):
         Max.debug = 0
 
+    @pytest.mark.parametrize("dtype", ["float64", "int64", "uint8", "bool"])
+    def test_initial(self, dtype):
+        x = tensor("x", shape=(None, 3), dtype=dtype)
+        f = function(
+            [x], [max(x, axis=0, initial=-np.inf), min(x, axis=0, initial=np.inf)]
+        )
+        if dtype == "float64":
+            lo, hi = -np.inf, np.inf
+        elif dtype == "bool":
+            lo, hi = False, True
+        else:
+            lo, hi = np.iinfo(dtype).min, np.iinfo(dtype).max
+        empty_max, empty_min = f(np.zeros((0, 3), dtype=dtype))
+        assert_array_equal(empty_max, np.full(3, lo, dtype=dtype))
+        assert_array_equal(empty_min, np.full(3, hi, dtype=dtype))
+
+        data = np.array([[1, 0, 1], [0, 0, 1]], dtype=dtype)
+        res_max, res_min = f(data)
+        assert_array_equal(res_max, data.max(axis=0))
+        assert_array_equal(res_min, data.min(axis=0))
+
+        with pytest.raises(ValueError):
+            max(x, initial=0)
+        with pytest.raises(ValueError):
+            min(x, initial=-np.inf)
+
     def test_scalar(self):
         for fct in [max, min]:
             n = as_tensor_variable(5.0)
diff --git a/tests/tensor/test_special.py b/tests/tensor/test_special.py
index 49ddbb15f4..317fa1508b 100644
--- a/tests/tensor/test_special.py
+++ b/tests/tensor/test_special.py
@@ -263,6 +263,13 @@ def test_logsumexp(shape, axis, keepdims):
     )
 
 
+def test_logsumexp_zero_size():
+    # https://github.com/pymc-devs/pytensor/issues/603
+    x = tensor("x", shape=(None, 2))
+    out = function([x], logsumexp(x))(np.zeros((0, 2), dtype=x.dtype))
+    assert out == -np.inf
+
+
 @pytest.mark.parametrize("mode", ["FAST_RUN", "FAST_COMPILE"])
 def test_logsumexp_logaddexp_stable_grad(mode):
     """Both helpers promise a stable forward, so their gradient must be stable too.

From af3036bb1f978588d3a3905aa753dd26d650ac47 Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Wed, 30 Sep 2026 02:35:40 +0530
Subject: [PATCH 02/11] Keep initial when chaining reductions, handle empty
 inputs in MLX, share initial_value

---
 pytensor/link/jax/dispatch/elemwise.py     | 16 +++-------------
 pytensor/link/jax/dispatch/math.py         |  3 +--
 pytensor/link/mlx/dispatch/math.py         | 11 ++++++++---
 pytensor/link/pytorch/dispatch/elemwise.py |  6 ------
 pytensor/tensor/math.py                    | 22 +++++++++++++---------
 pytensor/tensor/rewriting/math.py          | 13 ++++++++-----
 tests/tensor/rewriting/test_math.py        | 13 ++++++-------
 7 files changed, 39 insertions(+), 45 deletions(-)

diff --git a/pytensor/link/jax/dispatch/elemwise.py b/pytensor/link/jax/dispatch/elemwise.py
index c48e9ae295..586f589535 100644
--- a/pytensor/link/jax/dispatch/elemwise.py
+++ b/pytensor/link/jax/dispatch/elemwise.py
@@ -20,16 +20,6 @@ def elemwise_fn(*inputs):
     return elemwise_fn
 
 
-def careduce_initial(initial, dtype):
-    """Cast the `initial` of a `Max`/`Min` to a value representable in `dtype`."""
-    if initial is None or jnp.issubdtype(dtype, jnp.floating):
-        return initial
-    if dtype == jnp.bool_:
-        return initial > 0
-    info = jnp.iinfo(dtype)
-    return info.max if initial > 0 else info.min
-
-
 @jax_funcify.register(CAReduce)
 def jax_funcify_CAReduce(op, **kwargs):
     axis = op.axis
@@ -58,9 +48,9 @@ def careduce(x):
         if op_nfunc_spec:
             jax_op = getattr(jnp, op_nfunc_spec[0])
             if initial is not None:
-                return jax_op(
-                    x, axis=axis, initial=careduce_initial(initial, x.dtype)
-                ).astype(acc_dtype)
+                return jax_op(x, axis=axis, initial=op.initial_value(x.dtype)).astype(
+                    acc_dtype
+                )
             return jax_op(x, axis=axis).astype(acc_dtype)
 
         # The PyTensor `Op` didn't tell us which NumPy equivalent to use (or
diff --git a/pytensor/link/jax/dispatch/math.py b/pytensor/link/jax/dispatch/math.py
index a741612229..584bfc1a4e 100644
--- a/pytensor/link/jax/dispatch/math.py
+++ b/pytensor/link/jax/dispatch/math.py
@@ -2,7 +2,6 @@
 import numpy as np
 
 from pytensor.link.jax.dispatch import jax_funcify
-from pytensor.link.jax.dispatch.elemwise import careduce_initial
 from pytensor.tensor.math import Argmax, Dot, Max
 
 
@@ -22,7 +21,7 @@ def jax_funcify_Max(op, **kwargs):
     if initial is not None:
 
         def max(x):
-            return jnp.max(x, axis, initial=careduce_initial(initial, x.dtype))
+            return jnp.max(x, axis, initial=op.initial_value(x.dtype))
 
         return max
 
diff --git a/pytensor/link/mlx/dispatch/math.py b/pytensor/link/mlx/dispatch/math.py
index 800f4ee4de..a6ee064cd4 100644
--- a/pytensor/link/mlx/dispatch/math.py
+++ b/pytensor/link/mlx/dispatch/math.py
@@ -14,9 +14,6 @@ def dot(x, y):
 
 @mlx_funcify.register(Max)
 def mlx_funcify_Max(op, node=None, **kwargs):
-    if op.initial is not None:
-        raise NotImplementedError("Max with initial is not supported in MLX")
-
     def max_fn(x):
         axes = op.axis
         if axes is None:
@@ -26,6 +23,14 @@ def max_fn(x):
 
         keepdims = getattr(op, "keepdims", False)
 
+        if op.initial is not None:
+            # MLX shapes are concrete, so an empty reduction can be filled directly
+            reduced = range(x.ndim) if reduce_axes is None else reduce_axes
+            if any(x.shape[a] == 0 for a in reduced):
+                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 mx.max(x, axis=reduce_axes, keepdims=keepdims)
 
     return max_fn
diff --git a/pytensor/link/pytorch/dispatch/elemwise.py b/pytensor/link/pytorch/dispatch/elemwise.py
index fd5ccf0039..be7f091c40 100644
--- a/pytensor/link/pytorch/dispatch/elemwise.py
+++ b/pytensor/link/pytorch/dispatch/elemwise.py
@@ -105,9 +105,6 @@ def torch_any(x):
 
 @pytorch_funcify.register(Max)
 def pytorch_funcify_max(op, **kwargs):
-    if op.initial is not None:
-        raise NotImplementedError("Max with initial is not supported in PyTorch")
-
     def torch_max(x):
         if isinstance(op.axis, tuple):
             for d in sorted(op.axis, reverse=True):
@@ -121,9 +118,6 @@ def torch_max(x):
 
 @pytorch_funcify.register(Min)
 def pytorch_funcify_min(op, **kwargs):
-    if op.initial is not None:
-        raise NotImplementedError("Min with initial is not supported in PyTorch")
-
     def torch_min(x):
         if isinstance(op.axis, tuple):
             for d in sorted(op.axis, reverse=True):
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index 0e99fe1aab..1ff4695cb8 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -443,20 +443,24 @@ def __str__(self):
             return super().__str__()
         return f"{type(self).__name__}{{{self._axis_str()}, initial={self.initial}}}"
 
+    def initial_value(self, dtype):
+        """`initial` as a value of `dtype`: the dtype's bounds for ints and bools."""
+        dtype = np.dtype(dtype)
+        if self.initial is None or dtype.kind == "f":
+            return self.initial
+        if dtype.kind == "b":
+            return self.initial > 0
+        info = np.iinfo(dtype)
+        return info.max if self.initial > 0 else info.min
+
     def perform(self, node, inp, out):
         if self.initial is None:
             return super().perform(node, inp, out)
         [x] = inp
-        dtype = np.dtype(node.outputs[0].type.dtype)
-        if dtype.kind == "f":
-            initial = self.initial
-        elif dtype.kind == "b":
-            initial = self.initial > 0
-        else:
-            info = np.iinfo(dtype)
-            initial = info.max if self.initial > 0 else info.min
+        dtype = node.outputs[0].type.dtype
         out[0][0] = np.asarray(
-            self.ufunc.reduce(x, axis=self.axis, initial=initial), dtype=dtype
+            self.ufunc.reduce(x, axis=self.axis, initial=self.initial_value(dtype)),
+            dtype=dtype,
         )
 
     def pullback(self, inputs, outputs, output_grads):
diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index bd99c11a09..30ee7d4c3d 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -2116,9 +2116,12 @@ def local_reduce_chain(fgraph, node) -> list[TensorVariable] | None:
     if outer_op.scalar_op != inner_op.scalar_op:
         return None
 
-    # Merging would change what an empty inner or outer reduction returns
-    if getattr(outer_op, "initial", None) != getattr(inner_op, "initial", None):
-        return None
+    # Keeping `initial` from either side only turns an empty-axis error into it
+    clone_kwargs = {}
+    if hasattr(outer_op, "initial"):
+        for op in (outer_op, inner_op):
+            if getattr(op, "initial", None) is not None:
+                clone_kwargs["initial"] = op.initial
 
     outer_axis = outer_op.axis
     inner_axis = inner_op.axis
@@ -2126,7 +2129,7 @@ def local_reduce_chain(fgraph, node) -> list[TensorVariable] | None:
     # 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)
@@ -2139,7 +2142,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
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index af9015036d..bbba4401f7 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -2989,16 +2989,15 @@ def test_all(self):
         [new_out] = local_reduce_chain.transform(fg, out.owner)
         assert equal_computations([new_out], [x.all(axis=(0, 2))])
 
-    def test_max_initial(self):
+    @pytest.mark.parametrize("inner_initial", [None, -np.inf])
+    @pytest.mark.parametrize("outer_initial", [None, -np.inf])
+    def test_max_initial(self, inner_initial, outer_initial):
         x = tensor3()
-        out = x.max(axis=-1, initial=-np.inf).max(axis=0)
-        fg = FunctionGraph([x], [out], clone=False)
-        assert local_reduce_chain.transform(fg, out.owner) is None
-
-        out = x.max(axis=-1, initial=-np.inf).max(axis=0, initial=-np.inf)
+        out = x.max(axis=-1, initial=inner_initial).max(axis=0, initial=outer_initial)
         fg = FunctionGraph([x], [out], clone=False)
         [new_out] = local_reduce_chain.transform(fg, out.owner)
-        assert equal_computations([new_out], [x.max(axis=(0, 2), initial=-np.inf)])
+        initial = -np.inf if -np.inf in (inner_initial, outer_initial) else None
+        assert equal_computations([new_out], [x.max(axis=(0, 2), initial=initial)])
 
 
 class TestLocalSumProd:

From 13d85bb83f881fb38f5c6fba85ff5198bccc2397 Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Sun, 4 Oct 2026 20:00:08 +0530
Subject: [PATCH 03/11] Make initial a boolean prop, push max/min with initial
 through alloc

---
 pytensor/link/jax/dispatch/elemwise.py      |  4 +-
 pytensor/link/jax/dispatch/math.py          | 12 +----
 pytensor/link/mlx/dispatch/math.py          |  9 ++--
 pytensor/tensor/math.py                     | 51 ++++++++++-----------
 pytensor/tensor/rewriting/math.py           | 25 ++++++----
 pytensor/tensor/rewriting/uncanonicalize.py |  3 +-
 tests/tensor/rewriting/test_math.py         | 22 ++++++---
 7 files changed, 66 insertions(+), 60 deletions(-)

diff --git a/pytensor/link/jax/dispatch/elemwise.py b/pytensor/link/jax/dispatch/elemwise.py
index 586f589535..8be5f02b85 100644
--- a/pytensor/link/jax/dispatch/elemwise.py
+++ b/pytensor/link/jax/dispatch/elemwise.py
@@ -23,7 +23,7 @@ def elemwise_fn(*inputs):
 @jax_funcify.register(CAReduce)
 def jax_funcify_CAReduce(op, **kwargs):
     axis = op.axis
-    initial = getattr(op, "initial", None)
+    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)
@@ -47,7 +47,7 @@ def careduce(x):
 
         if op_nfunc_spec:
             jax_op = getattr(jnp, op_nfunc_spec[0])
-            if initial is not None:
+            if initial:
                 return jax_op(x, axis=axis, initial=op.initial_value(x.dtype)).astype(
                     acc_dtype
                 )
diff --git a/pytensor/link/jax/dispatch/math.py b/pytensor/link/jax/dispatch/math.py
index 584bfc1a4e..79a6876b82 100644
--- a/pytensor/link/jax/dispatch/math.py
+++ b/pytensor/link/jax/dispatch/math.py
@@ -16,19 +16,9 @@ def dot(x, y):
 @jax_funcify.register(Max)
 def jax_funcify_Max(op, **kwargs):
     axis = op.axis
-    initial = op.initial
-
-    if initial is not None:
-
-        def max(x):
-            return jnp.max(x, axis, initial=op.initial_value(x.dtype))
-
-        return max
 
     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
 
diff --git a/pytensor/link/mlx/dispatch/math.py b/pytensor/link/mlx/dispatch/math.py
index a6ee064cd4..14bae85032 100644
--- a/pytensor/link/mlx/dispatch/math.py
+++ b/pytensor/link/mlx/dispatch/math.py
@@ -23,13 +23,12 @@ def max_fn(x):
 
         keepdims = getattr(op, "keepdims", False)
 
-        if op.initial is not None:
+        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
-            if any(x.shape[a] == 0 for a in reduced):
-                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)
+            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 mx.max(x, axis=reduce_axes, keepdims=keepdims)
 
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index 1ff4695cb8..263339d664 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -369,8 +369,8 @@ 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 getattr(self, "initial", None) is not None:
-            # Empty reductions return the scalar_op identity, which is `initial`
+        if getattr(self, "initial", False):
+            # Empty reductions return the scalar_op identity
             return setup, alloc, loop, cast
 
         # We add an additional check for zero-sized dimensions (This seems like
@@ -415,22 +415,16 @@ class MaxAndMinCAReduce(NonZeroDimsCAReduce):
     Subclasses only need to bind the appropriate scalar ``Op`` (``maximum`` or
     ``minimum``) in their ``__init__`` and set ``nfunc_spec``.
 
-    ``initial`` can be ``None`` (reducing an empty axis raises) or the identity of
-    the scalar ``Op`` (``-inf`` for ``Max``, ``inf`` for ``Min``), in which case an
-    empty reduction returns it, like NumPy's ``initial``. For integer and boolean
-    inputs the lowest/highest value of the dtype is used.
+    With ``initial=True`` an empty reduction returns the identity of the scalar
+    ``Op`` (``-inf`` for ``Max``, ``inf`` for ``Min``) instead of raising, like
+    NumPy's ``initial``. For integer and boolean inputs the lowest/highest value of
+    the dtype is used.
     """
 
     __props__ = (*NonZeroDimsCAReduce.__props__, "initial")
 
-    def __init__(self, scalar_op, axis, initial=None):
-        if initial is not None:
-            initial = float(initial)
-            if initial != scalar_op.identity:
-                raise ValueError(
-                    f"initial must be None or {scalar_op.identity} for {type(self).__name__}, got {initial}"
-                )
-        self.initial = initial
+    def __init__(self, scalar_op, axis, initial=False):
+        self.initial = bool(initial)
         super().__init__(scalar_op, axis)
 
     def clone(self, **kwargs):
@@ -439,22 +433,25 @@ def clone(self, **kwargs):
         return type(self)(axis=axis, initial=initial)
 
     def __str__(self):
-        if self.initial is None:
+        if not self.initial:
             return super().__str__()
-        return f"{type(self).__name__}{{{self._axis_str()}, initial={self.initial}}}"
+        return f"{type(self).__name__}{{{self._axis_str()}, initial}}"
 
     def initial_value(self, dtype):
-        """`initial` as a value of `dtype`: the dtype's bounds for ints and bools."""
+        """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 self.initial is None or dtype.kind == "f":
-            return self.initial
+        if dtype.kind == "f":
+            return identity
         if dtype.kind == "b":
-            return self.initial > 0
+            return identity > 0
         info = np.iinfo(dtype)
-        return info.max if self.initial > 0 else info.min
+        return info.max if identity > 0 else info.min
 
     def perform(self, node, inp, out):
-        if self.initial is None:
+        if not self.initial:
             return super().perform(node, inp, out)
         [x] = inp
         dtype = node.outputs[0].type.dtype
@@ -514,14 +511,14 @@ def pushforward(self, inputs, outputs, tangents):
 class Max(MaxAndMinCAReduce):
     nfunc_spec = ("max", 1, 1)
 
-    def __init__(self, axis, initial=None):
+    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, initial=None):
+    def __init__(self, axis, initial=False):
         super().__init__(ps.minimum, axis, initial=initial)
 
 
@@ -549,7 +546,9 @@ def max(x, axis=None, keepdims=False, initial=None):
     a shape of 0.
 
     """
-    out = Max(axis=axis, initial=initial)(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)
@@ -582,7 +581,7 @@ def min(x, axis=None, keepdims=False, initial=None):
     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=initial)(x)
+        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, initial=max_initial)
diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index 30ee7d4c3d..6a24b885bd 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -2118,10 +2118,10 @@ def local_reduce_chain(fgraph, node) -> list[TensorVariable] | None:
 
     # Keeping `initial` from either side only turns an empty-axis error into it
     clone_kwargs = {}
-    if hasattr(outer_op, "initial"):
-        for op in (outer_op, inner_op):
-            if getattr(op, "initial", None) is not None:
-                clone_kwargs["initial"] = op.initial
+    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
@@ -2279,10 +2279,6 @@ def local_careduce_of_alloc(fgraph, node):
         case _:
             return None
 
-    # A broadcast axis may be empty at runtime, where the result is `initial`
-    if getattr(node.op, "initial", None) is not None:
-        return None
-
     ndim = len(shapes)
     axis = node.op.axis
     axis = tuple(range(ndim)) if axis is None else axis
@@ -2322,6 +2318,19 @@ def local_careduce_of_alloc(fgraph, node):
                 size = size.astype("float32")
             value = value * size if isinstance(node.op, Sum) else value**size
 
+    # With `initial`, a broadcast axis that is empty at runtime makes the result `initial`
+    if getattr(node.op, "initial", False):
+        broadcast_shapes = [
+            shapes[a] for a in axis if a < offset or value_bcast[a - offset]
+        ]
+        if broadcast_shapes:
+            dtype = node.outputs[0].dtype
+            value = switch(
+                eq(variadic_mul(*broadcast_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.
     value = value.astype(node.outputs[0].dtype)
diff --git a/pytensor/tensor/rewriting/uncanonicalize.py b/pytensor/tensor/rewriting/uncanonicalize.py
index b0bbbfff98..d80eaf11fd 100644
--- a/pytensor/tensor/rewriting/uncanonicalize.py
+++ b/pytensor/tensor/rewriting/uncanonicalize.py
@@ -64,9 +64,8 @@ def local_max_to_min(fgraph, node):
         ):
             neg_node = max.owner.inputs[0]
             if neg_node.owner and neg_node.owner.op == neg:
-                initial = getattr(max.owner.op, "initial", None)
                 new = Min(
-                    max.owner.op.axis, initial=None if initial is None else -initial
+                    max.owner.op.axis, initial=getattr(max.owner.op, "initial", False)
                 )(neg_node.owner.inputs[0])
                 return [copy_stack_trace(node.outputs[0], new)]
 
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index bbba4401f7..5e0c456912 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -3468,14 +3468,24 @@ def test_local_careduce_of_alloc_mixed_axes(self):
         result.assert_eval(v_val)
 
     def test_local_careduce_of_alloc_initial(self):
-        # The broadcast axis may be empty at runtime, where the result is `initial`
+        # A broadcast axis may be empty at runtime, where the result is `initial`
         v = vector("v")
         n = scalar("n", dtype="int64")
-        out = Max(axis=0, initial=-np.inf)(pt.alloc(v, n, 3))
-        f = function([v, n], out, mode=get_default_mode().including("specialize"))
-        np.testing.assert_array_equal(
-            f(np.ones(3, dtype=v.dtype), 0), np.full(3, -np.inf)
-        )
+        alloc_v = pt.alloc(v, n, 3)
+        outs = [
+            Max(axis=0, initial=True)(alloc_v),
+            Max(axis=None, initial=True)(alloc_v),
+        ]
+        f = function([v, n], outs, mode=get_default_mode().including("specialize"))
+        assert not any(isinstance(node.op, Alloc) for node in f.maker.fgraph.toposort())
+
+        v_val = np.array([1.0, 3.0, 2.0], dtype=v.dtype)
+        dropped, reduced = f(v_val, 0)
+        np.testing.assert_array_equal(dropped, np.full(3, -np.inf))
+        assert reduced == -np.inf
+        dropped, reduced = f(v_val, 2)
+        np.testing.assert_array_equal(dropped, v_val)
+        assert reduced == 3.0
 
     @pytest.mark.parametrize("reduce_op", [Max, Min, All, Any])
     def test_local_careduce_of_alloc_idempotent(self, reduce_op):

From e00d9b4edd44871e66a417c436985fd65dd0981f Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Sun, 11 Oct 2026 05:09:48 +0530
Subject: [PATCH 04/11] Move initial to NonZeroDimsCAReduce and address review

---
 pytensor/link/jax/dispatch/elemwise.py |  8 +--
 pytensor/tensor/math.py                | 83 +++++++++++++-------------
 pytensor/tensor/rewriting/math.py      | 35 +++++------
 tests/tensor/rewriting/test_math.py    | 17 ++++--
 4 files changed, 73 insertions(+), 70 deletions(-)

diff --git a/pytensor/link/jax/dispatch/elemwise.py b/pytensor/link/jax/dispatch/elemwise.py
index 8be5f02b85..76b8d3e146 100644
--- a/pytensor/link/jax/dispatch/elemwise.py
+++ b/pytensor/link/jax/dispatch/elemwise.py
@@ -47,11 +47,9 @@ def careduce(x):
 
         if op_nfunc_spec:
             jax_op = getattr(jnp, op_nfunc_spec[0])
-            if initial:
-                return jax_op(x, axis=axis, initial=op.initial_value(x.dtype)).astype(
-                    acc_dtype
-                )
-            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
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index 263339d664..a59c6057d0 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -365,11 +365,53 @@ 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 == "f":
+            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 getattr(self, "initial", False):
+        if self.initial:
             # Empty reductions return the scalar_op identity
             return setup, alloc, loop, cast
 
@@ -414,52 +456,13 @@ class MaxAndMinCAReduce(NonZeroDimsCAReduce):
 
     Subclasses only need to bind the appropriate scalar ``Op`` (``maximum`` or
     ``minimum``) in their ``__init__`` and set ``nfunc_spec``.
-
-    With ``initial=True`` an empty reduction returns the identity of the scalar
-    ``Op`` (``-inf`` for ``Max``, ``inf`` for ``Min``) instead of raising, like
-    NumPy's ``initial``. For integer and boolean inputs the lowest/highest value of
-    the dtype is used.
     """
 
-    __props__ = (*NonZeroDimsCAReduce.__props__, "initial")
-
-    def __init__(self, scalar_op, axis, initial=False):
-        self.initial = bool(initial)
-        super().__init__(scalar_op, axis)
-
     def clone(self, **kwargs):
         axis = kwargs.get("axis", self.axis)
         initial = kwargs.get("initial", self.initial)
         return type(self)(axis=axis, initial=initial)
 
-    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 == "f":
-            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 pullback(self, inputs, outputs, output_grads):
         # The strict-sense mathematical gradient of a maximum/minimum reduction
         # is not defined at points where the extremum is attained by more than
diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index 1f96f1f668..ac16d245f6 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -2308,28 +2308,23 @@ 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")
-            value = value * size if isinstance(node.op, Sum) else 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")
+        value = value * size if isinstance(node.op, Sum) else value**size
 
     # With `initial`, a broadcast axis that is empty at runtime makes the result `initial`
-    if getattr(node.op, "initial", False):
-        broadcast_shapes = [
-            shapes[a] for a in axis if a < offset or value_bcast[a - offset]
-        ]
-        if broadcast_shapes:
-            dtype = node.outputs[0].dtype
-            value = switch(
-                eq(variadic_mul(*broadcast_shapes), 0),
-                np.asarray(node.op.initial_value(dtype), dtype=dtype),
-                value,
-            )
+    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.
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 3ce1d528a0..ff2c7372c9 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -3472,19 +3472,26 @@ def test_local_careduce_of_alloc_initial(self):
         n = scalar("n", dtype="int64")
         alloc_v = pt.alloc(v, n, 3)
         outs = [
-            Max(axis=0, initial=True)(alloc_v),
-            Max(axis=None, initial=True)(alloc_v),
+            pt.max(alloc_v, axis=0, initial=-np.inf),
+            pt.max(alloc_v, axis=None, initial=-np.inf),
+            pt.max(alloc_v, axis=1, initial=-np.inf),
         ]
         f = function([v, n], outs, mode=get_default_mode().including("specialize"))
-        assert not any(isinstance(node.op, Alloc) for node in f.maker.fgraph.toposort())
+        # Only the 1d alloc of the axis=1 result may remain
+        assert not any(
+            isinstance(node.op, Alloc) and node.outputs[0].type.ndim == 2
+            for node in f.maker.fgraph.toposort()
+        )
 
         v_val = np.array([1.0, 3.0, 2.0], dtype=v.dtype)
-        dropped, reduced = f(v_val, 0)
+        dropped, reduced, kept = f(v_val, 0)
         np.testing.assert_array_equal(dropped, np.full(3, -np.inf))
         assert reduced == -np.inf
-        dropped, reduced = f(v_val, 2)
+        assert kept.shape == (0,)
+        dropped, reduced, kept = f(v_val, 2)
         np.testing.assert_array_equal(dropped, v_val)
         assert reduced == 3.0
+        np.testing.assert_array_equal(kept, [3.0, 3.0])
 
     @pytest.mark.parametrize("reduce_op", [Max, Min, All, Any])
     def test_local_careduce_of_alloc_idempotent(self, reduce_op):

From c1b8f53022202b325ef0188712a2bde5f5694295 Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Tue, 13 Oct 2026 04:12:10 +0530
Subject: [PATCH 05/11] Keep initial when local_reduce_broadcastable drops an
 axis

---
 pytensor/tensor/rewriting/math.py   | 13 ++-----------
 tests/tensor/rewriting/test_math.py | 12 ++++++++++++
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py
index ac16d245f6..166d88688b 100644
--- a/pytensor/tensor/rewriting/math.py
+++ b/pytensor/tensor/rewriting/math.py
@@ -2238,17 +2238,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)]
diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index ff2c7372c9..68f96834f8 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -115,6 +115,7 @@
     local_mul_canonizer,
     local_mul_switch_sink,
     local_neg_to_mul,
+    local_reduce_broadcastable,
     local_reduce_chain,
     local_sum_prod_of_mul_or_div,
     local_useless_abs,
@@ -3807,6 +3808,17 @@ def test_local_reduce_broadcast_some_1(self):
                 isinstance(node.op, CAReduce) for node in f.maker.fgraph.toposort()
             )
 
+    def test_local_reduce_broadcast_some_initial(self):
+        # Dropping the broadcastable axis must keep `initial`, the other axis can be empty
+        x = TensorType("float64", shape=(1, None))()
+        out = pt_max(x, axis=[0, 1], initial=-np.inf)
+        fg = FunctionGraph([x], [out], clone=False)
+        [new_out] = local_reduce_broadcastable.transform(fg, out.owner)
+        assert new_out.owner.op == Max(axis=(0,), initial=True)
+
+        f = function([x], out, mode="FAST_COMPILE")
+        assert f(np.zeros((1, 0))) == -np.inf
+
 
 class TestReduceJoin:
     def setup_method(self):

From b1f1e21224f599cd2d93bd58ac8a7d9794f3f17b Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Tue, 13 Oct 2026 04:12:10 +0530
Subject: [PATCH 06/11] Use the float identity for complex max/min initial

---
 pytensor/tensor/math.py      | 2 +-
 tests/tensor/test_special.py | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index a59c6057d0..371883c665 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -390,7 +390,7 @@ def initial_value(self, dtype):
             return None
         identity = self.scalar_op.identity
         dtype = np.dtype(dtype)
-        if dtype.kind == "f":
+        if dtype.kind in "fc":
             return identity
         if dtype.kind == "b":
             return identity > 0
diff --git a/tests/tensor/test_special.py b/tests/tensor/test_special.py
index 317fa1508b..8a15376b44 100644
--- a/tests/tensor/test_special.py
+++ b/tests/tensor/test_special.py
@@ -270,6 +270,14 @@ def test_logsumexp_zero_size():
     assert out == -np.inf
 
 
+def test_logsumexp_complex():
+    x = tensor("x", shape=(None,), dtype="complex128")
+    x_val = np.array([1 + 2j, 3 + 4j])
+    out = function([x], logsumexp(x), mode="FAST_COMPILE")(x_val)
+    # Only defined up to a multiple of 2*pi*j
+    np.testing.assert_allclose(np.exp(out), np.exp(x_val).sum())
+
+
 @pytest.mark.parametrize("mode", ["FAST_RUN", "FAST_COMPILE"])
 def test_logsumexp_logaddexp_stable_grad(mode):
     """Both helpers promise a stable forward, so their gradient must be stable too.

From 616ebf5936d1ec243ab4cb62c3fd0958def13858 Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Tue, 13 Oct 2026 04:12:11 +0530
Subject: [PATCH 07/11] Handle empty Min with initial in MLX

---
 pytensor/link/mlx/dispatch/math.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/pytensor/link/mlx/dispatch/math.py b/pytensor/link/mlx/dispatch/math.py
index 14bae85032..f392a2e5e5 100644
--- a/pytensor/link/mlx/dispatch/math.py
+++ b/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)
@@ -12,9 +12,8 @@ def dot(x, y):
     return dot
 
 
-@mlx_funcify.register(Max)
-def mlx_funcify_Max(op, node=None, **kwargs):
-    def max_fn(x):
+def _mlx_funcify_max_min(op, node, reduce_fn):
+    def max_min_fn(x):
         axes = op.axis
         if axes is None:
             reduce_axes = None
@@ -30,9 +29,19 @@ def max_fn(x):
             fill = op.initial_value(node.outputs[0].type.dtype)
             return mx.full(out_shape, fill, dtype=x.dtype)
 
-        return mx.max(x, axis=reduce_axes, keepdims=keepdims)
+        return reduce_fn(x, axis=reduce_axes, keepdims=keepdims)
+
+    return max_min_fn
+
+
+@mlx_funcify.register(Max)
+def mlx_funcify_Max(op, node=None, **kwargs):
+    return _mlx_funcify_max_min(op, node, mx.max)
+
 
-    return max_fn
+@mlx_funcify.register(Min)
+def mlx_funcify_Min(op, node=None, **kwargs):
+    return _mlx_funcify_max_min(op, node, mx.min)
 
 
 @mlx_funcify.register(Argmax)

From c72f722713788d31d91efc801af4571b880e54fc Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Tue, 13 Oct 2026 13:17:33 +0530
Subject: [PATCH 08/11] Use RewriteTester for the broadcastable-axis initial
 test

---
 tests/tensor/rewriting/test_math.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 68f96834f8..1286e16dce 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -3810,14 +3810,15 @@ def test_local_reduce_broadcast_some_1(self):
 
     def test_local_reduce_broadcast_some_initial(self):
         # Dropping the broadcastable axis must keep `initial`, the other axis can be empty
-        x = TensorType("float64", shape=(1, None))()
-        out = pt_max(x, axis=[0, 1], initial=-np.inf)
-        fg = FunctionGraph([x], [out], clone=False)
-        [new_out] = local_reduce_broadcastable.transform(fg, out.owner)
-        assert new_out.owner.op == Max(axis=(0,), initial=True)
-
-        f = function([x], out, mode="FAST_COMPILE")
-        assert f(np.zeros((1, 0))) == -np.inf
+        x = TensorType("float64", shape=(1, None))("x")
+        result = RewriteTester(
+            [x],
+            [pt_max(x, axis=[0, 1], initial=-np.inf)],
+            include=None,
+            custom_rewrite=local_reduce_broadcastable,
+        )
+        result.assert_graph(pt_max(x.dimshuffle(1), axis=0, initial=-np.inf))
+        result.assert_eval(np.zeros((1, 0)))
 
 
 class TestReduceJoin:

From 007c691c3fab6e8aeab9311ad3515bf56d67ab08 Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Tue, 13 Oct 2026 13:17:34 +0530
Subject: [PATCH 09/11] Register one MLX funcify for Max and Min

---
 pytensor/link/mlx/dispatch/math.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/pytensor/link/mlx/dispatch/math.py b/pytensor/link/mlx/dispatch/math.py
index f392a2e5e5..c43e16d888 100644
--- a/pytensor/link/mlx/dispatch/math.py
+++ b/pytensor/link/mlx/dispatch/math.py
@@ -12,7 +12,11 @@ def dot(x, y):
     return dot
 
 
-def _mlx_funcify_max_min(op, node, reduce_fn):
+@mlx_funcify.register(Max)
+@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:
@@ -34,16 +38,6 @@ def max_min_fn(x):
     return max_min_fn
 
 
-@mlx_funcify.register(Max)
-def mlx_funcify_Max(op, node=None, **kwargs):
-    return _mlx_funcify_max_min(op, node, mx.max)
-
-
-@mlx_funcify.register(Min)
-def mlx_funcify_Min(op, node=None, **kwargs):
-    return _mlx_funcify_max_min(op, node, mx.min)
-
-
 @mlx_funcify.register(Argmax)
 def mlx_funcify_Argmax(op, node=None, **kwargs):
     axis = op.axis

From 0387b745946edfedeeeee7712728d54754e8646b Mon Sep 17 00:00:00 2001
From: Anish Mehta <145196704+anishmehta24@users.noreply.github.com>
Date: Thu, 15 Oct 2026 03:15:58 +0530
Subject: [PATCH 10/11] Update tests/tensor/rewriting/test_math.py

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>
---
 tests/tensor/rewriting/test_math.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index 1ff7c6f6ce..cd874f906e 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -3835,7 +3835,7 @@ def test_local_reduce_broadcast_some_initial(self):
             include=None,
             custom_rewrite=local_reduce_broadcastable,
         )
-        result.assert_graph(pt_max(x.dimshuffle(1), axis=0, initial=-np.inf))
+        result.assert_graph(pt_max(x.squeeze(0), axis=0, initial=-np.inf))
         result.assert_eval(np.zeros((1, 0)))
 
 

From a2d69d39db661251be6eca96bbaa158b7ee800e2 Mon Sep 17 00:00:00 2001
From: Anish Mehta <mehtaanish111@gmail.com>
Date: Thu, 15 Oct 2026 03:16:17 +0530
Subject: [PATCH 11/11] Use RewriteTester in test_max_initial

---
 tests/tensor/rewriting/test_math.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py
index cd874f906e..f3a93cbb64 100644
--- a/tests/tensor/rewriting/test_math.py
+++ b/tests/tensor/rewriting/test_math.py
@@ -3002,12 +3002,14 @@ def test_all(self):
     @pytest.mark.parametrize("inner_initial", [None, -np.inf])
     @pytest.mark.parametrize("outer_initial", [None, -np.inf])
     def test_max_initial(self, inner_initial, outer_initial):
-        x = tensor3()
+        x = tensor3("x")
         out = x.max(axis=-1, initial=inner_initial).max(axis=0, initial=outer_initial)
-        fg = FunctionGraph([x], [out], clone=False)
-        [new_out] = local_reduce_chain.transform(fg, out.owner)
+        result = RewriteTester(
+            [x], [out], include=None, custom_rewrite=local_reduce_chain
+        )
         initial = -np.inf if -np.inf in (inner_initial, outer_initial) else None
-        assert equal_computations([new_out], [x.max(axis=(0, 2), initial=initial)])
+        result.assert_graph(x.max(axis=(0, 2), initial=initial))
+        result.assert_eval(np.random.default_rng(0).normal(size=(2, 3, 4)))
 
 
 class TestLocalSumProd: