code: def test_vectorize_softmax(op, constructor, core_axis, batch_axis):
      x = tensor(shape=(5, 5, 5, 5))
      batch_x = tensor(shape=(3, 5, 5, 5, 5))
  
      node = constructor(x, axis=core_axis).owner
      assert isinstance(node.op, op)
  
      new_node = vectorize_node(node, batch_x)
      if len(batch_axis) == 1:
          assert isinstance(new_node.op, op)
          assert (new_node.op.axis,) == batch_axis
      else:
          assert isinstance(new_node.op, Blockwise) and isinstance(
              new_node.op.core_op, op
          )
          assert new_node.op.core_op.axis == core_axis
  
imports: import numpy
  import pytest
  from scipy.special import beta
  from scipy.special import factorial
  from scipy.special import log_softmax
  from scipy.special import poch
  from scipy.special import softmax
  from pytensor.compile.function import function
  from pytensor.configdefaults import config
  from pytensor.graph.replace import vectorize_node
  from pytensor.tensor.blockwise import Blockwise
  from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad, beta, betaln, factorial, log_softmax, poch, softmax
  from pytensor.tensor.type import matrix, tensor, tensor3, tensor4, vector, vectors
  from tests import unittest_tools
  from tests.tensor.utils import random_ranged
label_desc: Use image node_id 1 for calling a global function with example usage: test_vectorize_softmax(op, constructor, core_axis, batch_axis) without return types