class_docstr: Context object for auto-batching multiple Python functions together.
  
  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:
  ```python
  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`.
class_docstr_tok: ["Context","object","for","auto-batching","multiple","Python","functions","together",".","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",":","``","`","python","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","`","."]
docstr: 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` (default) or
    `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` + 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: All functions called by the decorated function that need auto-batching
  must be decorated in the same `Context` before invoking any of them.
  
  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` objects describing the
      arguments at a particular call site, and must return a list of
      `instructions.Type` objects describing the values that call site will
      return.
  
  Returns:
    dec: A decorator that may be applied to a function with the given
      type signature to auto-batch it.
  
  Raises:
    ValueError: If the decorated function predictably cannot be auto-batched,
      e.g., name-clashing with another function already decorated in this
      `Context`.
code: def batch(self, 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` (default) or
        `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` + 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: All functions called by the decorated function that need auto-batching
      must be decorated in the same `Context` before invoking any of them.
  
      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` objects describing the
          arguments at a particular call site, and must return a list of
          `instructions.Type` objects describing the values that call site will
          return.
  
      Returns:
        dec: A decorator that may be applied to a function with the given
          type signature to auto-batch it.
  
      Raises:
        ValueError: If the decorated function predictably cannot be auto-batched,
          e.g., name-clashing with another function already decorated in this
          `Context`.
      """
      return functools.partial(
          self.batch_uncurried, type_inference=type_inference
      )
  
imports: import contextlib
  import functools
  import gast
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.auto_batching import allocation_strategy
  from tensorflow_probability.python.experimental.auto_batching import dsl
  from tensorflow_probability.python.experimental.auto_batching import gast_util
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import lowering
  from tensorflow_probability.python.experimental.auto_batching import stack_optimization
  from tensorflow_probability.python.experimental.auto_batching import stackless
  from tensorflow_probability.python.experimental.auto_batching import tf_backend
  from tensorflow_probability.python.experimental.auto_batching import type_inference
  from tensorflow_probability.python.experimental.auto_batching import virtual_machine
  from tensorflow.python.autograph.converters import break_statements
  from tensorflow.python.autograph.converters import continue_statements
  from tensorflow.python.autograph.converters import return_statements
  from tensorflow.python.autograph.core import converter
  from tensorflow.python.autograph.pyct import anno
  from tensorflow.python.autograph.pyct import inspect_utils
  from tensorflow.python.autograph.pyct import parser
  from tensorflow.python.autograph.pyct import qual_names
  from tensorflow.python.autograph.pyct import templates
  from tensorflow.python.autograph.pyct import transformer
  from tensorflow.python.autograph.pyct.common_transformers import anf
  from tensorflow.python.autograph.pyct.static_analysis import activity
label_desc: Use image node_id 4 for calling the Context obj's underlying member method code with example usage: obj.batch(type_inference) and returns: functools