Add `initial` to `max` and `min` by anishmehta24 · Pull Request #2444 · pymc-devs/pytensor

anishmehta24 · 2026-09-29T17:24:56Z

Description

Adds initial to pt.max/pt.min (and Max/Min), as discussed in #2439. For now it only accepts None or the identity (-inf for max, inf for min); the Op stores it as a boolean prop. With it set, reducing an empty axis returns initial (the dtype's lowest/highest value for ints and bools) instead of raising, the same in the Python, C, Numba and JAX backends. MLX checks the (concrete) input shape and fills empty reductions; PyTorch leaves it to torch.max, which raises on an empty axis.

Rewrites that could drop or change it:

local_reduce_chain keeps initial if either reduction has it
local_careduce_of_alloc still applies, and returns initial when a broadcast axis is empty at runtime
local_max_to_min and clone (used by vectorize) keep it

LogSumExp now uses max(..., initial=-inf), so logsumexp of an empty array is -inf in every backend.

While testing I found Numba returned True for max over an empty bool axis, because -inf was cast straight to bool. Both identity helpers now map it like the C backend does, and their cache versions are bumped.

Tests: test_initial in TestMinMax (float/int/uint/bool, empty and non-empty), logsumexp zero-size, the two rewrites, JAX and Numba empty cases. Ran the new tests under NUMBA, CVM and FAST_COMPILE.

Related Issue

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

ricardoV94 · 2026-09-29T17:28:51Z


 @mlx_funcify.register(Max)
 def mlx_funcify_Max(op, node=None, **kwargs):
+    if op.initial is not None:


this will break logsumexp completely for mlx? The dispatch can look at x shape, it's always statically known for mlx

Right, fixed. It checks x.shape now and returns the filled array when a reduced axis is empty, otherwise mx.max as before.

ricardoV94 · 2026-09-29T17:30:07Z

        return None

+    # Merging would change what an empty inner or outer reduction returns
+    if getattr(outer_op, "initial", None) != getattr(inner_op, "initial", None):


I think we can keep the one with more specific initial? Gotta think about it

I think that's safe: whichever side has initial, merging only turns an empty-axis error into initial and never changes a real value. Changed it to keep initial if either side has it, with tests for all four combinations.

ricardoV94 · 2026-09-29T17:30:49Z

        case _:
            return None

+    # A broadcast axis may be empty at runtime, where the result is `initial`


what did this rewrite do before to runtime empty maxes?

Before, max(alloc(v, n, 3), axis=0) with n=0 at runtime became just v, so an error (C/Python) or -inf (Numba) turned into v's values. With initial set, the right answer is -inf, so that would be a wrong value, not just a lost error. That's why it bails. It could bail only when the broadcast shape isn't known to be non-zero, if you'd prefer.

This bailing out seems wrong. Either we still reduce axis, so the original Op is fine.

Or we don't need the max/min at all (all reduced values were broadcasted). In that case if the alloc shapes are 0 we'd return initial, otherwise the value? A max/min on shapes 0 without initial is undefined, we can leave doing the same thing as before

Agreed, it doesn't bail anymore. With initial it returns initial when a broadcast axis is empty at runtime, otherwise the value; without initial nothing changes.

…are initial_value

ricardoV94 · 2026-10-02T11:29:45Z

+
+    if initial is not None:
+
+        def max(x):


does initial=None fail? If not we don't need two branches

It doesn't, so it's one branch now.

ricardoV94 · 2026-10-02T11:30:22Z

+        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):


x.size == 0?

Yes, switched to x.size == 0.

ricardoV94 · 2026-10-02T11:32:27Z

    if outer_op.scalar_op != inner_op.scalar_op:
        return None

+    # Keeping `initial` from either side only turns an empty-axis error into it


hmm, this is more complex than it should be. Maybe the default should be some sentinel Unset value, which can be converted to None. Then clone can distinguish "not specified, stay same" from override

With initial as a bool, clone keeps the current value unless it's passed, so no sentinel needed. The rewrite only passes initial=True when the inner op has it.

ricardoV94 · 2026-10-02T11:38:58Z

        case _:
            return None

+    # A broadcast axis may be empty at runtime, where the result is `initial`


This bailing out seems wrong. Either we still reduce axis, so the original Op is fine.

Or we don't need the max/min at all (all reduced values were broadcasted). In that case if the alloc shapes are 0 we'd return initial, otherwise the value? A max/min on shapes 0 without initial is undefined, we can leave doing the same thing as before

ricardoV94 · 2026-10-02T11:40:43Z

+    def __init__(self, scalar_op, axis, initial=None):
+        if initial is not None:
+            initial = float(initial)
+            if initial != scalar_op.identity:


if this is the only value we support the internal Op argument should be a boolean flag, not an arbitrary value. That seems fine to me.

Done, it's a boolean prop now. pt.max/pt.min still take initial=None or -inf/inf and check it.

# Conflicts: # pytensor/link/numba/dispatch/elemwise.py

ricardoV94 · 2026-10-10T17:02:15Z

+                return jax_op(x, axis=axis, initial=op.initial_value(x.dtype)).astype(
+                    acc_dtype
+                )
            return jax_op(x, axis=axis).astype(acc_dtype)


Can't this be op.initial_value(x.dtype)... if intial_value else None? (just one return?)

Yes, one return now (e00d9b4). initial goes in as a kwarg only when set, since jnp.all/jnp.any don't accept it.

ricardoV94

This is shaping up nice, left a few comments

ricardoV94 · 2026-10-10T17:06:27Z


+    # With `initial`, a broadcast axis that is empty at runtime makes the result `initial`
+    if getattr(node.op, "initial", False):
+        broadcast_shapes = [


this is called size_shapes in the branch above

Renamed. size_shapes is now computed once and used by both branches.

ricardoV94 · 2026-10-10T17:09:50Z

+    the dtype is used.
    """

+    __props__ = (*NonZeroDimsCAReduce.__props__, "initial")


shouldn't this be part of NonZeroDimsCareduce?

Moved it there: the prop, __init__, perform, initial_value and __str__ now live on NonZeroDimsCAReduce. MaxAndMinCAReduce only keeps clone.

ricardoV94 · 2026-10-10T17:13:23Z

+            Max(axis=0, initial=True)(alloc_v),
+            Max(axis=None, initial=True)(alloc_v),


I prefer if you test with the user facing max. Can we also include a third output along axis 1?

Switched to pt.max and added an axis=1 output. That one keeps a 1d alloc of the reduced vector, so the check now only rules out the 2d alloc.

ricardoV94 · 2026-10-12T16:58:19Z

bot found the following bugs:

[P2] Broadcast-axis removal drops initial. The rewrite (pytensor/tensor/rewriting/math.py:2250) constructs a fresh Op without preserving it. With shape=(1, None), x.max(axis=(0, 1), initial=-np.inf) raises on (1, 0) input in FAST_COMPILE and CVM. This also breaks empty logsumexp. Use node.op.clone(axis=new_axis).
[P2] Nonempty complex logsumexp now fails. The identity helper (pytensor/tensor/math.py:397) calls np.iinfo for complex dtypes. The previous inner graph evaluates [1+2j, 3+4j]; the new graph raises ValueError: Invalid integer data type 'c'. Handle complex identities before the integer branch.
[P2] MLX minimum reductions bypass the empty-input handler. The new handling (pytensor/link/mlx/dispatch/math.py:26) covers only Max. Float minima optimize to Min(initial=True), whose generic dispatch still calls mx.min without handling empty axes.

anishmehta24 · 2026-10-12T22:42:45Z

Fixed all three:

local_reduce_broadcastable now uses node.op.clone(axis=new_axis), so initial survives dropping the broadcastable axis (c1b8f53, with a test on the (1, 0) case).
Complex dtypes take the identity like floats (b1f1e21). Complex logsumexp works again in FAST_COMPILE. With numba it fails on main as well, so that part isn't from this PR.
MLX Min now goes through the same empty-input path as Max (616ebf5). I couldn't run the MLX tests, no Apple machine here.

ricardoV94 · 2026-10-13T07:34:45Z

                isinstance(node.op, CAReduce) for node in f.maker.fgraph.toposort()
            )

+    def test_local_reduce_broadcast_some_initial(self):


Use RewriteTest (gives you both the graph to check and the function to evaluate)

Done in c72f722, it checks the rewritten graph and evaluates both on a (1, 0) input.

ricardoV94 · 2026-10-13T07:37:00Z

-@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):


Nit, combine the two functions (decorate with both Ops). It's the same except one calls mx.max and the other mx.min

Done in 007c691, one function registered for both and it picks mx.max or mx.min from the op.

ricardoV94

Left two nits, let me know if you want to address. Otherwise good to go

# Conflicts: # pytensor/tensor/rewriting/math.py # tests/link/jax/test_math.py

anishmehta24 · 2026-10-13T22:52:23Z

Merged main to resolve the conflict with the new ProdWithoutZeros handling in local_careduce_of_alloc (kept both). Rewrite and max/min tests pass locally.

ricardoV94

last nits

ricardoV94 · 2026-10-14T13:13:58Z

+        fg = FunctionGraph([x], [out], clone=False)
+        [new_out] = local_reduce_chain.transform(fg, out.owner)
+        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)])


use RewriteTest

Done in a2d69d3

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Add initial to max and min

e3e9b6c

anishmehta24 mentioned this pull request Sep 29, 2026

Numba max/min over a zero-size axis returns -inf/inf instead of raising #2439

Closed

ricardoV94 reviewed Sep 29, 2026

View reviewed changes

Comment thread pytensor/link/pytorch/dispatch/elemwise.py Outdated

ricardoV94 reviewed Sep 29, 2026

View reviewed changes

Keep initial when chaining reductions, handle empty inputs in MLX, sh…

af3036b

…are initial_value

ricardoV94 reviewed Oct 2, 2026

View reviewed changes

anishmehta24 added 3 commits October 4, 2026 20:00

Make initial a boolean prop, push max/min with initial through alloc

13d85bb

Merge main into feat/max-min-initial

cf3fd8a

Merge main into feat/max-min-initial

8199171

# Conflicts: # pytensor/link/numba/dispatch/elemwise.py

ricardoV94 reviewed Oct 10, 2026

View reviewed changes

Move initial to NonZeroDimsCAReduce and address review

e00d9b4

anishmehta24 added 3 commits October 13, 2026 04:12

Keep initial when local_reduce_broadcastable drops an axis

c1b8f53

Use the float identity for complex max/min initial

b1f1e21

Handle empty Min with initial in MLX

616ebf5

ricardoV94 reviewed Oct 13, 2026

View reviewed changes

ricardoV94 approved these changes Oct 13, 2026

View reviewed changes

anishmehta24 added 3 commits October 13, 2026 13:17

Use RewriteTester for the broadcastable-axis initial test

c72f722

Register one MLX funcify for Max and Min

007c691

Merge remote-tracking branch 'origin/main' into feat/max-min-initial

d1c0c7e

# Conflicts: # pytensor/tensor/rewriting/math.py # tests/link/jax/test_math.py

ricardoV94 reviewed Oct 14, 2026

View reviewed changes

anishmehta24 and others added 2 commits October 15, 2026 03:15

Update tests/tensor/rewriting/test_math.py

0387b74

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Use RewriteTester in test_max_initial

a2d69d3

ricardoV94 approved these changes Oct 15, 2026

View reviewed changes

ricardoV94 merged commit ce6df07 into pymc-devs:main Oct 15, 2026
67 checks passed

ricardoV94 added enhancement New feature or request NumPy compatibility Op implementation labels Oct 15, 2026

		Max(axis=0, initial=True)(alloc_v),
		Max(axis=None, initial=True)(alloc_v),

Conversation

anishmehta24 commented Sep 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Oct 12, 2026

Uh oh!

anishmehta24 commented Oct 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

anishmehta24 commented Oct 13, 2026

Uh oh!

ricardoV94 left a comment

anishmehta24 commented Sep 29, 2026 •

edited

Loading