code: def shape_sequence_program(shape_sequence):
      """Program that writes into `answer` zeros having a sequence of shapes.
  
      This enables us to test that the final inferred shape is the broadcast of all
      intermediate shapes.
  
      Args:
        shape_sequence: The sequence of intermediate shapes.
  
      Returns:
        program: `instructions.Program` which returns an arbitrary value.
      """
      block_ops = []
  
      def op(shape, ans):
          return (np.zeros(shape, dtype=np.array(ans).dtype),)
  
      for shape in shape_sequence:
          # We use a partial instead of a lambda in order to capture a copy of shape.
          block_ops.append(
              instructions.prim_op(
                  ["ans"], ["ans"], functools.partial(op, shape)
              )
          )
      shape_seq_block = instructions.Block(
          block_ops, instructions.halt_op()
      )
      shape_seq_vars = {
          "ans": instructions.Type(None),
          instructions.pc_var: instructions.single_type(np.int64, ()),
      }
      return instructions.Program(
          instructions.ControlFlowGraph([shape_seq_block]),
          [],
          shape_seq_vars,
          ["ans"],
          ["ans"],
      )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import six
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.auto_batching import allocation_strategy
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import lowering
  from tensorflow_probability.python.experimental.auto_batching import numpy_backend
  from tensorflow_probability.python.experimental.auto_batching import test_programs
  from tensorflow_probability.python.experimental.auto_batching import tf_backend
  from tensorflow_probability.python.experimental.auto_batching import type_inference
  from tensorflow_probability.python.experimental.auto_batching import virtual_machine
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling a global function with example usage: shape_sequence_program(shape_sequence) and returns: instructions, np