blob: 6dc23e935928e7936b8732e1f9ec41215f5b5285 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
|
from sympy import *
from sympy.codegen.rewriting import ReplaceOptim
expand_pos_square = ReplaceOptim(
lambda e: e.is_Pow and e.exp.is_integer and e.exp == 2,
lambda p: UnevaluatedExpr(Mul(p.base, p.base, evaluate = False))
)
expand_neg_square = ReplaceOptim(
lambda e: e.is_Pow and e.exp.is_integer and e.exp == -2,
lambda p: UnevaluatedExpr(Mul(p.base, p.base, evaluate = False))
)
custom = [
(expand_pos_square, expand_pos_square),
(expand_neg_square, expand_neg_square)
] + cse_main.basic_optimizations
|