INNER CODE UNIT · Python

translate_vjp

HIPS/autograd · autograd/core.py:102

def translate_vjp(vjpfun, fun, argnum):
    if vjpfun is None:
        return lambda ans, *args, **kwargs: lambda g: vspace(args[argnum]).zeros()
    elif callable(vjpfun):
        return vjpfun
    else:
        raise TypeError(f"Bad VJP '{vjpfun}' for '{fun.__name__}'")


# -------------------- forward mode --------------------


def make_jvp(fun, x):
    def jvp(g):
        start_node = JVPNode.new_root(g)
        end_value, end_node = trace(start_node, fun, x)
        if end_node is None:
            return end_value, vspace(end_value).zeros()

View source record →