code: def _einsum_helper(input_shapes, output_shape, mesh_impl):
    
    input_shape_union = _shape_union(input_shapes)
    total_num_dims = input_shape_union.ndims
    # list of input shapes that contain all dimensions.
    full_shapes = [
        s for s in input_shapes + [output_shape] if s.ndims == total_num_dims]
    full_shape = full_shapes[0] if full_shapes else input_shape_union
    reduce_slice_fn, reduced_mesh_axes = _reduce_helper(
        full_shape, output_shape, mesh_impl.tensor_layout(full_shape))
    def einsum_slice_fn_naive(*slices):
      # naive einsum implementation where we broadcast all inputs to the full
      # shape, multiply componentwise, then reduce.
      return reduce_slice_fn(functools.reduce(tf.multiply, [
          _expand_dims(x, input_shape, full_shape)
          for x, input_shape in zip(slices, input_shapes)]))
    if full_shapes:
      # it is not wasteful of space to broadcast fully and then reduce.
      # this helps to avoid some inefficient GPU implementations.
      einsum_slice_fn = einsum_slice_fn_naive
    else:
      # call tf.einsum
      equation = _einsum_equation(input_shapes, output_shape)
      def einsum_slice_fn(*slices):
        if slices[0].dtype.is_floating:
          return mesh_impl.einsum(equation, *slices)
        else:
          return einsum_slice_fn_naive(*slices)
    return einsum_slice_fn, reduced_mesh_axes
docs: Returns slicewise function and reduced mesh dimensions.
  
    Assumes the output shape contains no new dimensions.
  
    Args:
      input_shapes: a list of Shapes
      output_shape: a Shape
      mesh_impl: a MeshImpl
    Returns:
      einsum_slice_fn: a function from tf.Tensors to tf.Tensor
      reduced_mesh_axes: a list of integers