code: def test_axis(self):
      input_arr = np.array(
          [
              [1, 2, 3],  # unit step
              [4, 5, 6],  # unit step
              [7, 8, 10],  # non-unit step
          ]
      )
      output = softmax(input_arr, axis=1)
  
      expected_output = np.array(
          [
              [0.09003057, 0.24472847, 0.66524096],  # unit step
              [0.09003057, 0.24472847, 0.66524096],  # unit step
              [0.04201007, 0.1141952, 0.84379473],  # non-unit step
          ]
      )
  
      assert np.allclose(output, expected_output)
  
file_path: cleanlab/tests/internal/test_numerics.py
imports: import pytest
  import numpy
  from cleanlab.internal.numerics import softmax
label_desc: Use image node_id 3 for calling the TestSoftmax obj's underlying member method code with example usage: obj.test_axis() without return types