tfp.experimental.auto_batching.virtual_machine.execute

Executes or stages a complete auto-batching VM program.

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.

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).

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.