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 for a previously declared function.
  
  Usually, one would use the `function` method to declare and define a
  function at the same time.  Explicit use of `define_function` is only useful
  for mutual recursion or controlling code order separately from the call
  graph.
  
  Example:
  ```python
  ab = dsl.ProgramBuilder()
  
  foo = ab.declare_function(...)
  
  with ab.function(...) as bar:
    ...
    ab.call(foo)
  
  with ab.define_function(foo):
    ...
    ab.call(bar)
  ```
  
  Function bodies appear in the compiled `instructions.Program` in order of
  definition, not declaration.
  
  Note:
  - The formal parameters are given by calling `ab.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 `ab.return_`,
    or the `Function` will be malformed.
  
  Args:
    function: The function (from `declare_function`) to define.
  
  Yields:
    function: The `function` being defined, by symmetry with the
      `context.function` method.
  
  Raises:
    ValueError: If invoked while defining a function, if the `function`
      argument has already been defined, or if the function definition does
      not end in a `return_`.
code: def define_function(self, function):
      """Registers a definition for a previously declared function.
  
      Usually, one would use the `function` method to declare and define a
      function at the same time.  Explicit use of `define_function` is only useful
      for mutual recursion or controlling code order separately from the call
      graph.
  
      Example:
      ```python
      ab = dsl.ProgramBuilder()
  
      foo = ab.declare_function(...)
  
      with ab.function(...) as bar:
        ...
        ab.call(foo)
  
      with ab.define_function(foo):
        ...
        ab.call(bar)
      ```
  
      Function bodies appear in the compiled `instructions.Program` in order of
      definition, not declaration.
  
      Note:
      - The formal parameters are given by calling `ab.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 `ab.return_`,
        or the `Function` will be malformed.
  
      Args:
        function: The function (from `declare_function`) to define.
  
      Yields:
        function: The `function` being defined, by symmetry with the
          `context.function` method.
  
      Raises:
        ValueError: If invoked while defining a function, if the `function`
          argument has already been defined, or if the function definition does
          not end in a `return_`.
      """
      if self._blocks is not None:
          raise ValueError("Nested function definitions not supported")
      msg = "Internal invariant violation"
      assert self._locals is None, msg
      assert self._pending_after_else_block is None, msg
      if function.graph is not None:
          raise ValueError(
              "Cannot redefine the function named {}.".format(
                  function.name
              )
          )
      self._functions.append(function)
      if function.name is not None:
          block_name = "enter_" + function.name
      else:
          block_name = "entry"
      self._blocks = [self._fresh_block(name=block_name)]
      self._locals = {}
      yield function
      if self._blocks[-1].terminator is None:
          msg = "Every function body must end in a return_."
          raise ValueError(msg)
      function.graph = inst.ControlFlowGraph(self._blocks)
      self._blocks = None
      self._locals = None
  
imports: import collections
  import contextlib
  import inspect
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 19 for calling the ProgramBuilder obj's underlying member method code with example usage: obj.define_function(function) without return types