def create_variable(self, name, alloc, type_, max_stack_depth, batch_size):
"""Returns an intialized Variable.
Args:
name: Name for the variable.
alloc: `VariableAllocation` for the variable.
type_: `instructions.TensorType` describing the sub-batch shape and dtype
of the variable being created.
max_stack_depth: Python `int`, the maximum stack depth to enforce.
batch_size: Python `int`, the number of parallel threads being executed.
Returns:
var: A new, initialized Variable object.
"""
del name
if alloc is instructions.VariableAllocation.NULL:
return instructions.NullVariable()
elif alloc is instructions.VariableAllocation.TEMPORARY:
return instructions.TemporaryVariable.empty()
else:
dtype, event_shape = type_
value = np.zeros([batch_size] + list(event_shape), dtype=dtype)
if alloc is instructions.VariableAllocation.REGISTER:
return RegisterNumpyVariable(value)
else:
return FullNumpyVariable(value,
_create_stack(max_stack_depth, value))
def create_variable(self, name, alloc, type_, max_stack_depth, batch_size):
"""Returns an intialized Variable.
Args:
name: Name for the variable.
alloc: `VariableAllocation` for the variable.
type_: `instructions.TensorType` describing the sub-batch shape and dtype
of the variable being created.
max_stack_depth: Scalar `int` `Tensor`, the maximum stack depth allocated.
batch_size: Scalar `int` `Tensor`, the number of parallel threads being
executed.
Returns:
var: A new, initialized Variable object.
"""
if alloc is instructions.VariableAllocation.NULL:
return instructions.NullVariable()
elif alloc is instructions.VariableAllocation.TEMPORARY:
return instructions.TemporaryVariable.empty()
else:
name = 'Variable' if name is None else 'VM.var_{}'.format(name)
dtype, event_shape = type_
with tf.name_scope('{}.initialize'.format(name)):
if (alloc is instructions.VariableAllocation.REGISTER
and tf.executing_eagerly()):
# Don't need to construct the empty value in Eager mode, because there
# is no tf.while_loop whose loop-carried state it would need to be.
# This is a substantial optimization for stackless mode, because that
# initializes variables on every function call, rather than just once.
value = (batch_size, dtype, event_shape)
else:
value = self.fill(0, batch_size, dtype, event_shape)
if alloc is instructions.VariableAllocation.REGISTER:
klass = RegisterTensorFlowVariable
extra = []
else:
klass = FullTensorFlowVariable
extra = [
_create_stack(max_stack_depth, value,
self._safety_checks)
]
class NamedVariable(klass):
"""Captures `name` to yield improved downstream TF op names."""
def _name(self):
return name
return NamedVariable(value, *extra)