gradient-script 0.2.0 → 0.3.1

package/dist/dsl/GradientChecker.js

@@ -4,6 +4,47 @@
  */
 import { Types } from './Types.js';
 import { expandBuiltIn, shouldExpand } from './Expander.js';
+/**
+ * Format gradient check results as a human-readable string
+ */
+export function formatGradCheckResult(result, funcName) {
+    const singularityCount = result.singularities.length;
+    const verifiedCount = result.totalChecks - result.errors.length - singularityCount;
+    if (result.passed) {
+        let msg = `✓ ${funcName}: ${verifiedCount} gradients verified (max error: ${result.maxError.toExponential(2)})`;
+        if (singularityCount > 0) {
+            msg += `\n  ⚠ ${singularityCount} singularities detected (both analytical and numerical produce NaN/Inf)`;
+        }
+        return msg;
+    }
+    const lines = [
+        `✗ ${funcName}: ${result.errors.length}/${result.totalChecks} gradients FAILED`
+    ];
+    // Group errors by parameter
+    const byParam = new Map();
+    for (const err of result.errors) {
+        const key = err.parameter;
+        if (!byParam.has(key))
+            byParam.set(key, []);
+        byParam.get(key).push(err);
+    }
+    for (const [param, errs] of byParam) {
+        if (errs.length === 1 && !errs[0].component) {
+            // Scalar parameter
+            const e = errs[0];
+            lines.push(`  ${param}: analytical=${e.analytical.toFixed(6)}, numerical=${e.numerical.toFixed(6)}, error=${e.error.toExponential(2)}`);
+        }
+        else {
+            // Structured parameter - show on one line if possible
+            const components = errs.map(e => `${e.component}:${e.error.toExponential(1)}`).join(', ');
+            lines.push(`  ${param}: {${components}}`);
+        }
+    }
+    if (singularityCount > 0) {
+        lines.push(`  ⚠ ${singularityCount} singularities also detected`);
+    }
+    return lines.join('\n');
+}
 /**
  * Gradient checker
  */
@@ -19,6 +60,9 @@ export class GradientChecker {
      */
     check(func, gradients, env, testPoint) {
         const errors = [];
+        const singularities = [];
+        let totalChecks = 0;
+        let maxError = 0;
         // For each parameter that has gradients
         for (const [paramName, gradient] of gradients.gradients.entries()) {
             const paramType = env.getOrThrow(paramName);
@@ -28,21 +72,21 @@ export class GradientChecker {
             }
             if (Types.isScalar(paramType)) {
                 // Scalar parameter
+                totalChecks++;
                 if (typeof paramValue !== 'number') {
                     throw new Error(`Expected scalar value for ${paramName}`);
                 }
                 const analytical = this.evaluateExpression(gradient, testPoint);
                 const numerical = this.numericalGradientScalar(func, testPoint, paramName);
-                const error = Math.abs(analytical - numerical);
-                const relativeError = Math.abs(error / (numerical + 1e-10));
-                if (error > this.tolerance && relativeError > this.tolerance) {
-                    errors.push({
-                        parameter: paramName,
-                        analytical,
-                        numerical,
-                        error,
-                        relativeError
-                    });
+                const checkResult = this.compareGradients(analytical, numerical, paramName);
+                if (checkResult.type === 'singularity') {
+                    singularities.push(checkResult.singularity);
+                }
+                else if (checkResult.type === 'error') {
+                    errors.push(checkResult.error);
+                }
+                else if (checkResult.error) {
+                    maxError = Math.max(maxError, checkResult.error.error);
                 }
             }
             else {
@@ -52,32 +96,74 @@ export class GradientChecker {
                 }
                 const structGrad = gradient;
                 for (const [comp, expr] of structGrad.components.entries()) {
+                    totalChecks++;
                     const analytical = this.evaluateExpression(expr, testPoint);
                     const numerical = this.numericalGradientComponent(func, testPoint, paramName, comp);
-                    const error = Math.abs(analytical - numerical);
-                    const relativeError = Math.abs(error / (numerical + 1e-10));
-                    if (error > this.tolerance && relativeError > this.tolerance) {
-                        errors.push({
-                            parameter: paramName,
-                            component: comp,
-                            analytical,
-                            numerical,
-                            error,
-                            relativeError
-                        });
+                    const checkResult = this.compareGradients(analytical, numerical, paramName, comp);
+                    if (checkResult.type === 'singularity') {
+                        singularities.push(checkResult.singularity);
+                    }
+                    else if (checkResult.type === 'error') {
+                        errors.push(checkResult.error);
+                    }
+                    else if (checkResult.error) {
+                        maxError = Math.max(maxError, checkResult.error.error);
                     }
                 }
             }
         }
-        const maxError = errors.length > 0 ? Math.max(...errors.map(e => e.error)) : 0;
         const meanError = errors.length > 0
             ? errors.reduce((sum, e) => sum + e.error, 0) / errors.length
             : 0;
         return {
             passed: errors.length === 0,
             errors,
+            singularities,
             maxError,
-            meanError
+            meanError,
+            totalChecks
+        };
+    }
+    /**
+     * Compare analytical and numerical gradients
+     * Distinguishes between: pass, error (mismatch), and singularity (both NaN/Inf)
+     */
+    compareGradients(analytical, numerical, parameter, component) {
+        const analyticalBad = !isFinite(analytical);
+        const numericalBad = !isFinite(numerical);
+        // Both produce NaN/Inf: singularity (not a bug)
+        if (analyticalBad && numericalBad) {
+            return {
+                type: 'singularity',
+                singularity: { parameter, component, analytical, numerical }
+            };
+        }
+        // One produces NaN/Inf, the other doesn't: actual bug
+        if (analyticalBad || numericalBad) {
+            return {
+                type: 'error',
+                error: {
+                    parameter,
+                    component,
+                    analytical,
+                    numerical,
+                    error: Infinity,
+                    relativeError: Infinity
+                }
+            };
+        }
+        // Both finite: compare values
+        const error = Math.abs(analytical - numerical);
+        const relativeError = Math.abs(error / (Math.abs(numerical) + 1e-10));
+        if (error > this.tolerance && relativeError > this.tolerance) {
+            return {
+                type: 'error',
+                error: { parameter, component, analytical, numerical, error, relativeError }
+            };
+        }
+        return {
+            type: 'pass',
+            error: { parameter, component, analytical, numerical, error, relativeError }
         };
     }
     /**

package/dist/dsl/Guards.d.ts

@@ -22,7 +22,7 @@ export declare function analyzeGuards(func: FunctionDef): GuardAnalysisResult;
 /**
  * Format guard analysis for display
  */
-export declare function formatGuardWarnings(result: GuardAnalysisResult): string;
+export declare function formatGuardWarnings(result: GuardAnalysisResult, asComments?: boolean): string;
 /**
  * Generate guard code snippets for common cases
  */

package/dist/dsl/Guards.js

@@ -154,21 +154,22 @@ function isSqExpression(expr) {
 /**
  * Format guard analysis for display
  */
-export function formatGuardWarnings(result) {
+export function formatGuardWarnings(result, asComments = false) {
     if (!result.hasIssues) {
         return '';
     }
+    const prefix = asComments ? '// ' : '';
     const lines = [];
-    lines.push('⚠️  EDGE CASE WARNINGS:');
-    lines.push('');
-    lines.push('The generated code may encounter edge cases that produce');
-    lines.push('NaN, Infinity, or incorrect gradients:');
-    lines.push('');
+    lines.push(`${prefix}⚠️  EDGE CASE WARNINGS:`);
+    lines.push(prefix);
+    lines.push(`${prefix}The generated code may encounter edge cases that produce`);
+    lines.push(`${prefix}NaN, Infinity, or incorrect gradients:`);
+    lines.push(prefix);
     // Show each guard individually with context
     for (const guard of result.guards) {
         const typeLabel = formatGuardType(guard.type);
         // Show location and variable if available
-        let location = '  •';
+        let location = `${prefix}  •`;
         if (guard.line) {
             location += ` Line ${guard.line}:`;
         }
@@ -176,13 +177,13 @@ export function formatGuardWarnings(result) {
             location += ` ${guard.variableName} =`;
         }
         lines.push(location);
-        lines.push(`    ${typeLabel}: ${guard.description}`);
-        lines.push(`    💡 Fix: ${guard.suggestion}`);
-        lines.push('');
+        lines.push(`${prefix}    ${typeLabel}: ${guard.description}`);
+        lines.push(`${prefix}    💡 Fix: ${guard.suggestion}`);
+        lines.push(prefix);
     }
-    lines.push('Add runtime checks or ensure inputs are within valid ranges.');
-    lines.push('Use --guards --epsilon 1e-10 to automatically emit epsilon guards.');
-    lines.push('');
+    lines.push(`${prefix}Add runtime checks or ensure inputs are within valid ranges.`);
+    lines.push(`${prefix}Use --guards --epsilon 1e-10 to automatically emit epsilon guards.`);
+    lines.push(prefix);
     return lines.join('\n');
 }
 function formatGuardType(type) {

package/dist/dsl/Inliner.d.ts

@@ -8,3 +8,8 @@ import { Expression, FunctionDef } from './AST.js';
  * Returns a new expression with all intermediate variables substituted
  */
 export declare function inlineIntermediateVariables(func: FunctionDef): Expression;
+/**
+ * Inline an expression using a substitution map
+ * Used to get the fully-expanded form of forward pass expressions
+ */
+export declare function inlineExpression(expr: Expression, substitutions: Map<string, Expression>): Expression;

package/dist/dsl/Inliner.js

@@ -39,3 +39,11 @@ export function inlineIntermediateVariables(func) {
     const transformer = new VariableSubstitutionTransformer(substitutions);
     return transformer.transform(func.returnExpr);
 }
+/**
+ * Inline an expression using a substitution map
+ * Used to get the fully-expanded form of forward pass expressions
+ */
+export function inlineExpression(expr, substitutions) {
+    const transformer = new VariableSubstitutionTransformer(substitutions);
+    return transformer.transform(expr);
+}

package/dist/dsl/Simplify.d.ts

@@ -15,3 +15,10 @@ export declare function simplifyGradients(gradients: Map<string, Expression | {
 }>): Map<string, Expression | {
     components: Map<string, Expression>;
 }>;
+/**
+ * Post-CSE simplification: applies rules that were intentionally skipped during
+ * initial simplification to avoid interfering with CSE.
+ *
+ * Specifically: a + a → 2 * a (now safe because CSE has already extracted temps)
+ */
+export declare function simplifyPostCSE(expr: Expression): Expression;

package/dist/dsl/Simplify.js

@@ -46,6 +46,9 @@ class Simplifier extends ExpressionTransformer {
                 return right;
             if (rightNum === 0)
                 return left;
+            // Note: a + a → 2 * a rules are intentionally NOT applied here
+            // because they flatten expression structure and interfere with CSE.
+            // The CSE pass will extract common subexpressions instead.
         }
         // Subtraction rules
         if (expr.operator === '-') {
@@ -57,6 +60,25 @@ class Simplifier extends ExpressionTransformer {
             if (expressionsEqual(left, right)) {
                 return { kind: 'number', value: 0 };
             }
+            // a - (-b) → a + b
+            if (right.kind === 'unary' && right.operator === '-') {
+                return this.transform({
+                    kind: 'binary',
+                    operator: '+',
+                    left,
+                    right: right.operand
+                });
+            }
+            // (-a) - (-b) → b - a
+            if (left.kind === 'unary' && left.operator === '-' &&
+                right.kind === 'unary' && right.operator === '-') {
+                return this.transform({
+                    kind: 'binary',
+                    operator: '-',
+                    left: right.operand,
+                    right: left.operand
+                });
+            }
         }
         // Multiplication rules
         if (expr.operator === '*') {
@@ -68,6 +90,50 @@ class Simplifier extends ExpressionTransformer {
                 return right;
             if (rightNum === 1)
                 return left;
+            // -1 * x → -x
+            if (leftNum === -1) {
+                return { kind: 'unary', operator: '-', operand: right };
+            }
+            // x * -1 → -x
+            if (rightNum === -1) {
+                return { kind: 'unary', operator: '-', operand: left };
+            }
+            // (-a) * (-b) → a * b
+            if (left.kind === 'unary' && left.operator === '-' &&
+                right.kind === 'unary' && right.operator === '-') {
+                return this.transform({
+                    kind: 'binary',
+                    operator: '*',
+                    left: left.operand,
+                    right: right.operand
+                });
+            }
+            // (-a) * b → -(a * b)
+            if (left.kind === 'unary' && left.operator === '-') {
+                return {
+                    kind: 'unary',
+                    operator: '-',
+                    operand: this.transform({
+                        kind: 'binary',
+                        operator: '*',
+                        left: left.operand,
+                        right
+                    })
+                };
+            }
+            // a * (-b) → -(a * b)
+            if (right.kind === 'unary' && right.operator === '-') {
+                return {
+                    kind: 'unary',
+                    operator: '-',
+                    operand: this.transform({
+                        kind: 'binary',
+                        operator: '*',
+                        left,
+                        right: right.operand
+                    })
+                };
+            }
             // (x / x) * y → y
             if (left.kind === 'binary' && left.operator === '/') {
                 if (expressionsEqual(left.left, left.right)) {
@@ -116,6 +182,42 @@ class Simplifier extends ExpressionTransformer {
             if (expressionsEqual(left, right)) {
                 return { kind: 'number', value: 1 };
             }
+            // (-a) / (-b) → a / b
+            if (left.kind === 'unary' && left.operator === '-' &&
+                right.kind === 'unary' && right.operator === '-') {
+                return this.transform({
+                    kind: 'binary',
+                    operator: '/',
+                    left: left.operand,
+                    right: right.operand
+                });
+            }
+            // (-a) / b → -(a / b)
+            if (left.kind === 'unary' && left.operator === '-') {
+                return {
+                    kind: 'unary',
+                    operator: '-',
+                    operand: this.transform({
+                        kind: 'binary',
+                        operator: '/',
+                        left: left.operand,
+                        right
+                    })
+                };
+            }
+            // a / (-b) → -(a / b)
+            if (right.kind === 'unary' && right.operator === '-') {
+                return {
+                    kind: 'unary',
+                    operator: '-',
+                    operand: this.transform({
+                        kind: 'binary',
+                        operator: '/',
+                        left,
+                        right: right.operand
+                    })
+                };
+            }
             // (a + a) / 2 → a
             if (rightNum === 2 && left.kind === 'binary' && left.operator === '+') {
                 if (expressionsEqual(left.left, left.right)) {
@@ -321,3 +423,37 @@ export function simplifyGradients(gradients) {
     }
     return simplified;
 }
+/**
+ * Post-CSE simplification: applies rules that were intentionally skipped during
+ * initial simplification to avoid interfering with CSE.
+ *
+ * Specifically: a + a → 2 * a (now safe because CSE has already extracted temps)
+ */
+export function simplifyPostCSE(expr) {
+    return new PostCSESimplifier().transform(expr);
+}
+class PostCSESimplifier extends ExpressionTransformer {
+    visitBinaryOp(expr) {
+        const left = this.transform(expr.left);
+        const right = this.transform(expr.right);
+        // a + a → 2 * a
+        if (expr.operator === '+' && expressionsEqual(left, right)) {
+            return {
+                kind: 'binary',
+                operator: '*',
+                left: { kind: 'number', value: 2 },
+                right: left
+            };
+        }
+        // Return simplified if no changes
+        if (left === expr.left && right === expr.right) {
+            return expr;
+        }
+        return {
+            kind: 'binary',
+            operator: expr.operator,
+            left,
+            right
+        };
+    }
+}

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

@@ -0,0 +1,23 @@
+/**
+ * Conversion between AST Expressions and E-Graph
+ */
+import { EGraph } from './EGraph.js';
+import { EClassId } from './ENode.js';
+import { Expression } from '../AST.js';
+/**
+ * Add an AST Expression to the e-graph, returning its e-class ID
+ */
+export declare function addExpression(egraph: EGraph, expr: Expression): EClassId;
+/**
+ * Add multiple expressions, returning a map of original keys to e-class IDs
+ */
+export declare function addExpressions<K extends string>(egraph: EGraph, expressions: Map<K, Expression>): Map<K, EClassId>;
+/**
+ * Add all gradients (Map<paramName, Map<component, Expression>>)
+ * Returns Map<paramName, Map<component, EClassId>>
+ */
+export declare function addGradients(egraph: EGraph, gradients: Map<string, Map<string, Expression>>): Map<string, Map<string, EClassId>>;
+/**
+ * Get all root e-class IDs from gradient structure
+ */
+export declare function getRootIds(gradientIds: Map<string, Map<string, EClassId>>): EClassId[];

package/dist/dsl/egraph/Convert.js

@@ -0,0 +1,84 @@
+/**
+ * Conversion between AST Expressions and E-Graph
+ */
+/**
+ * Add an AST Expression to the e-graph, returning its e-class ID
+ */
+export function addExpression(egraph, expr) {
+    switch (expr.kind) {
+        case 'number':
+            return egraph.add({ tag: 'num', value: expr.value });
+        case 'variable':
+            return egraph.add({ tag: 'var', name: expr.name });
+        case 'binary': {
+            const left = addExpression(egraph, expr.left);
+            const right = addExpression(egraph, expr.right);
+            switch (expr.operator) {
+                case '+':
+                    return egraph.add({ tag: 'add', children: [left, right] });
+                case '-':
+                    return egraph.add({ tag: 'sub', children: [left, right] });
+                case '*':
+                    return egraph.add({ tag: 'mul', children: [left, right] });
+                case '/':
+                    return egraph.add({ tag: 'div', children: [left, right] });
+                case '^':
+                case '**':
+                    return egraph.add({ tag: 'pow', children: [left, right] });
+            }
+        }
+        case 'unary': {
+            const operand = addExpression(egraph, expr.operand);
+            if (expr.operator === '-') {
+                return egraph.add({ tag: 'neg', child: operand });
+            }
+            // Unary + is identity
+            return operand;
+        }
+        case 'call': {
+            const args = expr.args.map(arg => addExpression(egraph, arg));
+            return egraph.add({ tag: 'call', name: expr.name, children: args });
+        }
+        case 'component': {
+            const object = addExpression(egraph, expr.object);
+            return egraph.add({ tag: 'component', object, field: expr.component });
+        }
+    }
+}
+/**
+ * Add multiple expressions, returning a map of original keys to e-class IDs
+ */
+export function addExpressions(egraph, expressions) {
+    const result = new Map();
+    for (const [key, expr] of expressions) {
+        result.set(key, addExpression(egraph, expr));
+    }
+    return result;
+}
+/**
+ * Add all gradients (Map<paramName, Map<component, Expression>>)
+ * Returns Map<paramName, Map<component, EClassId>>
+ */
+export function addGradients(egraph, gradients) {
+    const result = new Map();
+    for (const [paramName, components] of gradients) {
+        const componentIds = new Map();
+        for (const [comp, expr] of components) {
+            componentIds.set(comp, addExpression(egraph, expr));
+        }
+        result.set(paramName, componentIds);
+    }
+    return result;
+}
+/**
+ * Get all root e-class IDs from gradient structure
+ */
+export function getRootIds(gradientIds) {
+    const roots = [];
+    for (const components of gradientIds.values()) {
+        for (const id of components.values()) {
+            roots.push(id);
+        }
+    }
+    return roots;
+}

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

@@ -0,0 +1,93 @@
+/**
+ * E-Graph: Equality Graph for expression optimization
+ *
+ * An e-graph efficiently represents equivalence classes of expressions.
+ * It supports:
+ * - Adding expressions (returns e-class ID)
+ * - Merging e-classes (union)
+ * - Finding canonical e-class (find)
+ * - Rebuilding after merges (maintains congruence)
+ */
+import { ENode, EClassId } from './ENode.js';
+/**
+ * E-Class: An equivalence class of expressions
+ */
+export interface EClass {
+    id: EClassId;
+    nodes: Set<string>;
+    parents: Set<string>;
+}
+/**
+ * E-Graph: The main data structure
+ */
+export declare class EGraph {
+    private nextId;
+    private classes;
+    private parent;
+    private rank;
+    private hashcons;
+    private nodeStore;
+    private pending;
+    /**
+     * Find the canonical e-class ID (with path compression)
+     */
+    find(id: EClassId): EClassId;
+    /**
+     * Add an e-node to the e-graph, returning its e-class ID
+     * If the node already exists, returns the existing class
+     */
+    add(node: ENode): EClassId;
+    /**
+     * Merge two e-classes, returning the new canonical ID
+     */
+    merge(id1: EClassId, id2: EClassId): EClassId;
+    /**
+     * Rebuild the e-graph to restore congruence invariants
+     * Must be called after a batch of merges
+     */
+    rebuild(): void;
+    /**
+     * Repair an e-class after merges
+     */
+    private repair;
+    /**
+     * Find which e-class contains a node (by key)
+     */
+    private findClassForNode;
+    /**
+     * Canonicalize an e-node (update children to canonical IDs)
+     */
+    private canonicalize;
+    /**
+     * Get all e-class IDs
+     */
+    getClassIds(): EClassId[];
+    /**
+     * Get an e-class by ID
+     */
+    getClass(id: EClassId): EClass | undefined;
+    /**
+     * Get all e-nodes in an e-class
+     */
+    getNodes(classId: EClassId): ENode[];
+    /**
+     * Get the number of e-classes
+     */
+    get size(): number;
+    /**
+     * Get a node by its key
+     */
+    getNodeByKey(key: string): ENode | undefined;
+    /**
+     * Lookup e-class by node (if it exists)
+     */
+    lookup(node: ENode): EClassId | undefined;
+    /**
+     * Debug: print e-graph state
+     */
+    dump(): string;
+    /**
+     * Convert e-node to readable string
+     */
+    private nodeToString;
+}