code: def testIssue182(self):
      key = random.PRNGKey(0)
      init_fun, apply_fun = stax.Softmax
      input_shape = (10, 3)
      inputs = np.arange(30.0).astype("float32").reshape(input_shape)
  
      out_shape, params = init_fun(key, input_shape)
      out = apply_fun(params, inputs)
  
      assert out_shape == out.shape
      assert np.allclose(np.sum(np.asarray(out), -1), 1.0)
  
imports: from absl.testing import absltest
  import numpy
  from jax._src import test_util
  from jax import random
  from jax.example_libraries import stax
  from jax import dtypes
  from jax import config
label_desc: Use image node_id 14 for calling the StaxTest obj's underlying member method code with example usage: obj.testIssue182() without return types