code: def rng_op_spmd_fn(
      op,
      device_id,
      seed,
      shape,
      dtype,
      key,
      counter,
      alg,
      minval,
      maxval,
      op_version,
      device_index_fn,
      full_replicated=False,
  ):
      if full_replicated:
          # TODO(bfontain,xiejw): Consider to make this consistent with non-replicated
          # case. Seems very confusing.
          new_seed, new_key = seed, key
      else:
          # Runs on TF2 non-DTensor pure eager. This code should align the same
          # logic in RandomOpSPMDExpander.
          x_cord = device_id // 4
          y_cord = device_id % 4
          device_index = device_index_fn(x_cord, y_cord)
          device_id_seed = device_index * 65536 + 65521
          new_seed = gen_bitwise_ops.bitwise_xor(seed, device_id_seed)
          new_key = gen_bitwise_ops.bitwise_xor(
              key, math_ops.cast(device_id_seed, dtype=dtypes.uint64)
          )
      return _call_op(
          op=op,
          seed=new_seed,
          shape=shape,
          dtype=dtype,
          key=new_key,
          counter=counter,
          alg=alg,
          minval=minval,
          maxval=maxval,
          op_version=op_version,
      )
  
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 5 for calling a global function with example usage: rng_op_spmd_fn(op, device_id, seed, shape, dtype, key, counter, alg, minval, maxval, op_version, device_index_fn, full_replicated) and returns: _call_op