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.

def _compute_gradients(self, loss_fn, x, unused_optim_state):
    """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.
    """

    # 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)