def METHOD_NAME(self, objective_fn, *args, grad_fn=None, **kwargs):
    """Update trainable arguments with one step of the optimizer and return the corresponding
    objective function value prior to the step.
    Args:
        objective_fn (function): the objective function for optimization
        *args : variable length argument list for objective function
        grad_fn (function): optional gradient function of the
            objective function with respect to the variables ``*args``.
            If ``None``, the gradient function is computed automatically.
            Must return a ``tuple[array]`` with the same number of elements as ``*args``.
            Each array of the tuple should have the same shape as the corresponding argument.
        **kwargs : variable length of keyword arguments for the objective function
    Returns:
        tuple[list [array], float]: the new variable values :math:`x^{(t+1)}` and the objective
        function output prior to the step.
        If single arg is provided, list [array] is replaced by array.
    """
    g, forward = self.compute_grad(objective_fn, args, kwargs, grad_fn=grad_fn)
    new_args = self.apply_grad(g, args)
    if forward is None:
        forward = objective_fn(*args, **kwargs)
    # unwrap from list if one argument, cleaner return
    if len(new_args) == 1:
        return new_args[0], forward
    return new_args, forward