backup code

Author: lixinqi

graph_net/constraint_util.py

@@ -0,0 +1,41 @@
+from graph_net.dynamic_dim_constraints import DynamicDimConstraints
+from typing import Callable
+import copy
+
+
+def symbolize_data_input_dims(
+    dyn_dim_cstr: DynamicDimConstraints,
+    is_data_input: Callable[["input_var_name:str"], bool],
+    is_input_shape_valid: Callable[[DynamicDimConstraints], bool],
+) -> DynamicDimConstraints | None:
+    """
+    Symbolizes data input dimensions as much as possible.
+    Returns new DynamicDimConstraints if success.
+    Returns None if no symbolicable dim .
+    """
+    unqiue_dims = set()
+
+    def dumpy_filter_fn(input_name, input_idx, axis, dim):
+        unqiue_dims.add(dim)
+        # No symbolization because of returning True
+        return False
+
+    # Collect input dimensions into `unqiue_dims`
+    assert dyn_dim_cstr.symbolize(dumpy_filter_fn) is None
+    for picked_dim in unqiue_dims:
+        tmp_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
+
+        def filter_fn(input_name, input_idx, axis, dim):
+            return is_data_input(input_name) and dim == picked_dim
+
+        symbol = tmp_dyn_dim_cstr.symbolize(filter_fn)
+        if symbol is None:
+            continue
+        sym2example_value = {symbol: picked_dim + 1}
+        if not tmp_dyn_dim_cstr.check_delta_symbol2example_value(sym2example_value):
+            continue
+        tmp_dyn_dim_cstr.update_symbol2example_value(sym2example_value)
+        if not is_input_shape_valid(tmp_dyn_dim_cstr):
+            continue
+        dyn_dim_cstr = tmp_dyn_dim_cstr
+    return dyn_dim_cstr

graph_net/dynamic_dim_constraints.py

@@ -0,0 +1,220 @@
+import sys
+import sympy
+import importlib.util as imp
+from dataclasses import dataclass
+import copy
+from typing import Callable
+from collections import namedtuple
+
+
+@dataclass
+class DynamicDimConstraints:
+    kSymbolVarNamePrefix = "s"
+
+    symbols: list[sympy.Symbol]
+    kSymbols = "dynamic_dim_constraint_symbols"
+
+    # len(symbol2example_value) == len(symbols)
+    symbol2example_value: dict[sympy.Symbol, int]
+    kSymbol2ExampleValue = "dynamic_dim_constraint_symbol2example_value"
+
+    relations: list[sympy.Rel]
+    kRelations = "dynamic_dim_constraint_relations"
+
+    # len(input_shapes) equals number of Model.forward arguments
+    input_shapes: list[("var-name", tuple[sympy.Expr | int])]
+    kInputShapes = "dynamic_dim_constraint_input_shapes"
+
+    # len(input_shapes) equals number of Model.forward arguments
+    input_max_values: list[("var-name", tuple[sympy.Expr | int | None])]
+    kInputMaxValues = "dynamic_dim_constraint_input_max_values"
+
+    def symbolize(
+        self,
+        filter_fn=Callable[
+            ["input_name:str", "input_idx:int", "axis:int", "dim:int"], bool
+        ],
+    ) -> sympy.Symbol | None:
+        """
+        Returns created symbol.
+        """
+        InputDim = namedtuple("InputDim", ["input_idx", "axis", "dim"])
+        input_dims = [
+            InputDim(input_idx, axis, dim)
+            for input_idx, namedshape in enumerate(self.input_shapes)
+            for input_name, shape in [namedshape]
+            for axis, dim in enumerate(shape)
+            if isinstance(dim, int)
+            if filter_fn(input_name, input_idx, axis, dim)
+        ]
+        if len(input_dims) == 0:
+            return None
+        unique_dims = set(dim for input_dix, axis, dim in input_dims)
+        assert len(unique_dims) == 1
+        dim = list(unique_dims)[0]
+        new_sym = self._new_symbol(example_value=dim)
+        for input_dix, axis, _ in input_dims:
+            self.input_shapes[input_dix][1][axis] = new_sym
+        return new_sym
+
+    def _new_symbol(self, example_value):
+        max_existed_seq_no = max(
+            [
+                -1,
+                *(
+                    seq_no
+                    for symbol in self.symbols
+                    for seq_no in [int(symbol.name[1:])]
+                ),
+            ]
+        )
+        seq_no = max_existed_seq_no + 1
+        symbol = sympy.Symbol(f"{self.kSymbolVarNamePrefix}{seq_no}")
+        self.symbol2example_value[symbol] = example_value
+        self.symbols.append(symbol)
+        return symbol
+
+    def update_symbol2example_value(self, sym2example_value: dict):
+        self.symbol2example_value = self.merge_symbol2example_value(sym2example_value)
+        return self
+
+    def merge_symbol2example_value(self, sym2example_value: dict):
+        return {
+            k: v
+            for k, v in [*self.symbol2example_value.items(), *sym2example_value.items()]
+        }
+
+    def check_delta_symbol2example_value(self, sym2example_value: dict):
+        if len(sym2example_value) == 0:
+            return True
+
+        sym2example_value = self.merge_symbol2example_value(sym2example_value)
+
+        sym_exprs = [
+            *self._get_sym_exprs_from_input_shapes(),
+            *self._get_sym_exprs_from_input_max_values(),
+        ]
+
+        relations = [*self.relations, *(sym_expr > 0 for sym_expr in sym_exprs)]
+
+        return all(
+            relation.subs(sym2example_value) == sympy.true for relation in relations
+        )
+
+    def _get_sym_exprs_from_input_shapes(self):
+        yield from (
+            sym_dim
+            for name, shape in self.input_shapes
+            for sym_dim in shape
+            if isinstance(sym_dim, sympy.Expr)
+        )
+
+    def _get_sym_exprs_from_input_max_values(self):
+        yield from (
+            sym_value
+            for name, sym_value in self.input_max_values
+            if isinstance(sym_value, sympy.Expr)
+        )
+
+    def serialize_to_py_str(self):
+        symbols_definition = "\n".join(
+            f"{name} = Symbol('{name}')"
+            for symbol in self.symbols
+            for name in [symbol.name]
+        )
+        return f"""
+from sympy import Symbol, Expr, Rel, Eq
+
+{symbols_definition}
+
+{self.kSymbols} = {self.symbols}
+
+{self.kSymbol2ExampleValue} = {self.symbol2example_value}
+
+{self.kRelations} = {self.relations}
+
+{self.kInputShapes} = {self.input_shapes}
+
+{self.kInputMaxValues} = {self.input_max_values}
+        """
+
+    @classmethod
+    def unserialize_from_py_file(cls, filepath):
+        module = cls.load_module(filepath)
+        return cls(
+            symbols=cls.module_symbols(module),
+            symbol2example_value=cls.module_symbol2example_value(module),
+            relations=cls.module_relations(module),
+            input_shapes=cls.module_input_shapes(module),
+            input_max_values=cls.module_input_max_values(module),
+        )
+
+    @classmethod
+    def module_symbols(cls, module):
+        return cls.get_module_list_attr(module, cls.kSymbols)
+
+    @classmethod
+    def module_symbol2example_value(cls, module):
+        return cls.get_module_dict_attr(module, cls.kSymbol2ExampleValue)
+
+    @classmethod
+    def module_relations(cls, module):
+        return cls.get_module_list_attr(module, cls.kRelations)
+
+    @classmethod
+    def module_input_shapes(cls, module):
+        return cls.get_module_list_attr(module, cls.kInputShapes)
+
+    @classmethod
+    def module_input_max_values(cls, module):
+        return cls.get_module_list_attr(module, cls.kInputMaxValues)
+
+    @classmethod
+    def get_module_list_attr(cls, module, attr):
+        return cls.get_module_attr(module, attr, default=[])
+
+    @classmethod
+    def get_module_dict_attr(cls, module, attr):
+        return cls.get_module_attr(module, attr, default={})
+
+    @classmethod
+    def get_module_attr(cls, module, attr, default):
+        return getattr(module, attr) if hasattr(module, attr) else default
+
+    @classmethod
+    def load_module(cls, path, name="unamed"):
+        spec = imp.spec_from_file_location(name, path)
+        module = imp.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+
+
+if __name__ == "__main__":
+    cstr_code = """
+import sympy
+
+x = sympy.Symbol('x')
+y = sympy.Symbol('y')
+
+dynamic_dim_constraint_symbol2example_value = [(x, 2)]
+
+dynamic_dim_constraint_symbols = [x, y]
+
+dynamic_dim_constraint_relations = [x > 0]
+    """
+
+    import tempfile
+
+    with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", encoding="utf-8") as tmp:
+        tmp.write(cstr_code)
+        tmp.flush()
+        cstr = DynamicDimConstraints.unserialize_from_py_file(tmp.name)
+        print(cstr)
+        print(cstr.serialize_to_py_str())
+
+    with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", encoding="utf-8") as tmp:
+        tmp.write(cstr.serialize_to_py_str())
+        tmp.flush()
+        cstr = DynamicDimConstraints.unserialize_from_py_file(tmp.name)
+        print(cstr)
+        print(cstr.serialize_to_py_str())