code: def random(self, shape, tf_fn, kwargs):
      
      # 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
docs: 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