code: def test_row_wise_sparse_adam_output_grad(
      self,
      inputs,
      ITER,
      LR,
      beta1,
      beta2,
      epsilon,
      data_strategy,
      gc,
      dc,
  ):
      param, mom1, grad = inputs
      ITER = np.array([ITER], dtype=np.int64)
      LR = np.array([LR], dtype=np.float32)
  
      # Create a 1D row-wise average 2nd moment tensor.
      mom2 = data_strategy.draw(
          hu.tensor1d(
              min_len=param.shape[0],
              max_len=param.shape[0],
              elements=hu.elements_of_type(dtype=np.float32),
          )
      )
      mom2 = np.absolute(mom2)
  
      # Create an indexing array containing values which index into grad
      indices = data_strategy.draw(
          hu.tensor(
              max_dim=1,
              min_value=1,
              max_value=grad.shape[0],
              dtype=np.int64,
              elements=st.sampled_from(np.arange(grad.shape[0])),
          ),
      )
  
      # Note that unlike SparseAdam, RowWiseSparseAdam uses a moment
      # tensor that is strictly 1-dimensional and equal in length to the
      # first dimension of the parameters, so indices must also be
      # 1-dimensional.
      indices = indices.flatten()
  
      hypothesis.note("indices.shape: %s" % str(indices.shape))
  
      # Verify that the generated indices are unique
      hypothesis.assume(
          np.array_equal(np.unique(indices), np.sort(indices))
      )
  
      # Sparsify grad
      grad = grad[indices]
  
      op = core.CreateOperator(
          "RowWiseSparseAdam",
          ["param", "mom1", "mom2", "indices", "grad", "lr", "iter"],
          ["param", "mom1", "mom2", "output_grad"],
          beta1=beta1,
          beta2=beta2,
          epsilon=epsilon,
      )
  
      def ref_row_wise_sparse_output_grad(
          param,
          mom1,
          mom2,
          indices,
          grad,
          LR,
          ITER,
          beta1,
          beta2,
          epsilon,
          output_grad,
      ):
          param_out = np.copy(param)
          mom1_out = np.copy(mom1)
          mom2_out = np.copy(mom2)
          grad_out = np.copy(grad)
  
          for i, index in enumerate(indices):
              (
                  param_out[index],
                  mom1_out[index],
                  mom2_out[index],
                  grad_out[i],
              ) = self.ref_row_wise_adam(
                  param[index],
                  mom1[index],
                  mom2[index],
                  grad[i],
                  LR,
                  ITER,
                  beta1,
                  beta2,
                  epsilon,
                  output_grad,
              )
          return (param_out, mom1_out, mom2_out, grad_out)
  
      # Iter lives on the CPU
      input_device_options = {"iter": hu.cpu_do}
  
      self.assertDeviceChecks(
          dc,
          op,
          [param, mom1, mom2, indices, grad, LR, ITER],
          [0, 1, 2, 3],
          input_device_options=input_device_options,
      )
  
      self.assertReferenceChecks(
          gc,
          op,
          [param, mom1, mom2, indices, grad, LR, ITER],
          functools.partial(
              ref_row_wise_sparse_output_grad,
              beta1=beta1,
              beta2=beta2,
              epsilon=epsilon,
              output_grad=True,
          ),
          input_device_options=input_device_options,
      )
  
file_path: pytorch/caffe2/python/operator_test/adam_test.py
imports: import functools
  import hypothesis
  from hypothesis import given
  import hypothesis.strategies
  import numpy
  import unittest
  from caffe2.python import core, workspace
  import caffe2.python.hypothesis_test_util
label_desc: Use image node_id 10 for calling the TestAdam obj's underlying member method code with example usage: obj.test_row_wise_sparse_adam_output_grad(inputs, ITER, LR, beta1, beta2, epsilon, data_strategy, gc, dc) and returns: param_out,  mom1_out,  mom2_out,  grad_out