class_docstr: An auto-batching DSL context.
  
  Auto-batching DSL operations are methods on the `ProgramBuilder` object.  It's
  used like this:
  
  ```python
  ab = dsl.ProgramBuilder()
  
  def fib_type(arg_types):
    return arg_types[0]
  
  with ab.function(type_inference=fib_type) as fibonacci:
    n = ab.param('n')
    ab.var.cond = ab.primop(lambda n: n > 1)
    with ab.if_(ab.var.cond):
      ab.var.nm1 = ab.primop(lambda n: n - 1)
      ab.var.fibm1 = ab.call(fibonacci, [ab.var.nm1])
      ab.var.nm2 = ab.primop(lambda n: n - 2)
      ab.var.fibm2 = ab.call(fibonacci, [ab.var.nm2])
      ab.var.ans = ab.primop(lambda fibm1, fibm2: fibm1 + fibm2)
    with ab.else_():
      ab.var.ans = ab.const(1)
    ab.return_(ab.var.ans)
  
  prog = ab.program(main=fibonacci)
  # Now `prog` is a well-formed `instructions.Program`, and the context
  # `ab` is no longer needed.
  ```
  
  Note the sequence of method calls on `ProgramBuilder` corresponds to the
  source code of the `Program` being defined, not its runtime behavior.  This is
  because (a) functions are defined with a context manager (rather than a Python
  function) which executes its body immediately and exactly once; and (b)
  function call instructions (and primitive operations) are just recorded, not
  entered recursively.
docstr: Registers a definition of an auto-batchable function.
  
  Example:
  ```python
  ab = dsl.ProgramBuilder()
  
  with ab.function(...) as f:
    ab.param('n')
    ...
    ab.return_(...)
  ```
  
  Note:
  - The `as` clause (here `f`) binds an `instructions.Function` object
    representing the function being defined (see Yields).
  - The formal parameters are given by calling `param` inside the
    `with` block.
  - The body of the `with` block registers the body of the function being
    defined.
  - The last statement registered in the `with` block must be a call to
    `return_`, or the `Function` will be malformed.
  
  The `function` method is a shorthand of `declare_function` followed by
  `define_function`.  The example is equivalent to:
  
  ```python
  ab = dsl.ProgramBuilder()
  
  f = ab.declare_function(...)
  with ab.define_function(f):
    ab.param('n')
    ...
    ab.return_(...)
  ```
  
  Args:
    name: Optional string naming this function when the program is printed.
    type_inference: A Python callable giving the type signature of the
      function being defined.  The callable will be invoked with a single
      argument giving the list of `instruction.Type` objects describing
      the arguments at a particular call site, and must return a list
      of `instruction.Type` objects describing the values that call site
      will return.
  
  Raises:
    ValueError: If invoked while defining a function, or if the function
      definition does not end in a `return_`.
  
  Yields:
    function: An `instructions.Function` object representing the function
      being defined.  It can be passed to `call` to call it (including
      recursively).  Note that Python scopes `as` bindings to the definition
      enclosing the `with`, so a function thus bound can be referred to after
      its body as well.
code: def function(self, name=None, type_inference=None):
      """Registers a definition of an auto-batchable function.
  
      Example:
      ```python
      ab = dsl.ProgramBuilder()
  
      with ab.function(...) as f:
        ab.param('n')
        ...
        ab.return_(...)
      ```
  
      Note:
      - The `as` clause (here `f`) binds an `instructions.Function` object
        representing the function being defined (see Yields).
      - The formal parameters are given by calling `param` inside the
        `with` block.
      - The body of the `with` block registers the body of the function being
        defined.
      - The last statement registered in the `with` block must be a call to
        `return_`, or the `Function` will be malformed.
  
      The `function` method is a shorthand of `declare_function` followed by
      `define_function`.  The example is equivalent to:
  
      ```python
      ab = dsl.ProgramBuilder()
  
      f = ab.declare_function(...)
      with ab.define_function(f):
        ab.param('n')
        ...
        ab.return_(...)
      ```
  
      Args:
        name: Optional string naming this function when the program is printed.
        type_inference: A Python callable giving the type signature of the
          function being defined.  The callable will be invoked with a single
          argument giving the list of `instruction.Type` objects describing
          the arguments at a particular call site, and must return a list
          of `instruction.Type` objects describing the values that call site
          will return.
  
      Raises:
        ValueError: If invoked while defining a function, or if the function
          definition does not end in a `return_`.
  
      Yields:
        function: An `instructions.Function` object representing the function
          being defined.  It can be passed to `call` to call it (including
          recursively).  Note that Python scopes `as` bindings to the definition
          enclosing the `with`, so a function thus bound can be referred to after
          its body as well.
      """
      function = self.declare_function(name, type_inference)
      with self.define_function(function):
          yield function
  
imports: import collections
  import contextlib
  import inspect
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 17 for calling the ProgramBuilder obj's underlying member method code with example usage: obj.function(name, type_inference) without return types