code: def pea_nuts_program(latent_shape, choose_depth, step_state):
      """Synthetic program usable for benchmarking VM performance.
  
      This program is intended to resemble the control flow and scaling
      parameters of the NUTS algorithm, without any of the complexity.
      Hence the name.
  
      Each batch member looks like:
  
        state = ... # shape latent_shape
  
        def recur(depth, state):
          if depth > 1:
            state1 = recur(depth - 1, state)
            state2 = state1 + 1
            state3 = recur(depth - 1, state2)
            ans = state3 + 1
          else:
            ans = step_state(state)  # To simulate NUTS, something heavy
          return ans
  
        while count > 0:
          count = count - 1
          depth = choose_depth(count)
          state = recur(depth, state)
  
      Args:
        latent_shape: Python `tuple` of `int` giving the event shape of the
          latent state.
        choose_depth: Python `Tensor -> Tensor` callable.  The input
          `Tensor` will have shape `[batch_size]` (i.e., scalar event
          shape), and give the iteration of the outer while loop the
          thread is in.  The `choose_depth` function must return a `Tensor`
          of shape `[batch_size]` giving the depth, for each thread,
          to which to call `recur` in this iteration.
        step_state: Python `Tensor -> Tensor` callable.  The input and
          output `Tensor`s will have shape `[batch_size] + latent_shape`.
          This function is expected to update the state, and represents
          the "real work" versus which the VM overhead is being measured.
  
      Returns:
        program: `instructions.Program` that runs the above benchmark.
      """
      entry = instructions.Block()
      top_body = instructions.Block()
      finish_body = instructions.Block()
      enter_recur = instructions.Block()
      recur_body_1 = instructions.Block()
      recur_body_2 = instructions.Block()
      recur_body_3 = instructions.Block()
      recur_base_case = instructions.Block()
      # pylint: disable=bad-whitespace
      entry.assign_instructions(
          [
              instructions.prim_op(
                  ["count"], "cond", lambda count: count > 0
              ),  # cond = count > 0
              instructions.BranchOp(
                  "cond", top_body, instructions.halt()
              ),  # if cond
          ]
      )
      top_body.assign_instructions(
          [
              instructions.PopOp(["cond"]),  #   done with cond now
              instructions.prim_op(
                  ["count"], "ctm1", lambda count: count - 1
              ),  #   ctm1 = count - 1
              instructions.PopOp(["count"]),  #   done with count now
              instructions.push_op(
                  ["ctm1"], ["count"]
              ),  #   count = ctm1
              instructions.PopOp(["ctm1"]),  #   done with ctm1
              instructions.prim_op(
                  ["count"], "depth", choose_depth
              ),  #   depth = choose_depth(count)
              instructions.push_op(
                  ["depth", "state"], ["depth", "state"]
              ),  #   state = recur(depth, state)
              instructions.PopOp(
                  ["depth", "state"]
              ),  #     done with depth, state
              instructions.PushGotoOp(finish_body, enter_recur),
          ]
      )
      finish_body.assign_instructions(
          [
              instructions.push_op(["ans"], ["state"]),  #     ...
              instructions.PopOp(["ans"]),  #     pop callee's "ans"
              instructions.GotoOp(entry),  # end of while body
          ]
      )
      # Definition of recur begins here
      enter_recur.assign_instructions(
          [
              instructions.prim_op(
                  ["depth"], "cond1", lambda depth: depth > 0
              ),  # cond1 = depth > 0
              instructions.BranchOp(
                  "cond1", recur_body_1, recur_base_case
              ),  # if cond1
          ]
      )
      recur_body_1.assign_instructions(
          [
              instructions.PopOp(["cond1"]),  #   done with cond1 now
              instructions.prim_op(
                  ["depth"], "dm1", lambda depth: depth - 1
              ),  #   dm1 = depth - 1
              instructions.PopOp(["depth"]),  #   done with depth
              instructions.push_op(
                  ["dm1", "state"], ["depth", "state"]
              ),  #   state1 = recur(dm1, state)
              instructions.PopOp(["state"]),  #     done with state
              instructions.PushGotoOp(recur_body_2, enter_recur),
          ]
      )
      recur_body_2.assign_instructions(
          [
              instructions.push_op(["ans"], ["state1"]),  #     ...
              instructions.PopOp(["ans"]),  #     pop callee's "ans"
              instructions.prim_op(
                  ["state1"], "state2", lambda state: state + 1
              ),  #   state2 = state1 + 1
              instructions.PopOp(["state1"]),  #   done with state1
              instructions.push_op(
                  ["dm1", "state2"], ["depth", "state"]
              ),  #   state3 = recur(dm1, state2)
              instructions.PopOp(
                  ["dm1", "state2"]
              ),  #     done with dm1, state2
              instructions.PushGotoOp(recur_body_3, enter_recur),
          ]
      )
      recur_body_3.assign_instructions(
          [
              instructions.push_op(["ans"], ["state3"]),  #     ...
              instructions.PopOp(["ans"]),  #     pop callee's "ans"
              instructions.prim_op(
                  ["state3"], "ans", lambda state: state + 1
              ),  #   ans = state3 + 1
              instructions.PopOp(["state3"]),  #   done with state3
              instructions.IndirectGotoOp(),  #   return ans
          ]
      )
      recur_base_case.assign_instructions(
          [
              instructions.PopOp(
                  ["cond1", "depth"]
              ),  #   done with cond1, depth
              instructions.prim_op(
                  ["state"], "ans", step_state
              ),  #   ans = step_state(state)
              instructions.PopOp(["state"]),  #   done with state
              instructions.IndirectGotoOp(),  #   return ans
          ]
      )
  
      pea_nuts_graph = instructions.ControlFlowGraph(
          [
              entry,
              top_body,
              finish_body,
              enter_recur,
              recur_body_1,
              recur_body_2,
              recur_body_3,
              recur_base_case,
          ]
      )
  
      # pylint: disable=bad-whitespace
      pea_nuts_vars = {
          "count": instructions.single_type(np.int64, ()),
          "cond": instructions.single_type(np.bool_, ()),
          "cond1": instructions.single_type(np.bool_, ()),
          "ctm1": instructions.single_type(np.int64, ()),
          "depth": instructions.single_type(np.int64, ()),
          "dm1": instructions.single_type(np.int64, ()),
          "state": instructions.single_type(np.float32, latent_shape),
          "state1": instructions.single_type(np.float32, latent_shape),
          "state2": instructions.single_type(np.float32, latent_shape),
          "state3": instructions.single_type(np.float32, latent_shape),
          "ans": instructions.single_type(np.float32, latent_shape),
      }
  
      return instructions.Program(
          pea_nuts_graph, [], pea_nuts_vars, ["count", "state"], "state"
      )
  
file_path: probability/tensorflow_probability/python/experimental/auto_batching/test_programs.py
imports: import numpy
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 9 for calling a global function with example usage: pea_nuts_program(latent_shape, choose_depth, step_state) and returns: instructions