pytensor 2.18.4__tar.gz → 2.18.6__tar.gz

{pytensor-2.18.4 → pytensor-2.18.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.18.4
+Version: 2.18.6
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-12-15T10:33:30+0100",
+ "date": "2024-01-12T14:29:52+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "8ae14c20a3f1e57c3e7ba6b0f8bcb7ce4227e55f",
- "version": "2.18.4"
+ "full-revisionid": "0666fd564d2d67384051c5bb6e0c09e713e0329c",
+ "version": "2.18.6"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/graph/basic.py

@@ -1439,15 +1439,16 @@ def io_toposort(
         order = []
         while todo:
             cur = todo.pop()
-            # We suppose that all outputs are always computed
-            if cur.outputs[0] in computed:
+            if all(out in computed for out in cur.outputs):
                 continue
             if all(i in computed or i.owner is None for i in cur.inputs):
                 computed.update(cur.outputs)
                 order.append(cur)
             else:
                 todo.append(cur)
-                todo.extend(i.owner for i in cur.inputs if i.owner)
+                todo.extend(
+                    i.owner for i in cur.inputs if (i.owner and i not in computed)
+                )
         return order
 
     compute_deps = None

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/graph/replace.py

@@ -256,7 +256,7 @@ def vectorize_graph(
         import pytensor
         import pytensor.tensor as pt
 
-        from pytensor.graph import vectorize
+        from pytensor.graph import vectorize_graph
 
         # Original graph
         x = pt.vector("x")
@@ -264,7 +264,7 @@ def vectorize_graph(
 
         # Vectorized graph
         new_x = pt.matrix("new_x")
-        new_y = vectorize(y, replace={x: new_x})
+        new_y = vectorize_graph(y, replace={x: new_x})
 
         fn = pytensor.function([new_x], new_y)
         fn([[0, 1, 2], [2, 1, 0]])
@@ -277,7 +277,7 @@ def vectorize_graph(
         import pytensor
         import pytensor.tensor as pt
 
-        from pytensor.graph import vectorize
+        from pytensor.graph import vectorize_graph
 
         # Original graph
         x = pt.vector("x")
@@ -286,7 +286,7 @@ def vectorize_graph(
 
         # Vectorized graph
         new_x = pt.matrix("new_x")
-        [new_y1, new_y2] = vectorize([y1, y2], replace={x: new_x})
+        [new_y1, new_y2] = vectorize_graph([y1, y2], replace={x: new_x})
 
         fn = pytensor.function([new_x], [new_y1, new_y2])
         fn([[-10, 0, 10], [-11, 0, 11]])
@@ -306,6 +306,11 @@ def vectorize_graph(
         vect_inputs = [vect_vars.get(inp, inp) for inp in node.inputs]
         vect_node = vectorize_node(node, *vect_inputs)
         for output, vect_output in zip(node.outputs, vect_node.outputs):
+            if output in vect_vars:
+                # This can happen when some outputs of a multi-output node are given a replacement,
+                # while some of the remaining outputs are still needed in the graph.
+                # We make sure we don't overwrite the provided replacement with the newly vectorized output
+                continue
             vect_vars[output] = vect_output
 
     seq_vect_outputs = [vect_vars[out] for out in seq_outputs]

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/link/jax/dispatch/scalar.py

@@ -21,11 +21,14 @@ from pytensor.scalar.basic import (
     Sub,
 )
 from pytensor.scalar.math import (
+    BetaIncInv,
     Erf,
     Erfc,
     Erfcinv,
     Erfcx,
     Erfinv,
+    GammaIncCInv,
+    GammaIncInv,
     Iv,
     Ive,
     Log1mexp,
@@ -226,6 +229,20 @@ def jax_funcify_Second(op, **kwargs):
     return second
 
 
+@jax_funcify.register(GammaIncInv)
+def jax_funcify_GammaIncInv(op, **kwargs):
+    gammaincinv = try_import_tfp_jax_op(op, jax_op_name="igammainv")
+
+    return gammaincinv
+
+
+@jax_funcify.register(GammaIncCInv)
+def jax_funcify_GammaIncCInv(op, **kwargs):
+    gammainccinv = try_import_tfp_jax_op(op, jax_op_name="igammacinv")
+
+    return gammainccinv
+
+
 @jax_funcify.register(Erf)
 def jax_funcify_Erf(op, node, **kwargs):
     def erf(x):
@@ -250,6 +267,7 @@ def jax_funcify_Erfinv(op, **kwargs):
     return erfinv
 
 
+@jax_funcify.register(BetaIncInv)
 @jax_funcify.register(Erfcx)
 @jax_funcify.register(Erfcinv)
 def jax_funcify_from_tfp(op, **kwargs):

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/link/jax/dispatch/slinalg.py

@@ -1,7 +1,7 @@
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
-from pytensor.tensor.slinalg import Cholesky, Solve, SolveTriangular
+from pytensor.tensor.slinalg import BlockDiagonal, Cholesky, Solve, SolveTriangular
 
 
 @jax_funcify.register(Cholesky)
@@ -45,3 +45,11 @@ def jax_funcify_SolveTriangular(op, **kwargs):
         )
 
     return solve_triangular
+
+
+@jax_funcify.register(BlockDiagonal)
+def jax_funcify_BlockDiagonalMatrix(op, **kwargs):
+    def block_diag(*inputs):
+        return jax.scipy.linalg.block_diag(*inputs)
+
+    return block_diag

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/link/jax/linker.py

@@ -23,7 +23,7 @@ class JAXLinker(JITLinker):
         # Replace any shared RNG inputs so that their values can be updated in place
         # without affecting the original RNG container. This is necessary because
         # JAX does not accept RandomState/Generators as inputs, and they will have to
-        # be typyfied
+        # be tipyfied
         if shared_rng_inputs:
             warnings.warn(
                 f"The RandomType SharedVariables {shared_rng_inputs} will not be used "
@@ -52,9 +52,16 @@ class JAXLinker(JITLinker):
                 else:  # no break
                     raise ValueError()
                 input_storage[input_storage_idx] = new_inp_storage
+                # We need to change the order of the inputs of the FunctionGraph
+                # so that the new input is in the same position as to old one,
+                # to align with the storage_map. We hope this is safe!
+                old_inp_fgrap_index = fgraph.inputs.index(old_inp)
                 fgraph.remove_input(
-                    fgraph.inputs.index(old_inp), reason="JAXLinker.fgraph_convert"
+                    old_inp_fgrap_index,
+                    reason="JAXLinker.fgraph_convert",
                 )
+                fgraph.inputs.remove(new_inp)
+                fgraph.inputs.insert(old_inp_fgrap_index, new_inp)
 
         return jax_funcify(
             fgraph, input_storage=input_storage, storage_map=storage_map, **kwargs

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/link/numba/dispatch/slinalg.py

@@ -9,7 +9,7 @@ from scipy import linalg
 
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import numba_funcify
-from pytensor.tensor.slinalg import SolveTriangular
+from pytensor.tensor.slinalg import BlockDiagonal, SolveTriangular
 
 
 _PTR = ctypes.POINTER
@@ -273,3 +273,25 @@ def numba_funcify_SolveTriangular(op, node, **kwargs):
         return res
 
     return solve_triangular
+
+
+@numba_funcify.register(BlockDiagonal)
+def numba_funcify_BlockDiagonal(op, node, **kwargs):
+    dtype = node.outputs[0].dtype
+
+    # TODO: Why do we always inline all functions? It doesn't work with starred args, so can't use it in this case.
+    @numba_basic.numba_njit(inline="never")
+    def block_diag(*arrs):
+        shapes = np.array([a.shape for a in arrs], dtype="int")
+        out_shape = [int(s) for s in np.sum(shapes, axis=0)]
+        out = np.zeros((out_shape[0], out_shape[1]), dtype=dtype)
+
+        r, c = 0, 0
+        for arr, shape in zip(arrs, shapes):
+            rr, cc = shape
+            out[r : r + rr, c : c + cc] = arr
+            r += rr
+            c += cc
+        return out
+
+    return block_diag

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/scalar/math.py

@@ -733,6 +733,64 @@ class GammaIncC(BinaryScalarOp):
 gammaincc = GammaIncC(upgrade_to_float, name="gammaincc")
 
 
+class GammaIncInv(BinaryScalarOp):
+    """
+    Inverse to the regularized lower incomplete gamma function.
+    """
+
+    nfunc_spec = ("scipy.special.gammaincinv", 2, 1)
+
+    @staticmethod
+    def st_impl(k, x):
+        return scipy.special.gammaincinv(k, x)
+
+    def impl(self, k, x):
+        return GammaIncInv.st_impl(k, x)
+
+    def grad(self, inputs, grads):
+        (k, x) = inputs
+        (gz,) = grads
+        return [
+            grad_not_implemented(self, 0, k),
+            gz * exp(gammaincinv(k, x)) * gamma(k) * (gammaincinv(k, x) ** (1 - k)),
+        ]
+
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
+
+gammaincinv = GammaIncInv(upgrade_to_float, name="gammaincinv")
+
+
+class GammaIncCInv(BinaryScalarOp):
+    """
+    Inverse to the regularized upper incomplete gamma function.
+    """
+
+    nfunc_spec = ("scipy.special.gammainccinv", 2, 1)
+
+    @staticmethod
+    def st_impl(k, x):
+        return scipy.special.gammainccinv(k, x)
+
+    def impl(self, k, x):
+        return GammaIncCInv.st_impl(k, x)
+
+    def grad(self, inputs, grads):
+        (k, x) = inputs
+        (gz,) = grads
+        return [
+            grad_not_implemented(self, 0, k),
+            gz * -exp(gammainccinv(k, x)) * gamma(k) * (gammainccinv(k, x) ** (1 - k)),
+        ]
+
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
+
+gammainccinv = GammaIncCInv(upgrade_to_float, name="gammainccinv")
+
+
 def _make_scalar_loop(n_steps, init, constant, inner_loop_fn, name, loop_op=ScalarLoop):
     init = [as_scalar(x) if x is not None else None for x in init]
     constant = [as_scalar(x) for x in constant]
@@ -1648,6 +1706,43 @@ def betainc_grad(p, q, x, wrtp: bool):
     return grad
 
 
+class BetaIncInv(ScalarOp):
+    """
+    Inverse of the regularized incomplete beta function.
+    """
+
+    nfunc_spec = ("scipy.special.betaincinv", 3, 1)
+
+    def impl(self, a, b, x):
+        return scipy.special.betaincinv(a, b, x)
+
+    def grad(self, inputs, grads):
+        (a, b, x) = inputs
+        (gz,) = grads
+        return [
+            grad_not_implemented(self, 0, a),
+            grad_not_implemented(self, 0, b),
+            gz
+            * exp(betaln(a, b))
+            * ((1 - betaincinv(a, b, x)) ** (1 - b))
+            * (betaincinv(a, b, x) ** (1 - a)),
+        ]
+
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
+
+betaincinv = BetaIncInv(upgrade_to_float_no_complex, name="betaincinv")
+
+
+def betaln(a, b):
+    """
+    Beta function from gamma function.
+    """
+
+    return gammaln(a) + gammaln(b) - gammaln(a + b)
+
+
 class Hyp2F1(ScalarOp):
     """
     Gaussian hypergeometric function ``2F1(a, b; c; z)``.

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/sparse/basic.py

@@ -7,6 +7,7 @@ http://www-users.cs.umn.edu/~saad/software/SPARSKIT/paper.ps
 TODO: Automatic methods for determining best sparse format?
 
 """
+from typing import Literal
 from warnings import warn
 
 import numpy as np
@@ -47,6 +48,7 @@ from pytensor.tensor.math import (
     trunc,
 )
 from pytensor.tensor.shape import shape, specify_broadcastable
+from pytensor.tensor.slinalg import BaseBlockDiagonal, _largest_common_dtype
 from pytensor.tensor.type import TensorType
 from pytensor.tensor.type import continuous_dtypes as tensor_continuous_dtypes
 from pytensor.tensor.type import discrete_dtypes as tensor_discrete_dtypes
@@ -60,7 +62,6 @@ from pytensor.tensor.variable import (
 
 sparse_formats = ["csc", "csr"]
 
-
 """
 Types of sparse matrices to use for testing.
 
@@ -183,7 +184,6 @@ def as_sparse_variable(x, name=None, ndim=None, **kwargs):
 
 as_sparse = as_sparse_variable
 
-
 as_sparse_or_tensor_variable = as_symbolic
 
 
@@ -1800,7 +1800,7 @@ class SpSum(Op):
         return r
 
     def __str__(self):
-        return f"{self.__class__.__name__ }{{axis={self.axis}}}"
+        return f"{self.__class__.__name__}{{axis={self.axis}}}"
 
 
 def sp_sum(x, axis=None, sparse_grad=False):
@@ -2775,19 +2775,14 @@ class GreaterEqualSD(__ComparisonOpSD):
 
 greater_equal_s_d = GreaterEqualSD()
 
-
 eq = __ComparisonSwitch(equal_s_s, equal_s_d, equal_s_d)
 
-
 neq = __ComparisonSwitch(not_equal_s_s, not_equal_s_d, not_equal_s_d)
 
-
 lt = __ComparisonSwitch(less_than_s_s, less_than_s_d, greater_than_s_d)
 
-
 gt = __ComparisonSwitch(greater_than_s_s, greater_than_s_d, less_than_s_d)
 
-
 le = __ComparisonSwitch(less_equal_s_s, less_equal_s_d, greater_equal_s_d)
 
 ge = __ComparisonSwitch(greater_equal_s_s, greater_equal_s_d, less_equal_s_d)
@@ -2992,7 +2987,7 @@ class Remove0(Op):
         l = []
         if self.inplace:
             l.append("inplace")
-        return f"{self.__class__.__name__ }{{{', '.join(l)}}}"
+        return f"{self.__class__.__name__}{{{', '.join(l)}}}"
 
     def make_node(self, x):
         """
@@ -3291,6 +3286,7 @@ class TrueDot(Op):
     # Simplify code by splitting into DotSS and DotSD.
 
     __props__ = ()
+
     # The grad_preserves_dense attribute doesn't change the
     # execution behavior.  To let the optimizer merge nodes with
     # different values of this attribute we shouldn't compare it
@@ -4260,3 +4256,85 @@ class ConstructSparseFromList(Op):
 
 
 construct_sparse_from_list = ConstructSparseFromList()
+
+
+class SparseBlockDiagonal(BaseBlockDiagonal):
+    __props__ = (
+        "n_inputs",
+        "format",
+    )
+
+    def __init__(self, n_inputs: int, format: Literal["csc", "csr"] = "csc"):
+        super().__init__(n_inputs)
+        self.format = format
+
+    def make_node(self, *matrices):
+        matrices = self._validate_and_prepare_inputs(
+            matrices, as_sparse_or_tensor_variable
+        )
+        dtype = _largest_common_dtype(matrices)
+        out_type = matrix(format=self.format, dtype=dtype)
+
+        return Apply(self, matrices, [out_type])
+
+    def perform(self, node, inputs, output_storage, params=None):
+        dtype = node.outputs[0].type.dtype
+        output_storage[0][0] = scipy.sparse.block_diag(
+            inputs, format=self.format
+        ).astype(dtype)
+
+
+def block_diag(*matrices: TensorVariable, format: Literal["csc", "csr"] = "csc"):
+    r"""
+    Construct a block diagonal matrix from a sequence of input matrices.
+
+    Given the inputs `A`, `B` and `C`, the output will have these arrays arranged on the diagonal:
+
+    [[A, 0, 0],
+     [0, B, 0],
+     [0, 0, C]]
+
+    Parameters
+    ----------
+    A, B, C ... : tensors
+        Input tensors to form the block diagonal matrix. last two dimensions of the inputs will be used, and all
+        inputs should have at least 2 dimensins.
+
+        Note that the input matrices need not be sparse themselves, and will be automatically converted to the
+        requested format if they are not.
+
+    format: str, optional
+        The format of the output sparse matrix. One of 'csr' or 'csc'. Default is 'csr'. Ignored if sparse=False.
+
+    Returns
+    -------
+    out: sparse matrix tensor
+        Symbolic sparse matrix in the specified format.
+
+    Examples
+    --------
+    Create a sparse block diagonal matrix from two sparse 2x2 matrices:
+
+    ..code-block:: python
+        import numpy as np
+        from pytensor.sparse import block_diag
+        from scipy.sparse import csr_matrix
+
+        A = csr_matrix([[1, 2], [3, 4]])
+        B = csr_matrix([[5, 6], [7, 8]])
+        result_sparse = block_diag(A, B, format='csr', name='X')
+
+        print(result_sparse)
+        >>>  SparseVariable{csr,int32}
+
+        print(result_sparse.toarray().eval())
+        >>> array([[1, 2, 0, 0],
+        >>> [3, 4, 0, 0],
+        >>> [0, 0, 5, 6],
+        >>> [0, 0, 7, 8]])
+    """
+    if len(matrices) == 1:
+        return matrices
+
+    _sparse_block_diagonal = SparseBlockDiagonal(n_inputs=len(matrices), format=format)
+    return _sparse_block_diagonal(*matrices)

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/tensor/basic.py

@@ -43,7 +43,12 @@ from pytensor.tensor import (
     get_vector_length,
 )
 from pytensor.tensor.blockwise import Blockwise
-from pytensor.tensor.elemwise import DimShuffle, Elemwise, scalar_elemwise
+from pytensor.tensor.elemwise import (
+    DimShuffle,
+    Elemwise,
+    get_normalized_batch_axes,
+    scalar_elemwise,
+)
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.shape import (
     Shape,
@@ -3614,13 +3619,18 @@ def diagonal(a, offset=0, axis1=0, axis2=1):
 
 
 @_vectorize_node.register(ExtractDiag)
-def vectorize_extract_diag(op: ExtractDiag, node, batched_x):
-    batched_ndims = batched_x.type.ndim - node.inputs[0].type.ndim
+def vectorize_extract_diag(op: ExtractDiag, node, batch_x):
+    core_ndim = node.inputs[0].type.ndim
+    batch_ndim = batch_x.type.ndim - core_ndim
+    batch_axis1, batch_axis2 = get_normalized_batch_axes(
+        (op.axis1, op.axis2), core_ndim, batch_ndim
+    )
+
     return diagonal(
-        batched_x,
+        batch_x,
         offset=op.offset,
-        axis1=op.axis1 + batched_ndims,
-        axis2=op.axis2 + batched_ndims,
+        axis1=batch_axis1,
+        axis2=batch_axis2,
     ).owner
 
 
@@ -4269,6 +4279,25 @@ def take_along_axis(arr, indices, axis=0):
     return arr[_make_along_axis_idx(arr.shape, indices, axis)]
 
 
+def ix_(*args):
+    """
+    PyTensor np.ix_ analog
+
+    See numpy.lib.index_tricks.ix_ for reference
+    """
+    out = []
+    nd = len(args)
+    for k, new in enumerate(args):
+        if new is None:
+            out.append(slice(None))
+        new = as_tensor(new)
+        if new.ndim != 1:
+            raise ValueError("Cross index must be 1 dimensional")
+        new = new.reshape((1,) * k + (new.size,) + (1,) * (nd - k - 1))
+        out.append(new)
+    return tuple(out)
+
+
 __all__ = [
     "take_along_axis",
     "expand_dims",

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/tensor/blas.py

@@ -144,20 +144,20 @@ except ImportError as e:
 # If check_init_y() == True we need to initialize y when beta == 0.
 def check_init_y():
     if check_init_y._result is None:
-        if not have_fblas:
+        if not have_fblas:  # pragma: no cover
             check_init_y._result = False
-
-        y = float("NaN") * np.ones((2,))
-        x = np.ones((2,))
-        A = np.ones((2, 2))
-        gemv = _blas_gemv_fns[y.dtype]
-        gemv(1.0, A.T, x, 0.0, y, overwrite_y=True, trans=True)
-        check_init_y._result = np.isnan(y).any()
+        else:
+            y = float("NaN") * np.ones((2,))
+            x = np.ones((2,))
+            A = np.ones((2, 2))
+            gemv = _blas_gemv_fns[y.dtype]
+            gemv(1.0, A.T, x, 0.0, y, overwrite_y=True, trans=True)
+            check_init_y._result = np.isnan(y).any()
 
     return check_init_y._result
 
 
-check_init_y._result = None
+check_init_y._result = None  # type: ignore
 
 
 class Gemv(Op):

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/tensor/blas_scipy.py

@@ -19,17 +19,13 @@ if have_fblas:
 
 
 class ScipyGer(Ger):
-    def prepare_node(self, node, storage_map, compute_map, impl):
-        if impl == "py":
-            node.tag.local_ger = _blas_ger_fns[np.dtype(node.inputs[0].type.dtype)]
-
     def perform(self, node, inputs, output_storage):
         cA, calpha, cx, cy = inputs
         (cZ,) = output_storage
         # N.B. some versions of scipy (e.g. mine) don't actually work
         # in-place on a, even when I tell it to.
         A = cA
-        local_ger = node.tag.local_ger
+        local_ger = _blas_ger_fns[cA.dtype]
         if A.size == 0:
             # We don't have to compute anything, A is empty.
             # We need this special case because Numpy considers it

{pytensor-2.18.4 → pytensor-2.18.6}/pytensor/tensor/elemwise.py

@@ -1,6 +1,8 @@
 from copy import copy
+from typing import Union
 
 import numpy as np
+from numpy.core.numeric import normalize_axis_tuple
 
 import pytensor.tensor.basic
 from pytensor.configdefaults import config
@@ -1399,7 +1401,7 @@ class CAReduce(COp):
         # scalar inputs are treated as 1D regarding axis in this `Op`
         if axis is not None:
             try:
-                axis = np.core.numeric.normalize_axis_tuple(axis, ndim=max(1, inp_dims))
+                axis = normalize_axis_tuple(axis, ndim=max(1, inp_dims))
             except np.AxisError:
                 raise np.AxisError(axis, ndim=inp_dims)
 
@@ -1757,18 +1759,36 @@ def vectorize_dimshuffle(op: DimShuffle, node: Apply, x: TensorVariable) -> Appl
     return DimShuffle(input_broadcastable, new_order).make_node(x)
 
 
-@_vectorize_node.register(CAReduce)
-def vectorize_careduce(op: CAReduce, node: Apply, x: TensorVariable) -> Apply:
-    batched_ndims = x.type.ndim - node.inputs[0].type.ndim
-    if not batched_ndims:
-        return node.op.make_node(x)
-    axes = op.axis
-    # e.g., sum(matrix, axis=None) -> sum(tensor4, axis=(2, 3))
-    # e.g., sum(matrix, axis=0) -> sum(tensor4, axis=(2,))
-    if axes is None:
-        axes = list(range(node.inputs[0].type.ndim))
+def get_normalized_batch_axes(
+    core_axes: Union[None, int, tuple[int, ...]],
+    core_ndim: int,
+    batch_ndim: int,
+) -> tuple[int, ...]:
+    """Compute batch axes for a batched operation, from the core input ndim and axes.
+
+    e.g., sum(matrix, axis=None) -> sum(tensor4, axis=(2, 3))
+    batch_axes(None, 2, 4) -> (2, 3)
+
+    e.g., sum(matrix, axis=0) -> sum(tensor4, axis=(2,))
+    batch_axes(0, 2, 4) -> (2,)
+
+    e.g., sum(tensor3, axis=(0, -1)) -> sum(tensor4, axis=(1, 3))
+    batch_axes((0, -1), 3, 4) -> (1, 3)
+    """
+    if core_axes is None:
+        core_axes = tuple(range(core_ndim))
     else:
-        axes = list(axes)
-    new_axes = [axis + batched_ndims for axis in axes]
-    new_op = op.clone(axis=new_axes)
-    return new_op.make_node(x)
+        core_axes = normalize_axis_tuple(core_axes, core_ndim)
+    return tuple(core_axis + batch_ndim for core_axis in core_axes)
+
+
+@_vectorize_node.register(CAReduce)
+def vectorize_careduce(op: CAReduce, node: Apply, batch_x: TensorVariable) -> Apply:
+    core_ndim = node.inputs[0].type.ndim
+    batch_ndim = batch_x.type.ndim - core_ndim
+
+    if not batch_ndim:
+        return node.op.make_node(batch_x)
+
+    batch_axes = get_normalized_batch_axes(op.axis, core_ndim, batch_ndim)
+    return op.clone(axis=batch_axes).make_node(batch_x)