code: 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
  
file_path: pytensor/pytensor/link/jax/dispatch/elemwise.py
imports: import jax
  import jax.numpy
  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
label_desc: Use image node_id 5 for calling a global function with example usage: jax_funcify_SoftmaxGrad(op) and returns: softmax_grad, unknown