functions_on_tuples.py

from autograd.core import primitive
from autograd.container_types import make_tuple
from autograd.util import quick_grad_check
import autograd.numpy as np

@primitive
def _f(tupl):
    a, b = tupl
    return 2.*a + 3.*b

_f.defvjp(lambda g, ans, vs, gvs, tupl: (2.*g, 3.*g))

# for caller's convenience, ensure we pass in a boxed tuple rather than a tuple
# of boxes
def f(tupl):
    return _f(make_tuple(*tupl))

def g(x):
    return f((x, 2.*x))

quick_grad_check(g, 1.)


# here's a second version with a slightly different f signature
def f2(*args):
    return _f(make_tuple(*args))

def g2(x):
    return f2(x, 2.*x)

quick_grad_check(g, 1.)