def computation(self, results, *parameters):
"""
Adds a computation to the transformer.
Arguments:
results: Values to be computed
*parameters: Values to be set as arguments to evaluate
Returns:
Callable.
"""
return self.add_computation(computation(results, *parameters))
def computation(self, results, *parameters):
"""
Adds a computation to the transformer. In the case of not providing parameters
explicitly, the computation will keep using the old values for the parameters.
Arguments:
results: Values to be computed
*parameters: Values to be set as arguments to evaluate
Returns:
Callable.
"""
return self.add_computation(computation(results, *parameters))
def _transform_computations(self):
"""
Transform computation graphs to a form that can be run.
"""
# Run passes on the computation graphs
all_results = []
for comp in self.computations:
all_results.append(comp.computation)
all_ops = self.run_registered_graph_passes(all_results)
self.init_computation = \
self.add_computation(computation(doall(self.state_initialization_ops)).named('init'))
all_ops.append(self.init_computation.computation)
# Collect up all ops from the graph and obtain the init graph
all_ops = OrderedSet(Op.ordered_ops(all_ops))
def init_tensor_description(tensor_description):
if tensor_description.buffer is None:
tensor_description.buffer = self.device_buffer_storage(
tensor_description.base.tensor_size,
tensor_description.dtype,
tensor_description.name
)
self.device_buffers.add(tensor_description.buffer)
tensor_description.value = \
tensor_description.buffer.device_tensor(tensor_description)
for state in self.init_states:
init_tensor_description(state.tensor_description())
self.ops = Op.ordered_ops(all_ops)
for op in self.ops:
if op.is_tensor_op:
init_tensor_description(op.tensor_description())
self.start_transform_allocate()
for device_buffer in self.device_buffers:
device_buffer.transform_allocate()
self.finish_transform_allocate()
# Compile the computations now that we know their storage
for comp in self.computations:
comp.computation_name = \
self.transform_ordered_ops(Op.ordered_ops([comp.computation]),
name=comp.name)
self.finish_transform()
self.finalized = True