gradient-script 0.2.0 → 0.3.1

package/dist/dsl/CodeGen.js

@@ -2,11 +2,12 @@
  * Code generation for GradientScript DSL
  * Generates TypeScript/JavaScript code with gradient functions
  */
-import { simplifyGradients } from './Simplify.js';
+import { simplifyGradients, simplifyPostCSE } from './Simplify.js';
 import { ExpressionTransformer } from './ExpressionTransformer.js';
-import { eliminateCommonSubexpressionsStructured } from './CSE.js';
+import { optimizeWithEGraph } from './egraph/index.js';
 import { CodeGenError } from './Errors.js';
 import { serializeExpression } from './ExpressionUtils.js';
+import { inlineExpression } from './Inliner.js';
 function capitalize(str) {
     return str.charAt(0).toUpperCase() + str.slice(1);
 }
@@ -79,10 +80,12 @@ export class ExpressionCodeGen {
                         return left;
                     }
                     else if (exponent === 2) {
-                        return `${left} * ${left}`;
+                        // Wrap in parens because we're changing ^ (precedence 3) to * (precedence 2)
+                        // Without parens, a / b^2 would become a / b * b instead of a / (b * b)
+                        return `(${left} * ${left})`;
                     }
                     else if (exponent === 3) {
-                        return `${left} * ${left} * ${left}`;
+                        return `(${left} * ${left} * ${left})`;
                     }
                 }
             }
@@ -166,6 +169,10 @@ export class ExpressionCodeGen {
     }
     genUnary(expr) {
         const operand = this.generate(expr.operand);
+        // Parenthesize binary operands to avoid precedence bugs: -(a + b) not -a + b
+        if (expr.operand.kind === 'binary') {
+            return `${expr.operator}(${operand})`;
+        }
         return `${expr.operator}${operand}`;
     }
     genCall(expr) {
@@ -300,15 +307,162 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     // Forward pass - compute intermediate variables
     // Track which expressions are already computed for CSE reuse
     const forwardExpressionMap = new Map();
+    // Build substitution map for inlining (to match gradient expressions which are fully inlined)
+    const substitutionMap = new Map();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            substitutionMap.set(stmt.variable, stmt.expression);
+        }
+    }
+    // Collect forward variable names and expressions for optimization
+    const forwardVars = new Set();
+    const forwardVarExprs = new Map();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            forwardVars.add(stmt.variable);
+            forwardVarExprs.set(stmt.variable, stmt.expression);
+        }
+    }
+    // Optimize forward expressions with e-graph
+    let optimizedForwardExprs = forwardVarExprs;
+    let forwardCseTemps = new Map();
+    if (options.cse !== false && forwardVarExprs.size > 0) {
+        const forOptimizer = new Map();
+        forOptimizer.set('_forward', forwardVarExprs);
+        const forwardResult = optimizeWithEGraph(forOptimizer, { verbose: false });
+        // Rename forward temps to avoid conflicts with gradient temps (use _fwd prefix)
+        const rawTemps = forwardResult.intermediates;
+        const renameMap = new Map();
+        for (const oldName of rawTemps.keys()) {
+            const newName = oldName.replace('_tmp', '_fwd');
+            renameMap.set(oldName, newName);
+        }
+        // Apply renaming to temp definitions
+        for (const [oldName, expr] of rawTemps) {
+            const newName = renameMap.get(oldName);
+            forwardCseTemps.set(newName, renameTempRefs(expr, renameMap));
+        }
+        // Apply renaming to optimized expressions
+        optimizedForwardExprs = new Map();
+        for (const [varName, expr] of (forwardResult.gradients.get('_forward') || forwardVarExprs)) {
+            optimizedForwardExprs.set(varName, renameTempRefs(expr, renameMap));
+        }
+    }
+    // Helper to rename temp references in an expression
+    function renameTempRefs(expr, renameMap) {
+        if (expr.kind === 'variable') {
+            const newName = renameMap.get(expr.name);
+            return newName ? { kind: 'variable', name: newName } : expr;
+        }
+        else if (expr.kind === 'binary') {
+            return {
+                kind: 'binary',
+                operator: expr.operator,
+                left: renameTempRefs(expr.left, renameMap),
+                right: renameTempRefs(expr.right, renameMap)
+            };
+        }
+        else if (expr.kind === 'unary') {
+            return {
+                kind: 'unary',
+                operator: expr.operator,
+                operand: renameTempRefs(expr.operand, renameMap)
+            };
+        }
+        else if (expr.kind === 'call') {
+            return {
+                kind: 'call',
+                name: expr.name,
+                args: expr.args.map(a => renameTempRefs(a, renameMap))
+            };
+        }
+        else if (expr.kind === 'component') {
+            return {
+                kind: 'component',
+                object: renameTempRefs(expr.object, renameMap),
+                component: expr.component
+            };
+        }
+        return expr;
+    }
+    // Helper to find which forward vars an expression depends on
+    function findForwardVarDeps(expr) {
+        const deps = new Set();
+        function visit(e) {
+            if (e.kind === 'variable' && forwardVars.has(e.name)) {
+                deps.add(e.name);
+            }
+            else if (e.kind === 'binary') {
+                visit(e.left);
+                visit(e.right);
+            }
+            else if (e.kind === 'unary') {
+                visit(e.operand);
+            }
+            else if (e.kind === 'call') {
+                e.args.forEach(visit);
+            }
+            else if (e.kind === 'component') {
+                visit(e.object);
+            }
+        }
+        visit(expr);
+        return deps;
+    }
+    // Track which CSE temps need to be emitted after which forward var
+    const fwdTempAfterVar = new Map();
+    const fwdTempsBeforeAny = [];
+    for (const [tempName, tempExpr] of forwardCseTemps) {
+        const deps = findForwardVarDeps(tempExpr);
+        if (deps.size === 0) {
+            fwdTempsBeforeAny.push({ name: tempName, expr: tempExpr });
+        }
+        else {
+            let lastDep = '';
+            for (const stmt of func.body) {
+                if (stmt.kind === 'assignment' && deps.has(stmt.variable)) {
+                    lastDep = stmt.variable;
+                }
+            }
+            if (lastDep) {
+                if (!fwdTempAfterVar.has(lastDep)) {
+                    fwdTempAfterVar.set(lastDep, []);
+                }
+                fwdTempAfterVar.get(lastDep).push({ name: tempName, expr: tempExpr });
+            }
+        }
+    }
+    // Emit forward CSE temps that don't depend on forward vars
+    for (const { name: tempName, expr } of fwdTempsBeforeAny) {
+        const code = codegen.generate(expr);
+        if (format === 'typescript' || format === 'javascript') {
+            lines.push(`  const ${tempName} = ${code};`);
+        }
+        else if (format === 'python') {
+            lines.push(`  ${tempName} = ${code}`);
+        }
+        else if (format === 'csharp') {
+            lines.push(`    ${csharpFloatType} ${tempName} = ${code};`);
+        }
+    }
+    // Generate forward variable assignments with interleaved temps
     for (const stmt of func.body) {
         if (stmt.kind === 'assignment') {
             const varName = stmt.variable;
-            const generatedExpr = codegen.generate(stmt.expression);
+            const expr = optimizedForwardExprs.get(varName) || stmt.expression;
+            const generatedExpr = codegen.generate(expr);
             if (shouldTrackForForwardReuse(stmt.expression)) {
+                // Register the original expression
                 const exprKey = serializeExpression(stmt.expression);
                 if (!forwardExpressionMap.has(exprKey)) {
                     forwardExpressionMap.set(exprKey, varName);
                 }
+                // Also register the fully inlined form (this is what gradient expressions will have)
+                const inlinedExpr = inlineExpression(stmt.expression, substitutionMap);
+                const inlinedKey = serializeExpression(inlinedExpr);
+                if (inlinedKey !== exprKey && !forwardExpressionMap.has(inlinedKey)) {
+                    forwardExpressionMap.set(inlinedKey, varName);
+                }
             }
             if (format === 'typescript' || format === 'javascript') {
                 lines.push(`  const ${varName} = ${generatedExpr};`);
@@ -319,6 +473,20 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
             else if (format === 'csharp') {
                 lines.push(`    ${csharpFloatType} ${varName} = ${generatedExpr};`);
             }
+            // Emit any forward CSE temps that depend on this var
+            const tempsForVar = fwdTempAfterVar.get(varName) || [];
+            for (const { name: tempName, expr: tempExpr } of tempsForVar) {
+                const tempCode = codegen.generate(tempExpr);
+                if (format === 'typescript' || format === 'javascript') {
+                    lines.push(`  const ${tempName} = ${tempCode};`);
+                }
+                else if (format === 'python') {
+                    lines.push(`  ${tempName} = ${tempCode}`);
+                }
+                else if (format === 'csharp') {
+                    lines.push(`    ${csharpFloatType} ${tempName} = ${tempCode};`);
+                }
+            }
         }
     }
     // Compute output value - reuse forward pass variables if possible
@@ -372,77 +540,71 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     if (includeComments) {
         lines.push(`  ${comment} Gradients`);
     }
-    // Apply CSE if requested
+    // Apply e-graph optimization (CSE + algebraic simplification)
     const shouldApplyCSE = options.cse !== false; // Default to true
-    const cseIntermediates = new Map();
+    let cseIntermediates = new Map();
     if (shouldApplyCSE) {
-        // Collect all gradient expressions for CSE analysis
+        // Collect all gradient components into a single map for global optimization
+        // Include both structured and scalar gradients (scalar uses 'value' as component)
+        const allGradientComponents = new Map();
         for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
             if (isStructuredGradient(gradient)) {
-                const cseResult = eliminateCommonSubexpressionsStructured(gradient.components);
-                // Merge intermediates
-                for (const [name, expr] of cseResult.intermediates.entries()) {
-                    cseIntermediates.set(name, expr);
-                }
-                // Update gradient components with CSE-simplified versions
-                gradient.components = cseResult.components;
-            }
-        }
-        // Generate intermediate variables from CSE
-        if (cseIntermediates.size > 0) {
-            // Check if we should emit guards (opt-in)
-            const shouldEmitGuards = options.emitGuards === true;
-            const epsilon = options.epsilon || 1e-10;
-            // Identify potential denominators (sum of squares patterns)
-            const denominatorVars = new Set();
-            for (const [varName, expr] of cseIntermediates.entries()) {
-                const code = codegen.generate(expr);
-                // Check if this looks like a denominator (contains + and squared terms)
-                if (code.includes('+') && (code.includes('* ') || code.includes('Math.pow'))) {
-                    denominatorVars.add(varName);
-                }
+                allGradientComponents.set(paramName, gradient.components);
             }
-            for (const [varName, expr] of cseIntermediates.entries()) {
-                const code = codegen.generate(expr);
-                if (format === 'typescript' || format === 'javascript') {
-                    lines.push(`  const ${varName} = ${code};`);
-                }
-                else if (format === 'python') {
-                    lines.push(`  ${varName} = ${code}`);
-                }
-                else if (format === 'csharp') {
-                    lines.push(`    ${csharpFloatType} ${varName} = ${code};`);
-                }
+            else {
+                // Scalar gradient - wrap as a single 'value' component
+                allGradientComponents.set(paramName, new Map([['value', gradient]]));
+            }
+        }
+        // Run e-graph optimization globally across ALL gradient expressions
+        const globalCSE = optimizeWithEGraph(allGradientComponents, { verbose: false });
+        cseIntermediates = globalCSE.intermediates;
+        // Update gradient components with optimized versions
+        for (const [paramName, simplifiedComponents] of globalCSE.gradients.entries()) {
+            const gradient = gradientsToUse.gradients.get(paramName);
+            if (gradient && isStructuredGradient(gradient)) {
+                gradient.components = simplifiedComponents;
             }
-            // Emit epsilon guard if needed
-            if (shouldEmitGuards && denominatorVars.size > 0) {
-                lines.push('');
-                if (includeComments) {
-                    lines.push(`  ${comment} Guard against division by zero`);
+            else {
+                // Scalar gradient - unwrap from 'value' component
+                const optimizedExpr = simplifiedComponents.get('value');
+                if (optimizedExpr) {
+                    gradientsToUse.gradients.set(paramName, optimizedExpr);
                 }
-                for (const denom of denominatorVars) {
-                    if (format === 'typescript' || format === 'javascript') {
-                        lines.push(`  if (Math.abs(${denom}) < ${epsilon}) {`);
-                        lines.push(`    ${comment} Return zero gradients for degenerate case`);
-                        // Emit zero gradient structure
-                        const zeroGrads = [];
-                        for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
-                            if (isStructuredGradient(gradient)) {
-                                const components = Array.from(gradient.components.keys());
-                                const zeroStruct = components.map(c => `${c}: 0`).join(', ');
-                                zeroGrads.push(`d${paramName}: { ${zeroStruct} }`);
-                            }
-                            else {
-                                zeroGrads.push(`d${paramName}: 0`);
-                            }
-                        }
-                        lines.push(`    return { value, ${zeroGrads.join(', ')} };`);
-                        lines.push(`  }`);
-                    }
+            }
+        }
+        // Post-CSE simplification: apply rules that were skipped to avoid CSE interference
+        // Specifically: a + a → 2 * a (now safe because temps have been extracted)
+        for (const [varName, expr] of cseIntermediates) {
+            cseIntermediates.set(varName, simplifyPostCSE(expr));
+        }
+        for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
+            if (isStructuredGradient(gradient)) {
+                for (const [comp, expr] of gradient.components.entries()) {
+                    gradient.components.set(comp, simplifyPostCSE(expr));
                 }
             }
-            lines.push('');
+            else {
+                // Scalar gradient - apply post-CSE simplification
+                gradientsToUse.gradients.set(paramName, simplifyPostCSE(gradient));
+            }
+        }
+    }
+    // Generate CSE intermediate variables
+    if (cseIntermediates.size > 0) {
+        for (const [varName, expr] of cseIntermediates.entries()) {
+            const code = codegen.generate(expr);
+            if (format === 'typescript' || format === 'javascript') {
+                lines.push(`  const ${varName} = ${code};`);
+            }
+            else if (format === 'python') {
+                lines.push(`  ${varName} = ${code}`);
+            }
+            else if (format === 'csharp') {
+                lines.push(`    ${csharpFloatType} ${varName} = ${code};`);
+            }
         }
+        lines.push('');
     }
     for (const [paramName, gradient] of gradientsToUse.gradients.entries()) {
         // Use shorter names: du, dv instead of grad_u, grad_v
@@ -530,11 +692,12 @@ export function generateGradientFunction(func, gradients, env, options = {}) {
     return lines.join('\n');
 }
 /**
- * Generate the original forward function
+ * Generate the original forward function (with optional e-graph optimization)
  */
 export function generateForwardFunction(func, options = {}) {
     const format = options.format || 'typescript';
     const csharpFloatType = options.csharpFloatType || 'float';
+    const shouldOptimize = options.cse !== false; // Optimize by default
     const codegen = new ExpressionCodeGen(format, csharpFloatType);
     const lines = [];
     // Function signature
@@ -549,28 +712,127 @@ export function generateForwardFunction(func, options = {}) {
         const floatType = csharpFloatType;
         const params = func.parameters.map(p => {
             if (p.paramType && p.paramType.components) {
-                // Structured parameter - create a struct type name
                 return `${capitalize(p.name)}Struct ${p.name}`;
             }
             return `${floatType} ${p.name}`;
         }).join(', ');
-        // Generate struct definitions for structured parameters first (we'll prepend them later)
         lines.push(`public static ${floatType} ${capitalize(func.name)}(${params})`);
         lines.push('{');
     }
-    // Body
+    // Collect all forward variable names (for dependency tracking)
+    const forwardVars = new Set();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            forwardVars.add(stmt.variable);
+        }
+    }
+    // Collect expressions for optimization
+    const varExpressions = new Map();
+    for (const stmt of func.body) {
+        if (stmt.kind === 'assignment') {
+            varExpressions.set(stmt.variable, stmt.expression);
+        }
+    }
+    // Optimize with e-graph if enabled
+    let optimizedExprs = varExpressions;
+    let cseTemps = new Map();
+    if (shouldOptimize && varExpressions.size > 0) {
+        const forOptimizer = new Map();
+        forOptimizer.set('_forward', varExpressions);
+        const result = optimizeWithEGraph(forOptimizer, { verbose: false });
+        cseTemps = result.intermediates;
+        optimizedExprs = result.gradients.get('_forward') || varExpressions;
+    }
+    // Helper to find which forward vars an expression depends on
+    function findForwardVarDeps(expr) {
+        const deps = new Set();
+        function visit(e) {
+            if (e.kind === 'variable' && forwardVars.has(e.name)) {
+                deps.add(e.name);
+            }
+            else if (e.kind === 'binary') {
+                visit(e.left);
+                visit(e.right);
+            }
+            else if (e.kind === 'unary') {
+                visit(e.operand);
+            }
+            else if (e.kind === 'call') {
+                e.args.forEach(visit);
+            }
+            else if (e.kind === 'component') {
+                visit(e.object);
+            }
+        }
+        visit(expr);
+        return deps;
+    }
+    // Track which temps need to be emitted after which forward var
+    const tempAfterVar = new Map();
+    const tempsEmittedBeforeAny = [];
+    for (const [tempName, tempExpr] of cseTemps) {
+        const deps = findForwardVarDeps(tempExpr);
+        if (deps.size === 0) {
+            // Temp only depends on params - emit before any forward vars
+            tempsEmittedBeforeAny.push({ name: tempName, expr: tempExpr });
+        }
+        else {
+            // Find the last forward var this temp depends on
+            let lastDep = '';
+            for (const stmt of func.body) {
+                if (stmt.kind === 'assignment' && deps.has(stmt.variable)) {
+                    lastDep = stmt.variable;
+                }
+            }
+            if (lastDep) {
+                if (!tempAfterVar.has(lastDep)) {
+                    tempAfterVar.set(lastDep, []);
+                }
+                tempAfterVar.get(lastDep).push({ name: tempName, expr: tempExpr });
+            }
+        }
+    }
+    // Emit temps that don't depend on forward vars
+    for (const { name: tempName, expr } of tempsEmittedBeforeAny) {
+        const code = codegen.generate(expr);
+        if (format === 'typescript' || format === 'javascript') {
+            lines.push(`  const ${tempName} = ${code};`);
+        }
+        else if (format === 'python') {
+            lines.push(`  ${tempName} = ${code}`);
+        }
+        else if (format === 'csharp') {
+            lines.push(`    ${csharpFloatType} ${tempName} = ${code};`);
+        }
+    }
+    // Generate variable assignments with interleaved temps
     for (const stmt of func.body) {
         if (stmt.kind === 'assignment') {
             const varName = stmt.variable;
-            const expr = codegen.generate(stmt.expression);
+            const expr = optimizedExprs.get(varName) || stmt.expression;
+            const code = codegen.generate(expr);
             if (format === 'typescript' || format === 'javascript') {
-                lines.push(`  const ${varName} = ${expr};`);
+                lines.push(`  const ${varName} = ${code};`);
             }
             else if (format === 'python') {
-                lines.push(`  ${varName} = ${expr}`);
+                lines.push(`  ${varName} = ${code}`);
             }
             else if (format === 'csharp') {
-                lines.push(`    ${csharpFloatType} ${varName} = ${expr};`);
+                lines.push(`    ${csharpFloatType} ${varName} = ${code};`);
+            }
+            // Emit any temps that depend on this var
+            const tempsForVar = tempAfterVar.get(varName) || [];
+            for (const { name: tempName, expr: tempExpr } of tempsForVar) {
+                const tempCode = codegen.generate(tempExpr);
+                if (format === 'typescript' || format === 'javascript') {
+                    lines.push(`  const ${tempName} = ${tempCode};`);
+                }
+                else if (format === 'python') {
+                    lines.push(`  ${tempName} = ${tempCode}`);
+                }
+                else if (format === 'csharp') {
+                    lines.push(`    ${csharpFloatType} ${tempName} = ${tempCode};`);
+                }
             }
         }
     }

package/dist/dsl/ExpressionUtils.d.ts

@@ -54,10 +54,16 @@ export declare function containsVariable(expr: Expression, varName: string): boo
  */
 export declare function expressionDepth(expr: Expression): number;
 /**
- * Serializes an expression to canonical string representation.
- * Used for expression comparison and hashing (CSE, CodeGen forward reuse).
+ * Serializes an expression to structural string representation.
+ * Used for exact expression comparison - operand order matters.
  *
  * This ensures consistent string representation of expressions across different
  * parts of the codebase.
  */
 export declare function serializeExpression(expr: Expression): string;
+/**
+ * Serializes an expression to canonical form for CSE matching.
+ * Commutative operations (+ and *) have operands sorted lexicographically,
+ * so a*b and b*a produce the same canonical string.
+ */
+export declare function serializeCanonical(expr: Expression): string;

package/dist/dsl/ExpressionUtils.js

@@ -174,8 +174,8 @@ export function expressionDepth(expr) {
     }
 }
 /**
- * Serializes an expression to canonical string representation.
- * Used for expression comparison and hashing (CSE, CodeGen forward reuse).
+ * Serializes an expression to structural string representation.
+ * Used for exact expression comparison - operand order matters.
  *
  * This ensures consistent string representation of expressions across different
  * parts of the codebase.
@@ -197,3 +197,35 @@ export function serializeExpression(expr) {
             return `comp(${serializeExpression(expr.object)},${expr.component})`;
     }
 }
+/**
+ * Serializes an expression to canonical form for CSE matching.
+ * Commutative operations (+ and *) have operands sorted lexicographically,
+ * so a*b and b*a produce the same canonical string.
+ */
+export function serializeCanonical(expr) {
+    switch (expr.kind) {
+        case 'number':
+            return `num(${expr.value})`;
+        case 'variable':
+            return `var(${expr.name})`;
+        case 'binary': {
+            const leftStr = serializeCanonical(expr.left);
+            const rightStr = serializeCanonical(expr.right);
+            // For commutative operations, sort operands lexicographically
+            if (expr.operator === '+' || expr.operator === '*') {
+                const [first, second] = leftStr <= rightStr ? [leftStr, rightStr] : [rightStr, leftStr];
+                return `bin(${expr.operator},${first},${second})`;
+            }
+            // Non-commutative: preserve order
+            return `bin(${expr.operator},${leftStr},${rightStr})`;
+        }
+        case 'unary':
+            return `un(${expr.operator},${serializeCanonical(expr.operand)})`;
+        case 'call': {
+            const args = expr.args.map(arg => serializeCanonical(arg)).join(',');
+            return `call(${expr.name},${args})`;
+        }
+        case 'component':
+            return `comp(${serializeCanonical(expr.object)},${expr.component})`;
+    }
+}

package/dist/dsl/GradientChecker.d.ts

@@ -17,9 +17,25 @@ type NumValue = number | {
 export interface GradCheckResult {
     passed: boolean;
     errors: GradCheckError[];
+    singularities: GradCheckSingularity[];
     maxError: number;
     meanError: number;
+    totalChecks: number;
 }
+/**
+ * Singularity detected during gradient checking
+ * When both analytical and numerical produce NaN/Inf, it's a singularity, not a bug
+ */
+export interface GradCheckSingularity {
+    parameter: string;
+    component?: string;
+    analytical: number;
+    numerical: number;
+}
+/**
+ * Format gradient check results as a human-readable string
+ */
+export declare function formatGradCheckResult(result: GradCheckResult, funcName: string): string;
 export interface GradCheckError {
     parameter: string;
     component?: string;
@@ -39,6 +55,11 @@ export declare class GradientChecker {
      * Check gradients for a function
      */
     check(func: FunctionDef, gradients: GradientResult, env: TypeEnv, testPoint: Map<string, NumValue>): GradCheckResult;
+    /**
+     * Compare analytical and numerical gradients
+     * Distinguishes between: pass, error (mismatch), and singularity (both NaN/Inf)
+     */
+    private compareGradients;
     /**
      * Compute numerical gradient for scalar parameter using finite differences
      */