Matrix multiplication and Kronecker product.#

EGraph(_flatted_deps=[], _state=EGraphState(egraph=<builtins.EGraph object at 0x55c9c6eec570>, __egg_decls__=Declarations(_unnamed_functions=set(), _functions={'kron': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Matrix', args=()), TypeRefWithVars(name='Matrix', args=())), arg_names=('a', 'b'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Matrix', args=()), var_arg_type=None), builtin=False, egg_name='Kron', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, _constants={}, _classes={'Dim': ClassDecl(egg_name=None, type_vars=(), builtin=False, init=FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()),), arg_names=('value',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Dim', args=()), var_arg_type=None), builtin=False, egg_name='Lit', cost=None, default=None, on_merge=(), merge=None, unextractable=False), class_methods={'named': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='String', args=()),), arg_names=('name',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Dim', args=()), var_arg_type=None), builtin=False, egg_name='NamedDim', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, class_variables={}, methods={'__mul__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Dim', args=()), TypeRefWithVars(name='Dim', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Dim', args=()), var_arg_type=None), builtin=False, egg_name='Times', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, properties={}, preserved_methods={}), 'i64': ClassDecl(egg_name=None, type_vars=(), builtin=True, init=None, class_methods={}, class_variables={}, methods={'__add__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='+', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__sub__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='-', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__mul__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='*', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__truediv__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='/', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__mod__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='%', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__and__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='&', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__or__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='|', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__xor__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='^', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__lshift__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='<<', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__rshift__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='>>', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__invert__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()),), arg_names=('self',), arg_defaults=(None,), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='not-i64', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__lt__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Unit', args=()), var_arg_type=None), builtin=True, egg_name='<', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__gt__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Unit', args=()), var_arg_type=None), builtin=True, egg_name='>', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__le__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Unit', args=()), var_arg_type=None), builtin=True, egg_name='<=', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__ge__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Unit', args=()), var_arg_type=None), builtin=True, egg_name='>=', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'min': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='min', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'max': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='i64', args=()), var_arg_type=None), builtin=True, egg_name='max', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'to_string': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()),), arg_names=('self',), arg_defaults=(None,), return_type=TypeRefWithVars(name='String', args=()), var_arg_type=None), builtin=True, egg_name='to-string', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'bool_lt': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='bool-<', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'bool_gt': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='bool->', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'bool_le': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='bool-<=', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'bool_ge': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='i64', args=()), TypeRefWithVars(name='i64', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='bool->=', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, properties={}, preserved_methods={}), 'Unit': ClassDecl(egg_name='Unit', type_vars=(), builtin=True, init=None, class_methods={}, class_variables={}, methods={}, properties={}, preserved_methods={}), 'String': ClassDecl(egg_name=None, type_vars=(), builtin=True, init=None, class_methods={}, class_variables={}, methods={'replace': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='String', args=()), TypeRefWithVars(name='String', args=()), TypeRefWithVars(name='String', args=())), arg_names=('self', 'old', 'new'), arg_defaults=(None, None, None), return_type=TypeRefWithVars(name='String', args=()), var_arg_type=None), builtin=True, egg_name='replace', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, properties={}, preserved_methods={}), 'Bool': ClassDecl(egg_name='bool', type_vars=(), builtin=True, init=None, class_methods={}, class_variables={}, methods={'__invert__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Bool', args=()),), arg_names=('self',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='not', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__and__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Bool', args=()), TypeRefWithVars(name='Bool', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='and', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__or__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Bool', args=()), TypeRefWithVars(name='Bool', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='or', cost=None, default=None, on_merge=(), merge=None, unextractable=False), '__xor__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Bool', args=()), TypeRefWithVars(name='Bool', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='xor', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'implies': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Bool', args=()), TypeRefWithVars(name='Bool', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Bool', args=()), var_arg_type=None), builtin=True, egg_name='=>', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, properties={}, preserved_methods={}), 'Matrix': ClassDecl(egg_name='MExpr', type_vars=(), builtin=False, init=None, class_methods={'identity': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Dim', args=()),), arg_names=('dim',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Matrix', args=()), var_arg_type=None), builtin=False, egg_name='Id', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'named': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='String', args=()),), arg_names=('name',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Matrix', args=()), var_arg_type=None), builtin=False, egg_name='NamedMat', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, class_variables={}, methods={'__matmul__': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Matrix', args=()), TypeRefWithVars(name='Matrix', args=())), arg_names=('self', 'other'), arg_defaults=(None, None), return_type=TypeRefWithVars(name='Matrix', args=()), var_arg_type=None), builtin=False, egg_name='MMul', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'nrows': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Matrix', args=()),), arg_names=('self',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Dim', args=()), var_arg_type=None), builtin=False, egg_name='nrows', cost=None, default=None, on_merge=(), merge=None, unextractable=False), 'ncols': FunctionDecl(signature=FunctionSignature(arg_types=(TypeRefWithVars(name='Matrix', args=()),), arg_names=('self',), arg_defaults=(None,), return_type=TypeRefWithVars(name='Dim', args=()), var_arg_type=None), builtin=False, egg_name='ncols', cost=None, default=None, on_merge=(), merge=None, unextractable=False)}, properties={}, preserved_methods={})}, _rulesets={'': RulesetDecl(rules=[])}), rulesets={'': set()}, egg_fn_to_callable_refs=defaultdict(<class 'set'>, {'!=': {FunctionRef(name='!=')}, 'Times': {MethodRef(class_name='Dim', method_name='__mul__')}, 'Lit': {InitRef(class_name='Dim')}, '*': {MethodRef(class_name='i64', method_name='__mul__')}, 'nrows': {MethodRef(class_name='Matrix', method_name='nrows')}, 'Kron': {FunctionRef(name='kron')}, 'ncols': {MethodRef(class_name='Matrix', method_name='ncols')}, 'MMul': {MethodRef(class_name='Matrix', method_name='__matmul__')}, 'Id': {ClassMethodRef(class_name='Matrix', method_name='identity')}, 'NamedDim': {ClassMethodRef(class_name='Dim', method_name='named')}, 'NamedMat': {ClassMethodRef(class_name='Matrix', method_name='named')}}), callable_ref_to_egg_fn={FunctionRef(name='!='): '!=', MethodRef(class_name='Dim', method_name='__mul__'): 'Times', InitRef(class_name='Dim'): 'Lit', MethodRef(class_name='i64', method_name='__mul__'): '*', MethodRef(class_name='Matrix', method_name='nrows'): 'nrows', FunctionRef(name='kron'): 'Kron', MethodRef(class_name='Matrix', method_name='ncols'): 'ncols', MethodRef(class_name='Matrix', method_name='__matmul__'): 'MMul', ClassMethodRef(class_name='Matrix', method_name='identity'): 'Id', ClassMethodRef(class_name='Dim', method_name='named'): 'NamedDim', ClassMethodRef(class_name='Matrix', method_name='named'): 'NamedMat'}, type_ref_to_egg_sort={JustTypeRef(name='Dim', args=()): 'Dim', JustTypeRef(name='i64', args=()): 'i64', JustTypeRef(name='Matrix', args=()): 'MExpr', JustTypeRef(name='String', args=()): 'String'}, expr_to_egg_cache={VarDecl(name='a', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'a'), VarDecl(name='b', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'b'), VarDecl(name='c', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'c'), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='b', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='c', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'b'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'c')]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='a', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='b', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='c', is_let=False))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'a'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'b'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'c')])]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='a', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='b', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'a'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'b')]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='a', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='b', is_let=False))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='c', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'a'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'b')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'c')]), VarDecl(name='i', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'i'), CallDecl(callable=InitRef(class_name='Dim'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='i', is_let=False)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Lit', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'i')]), VarDecl(name='j', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'j'), CallDecl(callable=InitRef(class_name='Dim'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='j', is_let=False)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Lit', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'j')]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=InitRef(class_name='Dim'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='i', is_let=False)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=InitRef(class_name='Dim'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='j', is_let=False)),), bound_tp_params=()))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Lit', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'i')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Lit', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'j')])]), CallDecl(callable=MethodRef(class_name='i64', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='i', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='j', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), '*', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'i'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'j')]), CallDecl(callable=InitRef(class_name='Dim'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=CallDecl(callable=MethodRef(class_name='i64', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='i', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='i64', args=()), expr=VarDecl(name='j', is_let=False))), bound_tp_params=None)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Lit', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), '*', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'i'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'j')])]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='b', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='a', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'b'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'a')]), VarDecl(name='A', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), VarDecl(name='B', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)),), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)),), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)),), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)),), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')])]), VarDecl(name='n', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'n'), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=False)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'n')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=False)),), bound_tp_params=())),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'n')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=False)),), bound_tp_params=())),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'n')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=False)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'n')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=False)),), bound_tp_params=()))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'n')])]), VarDecl(name='C', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C'), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')])]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), VarDecl(name='D', is_let=False): Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'D'), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='D', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'D')]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='D', is_let=False))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'D')])]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='D', is_let=False))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'D')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=False))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='D', is_let=False))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'B')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'D')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=False)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'C')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='D', is_let=False)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), 'D')]), VarDecl(name='n', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%n'), LitDecl(value='n'): Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('n')), CallDecl(callable=ClassMethodRef(class_name='Dim', method_name='named'), args=(TypedExprDecl(tp=JustTypeRef(name='String', args=()), expr=LitDecl(value='n')),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'NamedDim', [Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('n'))]), VarDecl(name='m', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%m'), LitDecl(value='m'): Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('m')), CallDecl(callable=ClassMethodRef(class_name='Dim', method_name='named'), args=(TypedExprDecl(tp=JustTypeRef(name='String', args=()), expr=LitDecl(value='m')),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'NamedDim', [Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('m'))]), VarDecl(name='p', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%p'), LitDecl(value='p'): Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('p')), CallDecl(callable=ClassMethodRef(class_name='Dim', method_name='named'), args=(TypedExprDecl(tp=JustTypeRef(name='String', args=()), expr=LitDecl(value='p')),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'NamedDim', [Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('p'))]), VarDecl(name='A', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A'), LitDecl(value='A'): Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('A')), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='named'), args=(TypedExprDecl(tp=JustTypeRef(name='String', args=()), expr=LitDecl(value='A')),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'NamedMat', [Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('A'))]), VarDecl(name='B', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B'), LitDecl(value='B'): Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('B')), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='named'), args=(TypedExprDecl(tp=JustTypeRef(name='String', args=()), expr=LitDecl(value='B')),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'NamedMat', [Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('B'))]), VarDecl(name='C', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%C'), LitDecl(value='C'): Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('C')), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='named'), args=(TypedExprDecl(tp=JustTypeRef(name='String', args=()), expr=LitDecl(value='C')),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'NamedMat', [Lit(Span(SrcFile('<internal.egg>', None), 0, 0), String('C'))]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%C')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%C')]), VarDecl(name='ex1', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%ex1'), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=True)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%n')]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=True)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=True))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%n')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B')]), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='m', is_let=True)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%m')]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='m', is_let=True)),), bound_tp_params=()))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%m')])]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=True)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=True))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='m', is_let=True)),), bound_tp_params=()))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%n')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%m')])])]), VarDecl(name='rows', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%rows'), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='ex1', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%ex1')]), VarDecl(name='cols', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%cols'), CallDecl(callable=MethodRef(class_name='Matrix', method_name='ncols'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='ex1', is_let=True)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'ncols', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%ex1')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='nrows'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=True)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=True))), bound_tp_params=None)),), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'nrows', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%n')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B')])]), CallDecl(callable=MethodRef(class_name='Dim', method_name='__mul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='n', is_let=True)), TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='m', is_let=True))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Times', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%n'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%m')]), VarDecl(name='simple_ex1', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%simple_ex1'), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='B', is_let=True))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%B')]), VarDecl(name='ex2', is_let=True): Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%ex2'), CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='p', is_let=True)),), bound_tp_params=()): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%p')]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='p', is_let=True)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=True))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%p')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%C')]), CallDecl(callable=MethodRef(class_name='Matrix', method_name='__matmul__'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='p', is_let=True)),), bound_tp_params=())), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=True))), bound_tp_params=None)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=CallDecl(callable=ClassMethodRef(class_name='Matrix', method_name='identity'), args=(TypedExprDecl(tp=JustTypeRef(name='Dim', args=()), expr=VarDecl(name='m', is_let=True)),), bound_tp_params=()))), bound_tp_params=None))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'MMul', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%p')]), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%C')]), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A'), Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Id', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%m')])])]), CallDecl(callable=FunctionRef(name='kron'), args=(TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='A', is_let=True)), TypedExprDecl(tp=JustTypeRef(name='Matrix', args=()), expr=VarDecl(name='C', is_let=True))), bound_tp_params=None): Call(Span(SrcFile('<internal.egg>', None), 0, 0), 'Kron', [Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%A'), Var(Span(SrcFile('<internal.egg>', None), 0, 0), '%C')])}))

from __future__ import annotations

from egglog import *

egraph = EGraph()


class Dim(Expr):
    """
    A dimension of a matix.

    >>> Dim(3) * Dim.named("n")
    Dim(3) * Dim.named("n")
    """

    @method(egg_fn="Lit")
    def __init__(self, value: i64Like) -> None: ...

    @method(egg_fn="NamedDim")
    @classmethod
    def named(cls, name: StringLike) -> Dim:  # type: ignore[empty-body]
        ...

    @method(egg_fn="Times")
    def __mul__(self, other: Dim) -> Dim:  # type: ignore[empty-body]
        ...


a, b, c, n = vars_("a b c n", Dim)
i, j = vars_("i j", i64)
egraph.register(
    rewrite(a * (b * c)).to((a * b) * c),
    rewrite((a * b) * c).to(a * (b * c)),
    rewrite(Dim(i) * Dim(j)).to(Dim(i * j)),
    rewrite(a * b).to(b * a),
)


class Matrix(Expr, egg_sort="MExpr"):
    @method(egg_fn="Id")
    @classmethod
    def identity(cls, dim: Dim) -> Matrix:  # type: ignore[empty-body]
        """
        Create an identity matrix of the given dimension.
        """

    @method(egg_fn="NamedMat")
    @classmethod
    def named(cls, name: StringLike) -> Matrix:  # type: ignore[empty-body]
        """
        Create a named matrix.
        """

    @method(egg_fn="MMul")
    def __matmul__(self, other: Matrix) -> Matrix:  # type: ignore[empty-body]
        """
        Matrix multiplication.
        """

    @method(egg_fn="nrows")
    def nrows(self) -> Dim:  # type: ignore[empty-body]
        """
        Number of rows in the matrix.
        """

    @method(egg_fn="ncols")
    def ncols(self) -> Dim:  # type: ignore[empty-body]
        """
        Number of columns in the matrix.
        """


@function(egg_fn="Kron")
def kron(a: Matrix, b: Matrix) -> Matrix:  # type: ignore[empty-body]
    """
    Kronecker product of two matrices.

    https://en.wikipedia.org/wiki/Kronecker_product#Definition
    """


A, B, C, D = vars_("A B C D", Matrix)
egraph.register(
    # The dimensions of a kronecker product are the product of the dimensions
    rewrite(kron(A, B).nrows()).to(A.nrows() * B.nrows()),
    rewrite(kron(A, B).ncols()).to(A.ncols() * B.ncols()),
    # The dimensions of a matrix multiplication are the number of rows of the first
    # matrix and the number of columns of the second matrix.
    rewrite((A @ B).nrows()).to(A.nrows()),
    rewrite((A @ B).ncols()).to(B.ncols()),
    # The dimensions of an identity matrix are the input dimension
    rewrite(Matrix.identity(n).nrows()).to(n),
    rewrite(Matrix.identity(n).ncols()).to(n),
)
egraph.register(
    # Multiplication by an identity matrix is the same as the other matrix
    rewrite(Matrix.identity(n) @ A).to(A),
    rewrite(A @ Matrix.identity(n)).to(A),
    # Matrix multiplication is associative
    rewrite(A @ (B @ C)).to((A @ B) @ C),
    rewrite((A @ B) @ C).to(A @ (B @ C)),
    # Kronecker product is associative
    rewrite(kron(A, kron(B, C))).to(kron(kron(A, B), C)),
    rewrite(kron(kron(A, B), C)).to(kron(A, kron(B, C))),
    # Kronecker product distributes over matrix multiplication
    rewrite(kron(A @ C, B @ D)).to(kron(A, B) @ kron(C, D)),
    rewrite(kron(A, B) @ kron(C, D)).to(
        kron(A @ C, B @ D),
        # Only when the dimensions match
        eq(A.ncols()).to(C.nrows()),
        eq(B.ncols()).to(D.nrows()),
    ),
)
egraph.register(
    # demand rows and columns when we multiply matrices
    rule(eq(C).to(A @ B)).then(
        A.ncols(),
        A.nrows(),
        B.ncols(),
        B.nrows(),
    ),
    # demand rows and columns when we take the kronecker product
    rule(eq(C).to(kron(A, B))).then(
        A.ncols(),
        A.nrows(),
        B.ncols(),
        B.nrows(),
    ),
)


# Define a number of dimensions
n = egraph.let("n", Dim.named("n"))
m = egraph.let("m", Dim.named("m"))
p = egraph.let("p", Dim.named("p"))

# Define a number of matrices
A = egraph.let("A", Matrix.named("A"))
B = egraph.let("B", Matrix.named("B"))
C = egraph.let("C", Matrix.named("C"))

# Set each to be a square matrix of the given dimension
egraph.register(
    union(A.nrows()).with_(n),
    union(A.ncols()).with_(n),
    union(B.nrows()).with_(m),
    union(B.ncols()).with_(m),
    union(C.nrows()).with_(p),
    union(C.ncols()).with_(p),
)
# Create an example which should equal the kronecker product of A and B
ex1 = egraph.let("ex1", kron(Matrix.identity(n), B) @ kron(A, Matrix.identity(m)))
rows = egraph.let("rows", ex1.nrows())
cols = egraph.let("cols", ex1.ncols())

egraph.run(20)

egraph.check(eq(B.nrows()).to(m))
egraph.check(eq(kron(Matrix.identity(n), B).nrows()).to(n * m))

# Verify it matches the expected result
simple_ex1 = egraph.let("simple_ex1", kron(A, B))
egraph.check(eq(ex1).to(simple_ex1))

ex2 = egraph.let("ex2", kron(Matrix.identity(p), C) @ kron(A, Matrix.identity(m)))

egraph.run(10)
# Verify it is not simplified
egraph.check_fail(eq(ex2).to(kron(A, C)))
egraph

Total running time of the script: (0 minutes 0.031 seconds)

Gallery generated by Sphinx-Gallery