weifengpy

import itertools, time
from collections import defaultdict
from torch.distributed.tensor._ops._view_ops import view_groups, InputDim, Flatten, Split

def describe(rule):
    if isinstance(rule, InputDim):
        return 'InputDim'
    elif isinstance(rule, Flatten):
        return f'Flatten({",".join(describe(d) for d in rule.input_dims)})'
    elif isinstance(rule, Split):