class_docstr: Replicates input, parameters and compute over multiple TPUs.
  
  ``TpuReplicator`` is a TensorFlow "Distribution Strategy" implementing the
  programming model described in the TF-Replicator paper
  :cite:`buchlovsky2019tf` and TensorFlow RFC
  :cite:`buchlovsky2019distribution`. ``TpuReplicator`` enables data-parallel
  training across multiple TPUs on one or more machines, it supports
  :tf:`function`.
  
  To get started create a ``TpuReplicator`` instance:
  
      >>> replicator = snt.distribute.TpuReplicator()
  
  This provides a scope inside which any new :tf:`Variable`\ s will be
  replicated across all TPU cores:
  
      >>> with replicator.scope():
      ...    mod = snt.Linear(32)
  
  Additionally replicator provides utility functions to apply a module in
  parallel on multiple devices. First we need to define some computation that
  runs on each TPU. The "replica context" object provides us a way to
  communicate between replicas:
  
      >>> def forward():
      ...   # Compute a random output on each GPU.
      ...   x = tf.random.normal([8, 28 * 28])
      ...   y = mod(x)
      ...   # Synchronize the value of `y` between all GPUs.
      ...   ctx = tf.distribute.get_replica_context()
      ...   y = ctx.all_reduce("mean", y)
      ...   return y
  
  Finally we use the run API to apply ``forward`` in parallel on all TPU
  devices. This must be run as part of a :tf:`function` since ``TpuReplicator``
  uses XLA to compile and replicate our function to run in parallel over all
  TPU cores:
  
      >>> @tf.function(autograph=False)
      ... def all_forward():
      ...   return replicator.run(forward)
      >>> per_replica_y = all_forward()
class_docstr_tok: ["Replicates","input",",","parameters","and","compute","over","multiple","TPUs",".","``","TpuReplicator","``","is","a","TensorFlow","``","Distribution","Strategy","''","implementing","the","programming","model","described","in","the","TF-Replicator","paper",":","cite",":","`","buchlovsky2019tf","`","and","TensorFlow","RFC",":","cite",":","`","buchlovsky2019distribution","`",".","``","TpuReplicator","``","enables","data-parallel","training","across","multiple","TPUs","on","one","or","more","machines",",","it","supports",":","tf",":","`","function","`",".","To","get","started","create","a","``","TpuReplicator","``","instance",":",">",">",">","replicator","=","snt.distribute.TpuReplicator","(",")","This","provides","a","scope","inside","which","any","new",":","tf",":","`","Variable","`","\\","s","will","be","replicated","across","all","TPU","cores",":",">",">",">","with","replicator.scope","(",")",":","...","mod","=","snt.Linear","(","32",")","Additionally","replicator","provides","utility","functions","to","apply","a","module","in","parallel","on","multiple","devices",".","First","we","need","to","define","some","computation","that","runs","on","each","TPU",".","The","``","replica","context","''","object","provides","us","a","way","to","communicate","between","replicas",":",">",">",">","def","forward","(",")",":","...","#","Compute","a","random","output","on","each","GPU",".","...","x","=","tf.random.normal","(","[","8",",","28","*","28","]",")","...","y","=","mod","(","x",")","...","#","Synchronize","the","value","of","`","y","`","between","all","GPUs",".","...","ctx","=","tf.distribute.get_replica_context","(",")","...","y","=","ctx.all_reduce","(","``","mean","''",",","y",")","...","return","y","Finally","we","use","the","run","API","to","apply","``","forward","``","in","parallel","on","all","TPU","devices",".","This","must","be","run","as","part","of","a",":","tf",":","`","function","`","since","``","TpuReplicator","``","uses","XLA","to","compile","and","replicate","our","function","to","run","in","parallel","over","all","TPU","cores",":",">",">",">","@","tf.function","(","autograph=False",")","...","def","all_forward","(",")",":","...","return","replicator.run","(","forward",")",">",">",">","per_replica_y","=","all_forward","(",")"]
code: def scope(self):
      with contextlib.ExitStack() as stack:
          stack.enter_context(super().scope())
          stack.enter_context(
              tf.variable_creator_scope(replica_local_creator)
          )
          yield
  
file_path: sonnet/sonnet/src/distribute/replicator.py
imports: from typing import Callable, TypeVar
  from absl import logging
  import contextlib
  from sonnet.src import initializers
  import tensorflow
label_desc: Use image node_id 1 for calling the TpuReplicator obj's underlying member method code with example usage: obj.scope() without return types