gradient-script 0.2.0 → 0.3.1

package/dist/dsl/egraph/Rewriter.js

@@ -0,0 +1,131 @@
+/**
+ * Rewrite Engine for E-Graph Equality Saturation
+ *
+ * Applies rewrite rules until saturation (no new merges) or iteration limit.
+ */
+import { matchPattern, instantiatePattern } from './Pattern.js';
+/**
+ * Apply equality saturation to an e-graph
+ *
+ * Repeatedly applies rewrite rules until:
+ * - No new equivalences are discovered (saturated)
+ * - Max iterations reached
+ * - E-graph size limit exceeded
+ */
+export function saturate(egraph, rules, options = {}) {
+    const { maxIterations = 10, maxClassSize = 5000, maxMatchesPerIter = 5000, verbose = false } = options;
+    const stats = {
+        iterations: 0,
+        totalMatches: 0,
+        merges: 0,
+        saturated: false,
+        classCount: egraph.size
+    };
+    for (let iter = 0; iter < maxIterations; iter++) {
+        stats.iterations = iter + 1;
+        // Check size limit
+        if (egraph.size > maxClassSize) {
+            if (verbose) {
+                console.log(`[saturate] Size limit exceeded: ${egraph.size} > ${maxClassSize}`);
+            }
+            break;
+        }
+        // Collect all rule matches (with limit)
+        const matches = collectMatches(egraph, rules, maxMatchesPerIter);
+        stats.totalMatches += matches.length;
+        if (matches.length === 0) {
+            stats.saturated = true;
+            if (verbose) {
+                console.log(`[saturate] Saturated after ${iter + 1} iterations`);
+            }
+            break;
+        }
+        // Apply all matches
+        let mergesThisIter = 0;
+        for (const { rule, classId, subst } of matches) {
+            const rhsId = instantiatePattern(egraph, rule.rhs, subst);
+            const lhsCanon = egraph.find(classId);
+            const rhsCanon = egraph.find(rhsId);
+            if (lhsCanon !== rhsCanon) {
+                egraph.merge(lhsCanon, rhsCanon);
+                mergesThisIter++;
+            }
+        }
+        stats.merges += mergesThisIter;
+        // Rebuild to restore invariants
+        egraph.rebuild();
+        if (verbose) {
+            console.log(`[saturate] Iter ${iter + 1}: ${matches.length} matches, ${mergesThisIter} merges, ${egraph.size} classes`);
+        }
+        // If no merges happened, we're saturated
+        if (mergesThisIter === 0) {
+            stats.saturated = true;
+            break;
+        }
+    }
+    stats.classCount = egraph.size;
+    return stats;
+}
+/**
+ * Collect all rule matches across the e-graph
+ */
+function collectMatches(egraph, rules, maxMatches) {
+    const matches = [];
+    for (const classId of egraph.getClassIds()) {
+        for (const rule of rules) {
+            const substs = matchPattern(egraph, rule.lhs, classId);
+            for (const subst of substs) {
+                matches.push({ rule, classId, subst });
+                if (matches.length >= maxMatches) {
+                    // Hit limit - return what we have (prevents explosion)
+                    return matches;
+                }
+            }
+        }
+    }
+    return matches;
+}
+/**
+ * Apply a single rule once, returning number of merges
+ */
+export function applyRuleOnce(egraph, rule) {
+    const matches = collectMatches(egraph, [rule], 5000);
+    let merges = 0;
+    for (const { classId, subst } of matches) {
+        const rhsId = instantiatePattern(egraph, rule.rhs, subst);
+        const lhsCanon = egraph.find(classId);
+        const rhsCanon = egraph.find(rhsId);
+        if (lhsCanon !== rhsCanon) {
+            egraph.merge(lhsCanon, rhsCanon);
+            merges++;
+        }
+    }
+    egraph.rebuild();
+    return merges;
+}
+/**
+ * Apply rules in phases for better control
+ * E.g., apply core rules first, then algebra rules
+ */
+export function saturatePhased(egraph, phases, options = {}) {
+    const combinedStats = {
+        iterations: 0,
+        totalMatches: 0,
+        merges: 0,
+        saturated: true,
+        classCount: egraph.size
+    };
+    for (let phaseIdx = 0; phaseIdx < phases.length; phaseIdx++) {
+        const rules = phases[phaseIdx];
+        if (options.verbose) {
+            console.log(`[saturatePhased] Phase ${phaseIdx + 1}/${phases.length}: ${rules.length} rules`);
+        }
+        const phaseStats = saturate(egraph, rules, options);
+        combinedStats.iterations += phaseStats.iterations;
+        combinedStats.totalMatches += phaseStats.totalMatches;
+        combinedStats.merges += phaseStats.merges;
+        combinedStats.saturated = combinedStats.saturated && phaseStats.saturated;
+    }
+    combinedStats.classCount = egraph.size;
+    return combinedStats;
+}

package/dist/dsl/egraph/Rules.d.ts

@@ -0,0 +1,44 @@
+/**
+ * Algebraic Rewrite Rules for E-Graph Optimization
+ *
+ * Rules are organized by category:
+ * - Core: Always safe, essential (commutativity, identity, etc.)
+ * - Algebra: More powerful, can cause expansion (distribution, factoring)
+ * - Functions: Domain-specific (sqrt, trig, exp/log)
+ *
+ * Based on rules from Herbie (https://github.com/herbie-fp/herbie)
+ * and egg (https://egraphs-good.github.io/)
+ */
+import { Pattern } from './Pattern.js';
+/**
+ * A rewrite rule: if LHS matches, RHS is equivalent
+ */
+export interface Rule {
+    name: string;
+    lhs: Pattern;
+    rhs: Pattern;
+}
+/**
+ * Create a rule from pattern strings
+ */
+export declare function rule(name: string, lhs: string, rhs: string): Rule;
+/**
+ * Create bidirectional rules (both directions)
+ */
+export declare function biRule(name: string, a: string, b: string): Rule[];
+export declare const coreRules: Rule[];
+export declare const algebraRules: Rule[];
+export declare const functionRules: Rule[];
+/**
+ * All rules combined
+ */
+export declare const allRules: Rule[];
+/**
+ * Minimal rules for canonicalization only
+ * (no expansion, just normalization)
+ */
+export declare const canonRules: Rule[];
+/**
+ * Get rules by category
+ */
+export declare function getRuleSet(categories: ('core' | 'algebra' | 'function')[]): Rule[];

package/dist/dsl/egraph/Rules.js

@@ -0,0 +1,187 @@
+/**
+ * Algebraic Rewrite Rules for E-Graph Optimization
+ *
+ * Rules are organized by category:
+ * - Core: Always safe, essential (commutativity, identity, etc.)
+ * - Algebra: More powerful, can cause expansion (distribution, factoring)
+ * - Functions: Domain-specific (sqrt, trig, exp/log)
+ *
+ * Based on rules from Herbie (https://github.com/herbie-fp/herbie)
+ * and egg (https://egraphs-good.github.io/)
+ */
+import { parsePattern } from './Pattern.js';
+/**
+ * Create a rule from pattern strings
+ */
+export function rule(name, lhs, rhs) {
+    return {
+        name,
+        lhs: parsePattern(lhs),
+        rhs: parsePattern(rhs)
+    };
+}
+/**
+ * Create bidirectional rules (both directions)
+ */
+export function biRule(name, a, b) {
+    return [
+        rule(`${name}-l`, a, b),
+        rule(`${name}-r`, b, a)
+    ];
+}
+// =============================================================================
+// CORE RULES - Always safe, essential
+// =============================================================================
+export const coreRules = [
+    // === Commutativity ===
+    rule('comm-add', '(+ ?a ?b)', '(+ ?b ?a)'),
+    rule('comm-mul', '(* ?a ?b)', '(* ?b ?a)'),
+    // === Associativity ===
+    ...biRule('assoc-add', '(+ (+ ?a ?b) ?c)', '(+ ?a (+ ?b ?c))'),
+    ...biRule('assoc-mul', '(* (* ?a ?b) ?c)', '(* ?a (* ?b ?c))'),
+    // === Identity: addition ===
+    rule('add-0-l', '(+ 0 ?a)', '?a'),
+    rule('add-0-r', '(+ ?a 0)', '?a'),
+    // === Identity: subtraction ===
+    rule('sub-0', '(- ?a 0)', '?a'),
+    rule('0-sub', '(- 0 ?a)', '(neg ?a)'),
+    // === Identity: multiplication ===
+    rule('mul-1-l', '(* 1 ?a)', '?a'),
+    rule('mul-1-r', '(* ?a 1)', '?a'),
+    // === Identity: division ===
+    rule('div-1', '(/ ?a 1)', '?a'),
+    // === Identity: power ===
+    rule('pow-0', '(^ ?a 0)', '1'),
+    rule('pow-1', '(^ ?a 1)', '?a'),
+    // === Zero: multiplication ===
+    rule('mul-0-l', '(* 0 ?a)', '0'),
+    rule('mul-0-r', '(* ?a 0)', '0'),
+    // === Zero: division ===
+    rule('0-div', '(/ 0 ?a)', '0'),
+    // === Inverse: subtraction ===
+    rule('sub-self', '(- ?a ?a)', '0'),
+    // === Inverse: division ===
+    rule('div-self', '(/ ?a ?a)', '1'),
+    // === Double negation ===
+    rule('neg-neg', '(neg (neg ?a))', '?a'),
+    // === Negation with multiplication ===
+    rule('neg-mul-neg', '(* (neg ?a) (neg ?b))', '(* ?a ?b)'),
+    rule('mul-neg-1', '(* -1 ?a)', '(neg ?a)'),
+    rule('neg-to-mul', '(neg ?a)', '(* -1 ?a)'),
+];
+// =============================================================================
+// ALGEBRA RULES - More powerful, can cause expansion
+// =============================================================================
+export const algebraRules = [
+    // === Distribution (both directions) ===
+    ...biRule('dist-mul-add', '(* ?a (+ ?b ?c))', '(+ (* ?a ?b) (* ?a ?c))'),
+    ...biRule('dist-mul-sub', '(* ?a (- ?b ?c))', '(- (* ?a ?b) (* ?a ?c))'),
+    // === Negation propagation ===
+    rule('neg-add', '(neg (+ ?a ?b))', '(+ (neg ?a) (neg ?b))'),
+    rule('neg-sub', '(neg (- ?a ?b))', '(- ?b ?a)'),
+    rule('neg-mul-l', '(* (neg ?a) ?b)', '(neg (* ?a ?b))'),
+    rule('neg-mul-r', '(* ?a (neg ?b))', '(neg (* ?a ?b))'),
+    rule('neg-div-l', '(/ (neg ?a) ?b)', '(neg (/ ?a ?b))'),
+    rule('neg-div-r', '(/ ?a (neg ?b))', '(neg (/ ?a ?b))'),
+    // === Subtraction to addition ===
+    rule('sub-to-add', '(- ?a ?b)', '(+ ?a (neg ?b))'),
+    rule('add-neg-to-sub', '(+ ?a (neg ?b))', '(- ?a ?b)'),
+    // === Division to multiplication by inverse ===
+    rule('div-to-mul-inv', '(/ ?a ?b)', '(* ?a (inv ?b))'),
+    // === Inverse rules ===
+    rule('inv-1', '(inv 1)', '1'),
+    rule('inv-inv', '(inv (inv ?a))', '?a'),
+    rule('div-1-to-inv', '(/ 1 ?a)', '(inv ?a)'),
+    rule('mul-inv-cancel', '(* ?a (inv ?a))', '1'), // a * (1/a) = 1
+    rule('inv-mul', '(inv (* ?a ?b))', '(* (inv ?a) (inv ?b))'), // 1/(a*b) = (1/a)*(1/b)
+    rule('div-self-sq', '(/ ?a (* ?a ?a))', '(inv ?a)'), // a / (a*a) = 1/a
+    rule('div-sq-self', '(/ (* ?a ?a) ?a)', '?a'), // (a*a) / a = a
+    rule('div-self-pow2', '(/ ?a (^ ?a 2))', '(inv ?a)'), // a / a^2 = 1/a
+    rule('div-pow2-self', '(/ (^ ?a 2) ?a)', '?a'), // a^2 / a = a
+    // === Power rules ===
+    rule('pow-2', '(^ ?a 2)', '(* ?a ?a)'),
+    rule('sq-to-pow', '(* ?a ?a)', '(^ ?a 2)'),
+    // === Combining like terms ===
+    rule('add-same', '(+ ?a ?a)', '(* 2 ?a)'),
+    rule('sub-neg-same', '(- ?a (neg ?a))', '(* 2 ?a)'),
+];
+// =============================================================================
+// FUNCTION RULES - Domain-specific
+// =============================================================================
+export const functionRules = [
+    // === Sqrt ===
+    rule('sqrt-sq', '(* (sqrt ?a) (sqrt ?a))', '?a'),
+    rule('sqrt-mul', '(sqrt (* ?a ?b))', '(* (sqrt ?a) (sqrt ?b))'),
+    rule('sqrt-div', '(sqrt (/ ?a ?b))', '(/ (sqrt ?a) (sqrt ?b))'),
+    rule('sqrt-1', '(sqrt 1)', '1'),
+    rule('sqrt-0', '(sqrt 0)', '0'),
+    // === Abs ===
+    rule('abs-neg', '(abs (neg ?a))', '(abs ?a)'),
+    rule('abs-abs', '(abs (abs ?a))', '(abs ?a)'),
+    rule('abs-sq', '(abs (* ?a ?a))', '(* ?a ?a)'),
+    // === Trig ===
+    rule('sin-0', '(sin 0)', '0'),
+    rule('cos-0', '(cos 0)', '1'),
+    rule('sin-neg', '(sin (neg ?a))', '(neg (sin ?a))'),
+    rule('cos-neg', '(cos (neg ?a))', '(cos ?a)'),
+    // === Exp/Log ===
+    rule('exp-0', '(exp 0)', '1'),
+    rule('log-1', '(log 1)', '0'),
+    rule('exp-log', '(exp (log ?a))', '?a'),
+    rule('log-exp', '(log (exp ?a))', '?a'),
+    rule('log-mul', '(log (* ?a ?b))', '(+ (log ?a) (log ?b))'),
+    rule('log-div', '(log (/ ?a ?b))', '(- (log ?a) (log ?b))'),
+    rule('log-pow', '(log (^ ?a ?b))', '(* ?b (log ?a))'),
+];
+// =============================================================================
+// RULE SETS
+// =============================================================================
+/**
+ * All rules combined
+ */
+export const allRules = [
+    ...coreRules,
+    ...algebraRules,
+    ...functionRules,
+];
+/**
+ * Minimal rules for canonicalization only
+ * (no expansion, just normalization)
+ */
+export const canonRules = [
+    // Commutativity (for canonical ordering)
+    rule('comm-add', '(+ ?a ?b)', '(+ ?b ?a)'),
+    rule('comm-mul', '(* ?a ?b)', '(* ?b ?a)'),
+    // Identity removal
+    rule('add-0-l', '(+ 0 ?a)', '?a'),
+    rule('add-0-r', '(+ ?a 0)', '?a'),
+    rule('mul-1-l', '(* 1 ?a)', '?a'),
+    rule('mul-1-r', '(* ?a 1)', '?a'),
+    rule('sub-0', '(- ?a 0)', '?a'),
+    rule('div-1', '(/ ?a 1)', '?a'),
+    rule('pow-1', '(^ ?a 1)', '?a'),
+    // Zero
+    rule('mul-0-l', '(* 0 ?a)', '0'),
+    rule('mul-0-r', '(* ?a 0)', '0'),
+    rule('0-div', '(/ 0 ?a)', '0'),
+    rule('pow-0', '(^ ?a 0)', '1'),
+    // Inverse
+    rule('sub-self', '(- ?a ?a)', '0'),
+    rule('div-self', '(/ ?a ?a)', '1'),
+    // Double negation
+    rule('neg-neg', '(neg (neg ?a))', '?a'),
+    rule('neg-mul-neg', '(* (neg ?a) (neg ?b))', '(* ?a ?b)'),
+];
+/**
+ * Get rules by category
+ */
+export function getRuleSet(categories) {
+    const rules = [];
+    if (categories.includes('core'))
+        rules.push(...coreRules);
+    if (categories.includes('algebra'))
+        rules.push(...algebraRules);
+    if (categories.includes('function'))
+        rules.push(...functionRules);
+    return rules;
+}

package/dist/dsl/egraph/index.d.ts

@@ -0,0 +1,15 @@
+/**
+ * E-Graph Module for GradientScript
+ *
+ * Provides equality saturation-based optimization for gradient expressions.
+ * Can be used as an alternative to (or in addition to) the CSE module.
+ */
+export { EGraph } from './EGraph.js';
+export type { EClass } from './EGraph.js';
+export { EClassId, ENode, enodeKey, enodeChildren, enodeWithChildren } from './ENode.js';
+export { Pattern, Substitution, parsePattern, matchPattern, instantiatePattern, patternToString } from './Pattern.js';
+export { Rule, rule, biRule, coreRules, algebraRules, functionRules, allRules, canonRules, getRuleSet } from './Rules.js';
+export { saturate, saturatePhased, applyRuleOnce, SaturationStats, SaturationOptions } from './Rewriter.js';
+export { extractBest, extractWithCSE, ExtractionResult, CostModel, defaultCostModel } from './Extractor.js';
+export { addExpression, addExpressions, addGradients, getRootIds } from './Convert.js';
+export { optimizeWithEGraph, EGraphOptimizeResult } from './Optimizer.js';

package/dist/dsl/egraph/index.js

@@ -0,0 +1,21 @@
+/**
+ * E-Graph Module for GradientScript
+ *
+ * Provides equality saturation-based optimization for gradient expressions.
+ * Can be used as an alternative to (or in addition to) the CSE module.
+ */
+// Core e-graph
+export { EGraph } from './EGraph.js';
+export { enodeKey, enodeChildren, enodeWithChildren } from './ENode.js';
+// Pattern matching
+export { parsePattern, matchPattern, instantiatePattern, patternToString } from './Pattern.js';
+// Rewrite rules
+export { rule, biRule, coreRules, algebraRules, functionRules, allRules, canonRules, getRuleSet } from './Rules.js';
+// Saturation
+export { saturate, saturatePhased, applyRuleOnce } from './Rewriter.js';
+// Extraction
+export { extractBest, extractWithCSE, defaultCostModel } from './Extractor.js';
+// AST conversion
+export { addExpression, addExpressions, addGradients, getRootIds } from './Convert.js';
+// Main optimizer function
+export { optimizeWithEGraph } from './Optimizer.js';

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "gradient-script",
-  "version": "0.2.0",
+  "version": "0.3.1",
   "description": "Symbolic differentiation for structured types with a simple DSL",
   "type": "module",
   "main": "dist/index.js",

package/dist/dsl/CSE.d.ts

@@ -1,21 +0,0 @@
-/**
- * Common Subexpression Elimination (CSE)
- * Identifies repeated expressions and factors them out
- */
-import { Expression } from './AST.js';
-export interface CSEResult {
-    intermediates: Map<string, Expression>;
-    simplified: Expression;
-}
-export interface StructuredCSEResult {
-    intermediates: Map<string, Expression>;
-    components: Map<string, Expression>;
-}
-/**
- * Perform CSE on an expression
- */
-export declare function eliminateCommonSubexpressions(expr: Expression, minCount?: number): CSEResult;
-/**
- * Perform CSE on structured gradients (for structured types like {x, y})
- */
-export declare function eliminateCommonSubexpressionsStructured(components: Map<string, Expression>, minCount?: number): StructuredCSEResult;

package/dist/dsl/CSE.js

@@ -1,168 +0,0 @@
-/**
- * Common Subexpression Elimination (CSE)
- * Identifies repeated expressions and factors them out
- */
-import { ExpressionTransformer } from './ExpressionTransformer.js';
-import { serializeExpression } from './ExpressionUtils.js';
-/**
- * Perform CSE on an expression
- */
-export function eliminateCommonSubexpressions(expr, minCount = 2) {
-    const counter = new ExpressionCounter();
-    counter.count(expr);
-    const intermediates = new Map();
-    let varCounter = 0;
-    const subexprMap = new Map();
-    for (const [exprStr, count] of counter.counts.entries()) {
-        if (count >= minCount) {
-            const parsed = counter.expressions.get(exprStr);
-            if (parsed && shouldExtract(parsed)) {
-                const varName = `_tmp${varCounter++}`;
-                intermediates.set(varName, parsed);
-                subexprMap.set(exprStr, varName);
-            }
-        }
-    }
-    const simplified = substituteExpressions(expr, subexprMap, counter);
-    return { intermediates, simplified };
-}
-/**
- * Perform CSE on structured gradients (for structured types like {x, y})
- */
-export function eliminateCommonSubexpressionsStructured(components, minCount = 2) {
-    const counter = new ExpressionCounter();
-    for (const expr of components.values()) {
-        counter.count(expr);
-    }
-    const intermediates = new Map();
-    let varCounter = 0;
-    const subexprMap = new Map();
-    for (const [exprStr, count] of counter.counts.entries()) {
-        if (count >= minCount) {
-            const parsed = counter.expressions.get(exprStr);
-            if (parsed && shouldExtract(parsed)) {
-                const varName = `_tmp${varCounter++}`;
-                intermediates.set(varName, parsed);
-                subexprMap.set(exprStr, varName);
-            }
-        }
-    }
-    const simplifiedComponents = new Map();
-    for (const [comp, expr] of components.entries()) {
-        simplifiedComponents.set(comp, substituteExpressions(expr, subexprMap, counter));
-    }
-    return { intermediates, components: simplifiedComponents };
-}
-/**
- * Check if an expression should be extracted
- */
-function shouldExtract(expr) {
-    switch (expr.kind) {
-        case 'number':
-        case 'variable':
-            return false;
-        case 'component':
-            return expr.object.kind !== 'variable';
-        case 'unary':
-            return shouldExtract(expr.operand);
-        case 'binary':
-            return true;
-        case 'call':
-            return true;
-        default:
-            return false;
-    }
-}
-/**
- * Counts occurrences of subexpressions during traversal
- */
-class ExpressionCounter extends ExpressionTransformer {
-    counts = new Map();
-    expressions = new Map();
-    count(expr) {
-        this.transform(expr);
-    }
-    serialize(expr) {
-        return serializeExpression(expr);
-    }
-    recordExpression(expr) {
-        const key = this.serialize(expr);
-        const currentCount = this.counts.get(key) || 0;
-        this.counts.set(key, currentCount + 1);
-        if (!this.expressions.has(key)) {
-            this.expressions.set(key, expr);
-        }
-    }
-    visitNumber(node) {
-        this.recordExpression(node);
-        return node;
-    }
-    visitVariable(node) {
-        this.recordExpression(node);
-        return node;
-    }
-    visitBinaryOp(node) {
-        this.recordExpression(node);
-        return super.visitBinaryOp(node);
-    }
-    visitUnaryOp(node) {
-        this.recordExpression(node);
-        return super.visitUnaryOp(node);
-    }
-    visitFunctionCall(node) {
-        this.recordExpression(node);
-        return super.visitFunctionCall(node);
-    }
-    visitComponentAccess(node) {
-        this.recordExpression(node);
-        return super.visitComponentAccess(node);
-    }
-}
-/**
- * Substitute common subexpressions with variables
- */
-function substituteExpressions(expr, subexprMap, counter) {
-    const key = counter.serialize(expr);
-    if (subexprMap.has(key)) {
-        return {
-            kind: 'variable',
-            name: subexprMap.get(key)
-        };
-    }
-    let result = expr;
-    for (const [exprStr, varName] of subexprMap.entries()) {
-        const exprToReplace = counter.expressions.get(exprStr);
-        if (exprToReplace && counter.serialize(result) !== exprStr) {
-            result = substituteInExpression(result, exprToReplace, { kind: 'variable', name: varName }, counter);
-        }
-    }
-    return result;
-}
-/**
- * Transformer that substitutes a pattern with a replacement expression
- * Used for CSE optimization to replace repeated subexpressions with intermediate variables
- */
-class PatternSubstitutionTransformer extends ExpressionTransformer {
-    pattern;
-    replacement;
-    counter;
-    constructor(pattern, replacement, counter) {
-        super();
-        this.pattern = pattern;
-        this.replacement = replacement;
-        this.counter = counter;
-    }
-    transform(expr) {
-        if (this.counter.serialize(expr) === this.counter.serialize(this.pattern)) {
-            return this.replacement;
-        }
-        return super.transform(expr);
-    }
-}
-/**
- * Helper to substitute an expression pattern with a replacement
- */
-function substituteInExpression(expr, pattern, replacement, counter) {
-    const transformer = new PatternSubstitutionTransformer(pattern, replacement, counter);
-    return transformer.transform(expr);
-}