| 926 | |
| 927 | |
| 928 | def optimize_until(expr: Expr, stage: OptimizerStage) -> Expr: |
| 929 | result = expr |
| 930 | if stage == "logical": |
| 931 | return result |
| 932 | |
| 933 | # Simplify |
| 934 | expr = result.simplify() |
| 935 | if stage == "simplified-logical": |
| 936 | return expr |
| 937 | |
| 938 | # Manipulate Expression to make it more efficient |
| 939 | if dask.config.get("optimization.tune.active", True): |
| 940 | expr = expr.rewrite(kind="tune", rewritten={}) |
| 941 | if stage == "tuned-logical": |
| 942 | return expr |
| 943 | |
| 944 | # Lower |
| 945 | expr = expr.lower_completely() |
| 946 | if stage == "physical": |
| 947 | return expr |
| 948 | |
| 949 | # Simplify again |
| 950 | expr = expr.simplify() |
| 951 | if stage == "simplified-physical": |
| 952 | return expr |
| 953 | |
| 954 | # Final graph-specific optimizations |
| 955 | expr = expr.fuse() |
| 956 | if stage == "fused": |
| 957 | return expr |
| 958 | |
| 959 | raise ValueError(f"Stage {stage!r} not supported.") |
| 960 | |
| 961 | |
| 962 | class LLGExpr(Expr): |