class_docstr: Mesh implementation for TPU using SIMD and MPI operations.
class_docstr_tok: ["Mesh","implementation","for","TPU","using","SIMD","and","MPI","operations","."]
docstr: Call a random tf operation (e.g. random_uniform).
  
  Args:
    shape: a Shape
    tf_fn: a function such as tf.random.uniform
    kwargs: kwargs to pass to tf_fn, except for seed
  
  Returns:
    a LaidOutTensor
code: def random(self, shape, tf_fn, kwargs):
      """Call a random tf operation (e.g. random_uniform).
  
      Args:
        shape: a Shape
        tf_fn: a function such as tf.random.uniform
        kwargs: kwargs to pass to tf_fn, except for seed
  
      Returns:
        a LaidOutTensor
      """
      # TODO(noam): can we make things better with stateless_random?
      slice_shape = self.slice_shape(shape)
      x = tf_fn(slice_shape, **kwargs)
      # TPU does not have seeds enabled.  Sync up the
      # random choices by zeroing out all but the first core per group of
      # identical slices, then allreducing by group.
      layout = self.tensor_layout(shape)
      # we need to sync across these axes.
      mesh_axes = [
          i
          for i in xrange(self.ndims)
          if i not in layout.tensor_axis_to_mesh_axis
      ]
      multiplier = 1.0
      for axis in mesh_axes:
          multiplier *= tf.cast(
              tf.equal(self.laid_out_pcoord(axis).one_slice, 0), x.dtype
          )
      x *= multiplier
      x = self.LaidOutTensor([x])
      x = self.allreduce(x, mesh_axes, "SUM")
      return x
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import os
  import gin
  from mesh_tensorflow import ops_with_redefined_builtins
  from mesh_tensorflow import tpu_variables
  from mesh_tensorflow import utils
  from six.moves import xrange
  import tensorflow.compat.v1
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 15 for calling the SimdMeshImpl obj's underlying member method code with example usage: obj.random(shape, tf_fn, kwargs) and returns: x