def functional_gradient(self, f, wrt: tuple or list, get_output: bool):
        if get_output:
            @wraps(f)
            def aux_f(*args):
                output = f(*args)
                if isinstance(output, (tuple, list)) and len(output) == 1:
                    output = output[0]
                result = (output[0], output[1:]) if isinstance(output, (tuple, list)) else (output, None)
                if result[0].ndim > 0:
                    result = jnp.sum(result[0]), result[1]
                return result
            jax_grad_f = jax.value_and_grad(aux_f, argnums=wrt, has_aux=True)
            @wraps(f)
            def unwrap_outputs(*args):
                (loss, aux), grads = jax_grad_f(*args)
                return (loss, *aux, *grads) if aux is not None else (loss, *grads)
            return unwrap_outputs
        else:
            @wraps(f)
            def nonaux_f(*args):
                output = f(*args)
                result = output[0] if isinstance(output, (tuple, list)) else output
                if result.ndim > 0:
                    result = jnp.sum(result)
                return result
            return jax.grad(nonaux_f, argnums=wrt, has_aux=False)