TF 2.0 is out! Get hands-on practice at TF World, Oct 28-31. Use code TF20 for 20% off select passes. Register now

tfp.experimental.auto_batching.virtual_machine.execute

View source on GitHub

Executes or stages a complete auto-batching VM program.

Aliases:

tfp.experimental.auto_batching.virtual_machine.execute(
    program,
    args,
    max_stack_depth,
    backend,
    block_code_cache=None
)

Whether this executes or stages computation depends on whether the backend has an eager or deferred computation model.

The dimensions of the inputs and internal variables are split into one top batch dimension and an arbitrary number (here E) event dimensions. The event rank may be different for different inputs, outputs, and internal variables.

Args:

  • program: A instructions.Program to execute or stage.
  • args: Input values, a list of arrays, each of shape [batch_size, e1, ..., eE]. The batch size must be the same for all inputs. The other dimensions must agree with the declared shapes of the variables they will be stored in, but need not in general be the same as one another.
  • max_stack_depth: Python int. Maximum depth of stack to allocate.
  • backend: Object implementing required backend operations.
  • block_code_cache: Dict (allows cache to live across calls to vm.execute, or None (in which case a dict is created and used per call).

Returns:

  • results: A list of the output values. Each returned value is an array of shape [batch_size, e1, ..., eE]. The results are returned in the same order as the variables appear in program.out_vars.