code: def _initialize_environment(
      program, init_vals, max_stack_depth, backend
  ):
      """Construct initial environment.
  
      This pre-allocates Tensors (in the form of the backend's
      `Variable`s) of the right types and shapes to hold all the data that
      may be needed, and populates the variables holding the inputs with
      their input initial values.
  
      Args:
        program: The `Program` to be executed.
        init_vals: A dict mapping input variable names to patterns of Tensors of
          shape `[batch_size, e1, ..., eE]`.  These are the initial input
          values. Note that batch size is currently inferred from
          init_vals[0].shape[0]. To support arity-0 functions, we would need to
          either infer this from the program or accept it as an input argument.
        max_stack_depth: Python `int` giving the number of frames to preallocate on
          each stack Tensor.
        backend: Object implementing required backend operations.
  
      Returns:
        environment: The initial variable environment (including an initial program
          counter at instruction 0).
      """
      var_alloc = program.var_alloc
      var_defs = program.var_defs
      _check_initial_dtypes(var_defs, init_vals, backend)
      batch_size = inst.detect_batch_size(var_defs, init_vals, backend)
      environment = inst.Environment.initialize(
          backend, var_alloc, var_defs, max_stack_depth, batch_size
      )
      mask = backend.full_mask(batch_size)
      for name, val in six.iteritems(init_vals):
          environment[name] = environment.push(name, val, mask)
      return _initialize_pc(
          environment, program, batch_size, mask, backend
      )
  
file_path: probability/tensorflow_probability/python/experimental/auto_batching/virtual_machine.py
imports: import contextlib
  import functools
  from absl import logging
  import six
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 6 for calling a global function with example usage: _initialize_environment(program, init_vals, max_stack_depth, backend) and returns: _initialize_pc