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 primitive operation.
  
  Example:
  
  ```
  ab = dsl.ProgramBuilder()
  
  ab.var.five = ab.const(5)
  # Implicit output binding
  ab.var.ten = ab.primop(lambda five: five + five)
  # Explicit output binding
  ab.primop(lambda: (5, 10), vars_out=[ab.var.five, ab.var.ten])
  ```
  
  Args:
    f: A Python callable, the primitive operation to perform.  Can be
      an inline lambda expression in simple cases.  Must return a list or
      tuple of results, one for each intended output variable.
    vars_in: A list of Python strings, giving the auto-batched variables
      to pass into the callable when invoking it.  If absent, `primop`
      will try to infer it by inspecting the argument list of the callable
      and matching against variables bound in the local scope.
    vars_out: A pattern of Python strings, giving the auto-batched variable(s)
      to which to write the result of the callable.  Defaults to the empty
      list.
  
  Raises:
    ValueError: If the definition is invalid, if the primop references
      undefined auto-batched variables, or if auto-detection of input
      variables fails.
  
  Returns:
    op: An `instructions.PrimOp` instance representing this operation.  If one
      subsequently assigns this to a local, via `ProgramBuilder.var.foo = op`,
      that local becomes the output pattern.
code: def primop(self, f, vars_in=None, vars_out=None):
      """Records a primitive operation.
  
      Example:
  
      ```
      ab = dsl.ProgramBuilder()
  
      ab.var.five = ab.const(5)
      # Implicit output binding
      ab.var.ten = ab.primop(lambda five: five + five)
      # Explicit output binding
      ab.primop(lambda: (5, 10), vars_out=[ab.var.five, ab.var.ten])
      ```
  
      Args:
        f: A Python callable, the primitive operation to perform.  Can be
          an inline lambda expression in simple cases.  Must return a list or
          tuple of results, one for each intended output variable.
        vars_in: A list of Python strings, giving the auto-batched variables
          to pass into the callable when invoking it.  If absent, `primop`
          will try to infer it by inspecting the argument list of the callable
          and matching against variables bound in the local scope.
        vars_out: A pattern of Python strings, giving the auto-batched variable(s)
          to which to write the result of the callable.  Defaults to the empty
          list.
  
      Raises:
        ValueError: If the definition is invalid, if the primop references
          undefined auto-batched variables, or if auto-detection of input
          variables fails.
  
      Returns:
        op: An `instructions.PrimOp` instance representing this operation.  If one
          subsequently assigns this to a local, via `ProgramBuilder.var.foo = op`,
          that local becomes the output pattern.
      """
      self._prepare_for_instruction()
      if vars_out is None:
          vars_out = []
      if vars_in is None:
          # Deduce the intended variable names from the argument list of the callee.
          # Expected use case: the callee is an inline lambda expression.
          (
              args,
              varargs,
              varkw,
              _,
              kwonlyargs,
              _,
              _,
          ) = inspect.getfullargspec(f)
          vars_in = []
          for arg in args + kwonlyargs:
              if arg in self._locals:
                  vars_in.append(self._locals[arg])
              else:
                  raise ValueError(
                      "Auto-referencing unbound variable {}.".format(
                          arg
                      )
                  )
          if varargs is not None:
              raise ValueError("Varargs are not supported for primops")
          if varkw is not None:
              raise ValueError("kwargs are not supported for primops")
      for var in vars_in:
          if var not in self._var_defs:
              raise ValueError(
                  "Referencing undefined variable {}.".format(var)
              )
      prim = inst.prim_op(
          _str_list(vars_in), inst.pattern_map(str, vars_out), f
      )
      self._blocks[-1].instructions.append(prim)
      for var in inst.pattern_traverse(vars_out):
          self._mark_defined(var)
      return prim
  
imports: import collections
  import contextlib
  import inspect
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 12 for calling the ProgramBuilder obj's underlying member method code with example usage: obj.primop(f, vars_in, vars_out) and returns: prim