tfp.experimental.auto_batching.instructions.VariableAllocation

View source on GitHub

Class VariableAllocation

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

Aliases:

In general, a variable holding data with a longer or more complex lifetime will need a more expensive storage strategy.

Specifically, the four variable allocation strategies are: - NULL: Holds nothing. Drops writes, raises on reads. Useful for representing dummy variables that the user program never reads. - TEMPORARY: Holds one value per thread, but not across basic block boundaries. Only usable for temporaries that live in a single basic block, and thus never experience joins (or vm execution loop crossings). For such a variable, push just overwrites the whole Tensor; pop nulls the whole Tensor out. - REGISTER: Holds one value per thread, with no associated stack. Useful for representing temporaries that do not cross (recursive) function calls, but do span multiple basic blocks. For such a variable, push amounts to a where, with an optional runtime safety check for overwriting a defined value. - FULL: Holds a complete stack for each thread. Used as a last resort, when a stack is unavoidable.

The difference between register and temporary is that register is a [batch_size] + event_shape Tensor in the loop state of the toplevel while_loop, whereas temporary is represented as an empty tuple in the loop state, and only holds a Tensor during the execution of the virtual_machine._run_block call that uses it. Consequently, register updating involves a where, but writing to a temporary produces 0 TF ops. Also, in the (envisioned) gather-scatter batching mode, the temporary Tensor will automatically only hold data for the live threads, whereas reading and writing a register will still require the gathers and scatters.

__init__

View source

__init__(name)

Initialize self. See help(type(self)) for accurate signature.

Class Members

  • FULL
  • NULL
  • REGISTER
  • TEMPORARY