pytensor 3.0.0__tar.gz → 3.0.2__tar.gz

{pytensor-3.0.0/pytensor.egg-info → pytensor-3.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.0.0
+Version: 3.0.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/__init__.py

@@ -1,3 +1,6 @@
+import sys
+
+
 __docformat__ = "restructuredtext en"
 
 
@@ -14,7 +17,6 @@ from pytensor.configdefaults import config
 
 # isort: off
 from pytensor import tensor
-from pytensor import sparse
 from pytensor.compile import (
     In,
     Mode,
@@ -32,9 +34,24 @@ from pytensor.scan.basic import scan
 from pytensor.scan.views import map
 from pytensor.compile.builders import OpFromGraph
 from pytensor.link.jax.ops import wrap_jax
+from pytensor import _sparse_lazy
 # isort: on
 
 
+def __getattr__(name):
+    if name == "sparse":
+        # During pytensor.sparse's own import, submodules may do
+        # `import pytensor.sparse.X as Y` which probes pytensor.sparse via
+        # getattr before the parent attribute has been set. Return the
+        # partially-loaded module from sys.modules to avoid re-entry.
+        if "pytensor.sparse" in sys.modules:
+            return sys.modules["pytensor.sparse"]
+        import pytensor.sparse as sparse
+
+        return sparse
+    raise AttributeError(f"module 'pytensor' has no attribute {name!r}")
+
+
 # Some config variables are registered by submodules. Only after all those
 # imports were executed, we can warn about remaining flags provided by the user
 # through PYTENSOR_FLAGS.

pytensor-3.0.2/pytensor/_sparse_lazy.py

@@ -0,0 +1,31 @@
+"""Lazy registration of scipy.sparse handlers on pytensor's dispatchers.
+
+Imported by `pytensor/__init__.py` so `pytensor.sparse` doesn't have to be
+loaded eagerly at startup. The fallbacks match `type(x).__module__` against
+`scipy.sparse` as a string, so this module itself doesn't import scipy.sparse;
+the real handlers (and their scipy.sparse dependency) are pulled in only when
+an actual scipy.sparse value is passed to `as_symbolic` / `shared`.
+"""
+
+from pytensor.basic import _as_symbolic
+from pytensor.compile.sharedvalue import shared_constructor
+
+
+def _lazy_as_symbolic_sparse(x):
+    if type(x).__module__.startswith("scipy.sparse"):
+        from pytensor.sparse.basic import as_symbolic_sparse
+
+        return as_symbolic_sparse
+    return None
+
+
+def _lazy_shared_sparse(x):
+    if type(x).__module__.startswith("scipy.sparse"):
+        from pytensor.sparse.sharedvar import sparse_constructor
+
+        return sparse_constructor
+    return None
+
+
+_as_symbolic.register_lazy(_lazy_as_symbolic_sparse)  # type: ignore[attr-defined]
+shared_constructor.register_lazy(_lazy_shared_sparse)  # type: ignore[attr-defined]

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-05-01T19:05:54+0200",
+ "date": "2026-05-12T11:06:13+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "b3cfb34cb33c8e2283952f430fa657f2326c3328",
- "version": "3.0.0"
+ "full-revisionid": "a3e786e2a595d6b686b43686cc6587794b0c286d",
+ "version": "3.0.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/basic.py

@@ -2,6 +2,7 @@ from functools import singledispatch
 from typing import Any
 
 from pytensor.graph import Variable
+from pytensor.utils import add_lazy_dispatcher
 
 
 def as_symbolic(x: Any, name: str | None = None, **kwargs) -> Variable:
@@ -40,3 +41,6 @@ def _as_symbolic(x: Any, **kwargs) -> Variable:
     from pytensor.tensor import as_tensor_variable
 
     return as_tensor_variable(x, **kwargs)
+
+
+add_lazy_dispatcher(_as_symbolic)

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/compile/builders.py

@@ -11,6 +11,7 @@ from itertools import chain
 from typing import cast
 
 from pytensor.compile.maker import function
+from pytensor.compile.mode import get_mode
 from pytensor.compile.rebuild import rebuild_collect_shared
 from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.gradient import DisconnectedType, disconnected_type, grad, pushforward
@@ -917,7 +918,9 @@ class OpFromGraph(Op, HasInnerGraph):
         if getattr(self, "_fn", None) is not None:
             return self._fn
 
-        self._fn = function(self.inner_inputs, self.inner_outputs, **self.kwargs)
+        kwargs = self.kwargs.copy()
+        mode = get_mode(kwargs.pop("mode", None)).excluding("symbolic_op_recognition")
+        self._fn = function(self.inner_inputs, self.inner_outputs, mode=mode, **kwargs)
         self._fn.trust_input = True
 
         return self._fn
@@ -940,3 +943,64 @@ class OpFromGraph(Op, HasInnerGraph):
         # zip strict not specified because we are in a hot loop
         for output, variable in zip(outputs, variables):
             output[0] = variable
+
+
+class SymbolicOp(OpFromGraph):
+    r"""OpFromGraph subclass that builds the inner graph from input types.
+
+    Subclasses define the forward graph via :meth:`build_inner_graph` and
+    optionally override :meth:`pullback` / :meth:`pushforward`.
+
+    Override :meth:`filter_inputs` to coerce raw arguments (e.g. Python
+    scalars) into typed Variables at call sites.
+
+    Set the class attribute ``inline`` to control whether the inner graph is
+    inlined during compilation (default ``False``).
+    """
+
+    inline: bool = False
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if "__props__" in cls.__dict__:
+            # MetaType installs props-only __hash__ and __eq__ which ignores the inner graph
+            # override with fgraph-aware version
+            cls.__hash__ = OpFromGraph.__hash__
+            cls.__eq__ = OpFromGraph.__eq__
+
+    @staticmethod
+    def filter_inputs(*inputs):
+        return inputs
+
+    def build_inner_graph(self, *inputs) -> list[Variable]:
+        raise NotImplementedError
+
+    def __init__(self, input_types=None, **kwargs):
+        """Construct op for the given input Types.
+
+        When input_types is None, construction is deferred until the first
+        __call__, which inspects the actual input types and builds the graph.
+        """
+        for prop in getattr(type(self), "__props__", ()):
+            if prop in kwargs:
+                setattr(self, prop, kwargs.pop(prop))
+        self._init_kwargs = kwargs
+        if input_types is not None:
+            kwargs.setdefault("inline", type(self).inline)
+            kwargs.setdefault("strict", True)
+            dummy_inputs = [t() for t in input_types]
+            outputs = self.build_inner_graph(*dummy_inputs)
+            super().__init__(dummy_inputs, outputs, **kwargs)
+
+    def __call__(self, *inputs, **kwargs):
+        inputs = self.filter_inputs(*inputs)
+        input_types = tuple(inp.type for inp in inputs)
+
+        if hasattr(self, "fgraph") and input_types == tuple(self.input_types):
+            return super().__call__(*inputs, **kwargs)
+
+        init_kwargs = dict(self._init_kwargs)
+        for prop in getattr(type(self), "__props__", ()):
+            init_kwargs[prop] = getattr(self, prop)
+        op = type(self)(input_types=list(input_types), **init_kwargs)
+        return super(SymbolicOp, op).__call__(*inputs, **kwargs)

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/compile/sharedvalue.py

@@ -10,6 +10,7 @@ from pytensor.graph.basic import Variable
 from pytensor.graph.utils import add_tag_trace
 from pytensor.link.basic import Container
 from pytensor.link.c.type import generic
+from pytensor.utils import add_lazy_dispatcher
 
 
 if TYPE_CHECKING:
@@ -223,3 +224,6 @@ def shared_constructor(value, name=None, strict=False, allow_downcast=None, **kw
         allow_downcast=allow_downcast,
         name=name,
     )
+
+
+add_lazy_dispatcher(shared_constructor)

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/graph/features.py

@@ -510,8 +510,8 @@ class FullHistory(Feature):
         from pytensor.graph.features import FullHistory
         from pytensor.graph.rewriting.utils import rewrite_graph
 
-        x = pt.scalar("x")
-        out = pt.log(pt.exp(x) / pt.sum(pt.exp(x)))
+        x = pt.vector("x")
+        out = pt.log(pt.exp(x) / pt.sum(pt.exp(x), keepdims=True))
 
         fg = FunctionGraph(outputs=[out])
         history = FullHistory()
@@ -528,22 +528,24 @@ class FullHistory(Feature):
                 pytensor.dprint(history.next())
 
     .. testoutput::
-        Log [id A] 4
-         └─ True_div [id B] 3
-            ├─ Exp [id C] 2
+        Log [id A] 5
+         └─ True_div [id B] 4
+            ├─ Exp [id C] 3
             │  └─ x [id D]
-            └─ Sum{axes=None} [id E] 1
-               └─ Exp [id F] 0
-                  └─ x [id D]
+            └─ ExpandDims{axis=0} [id E] 2
+               └─ Sum{axes=None} [id F] 1
+                  └─ Exp [id G] 0
+                     └─ x [id D]
         >> MergeOptimizer
-        Log [id A] 3
-         └─ True_div [id B] 2
+        Log [id A] 4
+         └─ True_div [id B] 3
             ├─ Exp [id C] 0
             │  └─ x [id D]
-            └─ Sum{axes=None} [id E] 1
-               └─ Exp [id C] 0
-                  └─ ···
-        >> local_mul_canonizer
+            └─ ExpandDims{axis=0} [id E] 2
+               └─ Sum{axes=None} [id F] 1
+                  └─ Exp [id C] 0
+                     └─ ···
+        >> local_softmax_stabilize
         Log [id A] 1
          └─ Softmax{axis=None} [id B] 0
             └─ x [id C]
@@ -564,22 +566,24 @@ class FullHistory(Feature):
         Log [id A] 1
          └─ Softmax{axis=None} [id B] 0
             └─ x [id C]
-        >> local_mul_canonizer
-        Log [id A] 3
-         └─ True_div [id B] 2
+        >> local_softmax_stabilize
+        Log [id A] 4
+         └─ True_div [id B] 3
             ├─ Exp [id C] 0
             │  └─ x [id D]
-            └─ Sum{axes=None} [id E] 1
-               └─ Exp [id C] 0
-                  └─ ···
+            └─ ExpandDims{axis=0} [id E] 2
+               └─ Sum{axes=None} [id F] 1
+                  └─ Exp [id C] 0
+                     └─ ···
         >> MergeOptimizer
-        Log [id A] 4
-         └─ True_div [id B] 3
-            ├─ Exp [id C] 2
+        Log [id A] 5
+         └─ True_div [id B] 4
+            ├─ Exp [id C] 3
             │  └─ x [id D]
-            └─ Sum{axes=None} [id E] 1
-               └─ Exp [id F] 0
-                  └─ x [id D]
+            └─ ExpandDims{axis=0} [id E] 2
+               └─ Sum{axes=None} [id F] 1
+                  └─ Exp [id G] 0
+                     └─ x [id D]
 
 
     .. testcode::

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/jax/dispatch/elemwise.py

@@ -3,7 +3,7 @@ import jax.numpy as jnp
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
-from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
+from pytensor.tensor.special import LogSoftmax, Softmax
 
 
 @jax_funcify.register(Elemwise)
@@ -94,17 +94,6 @@ def jax_funcify_Softmax(op, **kwargs):
     return softmax
 
 
-@jax_funcify.register(SoftmaxGrad)
-def jax_funcify_SoftmaxGrad(op, **kwargs):
-    axis = op.axis
-
-    def softmax_grad(dy, sm):
-        dy_times_sm = dy * sm
-        return dy_times_sm - jnp.sum(dy_times_sm, axis=axis, keepdims=True) * sm
-
-    return softmax_grad
-
-
 @jax_funcify.register(LogSoftmax)
 def jax_funcify_LogSoftmax(op, **kwargs):
     axis = op.axis

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/jax/dispatch/linalg/solvers.py

@@ -82,8 +82,22 @@ def jax_funcify_ChoSolve(op, **kwargs):
 
 
 @jax_funcify.register(SolveSylvester)
-def jax_funcify_SolveSylsterer(op, **kwargs):
+def jax_funcify_SolveSylvester(op, **kwargs):
+    @jax.custom_vjp
     def solve_sylvester(a, b, c):
         return jax.scipy.linalg.solve_sylvester(a, b, c)
 
+    def _fwd(a, b, c):
+        x = jax.scipy.linalg.solve_sylvester(a, b, c)
+        return x, (a, b, x)
+
+    def _bwd(res, dx):
+        a, b, x = res
+        dc = jax.scipy.linalg.solve_sylvester(a.conj().T, b.conj().T, dx)
+        da = -dc @ x.conj().T
+        db = -x.conj().T @ dc
+        return da, db, dc
+
+    solve_sylvester.defvjp(_fwd, _bwd)
+
     return solve_sylvester

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/mlx/dispatch/elemwise.py

@@ -14,7 +14,7 @@ from pytensor.scalar.basic import (
     Mul,
 )
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
-from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
+from pytensor.tensor.special import LogSoftmax, Softmax
 
 
 @mlx_funcify.register(DimShuffle)
@@ -105,17 +105,6 @@ def mlx_funcify_Softmax(op, **kwargs):
     return softmax
 
 
-@mlx_funcify.register(SoftmaxGrad)
-def mlx_funcify_SoftmaxGrad(op, **kwargs):
-    axis = op.axis
-
-    def softmax_grad(dy, sm):
-        dy_times_sm = dy * sm
-        return dy_times_sm - mx.sum(dy_times_sm, axis=axis, keepdims=True) * sm
-
-    return softmax_grad
-
-
 @mlx_funcify.register(LogSoftmax)
 def mlx_funcify_LogSoftmax(op, **kwargs):
     axis = op.axis

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/numba/dispatch/compile_ops.py

@@ -50,7 +50,9 @@ def numba_deepcopy_tensor(x):
 
 
 @register_funcify_and_cache_key(OpFromGraph)
-def numba_funcify_OpFromGraph(op, node=None, **kwargs):
+def numba_funcify_OpFromGraph(
+    op, node=None, mode=NUMBA.excluding("symbolic_op_recognition"), **kwargs
+):
     _ = kwargs.pop("storage_map", None)
 
     # Apply inner rewrites
@@ -64,7 +66,7 @@ def numba_funcify_OpFromGraph(op, node=None, **kwargs):
         input_specs=input_specs,
         accept_inplace=True,
     )
-    NUMBA.optimizer(fgraph)
+    mode.optimizer(fgraph)
     output_specs = [Out(o, borrow=False) for o in fgraph.outputs]
     insert_deepcopy(fgraph, wrapped_inputs=input_specs, wrapped_outputs=output_specs)
     fgraph_fn, fgraph_cache_key = numba_funcify_and_cache_key(

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/numba/dispatch/cython_support.py

@@ -27,6 +27,7 @@ _C_TO_NUMPY: dict[str, DTypeLike] = {
     "long double": np.longdouble,
     "float complex": np.csingle,
     "double complex": np.cdouble,
+    "Py_ssize_t": np.intp,
 }
 
 
@@ -83,8 +84,10 @@ class Signature:
         raw_args = groups["args"]
 
         decl_expr = re.compile(
-            rb"\s*(?P<type>((long )|(unsigned )|(signed )|(double )|)"
-            rb"((double)|(float)|(int)|(short)|(char)|(long)|(bool)|(complex)))"
+            rb"\s*(?P<type>"
+            rb"((long )|(unsigned )|(signed )|(double )|)"
+            rb"((double)|(float)|(int)|(short)|(char)|(long)|(bool)|(complex))"
+            rb"|Py_ssize_t)"
             rb"(\s(?P<name>[\w_]*))?\s*"
         )
 

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/numba/dispatch/elemwise.py

@@ -36,13 +36,10 @@ from pytensor.scalar.basic import (
     Sub,
     TrueDiv,
     get_scalar_type,
-    maximum,
 )
-from pytensor.scalar.basic import add as add_as
 from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from pytensor.tensor.math import Argmax, Dot, MulWithoutZeros, Sum
-from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
 
 
 @singledispatch
@@ -505,122 +502,6 @@ def numba_funcify_DimShuffle(op: DimShuffle, node, **kwargs):
     return dimshuffle, cache_version
 
 
-@register_funcify_default_op_cache_key(Softmax)
-def numba_funcify_Softmax(op, node, **kwargs):
-    ndim = node.inputs[0].type.ndim
-    inp_dtype = node.inputs[0].type.numpy_dtype
-    axis = op.axis
-
-    if ndim > 1 and axis is not None:
-        reduce_max_py = create_multiaxis_reducer(
-            maximum,
-            identity=-np.inf,
-            axes=(axis,),
-            ndim=ndim,
-            out_dtype=inp_dtype,
-            keepdims=True,
-        )
-        reduce_sum_py = create_multiaxis_reducer(
-            add_as,
-            identity=0.0,
-            axes=(axis,),
-            ndim=ndim,
-            out_dtype=inp_dtype,
-            keepdims=True,
-        )
-
-        jit_fn = numba_basic.numba_njit(boundscheck=False)
-        reduce_max = jit_fn(reduce_max_py)
-        reduce_sum = jit_fn(reduce_sum_py)
-    else:
-        reduce_max = np.max
-        reduce_sum = np.sum
-
-    @numba_basic.numba_njit(boundscheck=False)
-    def softmax(x):
-        z = reduce_max(x)
-        e_x = np.exp(x - z)
-        w = reduce_sum(e_x)
-        sm = e_x / w
-        return sm
-
-    cache_version = 1
-    return softmax, cache_version
-
-
-@register_funcify_default_op_cache_key(SoftmaxGrad)
-def numba_funcify_SoftmaxGrad(op, node, **kwargs):
-    ndim = node.inputs[0].type.ndim
-    inp_dtype = node.inputs[0].type.numpy_dtype
-
-    axis = op.axis
-    if ndim > 1 and axis is not None:
-        reduce_sum_py = create_multiaxis_reducer(
-            add_as,
-            identity=0.0,
-            axes=(axis,),
-            ndim=ndim,
-            out_dtype=inp_dtype,
-            keepdims=True,
-        )
-
-        jit_fn = numba_basic.numba_njit(boundscheck=False)
-        reduce_sum = jit_fn(reduce_sum_py)
-    else:
-        reduce_sum = np.sum
-
-    @numba_basic.numba_njit(boundscheck=False)
-    def softmax_grad(dy, sm):
-        dy_times_sm = dy * sm
-        sum_dy_times_sm = reduce_sum(dy_times_sm)
-        dx = dy_times_sm - sum_dy_times_sm * sm
-        return dx
-
-    cache_version = 1
-    return softmax_grad, cache_version
-
-
-@register_funcify_default_op_cache_key(LogSoftmax)
-def numba_funcify_LogSoftmax(op, node, **kwargs):
-    ndim = node.inputs[0].type.ndim
-    inp_dtype = node.inputs[0].type.numpy_dtype
-    axis = op.axis
-
-    if ndim > 1 and axis is not None:
-        reduce_max_py = create_multiaxis_reducer(
-            maximum,
-            identity=-np.inf,
-            axes=(axis,),
-            ndim=ndim,
-            out_dtype=inp_dtype,
-            keepdims=True,
-        )
-        reduce_sum_py = create_multiaxis_reducer(
-            add_as,
-            identity=0.0,
-            axes=(axis,),
-            ndim=ndim,
-            out_dtype=inp_dtype,
-            keepdims=True,
-        )
-
-        jit_fn = numba_basic.numba_njit(boundscheck=False)
-        reduce_max = jit_fn(reduce_max_py)
-        reduce_sum = jit_fn(reduce_sum_py)
-    else:
-        reduce_max = np.max
-        reduce_sum = np.sum
-
-    @numba_basic.numba_njit(boundscheck=False)
-    def log_softmax(x):
-        xdev = x - reduce_max(x)
-        lsm = xdev - np.log(reduce_sum(np.exp(xdev)))
-        return lsm
-
-    cache_version = 1
-    return log_softmax, cache_version
-
-
 @register_funcify_default_op_cache_key(Argmax)
 def numba_funcify_Argmax(op, node, **kwargs):
     axis = op.axis

{pytensor-3.0.0 → pytensor-3.0.2}/pytensor/link/pytorch/dispatch/elemwise.py

@@ -6,7 +6,7 @@ from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
 from pytensor.scalar import ScalarLoop
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.math import All, Any, Max, Min, Prod, Sum
-from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
+from pytensor.tensor.special import LogSoftmax, Softmax
 
 
 @pytorch_funcify.register(Elemwise)
@@ -129,9 +129,34 @@ def pytorch_funcify_min(op, **kwargs):
     return torch_min
 
 
+def _pytorch_softmax_dispatch(torch_fn, axis):
+    if axis is None:
+
+        def fn(x):
+            return torch_fn(x.ravel(), dim=0).reshape(x.shape)
+
+    elif len(axis) == 1:
+
+        def fn(x):
+            return torch_fn(x, dim=axis[0])
+
+    else:
+
+        def fn(x):
+            orig_shape = x.shape
+            x = torch.movedim(x, axis, tuple(range(-len(axis), 0)))
+            unflatten_shape = x.shape[: -len(axis)]
+            x = x.reshape(*unflatten_shape, -1)
+            x = torch_fn(x, dim=-1)
+            x = x.reshape(*unflatten_shape, *[orig_shape[a] for a in axis])
+            x = torch.movedim(x, tuple(range(-len(axis), 0)), axis)
+            return x
+
+    return fn
+
+
 @pytorch_funcify.register(Softmax)
 def pytorch_funcify_Softmax(op, **kwargs):
-    axis = op.axis
     dtype = kwargs["node"].inputs[0].dtype
 
     if not dtype.startswith("float"):
@@ -139,18 +164,11 @@ def pytorch_funcify_Softmax(op, **kwargs):
             "Pytorch Softmax is not currently implemented for non-float types."
         )
 
-    def softmax(x):
-        if axis is not None:
-            return torch.softmax(x, dim=axis)
-        else:
-            return torch.softmax(x.ravel(), dim=0).reshape(x.shape)
-
-    return softmax
+    return _pytorch_softmax_dispatch(torch.softmax, op.axis)
 
 
 @pytorch_funcify.register(LogSoftmax)
 def pytorch_funcify_LogSoftmax(op, **kwargs):
-    axis = op.axis
     dtype = kwargs["node"].inputs[0].dtype
 
     if not dtype.startswith("float"):
@@ -158,24 +176,7 @@ def pytorch_funcify_LogSoftmax(op, **kwargs):
             "Pytorch LogSoftmax is not currently implemented for non-float types."
         )
 
-    def log_softmax(x):
-        if axis is not None:
-            return torch.log_softmax(x, dim=axis)
-        else:
-            return torch.log_softmax(x.ravel(), dim=0).reshape(x.shape)
-
-    return log_softmax
-
-
-@pytorch_funcify.register(SoftmaxGrad)
-def jax_funcify_SoftmaxGrad(op, **kwargs):
-    axis = op.axis
-
-    def softmax_grad(dy, sm):
-        dy_times_sm = dy * sm
-        return dy_times_sm - torch.sum(dy_times_sm, dim=axis, keepdim=True) * sm
-
-    return softmax_grad
+    return _pytorch_softmax_dispatch(torch.log_softmax, op.axis)
 
 
 def elemwise_ravel_fn(base_fn, op, node, **kwargs):