description: A token indicating how to allocate memory for an autobatched variable.
View source on GitHub |
A token indicating how to allocate memory for an autobatched variable.
tfp.experimental.auto_batching.instructions.VariableAllocation(
name
)
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.