Module: tf.contrib.tpu

Ops related to Tensor Processing Units.

Modules

profiler module: Stub file to maintain backwards compatibility.

Classes

class AsyncCheckpointSaverHook: Saves checkpoints every N steps or seconds.

class CrossShardOptimizer: An optimizer that averages gradients across TPU shards.

class DeviceAssignment: Mapping from logical cores in a computation to the physical TPU topology.

class InfeedQueue: A helper object to build a device infeed queue.

class InputPipelineConfig: Please see the definition of these values in TPUConfig.

class RunConfig: RunConfig with TPU support.

class TPUConfig: TPU related configuration required by TPUEstimator.

class TPUDistributionStrategy: The strategy to run Keras model on TPU.

class TPUEstimator: Estimator with TPU support.

class TPUEstimatorSpec: Ops and objects returned from a model_fn and passed to TPUEstimator.

class Topology: Describes a set of TPU devices.

Functions

batch_parallel(...): Shards computation along the batch dimension for parallel execution.

bfloat16_scope(...): Scope class for bfloat16 variables so that the model uses custom getter.

core(...): Returns the device name for a core in a replicated TPU computation.

cross_replica_sum(...): Sum the input tensor across replicas according to group_assignment.

device_assignment(...): Computes a device_assignment of a computation across a TPU topology.

export_estimator_savedmodel(...): Export Estimator trained model for TPU inference.

infeed_dequeue(...): A placeholder op for a value that will be fed into the computation.

infeed_dequeue_tuple(...): A placeholder op for values fed into the TPU simultaneously as a tuple.

infeed_enqueue(...): An op which feeds a single Tensor value into the computation.

infeed_enqueue_tuple(...): Feeds multiple Tensor values into the computation as an XLA tuple.

initialize_system(...): Initializes a distributed TPU system for use with TensorFlow.

keras_to_tpu_model(...): Copy model along with weights to the TPU. (deprecated)

outfeed_dequeue(...): Retrieves a single tensor from the computation outfeed.

outfeed_dequeue_tuple(...): Retrieve multiple values from the computation outfeed.

outfeed_enqueue(...): Enqueue a Tensor on the computation outfeed.

outfeed_enqueue_tuple(...): Enqueue multiple Tensor values on the computation outfeed.

outside_compilation(...): Builds part of a computation outside any current TPU replicate scope.

repeat(...): Builds a training loop that executes a fixed number of iterations.

replicate(...): Builds a graph operator that runs a replicated TPU computation.

rewrite(...): Rewrites computation for execution on a TPU system.

shard(...): Shards computation for parallel execution.

shutdown_system(...): Shuts down a running a distributed TPU system.

while_loop(...): Builds a training loop for TPUs.