API Reference

The public surface is small: two entry points cover almost all use, plus registration hooks for extending Tangent with new gradients and tensor types.

Differentiation

tangent.grad_util.grad = cached_grad(_grad_uncached) module-attribute

tangent.grad_util.autodiff = cached_autodiff(_autodiff_uncached) module-attribute

tangent.grad_util.vjp(func, wrt=(0,), optimized=True, check_dims=True, preserve_result=False, verbose=0)

Convenience function to produce vector-Jacobian products.

See autodiff for function arguments. Uses reverse-mode joint-motion autodiff to produce the VJP.

Source code in tangent/grad_util.py

def vjp(
    func: Callable[..., Any],
    wrt: Union[int, Tuple[int, ...]] = (0,),
    optimized: bool = True,
    check_dims: bool = True,
    preserve_result: bool = False,
    verbose: int = 0,
) -> Callable[..., Any]:
    """Convenience function to produce vector-Jacobian products.

    See `autodiff` for function arguments.
    Uses reverse-mode joint-motion autodiff to produce the VJP.
    """
    return _autodiff_uncached(
        func,
        wrt=wrt,
        motion='joint',
        mode='reverse',
        optimized=optimized,
        preserve_result=preserve_result,
        input_derivative=INPUT_DERIVATIVE.Required,
        check_dims=check_dims,
        verbose=verbose,
    )

tangent.grad_util.jvp(func, wrt=(0,), optimized=True, check_dims=True, preserve_result=False, verbose=0)

Convenience function to produce Jacobian-vector products.

See autodiff for function arguments. Uses forward-mode autodiff to produce the JVP.

Source code in tangent/grad_util.py

def jvp(
    func: Callable[..., Any],
    wrt: Union[int, Tuple[int, ...]] = (0,),
    optimized: bool = True,
    check_dims: bool = True,
    preserve_result: bool = False,
    verbose: int = 0,
) -> Callable[..., Any]:
    """Convenience function to produce Jacobian-vector products.

    See `autodiff` for function arguments.
    Uses forward-mode autodiff to produce the JVP.
    """
    return _autodiff_uncached(
        func,
        wrt=wrt,
        mode='forward',
        optimized=optimized,
        preserve_result=preserve_result,
        input_derivative=INPUT_DERIVATIVE.Required,
        check_dims=check_dims,
        verbose=verbose,
    )

Differentiable ODEs

tangent.ode.odeint(func, y0, ts, args=(), num_steps=10)

Integrate dy/dt = func(y, t, *args), returning y at each time in ts.

Parameters:

Name	Type	Description	Default
func		Dynamics func(y, t, *args) -> dy/dt, returning an array shaped like y.	required
y0		Initial state (array).	required
ts		1-D increasing array of output times; ts[0] is the initial time.	required
args		Tuple of extra arguments (the parameters theta) passed to func.	()
num_steps		Fixed RK4 sub-steps per output interval.	10

Returns:

Type	Description
	Array of shape (len(ts),) + y0.shape: the state at each time in ts.

Source code in tangent/ode.py

def odeint(func, y0, ts, args=(), num_steps=10):
    """Integrate dy/dt = func(y, t, *args), returning y at each time in `ts`.

    Args:
      func: Dynamics `func(y, t, *args) -> dy/dt`, returning an array shaped
          like `y`.
      y0: Initial state (array).
      ts: 1-D increasing array of output times; `ts[0]` is the initial time.
      args: Tuple of extra arguments (the parameters `theta`) passed to `func`.
      num_steps: Fixed RK4 sub-steps per output interval.

    Returns:
      Array of shape `(len(ts),) + y0.shape`: the state at each time in `ts`.
    """
    ts = numpy.asarray(ts, dtype=float)
    y0 = numpy.asarray(y0, dtype=float)
    out = numpy.empty((len(ts),) + y0.shape, dtype=float)
    out[0] = y0
    y = y0
    for i in range(len(ts) - 1):
        t0, t1 = ts[i], ts[i + 1]
        dt = (t1 - t0) / num_steps
        t = t0
        for _ in range(num_steps):
            y = _rk4_step(func, y, t, dt, args)
            t = t + dt
        out[i + 1] = y
    return out

Custom gradients

See the Custom Gradients guide for a walkthrough of when to use each of these.

tangent.custom_vjp.custom_vjp(func)

Mark func as differentiated by a user-registered VJP rule.

Returns func with two decorator attributes: func.defvjp registers the reverse-mode rule and func.defjvp the (optional) forward-mode rule. Until defvjp is called, differentiating a caller of func transforms its body as usual.

Source code in tangent/custom_vjp.py

def custom_vjp(func):
    """Mark `func` as differentiated by a user-registered VJP rule.

    Returns `func` with two decorator attributes: `func.defvjp` registers the
    reverse-mode rule and `func.defjvp` the (optional) forward-mode rule.
    Until `defvjp` is called, differentiating a caller of `func` transforms
    its body as usual.
    """
    params = _param_names(func)
    key = _key(func)

    def defvjp(bwd):
        from tangent import grads

        _custom_bwds[key] = bwd
        args = ', '.join(params)
        lines = ['def _adjoint_template(ans, %s):' % args]
        lines.append(
            "    _tng_customgrads = tangent.custom_bwd_call('%s', d[ans], ans, %s)" % (key, args)
        )
        if len(params) == 1:
            lines.append('    d[%s] = _tng_customgrads' % params[0])
        else:
            for i, p in enumerate(params):
                lines.append('    d[%s] = _tng_customgrads[%d]' % (p, i))
        template_fn = _make_template('\n'.join(lines) + '\n', '_adjoint_template')
        grads.adjoints[func] = template_fn
        grads.UNIMPLEMENTED_ADJOINTS.discard(func)
        return bwd

    def defjvp(jvp):
        from tangent import tangents

        _custom_jvps[key] = jvp
        args = ', '.join(params)
        dargs = ', '.join('d[%s]' % p for p in params)
        source = (
            'def _tangent_template(ans, %s):\n'
            "    d[ans] = tangent.custom_jvp_call('%s', ans, %s, %s)\n" % (args, key, args, dargs)
        )
        template_fn = _make_template(source, '_tangent_template')
        tangents.tangents[func] = template_fn
        tangents.UNIMPLEMENTED_TANGENTS.discard(func)
        return jvp

    func.defvjp = defvjp
    func.defjvp = defjvp
    return func

tangent.custom_vjp.register_adjoint(func)

Register a reverse-mode (adjoint) template for an existing callable.

Use this to teach Tangent the gradient of a function it cannot transform - typically a library or backend op you do not own. For your own Python functions, prefer custom_vjp, whose rule is plain Python.

The decorated function is an adjoint template written in Tangent's DSL: its first parameter is the primal output, the rest are the primal inputs (matched by position to the differentiated call), and it assigns each input's gradient into d[<param>] from the output gradient d[<output>]::

@tangent.register_adjoint(numpy.hypot)
def hypot_adjoint(z, x, y):
    d[x] = d[z] * x / z
    d[y] = d[z] * y / z

Registering a rule shadows any built-in rule for func. The template must have retrievable source (define it in a module, not the REPL), like any function Tangent differentiates.

Returns:

Type	Description
	A decorator that registers its argument as the adjoint template and
	returns it unchanged.

Source code in tangent/custom_vjp.py

def register_adjoint(func):
    """Register a reverse-mode (adjoint) template for an existing callable.

    Use this to teach Tangent the gradient of a function it cannot transform -
    typically a library or backend op you do not own. For your own Python
    functions, prefer `custom_vjp`, whose rule is plain Python.

    The decorated function is an adjoint *template* written in Tangent's DSL:
    its first parameter is the primal output, the rest are the primal inputs
    (matched by position to the differentiated call), and it assigns each
    input's gradient into `d[<param>]` from the output gradient `d[<output>]`::

        @tangent.register_adjoint(numpy.hypot)
        def hypot_adjoint(z, x, y):
            d[x] = d[z] * x / z
            d[y] = d[z] * y / z

    Registering a rule shadows any built-in rule for `func`. The template must
    have retrievable source (define it in a module, not the REPL), like any
    function Tangent differentiates.

    Returns:
      A decorator that registers its argument as the adjoint template and
      returns it unchanged.
    """
    from tangent import grads

    def decorator(template):
        grads.adjoints[func] = template
        grads.UNIMPLEMENTED_ADJOINTS.discard(func)
        return template

    return decorator

tangent.custom_vjp.register_tangent(func)

Register a forward-mode (tangent) template for an existing callable.

The forward-mode counterpart of register_adjoint. The decorated template receives the primal output followed by the primal inputs, and assigns the output's tangent into d[<output>] from the input tangents d[<input>]::

@tangent.register_tangent(numpy.hypot)
def hypot_tangent(z, x, y):
    d[z] = (x * d[x] + y * d[y]) / z

Returns:

Type	Description
	A decorator that registers its argument as the tangent template and
	returns it unchanged.

Source code in tangent/custom_vjp.py

def register_tangent(func):
    """Register a forward-mode (tangent) template for an existing callable.

    The forward-mode counterpart of `register_adjoint`. The decorated template
    receives the primal output followed by the primal inputs, and assigns the
    output's tangent into `d[<output>]` from the input tangents `d[<input>]`::

        @tangent.register_tangent(numpy.hypot)
        def hypot_tangent(z, x, y):
            d[z] = (x * d[x] + y * d[y]) / z

    Returns:
      A decorator that registers its argument as the tangent template and
      returns it unchanged.
    """
    from tangent import tangents

    def decorator(template):
        tangents.tangents[func] = template
        tangents.UNIMPLEMENTED_TANGENTS.discard(func)
        return template

    return decorator

tangent.custom_vjp.stop_gradient(x)

Identity in the primal; zero derivative in both modes.

Values flowing through stop_gradient are treated as constants by the differentiator - the standard way to freeze part of a computation.

Source code in tangent/custom_vjp.py

def stop_gradient(x):
    """Identity in the primal; zero derivative in both modes.

    Values flowing through `stop_gradient` are treated as constants by the
    differentiator - the standard way to freeze part of a computation.
    """
    return x

Working in notebooks and the REPL

tangent.capture.function(fn=None, *, source=None)

Decorator: capture fn's source now so Tangent can differentiate it later.

Usable bare (@tangent.function) or with an explicit source string (@tangent.function(source=...)) for functions whose source is not otherwise retrievable.

Source code in tangent/capture.py

def function(fn=None, *, source=None):
    """Decorator: capture `fn`'s source now so Tangent can differentiate it later.

    Usable bare (`@tangent.function`) or with an explicit source string
    (`@tangent.function(source=...)`) for functions whose source is not
    otherwise retrievable.
    """
    if fn is None:
        # Called with arguments: @tangent.function(source=...)
        def decorator(f):
            return function(f, source=source)

        return decorator

    if source is None:
        try:
            source = inspect.getsource(fn)
        except (OSError, TypeError):
            source = None
    if source is not None:
        _attach_source(fn, _strip_decorators(source))
    # If capture failed and no source was given, fall through unchanged: the
    # function still works if inspect.getsource succeeds at grad() time, and
    # otherwise raises the usual clear SourceCodeNotAvailableError.
    return fn

Shape checking

tangent.shape_check.check_shapes(func, *inputs)

Abstractly evaluate func on the given inputs' shapes.

Parameters:

Name	Type	Description	Default
func		The function to check (a primal, or a tangent.grad gradient).	required
*inputs		Example inputs - real arrays, or ShapedArrays, or scalars. Only their shapes and dtypes are used; no data is read.	()

Returns:

Type	Description
	The output ShapedArray (or structure of them).

Raises:

Type	Description
ShapeError	On a rank/broadcast/matmul/reshape mismatch, with the offending line of func's source in the message.

Source code in tangent/shape_check.py

def check_shapes(func, *inputs):
    """Abstractly evaluate `func` on the given inputs' shapes.

    Args:
      func: The function to check (a primal, or a `tangent.grad` gradient).
      *inputs: Example inputs - real arrays, or `ShapedArray`s, or scalars.
          Only their shapes and dtypes are used; no data is read.

    Returns:
      The output `ShapedArray` (or structure of them).

    Raises:
      ShapeError: On a rank/broadcast/matmul/reshape mismatch, with the
          offending line of `func`'s source in the message.
    """
    abstract = [_to_abstract(x) for x in inputs]
    try:
        return func(*abstract)
    except ShapeError as e:
        loc = _locate(e)
        if loc and loc not in str(e):
            raise ShapeError('%s\n    at %s' % (e, loc)) from None
        raise

Debugging

tangent.explain.explain(func, *args, wrt=(0,), out=print)

Show the primal, the generated adjoint, and a verified gradient.

Parameters:

Name	Type	Description	Default
func		The function to differentiate.	required
*args		Arguments to evaluate at.	()
wrt		Argument indices to differentiate with respect to.	(0,)
out		Where to print (pass lambda s: None to silence).	print

Returns:

Type	Description
	A dict with value, gradient, numeric_gradient, max_error,
	gradient_source, source_map, and (for straight-line functions)
	dataflow - the primal's data-flow graph and the set of variables that
	influence the output - and dead_inputs, differentiated arguments that
	provably do not affect the output.

Source code in tangent/explain.py

def explain(func, *args, wrt=(0,), out=print):
    """Show the primal, the generated adjoint, and a verified gradient.

    Args:
      func: The function to differentiate.
      *args: Arguments to evaluate at.
      wrt: Argument indices to differentiate with respect to.
      out: Where to print (pass `lambda s: None` to silence).

    Returns:
      A dict with `value`, `gradient`, `numeric_gradient`, `max_error`,
      `gradient_source`, `source_map`, and (for straight-line functions)
      `dataflow` - the primal's data-flow graph and the set of variables that
      influence the output - and `dead_inputs`, differentiated arguments that
      provably do not affect the output.
    """
    import tangent

    df = tangent.grad(func, wrt=wrt)
    gradient_source = getattr(df, '__tangent_source__', '<unavailable>')

    value = func(*args)
    gradient = df(*args)
    numeric = _central_difference(func, args, wrt)
    max_error = float(
        numpy.max(
            [
                numpy.max(numpy.abs(numpy.asarray(g) - numpy.asarray(n)))
                for g, n in zip(
                    gradient if isinstance(gradient, tuple) else (gradient,),
                    numeric if isinstance(numeric, list) else [numeric],
                )
            ]
        )
    )

    try:
        primal_source = inspect.getsource(func)
    except (OSError, TypeError):
        primal_source = '<source unavailable>'

    bar = '─' * 72
    out(bar)
    out('PRIMAL  %s%s' % (func.__name__, inspect.signature(func)))
    out(bar)
    out(primal_source.rstrip())
    out(bar)
    out('GRADIENT (generated source)')
    out(bar)
    out(gradient_source.rstrip())
    out(bar)

    # Data-flow graph and static influence analysis (straight-line functions).
    flow = _dataflow(func)
    dead_inputs = []
    if flow is not None:
        out('DATA FLOW (primal)')
        out(bar)
        for target, deps in flow['edges']:
            arrow = ('%s <- %s' % (target, ', '.join(deps))) if deps else ('%s <- ()' % target)
            mark = '' if target in flow['influential'] else '   [dead: never affects output]'
            out('  ' + arrow + mark)
        out('  return <- %s' % ', '.join(flow['out_names']))
        # A differentiated argument that the output does not depend on has a
        # structurally zero gradient - a static claim, not just "at this point".
        param_names = flow['params']
        for k in wrt:
            if k < len(param_names) and param_names[k] not in flow['influential']:
                dead_inputs.append(k)
                out(
                    'note: argument %d (%s) does not affect the output '
                    '(static data-flow).' % (k, param_names[k])
                )
        out(bar)

    out('value      = %r' % (value,))
    out('gradient   = %r' % (gradient,))
    out('fd check   = %r' % (numeric,))
    out('max |Δ|    = %.3g  %s' % (max_error, 'OK' if max_error < 1e-4 else 'MISMATCH'))
    grads_tuple = gradient if isinstance(gradient, tuple) else (gradient,)
    for k, g in zip(wrt, grads_tuple):
        if k not in dead_inputs and numpy.all(numpy.asarray(g) == 0):
            out(
                'note: argument %d has zero gradient here - it does not '
                'affect the output at this point.' % k
            )
    out(bar)

    return {
        'value': value,
        'gradient': gradient,
        'numeric_gradient': numeric,
        'max_error': max_error,
        'gradient_source': gradient_source,
        'source_map': source_map(df, func),
        'dataflow': flow,
        'dead_inputs': dead_inputs,
    }

tangent.explain.source_map(df, func=None)

Map generated-gradient lines back to primal statements.

Parameters:

Name	Type	Description	Default
df		A gradient function produced by tangent.grad/tangent.autodiff (it carries the generated source in __tangent_source__).	required
func		Optionally the primal function, to resolve primal line numbers in the user's source file. Defaults to the one recorded on df.	None

Returns:

Type	Description
	A list of dicts {line, code, primal, primal_line}: the 1-based line
	number and text in the generated source, the primal statement that line
	serves (None for scaffolding), and the primal's 1-based line number in
	the user's file (None when unresolvable).

Source code in tangent/explain.py

def source_map(df, func=None):
    """Map generated-gradient lines back to primal statements.

    Args:
      df: A gradient function produced by `tangent.grad`/`tangent.autodiff`
          (it carries the generated source in `__tangent_source__`).
      func: Optionally the primal function, to resolve primal line numbers in
          the user's source file. Defaults to the one recorded on `df`.

    Returns:
      A list of dicts `{line, code, primal, primal_line}`: the 1-based line
      number and text in the generated source, the primal statement that line
      serves (None for scaffolding), and the primal's 1-based line number in
      the user's file (None when unresolvable).
    """
    source = getattr(df, '__tangent_source__', None)
    if source is None:
        wrapped = getattr(df, '__wrapped__', None)
        source = getattr(wrapped, '__tangent_source__', None)
    if source is None:
        raise ValueError(
            'This function carries no generated source; pass a gradient '
            'produced by tangent.grad / tangent.autodiff.'
        )
    func = func or getattr(df, '__tangent_primal__', None)

    primal_lines = {}
    if func is not None:
        try:
            lines, start = inspect.getsourcelines(func)
            for offset, text in enumerate(lines):
                primal_lines[text.strip()] = start + offset
        except (OSError, TypeError):
            pass

    entries = []
    current_primal = None
    markers = ('# Grad of: ', '# Primal and tangent of: ', '# Beginning of forward pass')
    for lineno, text in enumerate(source.splitlines(), start=1):
        stripped = text.strip()
        for marker in markers[:2]:
            if stripped.startswith(marker):
                current_primal = stripped[len(marker) :]
                break
        entries.append(
            {
                'line': lineno,
                'code': text,
                'primal': current_primal,
                'primal_line': primal_lines.get(current_primal),
            }
        )
    return entries

Gradient surgery: insert_grad_of

with tangent.insert_grad_of(x) as dx: splices arbitrary code into the generated backward pass, with dx bound to the gradient of x at that point - Tangent's original killer feature for inspecting or modifying a gradient mid-flow. It survives the optimization passes (DCE/CSE), so the injected code is not eliminated:

import tangent

def f(x):
    y = x * x
    with tangent.insert_grad_of(y) as dy:
        print('gradient flowing into y:', dy)
        dy = dy * 0.9            # scale the gradient (surgery)
    return y * 3.0

df = tangent.grad(f)            # the `with` block runs during the backward pass

To call f directly (without differentiating) the block must be removed; the tangent decorator does that:

run_f = tangent.tangent(f)     # strips the insert_grad_of block so f runs normally

tangent.utils.insert_grad_of(var)

The context manager that allows insertion of arbitrary adjoint code.

This function can be used as a context manager e.g. with insert_grad_of(x) as dx to write code that will be inserted in the adjoint while having access to the gradients of certain variables.

This function is handled by reverse mode automatic differentiation, and shouldn't actually ever be called. If the user wants to use a function containing this context manager without taking the derivative, the tangent decorator should be used to remove it from the code.

Parameters:

Name	Type	Description	Default
var		The variable of which we want the gradient.	required

Returns:

Type	Description
	The gradient of this value.

Raises:

Type	Description
ValueError	If this context manager isn't removed using the tangent decorator and the code is actually run.

Source code in tangent/utils.py

def insert_grad_of(var):
    """The context manager that allows insertion of arbitrary adjoint code.

    This function can be used as a context manager e.g. `with insert_grad_of(x) as dx`
    to write code that will be inserted in the adjoint while having access to the
    gradients of certain variables.

    This function is handled by reverse mode automatic differentiation, and
    shouldn't actually ever be called. If the user wants to use a function
    containing this context manager without taking the derivative, the `tangent`
    decorator should be used to remove it from the code.

    Args:
      var: The variable of which we want the gradient.

    Returns:
      The gradient of this value.

    Raises:
      ValueError: If this context manager isn't removed using the `tangent`
          decorator and the code is actually run.
    """
    raise ValueError(
        'use the tangent decorator for functions containing the `with insert_grad_of` statement'
    )

Extending Tangent

Registering gradient templates

To teach Tangent the gradient of a library op it does not already know, use the public tangent.register_adjoint (reverse mode) and tangent.register_tangent (forward mode) decorators:

import numpy as np
import tangent

@tangent.register_adjoint(np.sinc)   # reverse rule: fills d[x] from d[y]
def asinc(y, x):
    d[x] = d[y] * (np.cos(np.pi * x) / x - y / x)

@tangent.register_tangent(np.sinc)   # forward rule: fills d[y] from d[x]
def tsinc(y, x):
    d[y] = d[x] * (np.cos(np.pi * x) / x - y / x)

Inside a template, d[v] denotes the derivative associated with variable v; the template's parameter names bind to the call site's result and arguments in order. (These wrap the lower-level tangent.grads.adjoint / tangent.tangents.tangent_ decorators that Tangent's own backend extensions use.) For your own Python functions, prefer custom_vjp, whose rule is plain Python.

Runtime type registries

New tensor types participate in gradient bookkeeping through the dispatch registries in tangent.utils:

tangent.utils.register_init_grad(t, init_grad_function)

Register a new gradient initializer.

Gradient initializers are used to initialize new adjoint and tangent variables. TODO: Link to the document explaining the overall terminology and mechanics.

Parameters:

Name	Type	Description	Default
t		A Python type object. The data type supported by the initializer.	required
init_grad_function		A unary function that takes an argument of type t and returns a zero object of the same size as the argument. For example, the gradient initializer for Numpy objects is zeros_like.	required

Source code in tangent/utils.py

def register_init_grad(t, init_grad_function):
    """Register a new gradient initializer.

    Gradient initializers are used to initialize new adjoint and tangent
    variables.
    TODO: Link to the document explaining the overall terminology and mechanics.

    Args:
      t: A Python type object. The data type supported by the initializer.
      init_grad_function: A unary function that takes an argument of type t
        and returns a zero object of the same size as the argument. For example,
        the gradient initializer for Numpy objects is zeros_like.
    """
    assert t not in grad_initializers
    grad_initializers[t] = (init_grad_function, True)

tangent.utils.register_add_grad(left_type, right_type, add_grad_function)

Register a new gradient adder supporting the given types.

Gradient adders are used to add (in the sense of arithmetic addition) intermediate adjoint and tangent variables. TODO: Link to the document explaining the overall terminology and mechanics.

Parameters:

Name	Type	Description	Default
left_type		A Python type object. The data type of the left operand supported by the adder.	required
right_type		A Python type object. The data type of the right operand supported by the adder.	required
add_grad_function		A binary function that takes two arguments, left and right, of the types left_type and right_type respectively, and returns their sum. For example, the gradient adder for Numpy objects is np.add.	required

Raises:

Type	Description
ValueError	If the given type pair was already registered.

Source code in tangent/utils.py

def register_add_grad(left_type, right_type, add_grad_function):
    """Register a new gradient adder supporting the given types.

    Gradient adders are used to add (in the sense of arithmetic addition)
    intermediate adjoint and tangent variables.
    TODO: Link to the document explaining the overall terminology and mechanics.

    Args:
      left_type: A Python type object. The data type of the left operand
        supported by the adder.
      right_type: A Python type object. The data type of the right operand
        supported by the adder.
      add_grad_function: A binary function that takes two arguments, left and
        right, of the types left_type and right_type respectively, and returns
        their sum. For example, the gradient adder for Numpy objects is np.add.

    Raises:
      ValueError: If the given type pair was already registered.
    """
    key = (left_type, right_type)
    if key in grad_adders:
        raise ValueError('Types %s already mapped to %s' % (key, grad_adders[key]))
    grad_adders[key] = add_grad_function

tangent.utils.register_unbroadcast(t, unbroadcaster_function)

Register a new unbroadcaster.

Unbroadcasters are used to undo broadcasting, e.g. np.eye(3) + 3 will broadcast 3 to np.shape(np.eye(3)). In the backward pass, we have to undo this.

Parameters:

Name	Type	Description	Default
t		A Python type object. The data type supported by the unbroadcaster.	required
unbroadcaster_function		A binary function that takes a first argument of type t, and a second argument that t needs to be unbroadcast to.	required

Source code in tangent/utils.py

def register_unbroadcast(t, unbroadcaster_function):
    """Register a new unbroadcaster.

    Unbroadcasters are used to undo broadcasting, e.g. np.eye(3) + 3
    will broadcast 3 to np.shape(np.eye(3)). In the backward pass, we have to
    undo this.

    Args:
      t: A Python type object. The data type supported by the
        unbroadcaster.
      unbroadcaster_function: A binary function that takes a first argument of
        type t, and a second argument that t needs to be unbroadcast to.
    """
    assert t not in unbroadcasters
    unbroadcasters[t] = unbroadcaster_function

tangent.utils.register_unreduce(t, unreducer_function)

Register a new unreducer.

Unreducers are used to undo reduction, e.g. np.sum(np.eye(3)) will reduce a (3,3) array to a scalar. In the backward pass, we have to undo this.

Parameters:

Name	Type	Description	Default
t		A Python type object. The data type supported by the unreducer.	required
unreducer_function		A function with the same signature as e.g. unreduce_array	required

Source code in tangent/utils.py

def register_unreduce(t, unreducer_function):
    """Register a new unreducer.

    Unreducers are used to undo reduction, e.g. np.sum(np.eye(3))
    will reduce a (3,3) array to a scalar. In the backward pass, we have to
    undo this.

    Args:
      t: A Python type object. The data type supported by the
        unreducer.
      unreducer_function: A function with the same signature
        as e.g. `unreduce_array`
    """
    assert t not in unreducers
    unreducers[t] = unreducer_function

tangent.utils.register_matmul_grad(t, grad_x_function, grad_y_function)

Register the partial gradients of the @ (matmul) operator for type t.

Parameters:

Name	Type	Description	Default
t		A Python type object. The array type the operands belong to.	required
grad_x_function		A ternary function fn(dz, x, y) returning d[z]/d[x] contributions for z = x @ y (including any rank-promotion cases).	required
grad_y_function		A ternary function fn(dz, x, y) returning d[z]/d[y].	required

Source code in tangent/utils.py

def register_matmul_grad(t, grad_x_function, grad_y_function):
    """Register the partial gradients of the `@` (matmul) operator for type t.

    Args:
      t: A Python type object. The array type the operands belong to.
      grad_x_function: A ternary function fn(dz, x, y) returning d[z]/d[x]
        contributions for z = x @ y (including any rank-promotion cases).
      grad_y_function: A ternary function fn(dz, x, y) returning d[z]/d[y].
    """
    matmul_grad_xs[t] = grad_x_function
    matmul_grad_ys[t] = grad_y_function

Frontend passes

tangent.passes.register_pass(pass_, before=None, after=None)

Insert an additional pass relative to a named existing pass.

This is the extension hook for backends and users: it splices a new Pass into the pipeline without editing the built-in sequence.

Parameters:

Name	Type	Description	Default
pass_		A Pass descriptor. Its name must not collide with a registered pass.	required
before		Name of the pass to insert immediately before.	None
after		Name of the pass to insert immediately after. Exactly one of before/after must be given.	None

Raises:

Type	Description
ValueError	On a name collision, a missing anchor, or if not exactly one of before/after is given.

Source code in tangent/passes.py

def register_pass(pass_, before=None, after=None):
    """Insert an additional pass relative to a named existing pass.

    This is the extension hook for backends and users: it splices a new `Pass`
    into the pipeline without editing the built-in sequence.

    Args:
      pass_: A `Pass` descriptor. Its name must not collide with a registered
          pass.
      before: Name of the pass to insert immediately before.
      after: Name of the pass to insert immediately after.
          Exactly one of `before`/`after` must be given.

    Raises:
      ValueError: On a name collision, a missing anchor, or if not exactly one
          of `before`/`after` is given.
    """
    if (before is None) == (after is None):
        raise ValueError('Pass exactly one of before= or after=')
    if any(p.name == pass_.name for p in _REGISTRY):
        raise ValueError('A pass named %r is already registered' % pass_.name)
    if before is not None:
        _REGISTRY.insert(_index_of(before), pass_)
    else:
        _REGISTRY.insert(_index_of(after) + 1, pass_)

Utilities

tangent.backend_status()

Return the load status of Tangent's optional backend extensions.

Returns:

Type	Description
	A dict mapping backend names (e.g. 'jax', 'torch') to one of
	'available', 'not installed', or 'broken: '.

Missing optional backends are logged at DEBUG level at import time; to see those messages, enable debug logging before importing tangent:

logging.getLogger('tangent').setLevel(logging.DEBUG)

Source code in tangent/__init__.py

def backend_status():
    """Return the load status of Tangent's optional backend extensions.

    Returns:
      A dict mapping backend names (e.g. 'jax', 'torch') to one of
      'available', 'not installed', or 'broken: <error message>'.

    Missing optional backends are logged at DEBUG level at import time; to see
    those messages, enable debug logging before importing tangent:

        logging.getLogger('tangent').setLevel(logging.DEBUG)
    """
    return dict(_backend_status)

tangent.utils.checkpoint()

Annotation: segment-checkpoint the loops inside this block.

Used as a context manager around a loop::

with tangent.checkpoint():
    for i in range(n):
        ...

In the primal it is a no-op. During reverse-mode differentiation it forces segment (sqrt-n) checkpointing for the loops directly inside it, regardless of the min_length threshold and without needing grad(..., checkpoint=True) - and, because the user opted in explicitly, it also applies to loops whose length is not a compile-time constant.

Source code in tangent/utils.py

def checkpoint():
    """Annotation: segment-checkpoint the loops inside this block.

    Used as a context manager around a loop::

        with tangent.checkpoint():
            for i in range(n):
                ...

    In the primal it is a no-op. During reverse-mode differentiation it forces
    segment (sqrt-n) checkpointing for the loops directly inside it,
    regardless of the `min_length` threshold and without needing
    `grad(..., checkpoint=True)` - and, because the user opted in explicitly,
    it also applies to loops whose length is not a compile-time constant.
    """
    return _CheckpointContext()

tangent.grads.UNIMPLEMENTED_ADJOINTS / tangent.tangents.UNIMPLEMENTED_TANGENTS hold the ops known to lack a derivative in each mode; calls to them raise the clean ReverseNotImplementedError / ForwardNotImplementedError.