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: A non-decorator version of `batch`, which see.
code: def batch_uncurried(self, function, type_inference):
      """A non-decorator version of `batch`, which see."""
      # TODO(axch): Friendly error message if `function` is not actually a
      # function.
      # TODO(axch): Friendly error message if `function` is a lambda and not a
      # `def`ed function.  Status quo will fail horribly on lambdas, because
      # the transformation requires the transformed functions be named.
      # Determine the name of the function being transformed
      name = function.__name__
      if name in self.function_names():
          msg = (
              "Cannot auto-batch the same-named function {} "
              "twice in the same `Context`."
          ).format(name)
          raise ValueError(msg)
      self._tag_function(function, type_inference)
  
      def batched(*args, **kwargs):
          """The batched function."""
          # Accepting kwargs here because Python 2.7 gets confused otherwise.
          # See, e.g.,
          # https://stackoverflow.com/questions/14003939/python-function-args-and-kwargs-with-other-specified-keyword-arguments
          # The whole strategy of adding arguments to the decorated function has
          # the unfortunate consequence that pylint uses the argument signature of
          # the decoratee to issue warnings, with the result that
          # disable=unexpected-keyword-arg needs to be added to use sites of
          # decorated functions, for example in `frontend_test.py`.
          # TODO(axch) Figure out how to avoid pylint warnings in callers.
          # Options:
          # - Add `backend` and `max_stack_depth` arguments to the `Context`
          #   constructor, and use those as defaults here.
          #   - Likely insufficient, because dynamic selection of the stack size
          #     seems important.
          # - Carry the same information in a context manager, used like
          #     with ctx.config(backend=TF_BACKEND, max_stack_depth=15):
          #       some_decorated_function(normal, arguments)
          # - Just invoke batched functions from the `Context` object:
          #     ctx.run(some_decorated_function, input_1, max_stack_depth=15)
          # - Maybe even make magic methods for the above?
          #     ctx.some_decorated_function(input_1, max_stack_depth=15)
          max_stack_depth = kwargs.pop("max_stack_depth", 15)
          backend = kwargs.pop("backend", TF_BACKEND)
          dry_run = kwargs.pop("dry_run", False)
          stackless = kwargs.pop("stackless", False)
          block_code_cache = kwargs.pop("block_code_cache", {})
          if self._in_dry_run:
              return _run_at_batch_size_one(backend, function, args)
          if dry_run:
              with self._dry_run():
                  return _run_at_batch_size_one(backend, function, args)
          if kwargs:
              msg = "Auto-batched function given unexpected keyword arguments {}."
              raise TypeError(msg.format(kwargs.keys()))
          sig = ab_type_inference.signature(
              self.program(main=name), args, backend
          )
          if stackless:
              compiled = self.program_compiled(
                  main=name, sig=sig, backend=backend
              )
              return st.execute(
                  compiled, backend, block_code_cache, *args
              )
          else:
              lowered = self.program_lowered(
                  main=name, sig=sig, backend=backend
              )
              return vm.execute(
                  lowered,
                  args,
                  max_stack_depth,
                  backend,
                  block_code_cache=block_code_cache,
              )
  
      return batched
  
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 5 for calling the Context obj's underlying member method code with example usage: obj.batch_uncurried(function, type_inference) and returns: batched, _run_at_batch_size_one, st, vm, _run_at_batch_size_one