code: def test_logical_device_assignment(self):
      strategy, num_replicas = get_tpu_strategy()
      with strategy.scope():
          v = variables.Variable(2.0)
          with strategy.extended.experimental_logical_device(1):
              w = variables.Variable(3.0)
  
      self.assertLen(
          strategy.experimental_local_results(v), num_replicas
      )
      self.assertLen(
          strategy.experimental_local_results(w), num_replicas
      )
      self.assertEqual(
          "/job:localhost/replica:0/task:0/device:TPU:0",
          strategy.experimental_local_results(v)[0].device,
      )
      self.assertEqual(
          "/job:localhost/replica:0/task:0/device:TPU:1",
          strategy.experimental_local_results(w)[0].device,
      )
  
      logical_devices = []
  
      @def_function.function
      def f(x):
          replica_ctx = distribute_lib.get_replica_context()
          with replica_ctx.experimental_logical_device(0):
              y = v * x
          with replica_ctx.experimental_logical_device(1):
              z = w * y
          logical_devices.append((y.device, z.device))
          return z
  
      result = strategy.run(f, args=(5.0,))
  
      self.assertEqual(
          [
              (
                  "/device:TPU_REPLICATED_CORE:0",
                  "/device:TPU_REPLICATED_CORE:1",
              )
          ],
          logical_devices,
      )
  
      with self.cached_session():
          self.evaluate(variables.global_variables_initializer())
          self.assertEqual(
              30.0 * num_replicas,
              self.evaluate(strategy.reduce("SUM", result, axis=None)),
          )
  
imports: import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.compiler.xla.experimental import xla_sharding
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import packed_distributed_variable
  from tensorflow.python.distribute import strategy_test_lib
  from tensorflow.python.distribute import tpu_replicated_variable
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute import tpu_values
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import remote
  from tensorflow.python.eager import test
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import summary_test_util
  from tensorflow.python.module import module
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import summary_ops_v2
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import flags
  from tensorflow.python.tpu import device_assignment
  from tensorflow.python.tpu import tpu_replication
label_desc: Use image node_id 2 for calling the TPUStrategyModelParallelismTest obj's underlying member method code with example usage: obj.test_logical_device_assignment() and returns: z