Apply name change(experimental_run_v2 -> run) for all callers.

PiperOrigin-RevId: 301919882
Change-Id: I14c6ed85bdf50d619d1bc572e0fbcc5f1821c70b

Author: kenfranko

examples/distributed_cifar10.ipynb

@@ -405,7 +405,7 @@
         "\n",
         "@tf.function\n",
         "def train_step(images, labels):\n",
-        "  per_replica_loss = strategy.experimental_run_v2(step, args=(images, labels))\n",
+        "  per_replica_loss = strategy.run(step, args=(images, labels))\n",
         "  return strategy.reduce(\"sum\", per_replica_loss, axis=None)\n",
         "\n",
         "def train_epoch(dataset):\n",
@@ -465,7 +465,7 @@
         "  total_correct = 0\n",
         "\n",
         "  for images, labels in cifar10_test_dist:\n",
-        "    per_replica_correct = strategy.experimental_run_v2(is_predicted, args=(images, labels))\n",
+        "    per_replica_correct = strategy.run(is_predicted, args=(images, labels))\n",
         "    total_correct += strategy.reduce(\"sum\", per_replica_correct, axis=0)\n",
         "\n",
         "  return tf.cast(total_correct, tf.float32) / num_cifar10_test_examples\n",

sonnet/src/conformance/checkpoint_test.py

@@ -195,7 +195,7 @@ def test_save_restore(self, golden, replicator_fn, use_function):
       variables = golden.create_all_variables(module)
 
     def forward():
-      per_replica = replicator.experimental_run_v2(
+      per_replica = replicator.run(
           lambda: golden.forward(module))
       return tree.map_structure(
           lambda args: tf.stack(replicator.unwrap(args), axis=0), per_replica)
@@ -350,7 +350,7 @@ def test_restore_on_create_in_replica_context(self, golden, replicator_fn,
       module = golden.create_module()
 
     def forward():
-      return replicator.experimental_run_v2(lambda: golden.forward(module))
+      return replicator.run(lambda: golden.forward(module))
 
     if use_function:
       forward = tf.function(forward)

sonnet/src/conformance/distribute_test.py

@@ -44,7 +44,7 @@ def test_variable_creation_in_replica_context(self, golden, replicator_fn):
     @tf.function
     def forward():
       step = lambda: golden.create_all_variables(mod)
-      return replicator.experimental_run_v2(step)
+      return replicator.run(step)
 
     # TODO(b/132329316) Remove when `xla.compile` allows tf.device(TPU).
     with tf.device(None):
@@ -83,7 +83,7 @@ def forward():
         state = core.initial_state(input_shape[0])
         return unroll(core, sequence, state)
 
-      return replicator.experimental_run_v2(forward)
+      return replicator.run(forward)
 
     # TpuReplicator doesn't support pure eager mode.
     if isinstance(replicator, snt_replicator.TpuReplicator):

sonnet/src/distribute/batch_norm_test.py

@@ -65,7 +65,7 @@ def foo():
       outputs = layer(inputs, True, False, scale, offset)
       return inputs, outputs
 
-    inputs, outputs = strategy.experimental_run_v2(foo)
+    inputs, outputs = strategy.run(foo)
     local_mean_metric = strategy.experimental_local_results(mean_metric.value)
     local_var_metric = strategy.experimental_local_results(var_metric.value)
     self.assertAllEqual(local_mean_metric[0].numpy(),
@@ -101,7 +101,7 @@ def compute():
         outputs = layer(inputs, True, False, scale, offset)
         return inputs, outputs
 
-      return strategy.experimental_run_v2(compute)
+      return strategy.run(compute)
     inputs, outputs = run()
 
     local_mean_metric = strategy.experimental_local_results(mean_metric.value)

sonnet/src/distribute/replicator.py

@@ -76,7 +76,7 @@ class Replicator(tf.distribute.MirroredStrategy):
   Finally we use the run API to apply ``forward`` in parallel on all accelerator
   devices:
 
-      >>> per_replica_y = replicator.experimental_run_v2(forward)
+      >>> per_replica_y = replicator.run(forward)
   """
 
   @contextlib.contextmanager
@@ -128,7 +128,7 @@ class TpuReplicator(tf.distribute.experimental.TPUStrategy):
 
       >>> @tf.function(autograph=False)
       ... def all_forward():
-      ...   return replicator.experimental_run_v2(forward)
+      ...   return replicator.run(forward)
       >>> per_replica_y = all_forward()
   """
 
@@ -139,6 +139,11 @@ def scope(self):
       stack.enter_context(tf.variable_creator_scope(replica_local_creator))
       yield
 
+# TODO(tomhennigan) Remove this once TF 2.2 is released.
+for cls in (Replicator, TpuReplicator):
+  if not hasattr(cls, "run"):
+    cls.run = cls.experimental_run_v2
+
 
 def create_variables_eagerly(f: Callable[..., T]) -> Callable[..., T]:
   """Wraps a function and attempts to create variables using eager mode.

sonnet/src/distribute/replicator_test.py

@@ -45,7 +45,7 @@ def _create_variable_in_replica_context(replicator):
   o = TrainableVariable()
 
   def create_var():
-    replicator.experimental_run_v2(o)
+    replicator.run(o)
 
   # TpuReplicator doesn't support pure eager mode.
   if isinstance(replicator, snt_replicator.TpuReplicator):
@@ -123,7 +123,7 @@ def test_assign(self, replicator_fn, method_name, value, cross_replica):
       # TpuReplicator doesn't support pure eager mode.
       if isinstance(replicator, snt_replicator.TpuReplicator):
         update_fn = tf.function(update_fn)
-      replicator.experimental_run_v2(update_fn)
+      replicator.run(update_fn)
     for component in v._values:
       self.assertAllEqual(component.read_value(), tf.ones_like(component))
 
@@ -144,7 +144,7 @@ def test_read_value(self, replicator_fn, cross_replica):
         read_value_fn = tf.function(v.read_value)
       else:
         read_value_fn = v.read_value
-      values = replicator.experimental_run_v2(read_value_fn)
+      values = replicator.run(read_value_fn)
       values = replicator.experimental_local_results(values)
     for component in v._values:
       for value in values:

sonnet/src/optimizers/optimizer_tests.py

@@ -114,7 +114,7 @@ def testUnsuppportedStrategyError(self):
     with self.assertRaisesRegexp(
         ValueError,
         "Sonnet optimizers are not compatible with `MirroredStrategy`"):
-      strategy.experimental_run_v2(lambda: optimizer.apply(updates, parameters))
+      strategy.run(lambda: optimizer.apply(updates, parameters))
 
 
 # NOTE: Avoiding ABCMeta because of metaclass conflict.