def get_n_trainable_vars(graph):
    """
    Gets number of trainable vars in a MTF model.

    :param graph: Mesh-Tensorflow graph
    :return: None
    """
    total_parameters = 0
    for variable in graph.trainable_variables:
      shape = variable.shape.dims
      variable_parameters = 1
      for dim in shape:
          variable_parameters *= dim.size
      total_parameters += variable_parameters
    print(f"\n\nN TRAINABLE VARS:\n{total_parameters:,}\n\n")
Gets number of trainable vars in a MTF model.

:param graph: Mesh-Tensorflow graph
:return: None
utils.py