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: Records a function return instruction.
  
  Example:
  ```python
  ab = dsl.ProgramBuilder()
  
  with ab.function(...) as f:
    ...
    ab.var.result = ...
    ab.return_(ab.var.result)
  ```
  
  A `return_` command must occur at the top level of the function definition
  (not inside any `if_`s), and must be the last statement therein.  You can
  always achieve this by assigning to a dedicated variable for the answer
  where you would otherwise return (and massaging your control flow).
  
  Args:
    vars_out: Pattern of Python strings giving the auto-batched variables to
      return.
  
  Raises:
    ValueError: If invoked more than once in a function body, or if trying to
      return variables that have not been written to.
code: def return_(self, vars_out):
      """Records a function return instruction.
  
      Example:
      ```python
      ab = dsl.ProgramBuilder()
  
      with ab.function(...) as f:
        ...
        ab.var.result = ...
        ab.return_(ab.var.result)
      ```
  
      A `return_` command must occur at the top level of the function definition
      (not inside any `if_`s), and must be the last statement therein.  You can
      always achieve this by assigning to a dedicated variable for the answer
      where you would otherwise return (and massaging your control flow).
  
      Args:
        vars_out: Pattern of Python strings giving the auto-batched variables to
          return.
  
      Raises:
        ValueError: If invoked more than once in a function body, or if trying to
          return variables that have not been written to.
      """
      # Assume the return_ call is at the top level, and the last statement in the
      # body.  If return_ is nested, the terminator may be overwritten
      # incorrectly.  If return_ is followed by something else, extra instructions
      # may get inserted before the return (becaue return_ doesn't set up a Block
      # to catch them).
      self._prepare_for_instruction()
      for var in inst.pattern_traverse(vars_out):
          if var not in self._var_defs:
              raise ValueError(
                  "Returning undefined variable {}.".format(var)
              )
      if self._functions[-1].vars_out:
          raise ValueError(
              "Function body must have exactly one return_ statement"
          )
      self._functions[-1].vars_out = inst.pattern_map(str, vars_out)
      self._blocks[-1].terminator = inst.halt_op()
  
imports: import collections
  import contextlib
  import inspect
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 16 for calling the ProgramBuilder obj's underlying member method code with example usage: obj.return_(vars_out) without return types