diff --git a/pytensor/link/jax/dispatch/elemwise.py b/pytensor/link/jax/dispatch/elemwise.py index dfdcd6d54e..76b8d3e146 100644 --- a/pytensor/link/jax/dispatch/elemwise.py +++ b/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 diff --git a/pytensor/link/jax/dispatch/math.py b/pytensor/link/jax/dispatch/math.py index bba4aca7d0..59cd5ed5d3 100644 --- a/pytensor/link/jax/dispatch/math.py +++ b/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 diff --git a/pytensor/link/mlx/dispatch/math.py b/pytensor/link/mlx/dispatch/math.py index 2f927634ad..c43e16d888 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) @@ -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) diff --git a/pytensor/link/numba/dispatch/elemwise.py b/pytensor/link/numba/dispatch/elemwise.py index 203c49fafd..8fa943f568 100644 --- a/pytensor/link/numba/dispatch/elemwise.py +++ b/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( ( diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index bf76ab97a1..371883c665 100644 --- a/pytensor/tensor/math.py +++ b/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() diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py index 7b3e2836a1..1e0a9dbaae 100644 --- a/pytensor/tensor/rewriting/math.py +++ b/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. diff --git a/pytensor/tensor/rewriting/uncanonicalize.py b/pytensor/tensor/rewriting/uncanonicalize.py index f9f4f950ee..d80eaf11fd 100644 --- a/pytensor/tensor/rewriting/uncanonicalize.py +++ b/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 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 b104969852..a483d9c027 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, ProdWithoutZeros, 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_prod_without_zeros(): x = matrix("x") outs = [ProdWithoutZeros(axis=None)(x), ProdWithoutZeros(axis=(1,))(x)] diff --git a/tests/link/numba/test_elemwise.py b/tests/link/numba/test_elemwise.py index d81522cfd3..c8957b47e5 100644 --- a/tests/link/numba/test_elemwise.py +++ b/tests/link/numba/test_elemwise.py @@ -426,6 +426,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 dd7f3e9fcd..f3a93cbb64 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, @@ -2998,6 +2999,18 @@ def test_all(self): [new_out] = local_reduce_chain.transform(fg, out.owner) assert equal_computations([new_out], [x.all(axis=(0, 2))]) + @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") + out = x.max(axis=-1, initial=inner_initial).max(axis=0, initial=outer_initial) + result = RewriteTester( + [x], [out], include=None, custom_rewrite=local_reduce_chain + ) + initial = -np.inf if -np.inf in (inner_initial, outer_initial) else None + 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: """Test sum/prod rewrites.""" @@ -3466,6 +3479,33 @@ 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): + # A broadcast axis may be empty at runtime, where the result is `initial` + v = vector("v") + n = scalar("n", dtype="int64") + alloc_v = pt.alloc(v, n, 3) + outs = [ + 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")) + # 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, kept = f(v_val, 0) + np.testing.assert_array_equal(dropped, np.full(3, -np.inf)) + assert reduced == -np.inf + 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): # Idempotent reductions (max/min/all/any) over the alloc's broadcast axes @@ -3788,6 +3828,18 @@ 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))("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.squeeze(0), axis=0, initial=-np.inf)) + result.assert_eval(np.zeros((1, 0))) + class TestReduceJoin: def setup_method(self): 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..8a15376b44 100644 --- a/tests/tensor/test_special.py +++ b/tests/tensor/test_special.py @@ -263,6 +263,21 @@ 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 + + +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.