code: def test_softmax(axis):
      x = matrix("x")
      x.tag.test_value = np.arange(6, dtype=config.floatX).reshape(2, 3)
      out = softmax(x, axis=axis)
      fgraph = FunctionGraph([x], [out])
      compare_jax_and_py(
          fgraph, [get_test_value(i) for i in fgraph.inputs]
      )
  
imports: import numpy
  import pytest
  import scipy.special
  import pytensor
  import pytensor.tensor
  from pytensor.compile import get_mode
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import get_test_value
  from pytensor.tensor import elemwise
  from pytensor.tensor.math import all
  from pytensor.tensor.math import prod
  from pytensor.tensor.math import sum
  from pytensor.tensor.special import SoftmaxGrad, log_softmax, softmax
  from pytensor.tensor.type import matrix, tensor, vector, vectors
  from tests.link.jax.test_basic import compare_jax_and_py
  from tests.tensor.test_elemwise import TestElemwise
label_desc: Use image node_id 4 for calling a global function with example usage: test_softmax(axis) without return types