Spaces:
Running
Running
| from torch.fx.experimental.graph_gradual_typechecker import Refine | |
| from torch.fx.tensor_type import TensorType | |
| from torch.fx.experimental.unification import Var, unify # type: ignore[attr-defined] | |
| def infer_symbolic_types_single_pass(traced): | |
| """ | |
| Calls our symbolic inferencer once. | |
| """ | |
| r = Refine(traced) | |
| r.refine() | |
| mgu = unify_eq(r.constraints) | |
| substitute_all_types(traced.graph, mgu) | |
| def infer_symbolic_types(traced): | |
| """ | |
| Calls our symbolic inferencer twice. | |
| This is useful when one pass is not enough | |
| to infer all the information such as the case | |
| for braodcasting. | |
| """ | |
| r = Refine(traced) | |
| r.refine() | |
| mgu = unify_eq(r.constraints) | |
| substitute_all_types(traced.graph, mgu) | |
| r = Refine(traced) | |
| r.refine() | |
| mgu = unify_eq(r.constraints) | |
| substitute_all_types(traced.graph, mgu) | |
| r.symbolic_relations() | |
| def convert_eq(list_of_eq): | |
| """ | |
| Convert equality constraints in the right format | |
| to be used by unification library. | |
| """ | |
| lhs = [] | |
| rhs = [] | |
| for eq in list_of_eq: | |
| lhs.append(eq.lhs) | |
| rhs.append(eq.rhs) | |
| return tuple(lhs), tuple(rhs) | |
| def unify_eq(list_of_eq): | |
| """ | |
| Apply unification to a set of | |
| equality constraints | |
| """ | |
| lhs, rhs = convert_eq(list_of_eq) | |
| return unify(lhs, rhs) | |
| def substitute_solution_one_type(mapping, t): | |
| """ | |
| Apply the most general unifier to a type | |
| """ | |
| if isinstance(t, Var): | |
| if t in mapping.keys(): | |
| return mapping[t] | |
| else: | |
| return t | |
| elif isinstance(t, TensorType): | |
| new_type = [] | |
| for typ in t.__args__: | |
| if typ in mapping.keys(): | |
| new_type.append(mapping[typ]) | |
| else: | |
| new_type.append(typ) | |
| return TensorType(tuple(new_type)) | |
| elif isinstance(t, list): | |
| new_type = [] | |
| for typ in t: | |
| new_type.append(substitute_solution_one_type(mapping, typ)) | |
| return new_type | |
| elif isinstance(t, tuple): | |
| new_type = [] | |
| for typ in t: | |
| new_type.append(substitute_solution_one_type(mapping, typ)) | |
| return tuple(new_type) | |
| else: | |
| return t | |
| def substitute_all_types(graph, mapping): | |
| """ | |
| Apply the most general unifier to all types in a graph | |
| till reaching a fixed point. If the input and output graph | |
| are the same, we converge. | |
| """ | |
| flag = True | |
| while flag: | |
| flag = False | |
| for k in mapping: | |
| old_mapping_val = mapping[k] | |
| if mapping[k] in mapping.keys(): | |
| new_key = mapping[k] | |
| mapping[k] = mapping[new_key] | |
| if old_mapping_val != mapping[k]: | |
| flag = True | |
| for n in graph.nodes: | |
| n.type = substitute_solution_one_type(mapping, n.type) | |
| def check_for_type_equality(g1, g2): | |
| """ | |
| A check equality to be used in fixed points. | |
| We do not use graph equality but instead type | |
| equality. | |
| """ | |
| for n, m in zip(g1.nodes, g2.nodes): | |
| if n.type != m.type: | |
| return False | |
| return True | |