1from typing import Any, Callable, Dict, Optional 2 3from torch.utils._pytree import Context, TreeSpec 4 5 6def reorder_kwargs(user_kwargs: Dict[str, Any], spec: TreeSpec) -> Dict[str, Any]: 7 """Reorder user-provided kwargs to match the order in `spec`. `spec` is 8 expected to be the in_spec of an exported program, i.e. the spec that 9 results from flattening `(args, kwargs)`. 10 11 We need this to provide consistent input ordering, such so that users can 12 pass in foo(a=a, b=b) OR foo(b=b, a=a) and receive the same result. 13 """ 14 # Make sure that the spec is actually shaped like (args, kwargs) 15 assert spec.type is tuple 16 assert spec.num_children == 2 17 kwargs_spec = spec.children_specs[1] 18 assert kwargs_spec.type is dict 19 20 if set(user_kwargs) != set(kwargs_spec.context): 21 raise ValueError( 22 f"kwarg key mismatch: " 23 f"Got {list(user_kwargs)} but expected {kwargs_spec.context}" 24 ) 25 26 reordered_kwargs = {} 27 for kw in kwargs_spec.context: 28 reordered_kwargs[kw] = user_kwargs[kw] 29 30 return reordered_kwargs 31 32 33def is_equivalent( 34 spec1: TreeSpec, 35 spec2: TreeSpec, 36 equivalence_fn: Callable[[Optional[type], Context, Optional[type], Context], bool], 37) -> bool: 38 """Customizable equivalence check for two TreeSpecs. 39 40 Arguments: 41 spec1: The first TreeSpec to compare 42 spec2: The second TreeSpec to compare 43 equivalence_fn: A function to determine the equivalence of two 44 TreeSpecs by examining their types and contexts. It will be called like: 45 46 equivalence_fn(spec1.type, spec1.context, spec2.type, spec2.context) 47 48 This function will be applied recursively to all children. 49 50 Returns: 51 True if the two TreeSpecs are equivalent, False otherwise. 52 """ 53 if not equivalence_fn(spec1.type, spec1.context, spec2.type, spec2.context): 54 return False 55 56 # Recurse on children 57 if len(spec1.children_specs) != len(spec2.children_specs): 58 return False 59 60 for child_spec1, child_spec2 in zip(spec1.children_specs, spec2.children_specs): 61 if not is_equivalent(child_spec1, child_spec2, equivalence_fn): 62 return False 63 64 return True 65