# tensorplay.func

Source: https://www.tensorplay.cn/docs/func.html

tensorplay.func is a library of composable function transforms.

- A “function transform” is a higher-order function that accepts a numerical function and returns a new function that computes a different quantity.

- It has auto-differentiation transforms ([grad()](/docs/generated/tensorplay.func.grad.html#tensorplay.func.grad) returns a function that computes the gradient of f), a vectorization/batching transform ([vmap()](/docs/generated/tensorplay.func.vmap.html#tensorplay.func.vmap) returns a function that computes f over batches of inputs), and others.

- The transforms compose with each other: vmap(grad(f)) computes per-sample gradients, jacrev(jacrev(f)) a Hessian, and vmap(vmap(f)) maps over two independent batch dimensions.

## Why composable function transforms?

A number of use cases are awkward to express with the ordinary
tensor-and-module API alone:

- computing per-sample gradients (or other per-sample quantities)

- running ensembles of models on a single machine

- efficiently computing Jacobians and Hessians

Composing [vmap()](/docs/generated/tensorplay.func.vmap.html#tensorplay.func.vmap), [grad()](/docs/generated/tensorplay.func.grad.html#tensorplay.func.grad), [vjp()](/docs/generated/tensorplay.func.vjp.html#tensorplay.func.vjp), and [jvp()](/docs/generated/tensorplay.func.jvp.html#tensorplay.func.jvp) covers
all of them without designing a separate subsystem for each use case.

Transforms operate on plain callables, so a module must first be turned into
a function of its state: [functional_call()](/docs/generated/tensorplay.func.functional_call.html#tensorplay.func.functional_call) runs a module with
supplied parameters and buffers, and [stack_module_state()](/docs/generated/tensorplay.func.stack_module_state.html#tensorplay.func.stack_module_state) stacks the
state of an ensemble so [vmap()](/docs/generated/tensorplay.func.vmap.html#tensorplay.func.vmap) can map over its members.

> **Note**
>
> Operator coverage under the transforms is not complete: an operator without a batching rule raises NotImplementedError naming the operator, and a few known rough edges are collected on the UX limitations page.

- [tensorplay.func whirlwind tour](/docs/func.whirlwind_tour.html)

- [tensorplay.func API reference](/docs/func.api.html)

- [UX limitations](/docs/func.ux_limitations.html)

- [Patching batch norm](/docs/func.batch_norm.html)