code: def test_storm_sparse(
      self, inputs, grad_sq_sum, lr, momentum, beta, gc, dc
  ):
      param, moment, grad = inputs
      grad_sq_sum = np.array([grad_sq_sum], dtype=np.float32)
      lr = np.array([lr], dtype=np.float32)
  
      # Create an indexing array containing values that are lists of indices,
      # which index into grad
      indices = np.random.choice(
          np.arange(grad.shape[0]),
          size=np.random.randint(grad.shape[0]),
          replace=False,
      )
  
      # Sparsify grad
      grad = grad[indices]
  
      op = core.CreateOperator(
          "SparseStorm",
          ["param", "moment", "grad_sq_sum", "grad", "indices", "lr"],
          ["param", "moment", "grad_sq_sum"],
          momentum=momentum,
          beta=beta,
          device_option=gc,
      )
  
      def ref_sparse(
          param, moment, grad_sq_sum, grad, indices, lr, momentum, beta
      ):
          param_out = np.copy(param)
          moment_out = np.copy(moment)
          grad_sq_sum_out = np.copy(grad_sq_sum)
  
          grad_sq_sum_out = grad_sq_sum + np.sum(grad * grad)
          nlr = lr * np.power(beta + grad_sq_sum_out, -1.0 / 3.0)
          alpha = momentum * np.square(nlr)
          for i, index in enumerate(indices):
              gi = grad[i]
              moment_out[index] = gi + (1 - alpha) * (
                  moment[index] - gi
              )
              param_out[index] = param[index] + nlr * moment_out[index]
  
          return (
              param_out.astype(np.float32),
              moment_out.astype(np.float32),
              grad_sq_sum_out.astype(np.float32),
          )
  
      self.assertReferenceChecks(
          gc,
          op,
          [param, moment, grad_sq_sum, grad, indices, lr],
          functools.partial(ref_sparse, momentum=momentum, beta=beta),
      )
  
file_path: pytorch/caffe2/python/operator_test/storm_test.py
imports: import functools
  from hypothesis import given, settings, HealthCheck
  import hypothesis.strategies
  import numpy
  from caffe2.python import core
  import caffe2.python.hypothesis_test_util
label_desc: Use image node_id 2 for calling the TestStorm obj's underlying member method code with example usage: obj.test_storm_sparse(inputs, grad_sq_sum, lr, momentum, beta, gc, dc) and returns: param_out,  moment_out,  grad_sq_sum_out