def get_operation_mtf_dimension_names(self, operation_name):
    """The Mesh TensorFlow dimensions associated with an operation.

    Args:
      operation_name: a string, name of an operation in the graph.

    Returns:
      a set(string), the names of Mesh TensorFlow dimensions.
    """
    mtf_dimension_names = set()
    for tensor_name in self.get_operation_input_names(operation_name):
      mtf_dimension_names.update(self.get_tensor_mtf_dimension_names(
          tensor_name))
    for tensor_name in self.get_operation_output_names(operation_name):
      mtf_dimension_names.update(self.get_tensor_mtf_dimension_names(
          tensor_name))
    return mtf_dimension_names
The Mesh TensorFlow dimensions associated with an operation.

    Args:
      operation_name: a string, name of an operation in the graph.

    Returns:
      a set(string), the names of Mesh TensorFlow dimensions.