pytensor 2.17.4__tar.gz → 2.18.0__tar.gz

{pytensor-2.17.4 → pytensor-2.18.0}/PKG-INFO

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

{pytensor-2.17.4 → pytensor-2.18.0}/doc/library/tensor/random/basic.rst

@@ -82,9 +82,6 @@ PyTensor can produce :class:`RandomVariable`\s that draw samples from many diffe
 .. autoclass:: pytensor.tensor.random.basic.CategoricalRV
    :members: __call__
 
-.. autoclass:: pytensor.tensor.random.basic.ChiSquareRV
-   :members: __call__
-
 .. autoclass:: pytensor.tensor.random.basic.DirichletRV
    :members: __call__
 

{pytensor-2.17.4 → pytensor-2.18.0}/pyproject.toml

@@ -2,7 +2,7 @@
 requires = [
   "setuptools>=48.0.0",
   "cython",
-  "numpy>=1.17.0,<1.26",
+  "numpy>=1.17.0",
   "versioneer[toml]==0.28",
 ]
 build-backend = "setuptools.build_meta"

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-11-11T07:10:07+0000",
+ "date": "2023-11-20T14:04:50+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "0f802ab2ddb5f960218f88d1e984417cdbe94bc8",
- "version": "2.17.4"
+ "full-revisionid": "eb552eef7a8381e47a38ce0694244700693cb978",
+ "version": "2.18.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/gradient.py

@@ -101,7 +101,7 @@ def grad_undefined(op, x_pos, x, comment=""):
     return (
         NullType(
             "This variable is Null because the grad method for "
-            f"input {x_pos} ({x}) of the {op} op is not implemented. {comment}"
+            f"input {x_pos} ({x}) of the {op} op is undefined. {comment}"
         )
     )()
 

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/graph/__init__.py

@@ -9,7 +9,7 @@ from pytensor.graph.basic import (
     clone,
     ancestors,
 )
-from pytensor.graph.replace import clone_replace, graph_replace, vectorize
+from pytensor.graph.replace import clone_replace, graph_replace, vectorize_graph
 from pytensor.graph.op import Op
 from pytensor.graph.type import Type
 from pytensor.graph.fg import FunctionGraph

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/graph/replace.py

@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Iterable, Mapping, Sequence
 from functools import partial, singledispatch
 from typing import Optional, Union, cast, overload
@@ -203,7 +204,7 @@ def graph_replace(
 
 
 @singledispatch
-def _vectorize_node(op: Op, node: Apply, *bached_inputs) -> Apply:
+def _vectorize_node(op: Op, node: Apply, *batched_inputs) -> Apply:
     # Default implementation is provided in pytensor.tensor.blockwise
     raise NotImplementedError
 
@@ -214,8 +215,12 @@ def vectorize_node(node: Apply, *batched_inputs) -> Apply:
     return _vectorize_node(op, node, *batched_inputs)
 
 
+def _vectorize_not_needed(op, node, *batched_inputs):
+    return op.make_node(*batched_inputs)
+
+
 @overload
-def vectorize(
+def vectorize_graph(
     outputs: Variable,
     replace: Mapping[Variable, Variable],
 ) -> Variable:
@@ -223,14 +228,14 @@ def vectorize(
 
 
 @overload
-def vectorize(
+def vectorize_graph(
     outputs: Sequence[Variable],
     replace: Mapping[Variable, Variable],
 ) -> Sequence[Variable]:
     ...
 
 
-def vectorize(
+def vectorize_graph(
     outputs: Union[Variable, Sequence[Variable]],
     replace: Mapping[Variable, Variable],
 ) -> Union[Variable, Sequence[Variable]]:
@@ -309,3 +314,8 @@ def vectorize(
     else:
         [vect_output] = seq_vect_outputs
         return vect_output
+
+
+def vectorize(*args, **kwargs):
+    warnings.warn("vectorize was renamed to vectorize_graph", UserWarning)
+    return vectorize_graph(*args, **kwargs)

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/link/c/cmodule.py

@@ -2731,6 +2731,15 @@ def default_blas_ldflags():
             shell=True,
         )
         (stdout, stderr) = p.communicate(input=b"")
+        if p.returncode != 0:
+            warnings.warn(
+                "Pytensor cxx failed to communicate its search dirs. As a consequence, "
+                "it might not be possible to automatically determine the blas link flags to use.\n"
+                f"Command that was run: {config.cxx} -print-search-dirs\n"
+                f"Output printed to stderr: {stderr.decode(sys.stderr.encoding)}"
+            )
+            return []
+
         maybe_lib_dirs = [
             [pathlib.Path(p).resolve() for p in line[len("libraries: =") :].split(":")]
             for line in stdout.decode(sys.stdout.encoding).splitlines()
@@ -2773,6 +2782,9 @@ def default_blas_ldflags():
         else:
             raise RuntimeError(f"Supplied flags {flags} failed to compile")
 
+    # If no compiler is available we default to empty ldflags
+    if not config.cxx:
+        return ""
     _std_lib_dirs = std_lib_dirs()
     if len(_std_lib_dirs) > 0:
         rpath = _std_lib_dirs[0]

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/link/jax/dispatch/scalar.py

@@ -20,7 +20,17 @@ from pytensor.scalar.basic import (
     Second,
     Sub,
 )
-from pytensor.scalar.math import Erf, Erfc, Erfcinv, Erfcx, Erfinv, Iv, Log1mexp, Psi
+from pytensor.scalar.math import (
+    Erf,
+    Erfc,
+    Erfcinv,
+    Erfcx,
+    Erfinv,
+    Iv,
+    Log1mexp,
+    Psi,
+    TriGamma,
+)
 
 
 def try_import_tfp_jax_op(op: ScalarOp, jax_op_name: Optional[str] = None) -> Callable:
@@ -37,7 +47,7 @@ def try_import_tfp_jax_op(op: ScalarOp, jax_op_name: Optional[str] = None) -> Ca
     return typing.cast(Callable, getattr(tfp_jax_math, jax_op_name))
 
 
-def check_if_inputs_scalars(node):
+def all_inputs_are_scalar(node):
     """Check whether all the inputs of an `Elemwise` are scalar values.
 
     `jax.lax` or `jax.numpy` functions systematically return `TracedArrays`,
@@ -62,56 +72,68 @@ def check_if_inputs_scalars(node):
 
 @jax_funcify.register(ScalarOp)
 def jax_funcify_ScalarOp(op, node, **kwargs):
-    func_name = op.nfunc_spec[0]
+    """Return JAX function that implements the same computation as the Scalar Op.
+
+    This dispatch is expected to return a JAX function that works on Array inputs as Elemwise does,
+    even though it's dispatched on the Scalar Op.
+    """
 
     # We dispatch some PyTensor operators to Python operators
     # whenever the inputs are all scalars.
-    are_inputs_scalars = check_if_inputs_scalars(node)
-    if are_inputs_scalars:
-        elemwise = elemwise_scalar(op)
-        if elemwise is not None:
-            return elemwise
+    if all_inputs_are_scalar(node):
+        jax_func = jax_funcify_scalar_op_via_py_operators(op)
+        if jax_func is not None:
+            return jax_func
+
+    nfunc_spec = getattr(op, "nfunc_spec", None)
+    if nfunc_spec is None:
+        raise NotImplementedError(f"Dispatch not implemented for Scalar Op {op}")
 
+    func_name = nfunc_spec[0]
     if "." in func_name:
-        jnp_func = functools.reduce(getattr, [jax] + func_name.split("."))
+        jax_func = functools.reduce(getattr, [jax] + func_name.split("."))
     else:
-        jnp_func = getattr(jnp, func_name)
-
-    if hasattr(op, "nfunc_variadic"):
-        # These are special cases that handle invalid arities due to the broken
-        # PyTensor `Op` type contract (e.g. binary `Op`s that also function as
-        # their own variadic counterparts--even when those counterparts already
-        # exist as independent `Op`s).
-        jax_variadic_func = getattr(jnp, op.nfunc_variadic)
-
-        def elemwise(*args):
-            if len(args) > op.nfunc_spec[1]:
-                return jax_variadic_func(
-                    jnp.stack(jnp.broadcast_arrays(*args), axis=0), axis=0
-                )
-            else:
-                return jnp_func(*args)
-
-        return elemwise
-    else:
-        return jnp_func
+        jax_func = getattr(jnp, func_name)
+
+    if len(node.inputs) > op.nfunc_spec[1]:
+        # Some Scalar Ops accept multiple number of inputs, behaving as a variadic function,
+        # even though the base Op from `func_name` is specified as a binary Op.
+        # This happens with `Add`, which can work as a `Sum` for multiple scalars.
+        jax_variadic_func = getattr(jnp, op.nfunc_variadic, None)
+        if not jax_variadic_func:
+            raise NotImplementedError(
+                f"Dispatch not implemented for Scalar Op {op} with {len(node.inputs)} inputs"
+            )
+
+        def jax_func(*args):
+            return jax_variadic_func(
+                jnp.stack(jnp.broadcast_arrays(*args), axis=0), axis=0
+            )
+
+    return jax_func
 
 
 @functools.singledispatch
-def elemwise_scalar(op):
+def jax_funcify_scalar_op_via_py_operators(op):
+    """Specialized JAX dispatch for Elemwise operations where all inputs are Scalar arrays.
+
+    Scalar (constant) arrays in the JAX backend get lowered to the native types (int, floats),
+    which can perform better with Python operators, and more importantly, avoid upcasting to array types
+    not supported by some JAX functions.
+    """
     return None
 
 
-@elemwise_scalar.register(Add)
-def elemwise_scalar_add(op):
+@jax_funcify_scalar_op_via_py_operators.register(Add)
+def jax_funcify_scalar_Add(op):
     def elemwise(*inputs):
         return sum(inputs)
 
     return elemwise
 
 
-@elemwise_scalar.register(Mul)
-def elemwise_scalar_mul(op):
+@jax_funcify_scalar_op_via_py_operators.register(Mul)
+def jax_funcify_scalar_Mul(op):
     import operator
     from functools import reduce
 
@@ -121,24 +143,24 @@ def elemwise_scalar_mul(op):
     return elemwise
 
 
-@elemwise_scalar.register(Sub)
-def elemwise_scalar_sub(op):
+@jax_funcify_scalar_op_via_py_operators.register(Sub)
+def jax_funcify_scalar_Sub(op):
     def elemwise(x, y):
         return x - y
 
     return elemwise
 
 
-@elemwise_scalar.register(IntDiv)
-def elemwise_scalar_intdiv(op):
+@jax_funcify_scalar_op_via_py_operators.register(IntDiv)
+def jax_funcify_scalar_IntDiv(op):
     def elemwise(x, y):
         return x // y
 
     return elemwise
 
 
-@elemwise_scalar.register(Mod)
-def elemwise_scalar_mod(op):
+@jax_funcify_scalar_op_via_py_operators.register(Mod)
+def jax_funcify_scalar_Mod(op):
     def elemwise(x, y):
         return x % y
 
@@ -263,6 +285,14 @@ def jax_funcify_Psi(op, node, **kwargs):
     return psi
 
 
+@jax_funcify.register(TriGamma)
+def jax_funcify_TriGamma(op, node, **kwargs):
+    def tri_gamma(x):
+        return jax.scipy.special.polygamma(1, x)
+
+    return tri_gamma
+
+
 @jax_funcify.register(Softplus)
 def jax_funcify_Softplus(op, **kwargs):
     def softplus(x):

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/link/numba/dispatch/random.py

@@ -195,7 +195,6 @@ def {sized_fn_name}({random_fn_input_names}):
 @numba_funcify.register(aer.NormalRV)
 @numba_funcify.register(aer.LogNormalRV)
 @numba_funcify.register(aer.GammaRV)
-@numba_funcify.register(aer.ChiSquareRV)
 @numba_funcify.register(aer.ParetoRV)
 @numba_funcify.register(aer.GumbelRV)
 @numba_funcify.register(aer.ExponentialRV)

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/scalar/math.py

@@ -13,7 +13,7 @@ import scipy.special
 import scipy.stats
 
 from pytensor.configdefaults import config
-from pytensor.gradient import grad_not_implemented
+from pytensor.gradient import grad_not_implemented, grad_undefined
 from pytensor.scalar.basic import BinaryScalarOp, ScalarOp, UnaryScalarOp
 from pytensor.scalar.basic import abs as scalar_abs
 from pytensor.scalar.basic import (
@@ -473,8 +473,12 @@ class TriGamma(UnaryScalarOp):
     def impl(self, x):
         return TriGamma.st_impl(x)
 
-    def grad(self, inputs, outputs_gradients):
-        raise NotImplementedError()
+    def L_op(self, inputs, outputs, outputs_gradients):
+        (x,) = inputs
+        (g_out,) = outputs_gradients
+        if x in complex_types:
+            raise NotImplementedError("gradient not implemented for complex types")
+        return [g_out * polygamma(2, x)]
 
     def c_support_code(self, **kwargs):
         # The implementation has been copied from
@@ -541,7 +545,52 @@ class TriGamma(UnaryScalarOp):
         raise NotImplementedError("only floating point is implemented")
 
 
-tri_gamma = TriGamma(upgrade_to_float, name="tri_gamma")
+# Scipy polygamma does not support complex inputs: https://github.com/scipy/scipy/issues/7410
+tri_gamma = TriGamma(upgrade_to_float_no_complex, name="tri_gamma")
+
+
+class PolyGamma(BinaryScalarOp):
+    """Polygamma function of order n evaluated at x.
+
+    It corresponds to the (n+1)th derivative of the log gamma function.
+
+    TODO: Because the first input is discrete and the output is continuous,
+     the default elemwise inplace won't work, as it always tries to store the results in the first input.
+    """
+
+    nfunc_spec = ("scipy.special.polygamma", 2, 1)
+
+    @staticmethod
+    def output_types_preference(n_type, x_type):
+        if n_type not in discrete_types:
+            raise TypeError(
+                f"Polygamma order parameter must be discrete, got {n_type} dtype"
+            )
+        # Scipy doesn't support it
+        return upgrade_to_float_no_complex(x_type)
+
+    @staticmethod
+    def st_impl(n, x):
+        return scipy.special.polygamma(n, x)
+
+    def impl(self, n, x):
+        return PolyGamma.st_impl(n, x)
+
+    def L_op(self, inputs, outputs, output_gradients):
+        (n, x) = inputs
+        (g_out,) = output_gradients
+        if x in complex_types:
+            raise NotImplementedError("gradient not implemented for complex types")
+        return [
+            grad_undefined(self, 0, n),
+            g_out * self(n + 1, x),
+        ]
+
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
+
+polygamma = PolyGamma(name="polygamma")
 
 
 class Chi2SF(BinaryScalarOp):

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/tensor/__init__.py

@@ -148,6 +148,7 @@ from pytensor.tensor.subtensor import *  # noqa
 from pytensor.tensor.type import *  # noqa
 from pytensor.tensor.type_other import *  # noqa
 from pytensor.tensor.variable import TensorConstant, TensorVariable  # noqa
+from pytensor.tensor.functional import vectorize  # noqa
 
 # Allow accessing numpy constants from pytensor.tensor
 from numpy import e, euler_gamma, inf, infty, nan, newaxis, pi  # noqa

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/tensor/blockwise.py

@@ -1,4 +1,3 @@
-import re
 from collections.abc import Sequence
 from typing import Any, Optional, cast
 
@@ -9,53 +8,22 @@ from pytensor.gradient import DisconnectedType
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.null_type import NullType
 from pytensor.graph.op import Op
-from pytensor.graph.replace import _vectorize_node, vectorize
+from pytensor.graph.replace import (
+    _vectorize_node,
+    _vectorize_not_needed,
+    vectorize_graph,
+)
 from pytensor.tensor import as_tensor_variable
 from pytensor.tensor.shape import shape_padleft
 from pytensor.tensor.type import continuous_dtypes, discrete_dtypes, tensor
-from pytensor.tensor.utils import broadcast_static_dim_lengths, import_func_from_string
+from pytensor.tensor.utils import (
+    _parse_gufunc_signature,
+    broadcast_static_dim_lengths,
+    import_func_from_string,
+)
 from pytensor.tensor.variable import TensorVariable
 
 
-# TODO: Implement vectorize helper to batch whole graphs (similar to what Blockwise does for the grad)
-
-# Copied verbatim from numpy.lib.function_base
-# https://github.com/numpy/numpy/blob/f2db090eb95b87d48a3318c9a3f9d38b67b0543c/numpy/lib/function_base.py#L1999-L2029
-_DIMENSION_NAME = r"\w+"
-_CORE_DIMENSION_LIST = "(?:{0:}(?:,{0:})*)?".format(_DIMENSION_NAME)
-_ARGUMENT = rf"\({_CORE_DIMENSION_LIST}\)"
-_ARGUMENT_LIST = "{0:}(?:,{0:})*".format(_ARGUMENT)
-_SIGNATURE = "^{0:}->{0:}$".format(_ARGUMENT_LIST)
-
-
-def _parse_gufunc_signature(signature):
-    """
-    Parse string signatures for a generalized universal function.
-
-    Arguments
-    ---------
-    signature : string
-        Generalized universal function signature, e.g., ``(m,n),(n,p)->(m,p)``
-        for ``np.matmul``.
-
-    Returns
-    -------
-    Tuple of input and output core dimensions parsed from the signature, each
-    of the form List[Tuple[str, ...]].
-    """
-    signature = re.sub(r"\s+", "", signature)
-
-    if not re.match(_SIGNATURE, signature):
-        raise ValueError(f"not a valid gufunc signature: {signature}")
-    return tuple(
-        [
-            tuple(re.findall(_DIMENSION_NAME, arg))
-            for arg in re.findall(_ARGUMENT, arg_list)
-        ]
-        for arg_list in signature.split("->")
-    )
-
-
 def safe_signature(
     core_inputs: Sequence[Variable],
     core_outputs: Sequence[Variable],
@@ -73,17 +41,6 @@ def safe_signature(
     return f"{inputs_sig}->{outputs_sig}"
 
 
-@_vectorize_node.register(Op)
-def vectorize_node_fallback(op: Op, node: Apply, *bached_inputs) -> Apply:
-    if hasattr(op, "gufunc_signature"):
-        signature = op.gufunc_signature
-    else:
-        # TODO: This is pretty bad for shape inference and merge optimization!
-        #  Should get better as we add signatures to our Ops
-        signature = safe_signature(node.inputs, node.outputs)
-    return cast(Apply, Blockwise(op, signature=signature).make_node(*bached_inputs))
-
-
 class Blockwise(Op):
     """Generalizes a core `Op` to work with batched dimensions.
 
@@ -274,7 +231,7 @@ class Blockwise(Op):
 
             core_igrads = self.core_op.L_op(core_inputs, core_outputs, core_ograds)
 
-        igrads = vectorize(
+        igrads = vectorize_graph(
             [core_igrad for core_igrad in core_igrads if core_igrad is not None],
             replace=dict(
                 zip(core_inputs + core_outputs + core_ograds, inputs + outputs + ograds)
@@ -397,6 +354,15 @@ class Blockwise(Op):
             return self.name
 
 
-@_vectorize_node.register(Blockwise)
-def vectorize_not_needed(op, node, *batch_inputs):
-    return op.make_node(*batch_inputs)
+@_vectorize_node.register(Op)
+def vectorize_node_fallback(op: Op, node: Apply, *bached_inputs) -> Apply:
+    if hasattr(op, "gufunc_signature"):
+        signature = op.gufunc_signature
+    else:
+        # TODO: This is pretty bad for shape inference and merge optimization!
+        #  Should get better as we add signatures to our Ops
+        signature = safe_signature(node.inputs, node.outputs)
+    return cast(Apply, Blockwise(op, signature=signature).make_node(*bached_inputs))
+
+
+_vectorize_node.register(Blockwise, _vectorize_not_needed)

{pytensor-2.17.4 → pytensor-2.18.0}/pytensor/tensor/elemwise.py

@@ -7,7 +7,7 @@ from pytensor.configdefaults import config
 from pytensor.gradient import DisconnectedType
 from pytensor.graph.basic import Apply
 from pytensor.graph.null_type import NullType
-from pytensor.graph.replace import _vectorize_node
+from pytensor.graph.replace import _vectorize_node, _vectorize_not_needed
 from pytensor.graph.utils import MethodNotDefined
 from pytensor.link.c.basic import failure_code
 from pytensor.link.c.op import COp, ExternalCOp, OpenMPOp
@@ -22,7 +22,6 @@ from pytensor.scalar.basic import transfer_type, upcast
 from pytensor.tensor import elemwise_cgen as cgen
 from pytensor.tensor import get_vector_length
 from pytensor.tensor.basic import _get_vector_length, as_tensor_variable
-from pytensor.tensor.blockwise import vectorize_not_needed
 from pytensor.tensor.type import (
     TensorType,
     continuous_dtypes,
@@ -1741,7 +1740,7 @@ def _get_vector_length_Elemwise(op, var):
     raise ValueError(f"Length of {var} cannot be determined")
 
 
-_vectorize_node.register(Elemwise, vectorize_not_needed)
+_vectorize_node.register(Elemwise, _vectorize_not_needed)
 
 
 @_vectorize_node.register(DimShuffle)

pytensor-2.18.0/pytensor/tensor/functional.py

@@ -0,0 +1,125 @@
+from typing import Callable, Optional
+
+from pytensor.graph import vectorize_graph
+from pytensor.tensor import TensorVariable
+from pytensor.tensor.utils import _parse_gufunc_signature
+
+
+def vectorize(func: Callable, signature: Optional[str] = None) -> Callable:
+    """Create a vectorized version of a python function that takes TensorVariables as inputs and outputs.
+
+    Similar to numpy.vectorize. See respective docstrings for more details.
+
+    Parameters
+    ----------
+    func: Callable
+        Function that creates the desired outputs from TensorVariable inputs with the core dimensions.
+    signature: str, optional
+        Generalized universal function signature, e.g., (m,n),(n)->(m) for vectorized matrix-vector multiplication.
+        If not provided, it is assumed all inputs have scalar core dimensions. Unlike numpy, the outputs
+        can have arbitrary shapes when the signature is not provided.
+
+    Returns
+    -------
+    vectorized_func: Callable
+        Callable that takes TensorVariables with arbitrarily batched dimensions on the left
+        and returns variables whose graphs correspond to the vectorized expressions of func.
+
+    Notes
+    -----
+    Unlike numpy.vectorize, the equality of core dimensions implied by the signature is not explicitly asserted.
+
+    To vectorize an existing graph, use `pytensor.graph.replace.vectorize_graph` instead.
+
+
+    Examples
+    --------
+    .. code-block:: python
+
+        import pytensor
+        import pytensor.tensor as pt
+
+        def func(x):
+            return pt.exp(x) / pt.sum(pt.exp(x))
+
+        vec_func = pt.vectorize(func, signature="(a)->(a)")
+
+        x = pt.matrix("x")
+        y = vec_func(x)
+
+        fn = pytensor.function([x], y)
+        fn([[0, 1, 2], [2, 1, 0]])
+        # array([[0.09003057, 0.24472847, 0.66524096],
+        #        [0.66524096, 0.24472847, 0.09003057]])
+
+
+    .. code-block:: python
+
+        import pytensor
+        import pytensor.tensor as pt
+
+        def func(x):
+            return x[0], x[-1]
+
+        vec_func = pt.vectorize(func, signature="(a)->(),()")
+
+        x = pt.matrix("x")
+        y1, y2 = vec_func(x)
+
+        fn = pytensor.function([x], [y1, y2])
+        fn([[-10, 0, 10], [-11, 0, 11]])
+        # [array([-10., -11.]), array([10., 11.])]
+
+    """
+
+    def inner(*inputs):
+        if signature is None:
+            # Assume all inputs are scalar
+            inputs_sig = [()] * len(inputs)
+        else:
+            inputs_sig, outputs_sig = _parse_gufunc_signature(signature)
+            if len(inputs) != len(inputs_sig):
+                raise ValueError(
+                    f"Number of inputs does not match signature: {signature}"
+                )
+
+        # Create dummy core inputs by stripping the batched dimensions of inputs
+        core_inputs = []
+        for input, input_sig in zip(inputs, inputs_sig):
+            if not isinstance(input, TensorVariable):
+                raise TypeError(
+                    f"Inputs to vectorize function must be TensorVariable, got {type(input)}"
+                )
+
+            if input.ndim < len(input_sig):
+                raise ValueError(
+                    f"Input {input} has less dimensions than signature {input_sig}"
+                )
+            if len(input_sig):
+                core_shape = input.type.shape[-len(input_sig) :]
+            else:
+                core_shape = ()
+
+            core_input = input.type.clone(shape=core_shape)(name=input.name)
+            core_inputs.append(core_input)
+
+        # Call function on dummy core inputs
+        core_outputs = func(*core_inputs)
+        if core_outputs is None:
+            raise ValueError("vectorize function returned no outputs")
+
+        if signature is not None:
+            if isinstance(core_outputs, (list, tuple)):
+                n_core_outputs = len(core_outputs)
+            else:
+                n_core_outputs = 1
+            if n_core_outputs != len(outputs_sig):
+                raise ValueError(
+                    f"Number of outputs does not match signature: {signature}"
+                )
+
+        # Vectorize graph by replacing dummy core inputs by original inputs
+        outputs = vectorize_graph(core_outputs, replace=dict(zip(core_inputs, inputs)))
+        return outputs
+
+    return inner