class_docstr: Softmax test class.
class_docstr_tok: ["Softmax","test","class","."]
docstr: Test softmax.
docstr_tok: ["Test","softmax","."]
code: def test_softmax(
      self, hexagon_session: Session, dtype, shape, softmax_operation
  ):
      """Test softmax."""
      if dtype == "float16":
          pytest.xfail("float16 is not supported.")
  
      a_tensor = te.placeholder(shape, dtype=dtype, name="a_tensor")
  
      topi_op = OPERATOR_CONFIGS[softmax_operation]["topi"]
      b_tensor = topi_op(a_tensor, axis=1)
  
      def get_ref_data(shape):
          ref_func = tvm.topi.testing.softmax_python
          a_np = np.random.uniform(size=shape).astype(dtype)
  
          if len(shape) == 2:
              b_np = ref_func(a_np)
          elif len(shape) == 4:
              _, c, height, width = a_np.shape
              a_np_2d = a_np.transpose(0, 2, 3, 1).reshape(
                  height * width, c
              )
              b_np_2d = tvm.topi.testing.softmax_python(a_np_2d)
              b_np = b_np_2d.reshape(1, height, width, c).transpose(
                  0, 3, 1, 2
              )
  
          return a_np, b_np
  
      # get the test data
      a_np, b_np = get_ref_data(shape)
  
      with tvm.target.Target(get_hexagon_target("v68")):
          fschedule = topi.hexagon.schedule_softmax
          s = fschedule(b_tensor)
  
      func = tvm.build(
          s,
          [a_tensor, b_tensor],
          get_hexagon_target("v68"),
          name="softmax",
      )
      mod = hexagon_session.load_module(func)
  
      dev = hexagon_session.device
      a = tvm.nd.array(a_np, dev)
      b = tvm.nd.array(
          np.zeros(
              get_const_tuple(b_tensor.shape), dtype=b_tensor.dtype
          ),
          dev,
      )
      mod["softmax"](a, b)
  
      tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
  
file_path: tvm/tests/python/contrib/test_hexagon/topi/test_softmax.py
imports: import numpy
  import pytest
  import tvm
  import tvm.testing
  from tvm import topi
  from tvm import te
  from tvm.contrib.hexagon.session import Session
  import tvm.topi.testing
  from tvm.topi.utils import get_const_tuple
  from ..infrastructure import get_hexagon_target
label_desc: Use image node_id 2 for calling the TestSoftmax obj's underlying member method code with example usage: obj.test_softmax(hexagon_session, dtype, shape, softmax_operation) and returns: a_np,  b_np