Analogous to tf.Optimizer.minimize

    :param loss_fn: tf Tensor, representing the loss to minimize
    :param x: list of Tensor, analogous to tf.Optimizer's var_list
    :param optim_state: A possibly nested dict, containing any optimizer state.

    Returns:
      new_x: list of Tensor, updated version of `x`
      new_optim_state: dict, updated version of `optim_state`

def minimize(self, loss_fn, x, optim_state):
    """
    Analogous to tf.Optimizer.minimize

    :param loss_fn: tf Tensor, representing the loss to minimize
    :param x: list of Tensor, analogous to tf.Optimizer's var_list
    :param optim_state: A possibly nested dict, containing any optimizer state.

    Returns:
      new_x: list of Tensor, updated version of `x`
      new_optim_state: dict, updated version of `optim_state`
    """
    grads = self._compute_gradients(loss_fn, x, optim_state)
    return self._apply_gradients(grads, x, optim_state)