code: def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
      
      val = [i for i, mesh_dimension in enumerate(mesh_shape)
             if (tensor_dimension.name, mesh_dimension.name) in self._pairs]
      if len(val) > 1:
        raise ValueError(
            "Tensor dimension maps to multiple mesh dimensions"
            " tensor_dimension=%s mesh_shape=%s layout=%s"
            % (tensor_dimension, mesh_shape, self._pairs))
      return val[0] if val else None
docs: Mesh axis associated with tensor dimension (or None).
  
      Args:
        tensor_dimension: Dimension.
        mesh_shape: Shape.
  
      Returns:
        Integer or None.
  
      Raises:
        ValueError: If one Tensor dimension maps to two mesh dimensions.