tfp.experimental.auto_batching.virtual_machine.execute

Executes or stages a complete auto-batching VM program.

View aliases

Main aliases

tfp.experimental.auto_batching.frontend.vm.execute

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.