Module: tfp.experimental.auto_batching.instructions

Instruction language for auto-batching virtual machine.


class Block: A basic block.

class BranchOp: A conditional.

class ControlFlowGraph: A control flow graph (CFG).

class Function: A function subject to auto-batching, callable with FunctionCallOp.

class FunctionCallOp: Call a Function.

class GotoOp: An unconditional jump.

class IndirectGotoOp: Jump to the address in the reserved "program counter" variable.

class PopOp: Restore the given variables from their stacks.

class PrimOp: An arbitrary already-batched computation, a 'primitive operation'.

class Program: An auto-batchable program.

class PushGotoOp: Save an address for IndirectGotoOp and unconditionally jump to another.

class TensorType: TensorType(dtype, shape)

class Type: Type(tensors,)

class VariableAllocation: A token indicating how to allocate memory for an autobatched variable.


extract_referenced_variables(...): Extracts a set of the variable names referenced by the node in question.

halt_op(...): Returns a control transfer Op that means "exit this graph".

interpret(...): Interprets a program in this instruction language and returns the result.

push_op(...): Returns an Op that pushes values from vars_in into vars_out.

pc_var '__program_counter__'