Args:

computation: A Python function that builds a computation to apply to the
input. If the function takes n inputs, 'inputs' should be a list of n
tensors.

computation may return a list of operations and tensors. Tensors must
come before operations in the returned list. The return value of
rewrite is a list of tensors corresponding to the tensors from the
output of computation.

All Operations returned from computation will be executed when
evaluating any of the returned output tensors.

inputs: A list of input tensors or None (equivalent to an empty list).

infeed_queue: If not None, the InfeedQueue from which to append a tuple
of arguments as inputs to computation.

device_assignment: if not None, a DeviceAssignment describing the
mapping between logical cores in the computation with physical cores in
the TPU topology. May be omitted for a single-core computation, in which
case the core attached to task 0, TPU device 0 is used.