code: def _compute_gradients(self, loss_fn, x, unused_optim_state):
      
  
      # Assumes `x` is a list,
      # and contains a tensor representing a batch of images
      assert len(x) == 1 and isinstance(x, list), \
          'x should be a list and contain only one image tensor'
      x = x[0]
      loss = reduce_mean(loss_fn(x), axis=0)
      return tf.gradients(loss, x)
docs: Compute a new value of `x` to minimize `loss_fn`.
  
      Args:
          loss_fn: a callable that takes `x`, a batch of images, and returns
              a batch of loss values. `x` will be optimized to minimize
              `loss_fn(x)`.
          x: A list of Tensors, the values to be updated. This is analogous
              to the `var_list` argument in standard TF Optimizer.
          unused_optim_state: A (possibly nested) dict, containing any state
              info needed for the optimizer.
  
      Returns:
          new_x: A list of Tensors, the same length as `x`, which are updated
          new_optim_state: A dict, with the same structure as `optim_state`,
              which have been updated.