class_docstr: python reference implemenation of sgd
docstr: Perform an optimization step using gradients and states.
  
  Parameters
  ----------
  indices : list of int
      List of unique indices of the parameters into the individual learning rates
      and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
      and `set_wd_mult()`, respectively.
  weights : list of NDArray
      List of parameters to be updated.
  grads : list of NDArray
      List of gradients of the objective with respect to this parameter.
  states : List of any obj
      List of state returned by `create_state()`.
docstr_tok: ["Perform","an","optimization","step","using","gradients","and","states",".","Parameters","--","--","--","--","--","indices",":","list","of","int","List","of","unique","indices","of","the","parameters","into","the","individual","learning","rates","and","weight","decays",".","Learning","rates","and","weight","decay","may","be","set","via","`","set_lr_mult","(",")","`","and","`","set_wd_mult","(",")","`",",","respectively",".","weights",":","list","of","NDArray","List","of","parameters","to","be","updated",".","grads",":","list","of","NDArray","List","of","gradients","of","the","objective","with","respect","to","this","parameter",".","states",":","List","of","any","obj","List","of","state","returned","by","`","create_state","(",")","`","."]
code: def step(self, indices, weights, grads, states):
      """Perform an optimization step using gradients and states.
  
      Parameters
      ----------
      indices : list of int
          List of unique indices of the parameters into the individual learning rates
          and weight decays. Learning rates and weight decay may be set via `set_lr_mult()`
          and `set_wd_mult()`, respectively.
      weights : list of NDArray
          List of parameters to be updated.
      grads : list of NDArray
          List of gradients of the objective with respect to this parameter.
      states : List of any obj
          List of state returned by `create_state()`.
      """
      for index, weight, grad, state in zip(
          indices, weights, grads, states
      ):
          lr = self._get_lr(index)
          wd = self._get_wd(index)
          self._update_count(index)
          num_rows = weight.shape[0]
          if self.momentum == 0.0:
              # Update on a per row basis, skip all-zero rows
              for row in range(num_rows):
                  grad_row = grad[row].asnumpy()
                  all_zeros = mx.test_utils.almost_equal(
                      grad_row, np.zeros_like(grad_row)
                  )
                  if all_zeros:
                      continue
                  grad[row] *= self.rescale_grad
                  if self.clip_gradient is not None:
                      grad[row] = mx.nd.clip(
                          grad[row],
                          -self.clip_gradient,
                          self.clip_gradient,
                      )
                  grad[row] += wd * weight[row]
                  weight[row] -= lr * grad[row]
          else:
              mom = state
              for row in range(num_rows):
                  grad_row = grad[row].asnumpy()
                  all_zeros = mx.test_utils.almost_equal(
                      grad_row, np.zeros_like(grad_row)
                  )
                  if all_zeros:
                      continue
                  grad[row] *= self.rescale_grad
                  if self.clip_gradient is not None:
                      grad[row] = mx.nd.clip(
                          grad[row],
                          -self.clip_gradient,
                          self.clip_gradient,
                      )
                  grad[row] += wd * weight[row]
                  mom[row] *= self.momentum
                  mom[row] -= lr * grad[row]
                  weight[row] += mom[row]
  
imports: import itertools
  import numpy
  import itertools
  import mxnet
  import mxnet.lr_scheduler
  from mxnet import gluon
  import unittest
  import pytest
  import math
  from mxnet.test_utils import *
  from common import retry, xfail_when_nonstandard_decimal_separator
label_desc: Use image node_id 3 for calling the PySparseSGD obj's underlying member method code with example usage: obj.step(indices, weights, grads, states) without return types