code: def jax_funcify_Softmax(op, **kwargs):
      axis = op.axis
  
      def softmax(x):
          return jax.nn.softmax(x, axis=axis)
  
      return softmax
  
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 4 for calling a global function with example usage: jax_funcify_Softmax(op) and returns: softmax, jax