<br />

|---|
| [![](https://www.tensorflow.org/images/GitHub-Mark-32px.png) View source on GitHub](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L427-L754) |

Context object for auto-batching multiple Python functions together.

#### View aliases


**Main aliases**

[`tfp.experimental.auto_batching.frontend.Context`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/Context)

<br />

    tfp.experimental.auto_batching.Context()

> [!WARNING]
> **Warning:** This is alpha-grade software. The exact subset of Python that can be successfully auto-batched is ill-specified and subject to change. Furthermore, the errors elicited by stepping outside this subset are likely to be confusing and misleading. Use at your own risk.

#### Usage:

    ctx = frontend.Context()

    @ctx.batch(type_inference=lambda ...)
    def my_single_example_function_1(args):
      ...

    @ctx.batch(type_inference=lambda ...)
    def my_single_example_function_2(args):
      ...

    # etc

Then calling any of the decorated functions will execute a batch computation.
The decorated functions may call each other, including mutually recursively.
See also the `batch` method.

#### Limitations:

- You must explicitly decorate every function to be auto-batched.
- All calls to them must call them by name (no higher-order auto-batching).
- Auto-batched functions must be defined with `def`, not `lambda`.

## Methods

### `batch`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L486-L542)

    batch(
        type_inference
    )

Decorates one function to auto-batch.

The decorated function will run in batch. It accepts all the same
arguments, except:

- All arguments must have an additional leading dimension for the batch. (By special dispensation, scalar inputs are promoted to shape \[1\], which then leads to broadcasting.)
- All the arguments' sizes in the batch dimension must be the same, or 1. The latter are broadcast.
- The returned value will also have a leading batch dimension, and will have the same size.
- The batched function accepts an additional `bool` keyword argument `dry_run`. If present and `True`, just calls the unbatched version, circumventing the auto-batching system. This can be useful for debugging the program subject to auto-batching.
- The batched function accepts an additional `bool` keyword argument `stackless`. If present and `True`, invokes the stackless version of the auto-batching system. This can be useful for avoiding stack maintenance overhead; but in general, it will recover less batching, and not work in graph-mode TensorFlow.
- The batched function accepts an additional `int` keyword argument `max_stack_depth` specifying the maximum stack depth (default 15). Ignored in stackless execution.
- The batched function accepts an additional keyword argument `backend` specifying the backend to use. Must be an instance of [`auto_batching.TensorFlowBackend`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/TensorFlowBackend) (default) or [`auto_batching.NumpyBackend`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/NumpyBackend).
- The batched function accepts an additional keyword argument `block_code_cache`, a dict which allows the caching of basic block rewrites (i.e. [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function) + XLA) to live across calls to the autobatched function. The default value of `None` results in caching only within a given call to the batched function. Currently, stackless autobatching ignores the cache completely.

> [!NOTE]
> **Note:** All functions called by the decorated function that need auto-batching must be decorated in the same `Context` before invoking any of them.

<br />

<br />

<br />

| Args ||
|---|---|
| `type_inference` | A Python callable giving the type signature of the function being auto-batched. The callable will be invoked with a single argument giving the list of [`instructions.Type`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/Type) objects describing the arguments at a particular call site, and must return a list of [`instructions.Type`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/Type) objects describing the values that call site will return. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| `dec` | A decorator that may be applied to a function with the given type signature to auto-batch it. |

<br />

<br />

<br />

<br />

| Raises ||
|---|---|
| `ValueError` | If the decorated function predictably cannot be auto-batched, e.g., name-clashing with another function already decorated in this `Context`. |

<br />

### `batch_uncurried`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L544-L602)

    batch_uncurried(
        function, type_inference
    )

A non-decorator version of `batch`, which see.

### `function_names`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L604-L605)

    function_names()

### `lowered_for_args`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L750-L754)

    lowered_for_args(
        name, args, backend
    )

Helper for calling program_lowered that computes the type signature.

### `module`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L607-L635)

    module()

Constructs an `instructions.Module` for this `Context`.

<br />

<br />

<br />

| Returns ||
|---|---|
| `module` | An `instructions.Module` representing the batched computation defined by all the functions decorated with `batch` in this `Context` so far. |

<br />

### `program`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L637-L655)

    program(
        main
    )

Constructs an [`instructions.Program`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/instructions/Program) for this `Context`.

This is a helper method, equivalent to `self.module().program(main)`.

<br />

<br />

<br />

| Args ||
|---|---|
| `main` | Python string name of the function that should be the entry point. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| `prog` | An [`instructions.Program`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/instructions/Program) representing the batched computation defined by all the functions decorated with `batch` in this `Context` so far. Suitable for downstream compilation with other passes in `auto_batching`. |

<br />

<br />

<br />

<br />

| Raises ||
|---|---|
| `ValueError` | If the intended `main` function was not decorated with `batch`. |

<br />

### `program_compiled`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L657-L700)

    program_compiled(
        main, sig=None, backend=None
    )

Constructs a compiled [`instructions.Program`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/instructions/Program) for this `Context`.

This constructs the program with `self.program(main)`, and the performs type
inference and optimization, to emit a result that can be executed by the
stackless auto-batching VM.

The point of having this as a method in its own right is that it caches the
compilation on the types of the arguments.

If either `sig` or `backend` are omitted or `None`, type inference is
skipped. The result is not executable, but it can be enlightening to
inspect.

<br />

<br />

<br />

| Args ||
|---|---|
| `main` | Python string name of the function that should be the entry point. |
| `sig` | A `list` of (patterns of) [`instructions.TensorType`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/TensorType) aligned with the formal parameters to `main`. |
| `backend` | Backend implementation. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| `prog` | An [`instructions.Program`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/instructions/Program) representing the batched computation defined by all the functions decorated with `batch` in this `Context` so far. Suitable for execution or staging on real data by the auto-batching VM. |

<br />

### `program_lowered`

[View source](https://github.com/tensorflow/probability/blob/v0.25.0/tensorflow_probability/python/experimental/auto_batching/frontend.py#L702-L748)

    program_lowered(
        main, sig=None, backend=None
    )

Constructs a lowered [`instructions.Program`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/instructions/Program) for this `Context`.

This constructs the program with `self.program(main)`, and the performs type
inference, optimization, and lowering, to emit a result that can be executed
(or staged) by the auto-batching VM.

The point of having this as a method in its own right is that it caches the
compilation on the types of the arguments.

If either `sig` or `backend` are omitted or `None`, type inference is
skipped. The result is not executable, but it can be enlightening to
inspect.

<br />

<br />

<br />

| Args ||
|---|---|
| `main` | Python string name of the function that should be the entry point. |
| `sig` | A `list` of (patterns of) [`instructions.TensorType`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/TensorType) aligned with the formal parameters to `main`. |
| `backend` | Backend implementation. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| `prog` | An [`instructions.Program`](https://www.tensorflow.org/probability/api_docs/python/tfp/experimental/auto_batching/instructions/Program) representing the batched computation defined by all the functions decorated with `batch` in this `Context` so far. Suitable for execution or staging on real data by the auto-batching VM. |

<br />