Compute gradient estimates using SPSA.

def _compute_gradients(self, loss_fn, x, unused_optim_state):
    """Compute gradient estimates using SPSA."""
    # Assumes `x` is a list, containing a [1, H, W, C] image
    # If static batch dimension is None, tf.reshape to batch size 1
    # so that static shape can be inferred
    assert len(x) == 1
    static_x_shape = x[0].get_shape().as_list()
    if static_x_shape[0] is None:
      x[0] = tf.reshape(x[0], [1] + static_x_shape[1:])
    assert x[0].get_shape().as_list()[0] == 1
    x = x[0]
    x_shape = x.get_shape().as_list()

    def body(i, grad_array):
      delta = self._delta
      delta_x = self._get_delta(x, delta)
      delta_x = tf.concat([delta_x, -delta_x], axis=0)
      loss_vals = tf.reshape(
          loss_fn(x + delta_x),
          [2 * self._num_samples] + [1] * (len(x_shape) - 1))
      avg_grad = reduce_mean(loss_vals * delta_x, axis=0) / delta
      avg_grad = tf.expand_dims(avg_grad, axis=0)
      new_grad_array = grad_array.write(i, avg_grad)
      return i + 1, new_grad_array

    def cond(i, _):
      return i < self._num_iters

    _, all_grads = tf.while_loop(
        cond,
        body,
        loop_vars=[
            0, tf.TensorArray(size=self._num_iters, dtype=tf_dtype)
        ],
        back_prop=False,
        parallel_iterations=1)
    avg_grad = reduce_sum(all_grads.stack(), axis=0)
    return [avg_grad]