code: def check_ste(net_type_str, w_init, hybridize, in_data, device=None):
      device = device or default_device()
  
      net = eval(net_type_str)(w_init=w_init)
      if hybridize:
          net.hybridize()
      # Init
      net.initialize(mx.init.Constant([w_init]), device=device)
  
      # Test:
      in_data = in_data.to_device(device)
      with mx.autograd.record():
          out = net(in_data)
      assert all(out == net.expected_output(in_data, w_init)), (
          net_type_str
          + " output is "
          + str(out)
          + ", but"
          + " expected "
          + str(net.expected_output(in_data, w_init))
      )
  
      out.backward()
      assert all(net.w.grad() == net.expected_grads(in_data, w_init)), (
          net_type_str
          + " w grads are "
          + str(net.w.grad())
          + " but expected "
          + str(net.expected_grads(in_data, w_init))
      )
      with mx.autograd.record():
          out = net(in_data)
      assert all(out == net.expected_output(in_data, w_init)), (
          net_type_str
          + " output is "
          + str(out)
          + ", but"
          + " expected "
          + str(net.expected_output(in_data, w_init))
      )
      out.backward()
      assert all(net.w.grad() == net.expected_grads(in_data, w_init)), (
          net_type_str
          + " w grads are "
          + str(net.w.grad())
          + " but expected "
          + str(net.expected_grads(in_data, w_init))
      )
  
file_path: mxnet/tests/python/unittest/test_contrib_stes_op.py
imports: from common import xfail_when_nonstandard_decimal_separator
  import mxnet
  from mxnet import nd, autograd, gluon
  from mxnet.test_utils import default_device
label_desc: Use image node_id 1 for calling a global function with example usage: check_ste(net_type_str, w_init, hybridize, in_data, device) without return types