code: def rng_op_spmd(
      op,
      device_id,
      seed,
      shape,
      dtype,
      key,
      counter,
      alg,
      minval,
      maxval,
      op_version,
      device_index_fn,
      full_replicated=False,
      is_tpu=False,
  ):
      if not is_tpu:
          return rng_op_spmd_fn(
              op,
              device_id,
              seed,
              shape,
              dtype,
              key,
              counter,
              alg,
              minval,
              maxval,
              op_version,
              device_index_fn,
              full_replicated=full_replicated,
          )
  
      # As of 2021-April, TPU eager and multi-device function produce different
      # stateless rng results compared with bridge compiled function. As DTensor
      # uses bridge to lower TPU function by default, we need to create a
      # TPUStrategy for single core and invoke `run` on it.
      @polymorphic_function.function
      def tpu_fn(device_id, seed):
          return rng_op_spmd_fn(
              op,
              device_id,
              seed,
              shape,
              dtype,
              key,
              counter,
              alg,
              minval,
              maxval,
              op_version,
              device_index_fn,
              full_replicated=full_replicated,
          )
  
      return get_tpu_strategy().run(tpu_fn, args=(device_id, seed))
  
imports: from absl.testing import parameterized
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.dtensor.python.tests import test_util_ops
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver
  from tensorflow.python.eager import remote
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import gen_bitwise_ops
  from tensorflow.python.ops import gen_stateful_random_ops
  from tensorflow.python.ops import gen_stateless_random_ops_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import device_assignment
label_desc: Use image node_id 4 for calling a global function with example usage: rng_op_spmd(op, device_id, seed, shape, dtype, key, counter, alg, minval, maxval, op_version, device_index_fn, full_replicated, is_tpu) and returns: get_tpu_strategy, rng_op_spmd_fn, rng_op_spmd_fn