gradient-script 0.1.0 → 0.3.0

package/dist/dsl/Inliner.js

@@ -6,8 +6,9 @@ import { ExpressionTransformer } from './ExpressionTransformer.js';
 /**
  * Expression transformer that substitutes variables from a substitution map
  * Handles recursive inlining by reprocessing substituted expressions
+ * Used for inlining intermediate variables to eliminate assignments
  */
-class SubstitutionTransformer extends ExpressionTransformer {
+class VariableSubstitutionTransformer extends ExpressionTransformer {
     substitutions;
     constructor(substitutions) {
         super();
@@ -35,6 +36,14 @@ export function inlineIntermediateVariables(func) {
         }
     }
     // Use transformer to inline all variables
-    const transformer = new SubstitutionTransformer(substitutions);
+    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/Lexer.js

@@ -2,6 +2,7 @@
  * Lexer for GradientScript DSL
  * Tokenizes input with support for ∇, ^, **, and structured types
  */
+import { ParseError } from './Errors.js';
 export var TokenType;
 (function (TokenType) {
     // Literals
@@ -126,7 +127,8 @@ export class Lexer {
                 this.advance();
                 return { type: TokenType.NEWLINE, value: '\n', line, column };
         }
-        throw new Error(`Unexpected character '${char}' at line ${line}, column ${column}`);
+        // Create a more helpful error for common mistakes
+        throw new ParseError(`Unexpected character '${char}'`, line, column, char);
     }
     number() {
         const line = this.line;

package/dist/dsl/Parser.js

@@ -123,13 +123,15 @@ export class Parser {
      * variable = expression
      */
     assignment() {
-        const variable = this.consume(TokenType.IDENTIFIER, 'Expected variable name').value;
+        const varToken = this.consume(TokenType.IDENTIFIER, 'Expected variable name');
+        const variable = varToken.value;
         this.consume(TokenType.EQUALS, "Expected '=' in assignment");
         const expression = this.expression();
         return {
             kind: 'assignment',
             variable,
-            expression
+            expression,
+            loc: { line: varToken.line, column: varToken.column }
         };
     }
     /**
@@ -161,13 +163,15 @@ export class Parser {
     multiplicative() {
         let expr = this.power();
         while (this.match(TokenType.MULTIPLY, TokenType.DIVIDE)) {
-            const operator = this.previous().value;
+            const opToken = this.previous();
+            const operator = opToken.value;
             const right = this.power();
             expr = {
                 kind: 'binary',
                 operator,
                 left: expr,
-                right
+                right,
+                loc: { line: opToken.line, column: opToken.column }
             };
         }
         return expr;
@@ -213,6 +217,7 @@ export class Parser {
         while (true) {
             if (this.match(TokenType.LPAREN)) {
                 // Function call
+                const startLoc = expr.loc || { line: this.previous().line, column: this.previous().column };
                 const args = this.argumentList();
                 this.consume(TokenType.RPAREN, "Expected ')' after arguments");
                 if (expr.kind !== 'variable') {
@@ -222,7 +227,8 @@ export class Parser {
                 expr = {
                     kind: 'call',
                     name: expr.name,
-                    args
+                    args,
+                    loc: startLoc
                 };
             }
             else if (this.match(TokenType.DOT)) {

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,89 @@ 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)) {
+                    return left.left;
+                }
+            }
+            // (a + a) / (2 * b) → a / b
+            if (right.kind === 'binary' && right.operator === '*') {
+                const rightLeft = right.left;
+                const rightRight = right.right;
+                const rightLeftNum = isNumber(rightLeft) ? rightLeft.value : null;
+                if (rightLeftNum === 2 && left.kind === 'binary' && left.operator === '+') {
+                    if (expressionsEqual(left.left, left.right)) {
+                        return {
+                            kind: 'binary',
+                            operator: '/',
+                            left: left.left,
+                            right: rightRight
+                        };
+                    }
+                    // (-1 * a + a * -1) / (2 * b) → -a / b
+                    const leftLeft = left.left;
+                    const leftRight = left.right;
+                    if (leftLeft.kind === 'binary' && leftLeft.operator === '*' &&
+                        leftRight.kind === 'binary' && leftRight.operator === '*') {
+                        const ll_left = leftLeft.left;
+                        const ll_right = leftLeft.right;
+                        const lr_left = leftRight.left;
+                        const lr_right = leftRight.right;
+                        const ll_leftNum = isNumber(ll_left) ? ll_left.value : null;
+                        const lr_rightNum = isNumber(lr_right) ? lr_right.value : null;
+                        // (-1 * a) + (a * -1)
+                        if (ll_leftNum === -1 && lr_rightNum === -1 && expressionsEqual(ll_right, lr_left)) {
+                            return {
+                                kind: 'unary',
+                                operator: '-',
+                                operand: {
+                                    kind: 'binary',
+                                    operator: '/',
+                                    left: ll_right,
+                                    right: rightRight
+                                }
+                            };
+                        }
+                    }
+                }
+            }
         }
         // Power rules
         if (expr.operator === '^' || expr.operator === '**') {
@@ -274,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;
+}