code: def test_gumbel_softmax(shape):
      # Here, we just verify that it will generate one-hot vectors and will have gradient
      logits = mx.np.random.uniform(-2, -1, shape)
      ret = gumbel_softmax(logits)
      assume_allones = (ret == 1).sum(axis=-1).asnumpy()
      assert_allclose(assume_allones, np.ones_like(assume_allones))
  
file_path: gluon-nlp/tests/test_op.py
imports: import numpy
  from numpy.testing import assert_allclose
  import mxnet
  from mxnet import gluon
  from scipy.stats import ks_2samp
  import pytest
  from gluonnlp.op import *
label_desc: Use image node_id 4 for calling a global function with example usage: test_gumbel_softmax(shape) without return types