diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..8fff6a9 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,11 @@ +# Cargo configuration for petite-ad. +# +# The mold linker flag below speeds up local builds substantially. +# If mold is not installed, the flag is silently ignored by the linker. +# Remove or adjust the rustflags if needed for your machine. + +[target.x86_64-unknown-linux-gnu] +rustflags = ["-C", "link-arg=-fuse-ld=mold"] + +[registries.crates-io] +protocol = "sparse" # Faster crate downloads diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..06cd551 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,104 @@ +name: CI + +on: + push: + branches: [main, master, develop] + pull_request: + branches: [main, master] + +env: + CARGO_TERM_COLOR: always + +jobs: + test: + name: Test + runs-on: ubuntu-latest + strategy: + matrix: + rust: [stable, nightly] + + steps: + - uses: actions/checkout@v4 + + - name: Setup mold linker + uses: rui314/setup-mold@v1 + + - name: Install Rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: ${{ matrix.rust }} + components: rustfmt, clippy + + - name: Setup Rust cache + uses: Swatinem/rust-cache@v2 + with: + key: ${{ matrix.rust }} + + - name: Check formatting + run: cargo fmt -- --check + + - name: Run clippy + run: cargo clippy --all-targets --all-features -- -D warnings + + - name: Run tests + run: cargo test --verbose --all-features + + - name: Run examples + run: cargo test --examples --all-features + + - name: Run benchmarks (check only) + run: cargo bench --no-run + + coverage: + name: Code Coverage + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Setup mold linker + uses: rui314/setup-mold@v1 + + - name: Install Rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + + - name: Setup Rust cache + uses: Swatinem/rust-cache@v2 + + - name: Install cargo-tarpaulin + run: cargo install cargo-tarpaulin + + - name: Generate coverage report + run: cargo tarpaulin --lib --out Xml + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v4 + with: + file: ./cobertura.xml + fail_ci_if_error: false + + security: + name: Security audit + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Setup mold linker + uses: rui314/setup-mold@v1 + + - name: Install Rust + uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + + - name: Install cargo-audit + run: cargo install cargo-audit + + - name: Generate temporary lockfile + run: cargo generate-lockfile + + - name: Run security audit + run: cargo audit diff --git a/.gitignore b/.gitignore index ad67955..2f60cf6 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,5 @@ target # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +Cargo.lock diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 59039be..d50c929 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,65 +1,61 @@ -# Pre-commit hooks for Rust project +# Pre-commit hooks for petite-ad # See https://pre-commit.com for more information -# Run `pre-commit install` to install these hooks +# Install: pip install pre-commit +# Setup: pre-commit install +# Run manually: prek run --all-files # alias for: pre-commit run --all-files repos: - # General hooks - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 - hooks: - - id: trailing-whitespace - - id: end-of-file-fixer - - id: check-yaml - - id: check-toml - - id: check-added-large-files - args: ['--maxkb=1000'] - - id: check-merge-conflict - - id: check-case-conflict - - id: mixed-line-ending - args: ['--fix=lf'] - - # Rust specific hooks using local cargo commands - repo: local hooks: + # Run cargo fmt to automatically format Rust files. - id: cargo-fmt name: cargo fmt entry: cargo fmt language: system + types: [rust] pass_filenames: false - - id: cargo-check - name: cargo check - entry: cargo check - language: system - args: ['--all-targets', '--all-features'] - pass_filenames: false - - # Clippy linting - - repo: local - hooks: - - id: clippy + # Run cargo clippy to catch common mistakes. + - id: cargo-clippy name: cargo clippy - entry: cargo clippy + entry: cargo clippy --all-targets --all-features -- -D warnings language: system - args: ['--all-targets', '--all-features', '--', '-D', 'warnings'] + types: [rust] pass_filenames: false - # Run tests - - repo: local - hooks: + # Run cargo test to ensure all tests pass. - id: cargo-test name: cargo test - entry: cargo test + entry: cargo test --all-features language: system - args: ['--all-features'] + types: [rust] pass_filenames: false - # Check documentation - - repo: local + # Generic pre-commit hooks for non-Rust files. + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 hooks: - - id: cargo-doc - name: cargo doc - entry: cargo doc - language: system - args: ['--no-deps', '--all-features'] - pass_filenames: false + # Check for accidentally committed large files. + - id: check-added-large-files + args: ["--maxkb=10000"] + + # Check config file syntax. + - id: check-yaml + - id: check-json + - id: check-toml + + # Fix whitespace-only file hygiene issues. + - id: trailing-whitespace + exclude: ^(tests/|fixtures/) + - id: end-of-file-fixer + + # Detect conflict markers and case-insensitive path collisions. + - id: check-merge-conflict + - id: check-case-conflict + + # Normalize line endings. + - id: mixed-line-ending + args: ["--fix=lf"] + + # Check that scripts with shebangs are executable. + - id: check-executables-have-shebangs diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..0667f70 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,209 @@ +# AGENTS.md + +Guidance for agentic coding agents working with petite-ad. + +## Quick Start + +**Build directory**: `~/.cache/rust-build/` (cached, never mention in docs) +**Edition**: Rust 2021 | **License**: MIT OR Apache-2.0 + +### Essential Commands + +```bash +# Build & check +cargo build --release # Build optimized release +cargo check # Quick syntax check +cargo fmt # Auto-format code +cargo clippy --all-targets --all-features -- -D warnings # Lint with strict warnings + +# Testing +cargo test # Run all tests +cargo test test_name -- --nocapture # Run specific test with output +cargo test --lib # Library tests only +cargo tarpaulin --lib --out Stdout # Generate coverage report + +# Benchmarking +cargo bench # Run benchmarks +cargo bench --no-run # Check benchmarks compile without running +``` + +## Project Structure + +**Core Library** (`src/lib.rs`): `MonoAD`, `MultiAD`, `MonoAD2RR`, `MonoAD2FR`, `MonoAD2RF`, `GraphBuilder` + +**Modules**: +- `src/mono/` - Single-variable automatic differentiation +- `src/multi/` - Multi-variable automatic differentiation +- `src/error.rs` - Error types (`AutodiffError`) +- `src/macros.rs` - Macros (`mono_ops!`, `multi_ops!`) + +**Tests**: Unit tests in `src/*/tests.rs`, comprehensive tests in `src/tests_comprehensive.rs` + +**Benchmarks**: `benches/compute_benchmark.rs`, `benches/hessian_benchmark.rs` + +## Code Style & Conventions + +### Imports & Organization + +- **Standard library first**, then external crates, then internal modules (separated by blank lines) +- Use explicit imports for clarity +- Re-export public types in `lib.rs` for API simplicity +- Use `crate::Result` for error-returning functions + +### Naming & Types + +- **Functions**: `snake_case`, methods starting with action verbs (`compute_`, `forward_`, `backward_`) +- **Types**: `PascalCase` for structs/enums/traits +- **Constants**: `SCREAMING_SNAKE_CASE` +- **Private helpers**: Prefix internal functions with underscore if needed +- **Type annotations**: Always explicit for public functions and struct fields + +### Error Handling + +- Use `crate::Result` (alias for `std::result::Result`) +- Use `?` operator for propagation +- Document error conditions in doc comments + +### Performance + +- Use `#[inline]` and `#[inline(always)]` for hot functions +- Pre-allocate vectors with `with_capacity()` when size is known +- Capture computed values in closures to avoid recomputation +- Use mold linker for faster builds (configured in `.cargo/config.toml`) + +### Comments & Documentation + +- **Public items**: Full doc comments `///` with examples +- **Complex logic**: Explain the "why" not the "what" +- **Inline comments**: Use `//` for clarification, keep minimal +- **Tests**: Include usage examples in doc comments for public APIs + +### Formatting + +- 4-space indentation (enforced by `cargo fmt`) +- Max line length: ~100 characters (soft limit) +- One statement per line; use temporary variables for readability + +## Pre-commit Hooks + +All commits run: `cargo fmt --check` → `cargo clippy -- -D warnings` → `cargo test --all-features` + +### ⚠️ CRITICAL: Never use `--no-verify` + +**❌ NEVER use `git commit --no-verify`** - This bypasses all pre-commit checks and will likely cause CI failures. + +If you used `--no-verify` and CI fails: +1. Fix the issues (usually formatting: run `cargo fmt`) +2. Re-stage files: `git add .` +3. Commit **without** `--no-verify`: `git commit -m "fix: ..."` +4. Push and verify CI passes + +### Common Pre-commit Failures + +| Failure | Cause | Fix | +|---------|-------|-----| +| `cargo fmt` | Code formatting | Run `cargo fmt` then re-stage | +| `cargo clippy` | Lint warnings | Fix warnings or run `cargo clippy --fix` | +| `cargo test` | Test failures | Fix failing tests | + +## GitHub CI + +The CI workflow (`.github/workflows/ci.yml`) runs: +- **Test matrix**: Rust stable and nightly +- **Coverage**: Code coverage with cargo-tarpaulin +- **Security audit**: cargo-audit for vulnerability checking + +### Post-Commit CI Verification Checklist + +After pushing any commit, agents MUST verify CI passes: + +- [ ] Push commit and note the run ID +- [ ] Run `gh run watch --repo leafyoung/petite-ad --exit-status` until completion +- [ ] If CI fails, run `gh run view --log-failed --repo leafyoung/petite-ad` to diagnose +- [ ] Fix any issues and re-push until all jobs pass +- [ ] Only proceed to next task after CI is green + +### GitHub CLI Commands + +```bash +# View CI run status (get run ID from push output or list) +gh run list --repo leafyoung/petite-ad --limit 3 + +# Watch CI run progress with live updates +gh run watch --repo leafyoung/petite-ad --exit-status + +# If CI fails, view detailed failure logs +gh run view --repo leafyoung/petite-ad --log-failed + +# View specific job logs +gh run view --repo leafyoung/petite-ad --job +``` + +## Key Patterns in Codebase + +### Forward Pass + +```rust +// ✅ Good: Pre-allocate with capacity, inline hot functions +#[inline] +pub fn compute(exprs: &[MonoAD], x: f64) -> f64 { + let mut value = x; + for expr in exprs { + value = expr.forward(value); + } + value +} +``` + +### Backward Pass + +```rust +// ✅ Good: Capture computed values in closures +MonoAD::Sin => { + let y = x.sin(); + let x_cos = x.cos(); // Capture to avoid recomputation + let grad = Box::new(move |dy: f64| -> f64 { dy * x_cos }); + (y, grad) +} +``` + +### Error Handling + +```rust +// ✅ Good: Use Result type alias and ? operator +pub fn compute(exprs: &[(MultiAD, Vec)], inputs: &[f64]) -> Result { + let mut values: Vec = inputs.to_vec(); + // ... operations with ? for error propagation + Ok(values.last().copied().unwrap_or(0.0)) +} +``` + +## Important Constraints + +- ❌ Never: `cargo clean` (breaks build cache) +- ❌ Never: Use `git commit --no-verify` (bypasses pre-commit checks) +- ❌ Never: Commit `Cargo.lock` if this is a library (only for binaries) +- ✅ Always: Use `cargo test` before committing +- ✅ Always: Run `cargo fmt` on modified files +- ✅ Always: Verify CI passes after pushing + +## Testing Strategy + +**Unit tests** (`#[test]`): Logic validation in `src/*/tests.rs` +**Comprehensive tests** (`src/tests_comprehensive.rs`): Edge cases and coverage +**Benchmarks** (`benches/`): Performance regression testing +**Doc tests**: Examples in `///` doc comments + +Example: `cargo test test_mono_neg_forward -- --nocapture` + +## Coverage + +Target: >85% code coverage (currently ~87%) + +```bash +# Generate coverage report +cargo tarpaulin --lib --out Stdout + +# Generate HTML report +cargo tarpaulin --lib --out Html +``` diff --git a/Cargo.lock b/Cargo.lock deleted file mode 100644 index f828aec..0000000 --- a/Cargo.lock +++ /dev/null @@ -1,606 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 4 - -[[package]] -name = "aho-corasick" -version = "1.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" -dependencies = [ - "memchr", -] - -[[package]] -name = "alloca" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" -dependencies = [ - "cc", -] - -[[package]] -name = "anes" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" - -[[package]] -name = "anstyle" -version = "1.0.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78" - -[[package]] -name = "autocfg" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" - -[[package]] -name = "bumpalo" -version = "3.19.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" - -[[package]] -name = "cast" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" - -[[package]] -name = "cc" -version = "1.2.52" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd4932aefd12402b36c60956a4fe0035421f544799057659ff86f923657aada3" -dependencies = [ - "find-msvc-tools", - "shlex", -] - -[[package]] -name = "cfg-if" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" - -[[package]] -name = "ciborium" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" -dependencies = [ - "ciborium-io", - "ciborium-ll", - "serde", -] - -[[package]] -name = "ciborium-io" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" - -[[package]] -name = "ciborium-ll" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" -dependencies = [ - "ciborium-io", - "half", -] - -[[package]] -name = "clap" -version = "4.5.54" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6e6ff9dcd79cff5cd969a17a545d79e84ab086e444102a591e288a8aa3ce394" -dependencies = [ - "clap_builder", -] - -[[package]] -name = "clap_builder" -version = "4.5.54" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa42cf4d2b7a41bc8f663a7cab4031ebafa1bf3875705bfaf8466dc60ab52c00" -dependencies = [ - "anstyle", - "clap_lex", -] - -[[package]] -name = "clap_lex" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1d728cc89cf3aee9ff92b05e62b19ee65a02b5702cff7d5a377e32c6ae29d8d" - -[[package]] -name = "criterion" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d883447757bb0ee46f233e9dc22eb84d93a9508c9b868687b274fc431d886bf" -dependencies = [ - "alloca", - "anes", - "cast", - "ciborium", - "clap", - "criterion-plot", - "itertools", - "num-traits", - "oorandom", - "page_size", - "plotters", - "rayon", - "regex", - "serde", - "serde_json", - "tinytemplate", - "walkdir", -] - -[[package]] -name = "criterion-plot" -version = "0.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed943f81ea2faa8dcecbbfa50164acf95d555afec96a27871663b300e387b2e4" -dependencies = [ - "cast", - "itertools", -] - -[[package]] -name = "crossbeam-deque" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" -dependencies = [ - "crossbeam-epoch", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.9.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" -dependencies = [ - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-utils" -version = "0.8.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" - -[[package]] -name = "crunchy" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" - -[[package]] -name = "either" -version = "1.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" - -[[package]] -name = "find-msvc-tools" -version = "0.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f449e6c6c08c865631d4890cfacf252b3d396c9bcc83adb6623cdb02a8336c41" - -[[package]] -name = "half" -version = "2.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" -dependencies = [ - "cfg-if", - "crunchy", - "zerocopy", -] - -[[package]] -name = "itertools" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" -dependencies = [ - "either", -] - -[[package]] -name = "itoa" -version = "1.0.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" - -[[package]] -name = "js-sys" -version = "0.3.83" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" -dependencies = [ - "once_cell", - "wasm-bindgen", -] - -[[package]] -name = "libc" -version = "0.2.180" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bcc35a38544a891a5f7c865aca548a982ccb3b8650a5b06d0fd33a10283c56fc" - -[[package]] -name = "memchr" -version = "2.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" - -[[package]] -name = "num-traits" -version = "0.2.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" -dependencies = [ - "autocfg", -] - -[[package]] -name = "once_cell" -version = "1.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" - -[[package]] -name = "oorandom" -version = "11.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" - -[[package]] -name = "page_size" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" -dependencies = [ - "libc", - "winapi", -] - -[[package]] -name = "petite-ad" -version = "0.1.2" -dependencies = [ - "criterion", -] - -[[package]] -name = "plotters" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" -dependencies = [ - "num-traits", - "plotters-backend", - "plotters-svg", - "wasm-bindgen", - "web-sys", -] - -[[package]] -name = "plotters-backend" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" - -[[package]] -name = "plotters-svg" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" -dependencies = [ - "plotters-backend", -] - -[[package]] -name = "proc-macro2" -version = "1.0.105" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "535d180e0ecab6268a3e718bb9fd44db66bbbc256257165fc699dadf70d16fe7" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "quote" -version = "1.0.43" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc74d9a594b72ae6656596548f56f667211f8a97b3d4c3d467150794690dc40a" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "rayon" -version = "1.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" -dependencies = [ - "either", - "rayon-core", -] - -[[package]] -name = "rayon-core" -version = "1.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" -dependencies = [ - "crossbeam-deque", - "crossbeam-utils", -] - -[[package]] -name = "regex" -version = "1.12.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4" -dependencies = [ - "aho-corasick", - "memchr", - "regex-automata", - "regex-syntax", -] - -[[package]] -name = "regex-automata" -version = "0.4.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5276caf25ac86c8d810222b3dbb938e512c55c6831a10f3e6ed1c93b84041f1c" -dependencies = [ - "aho-corasick", - "memchr", - "regex-syntax", -] - -[[package]] -name = "regex-syntax" -version = "0.8.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" - -[[package]] -name = "rustversion" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" - -[[package]] -name = "same-file" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" -dependencies = [ - "winapi-util", -] - -[[package]] -name = "serde" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" -dependencies = [ - "serde_core", - "serde_derive", -] - -[[package]] -name = "serde_core" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" -dependencies = [ - "serde_derive", -] - -[[package]] -name = "serde_derive" -version = "1.0.228" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "serde_json" -version = "1.0.149" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" -dependencies = [ - "itoa", - "memchr", - "serde", - "serde_core", - "zmij", -] - -[[package]] -name = "shlex" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" - -[[package]] -name = "syn" -version = "2.0.114" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "tinytemplate" -version = "1.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" -dependencies = [ - "serde", - "serde_json", -] - -[[package]] -name = "unicode-ident" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" - -[[package]] -name = "walkdir" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" -dependencies = [ - "same-file", - "winapi-util", -] - -[[package]] -name = "wasm-bindgen" -version = "0.2.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" -dependencies = [ - "cfg-if", - "once_cell", - "rustversion", - "wasm-bindgen-macro", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-macro" -version = "0.2.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" -dependencies = [ - "quote", - "wasm-bindgen-macro-support", -] - -[[package]] -name = "wasm-bindgen-macro-support" -version = "0.2.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" -dependencies = [ - "bumpalo", - "proc-macro2", - "quote", - "syn", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-shared" -version = "0.2.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "web-sys" -version = "0.3.83" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" -dependencies = [ - "js-sys", - "wasm-bindgen", -] - -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-util" -version = "0.1.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" -dependencies = [ - "windows-sys", -] - -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - -[[package]] -name = "windows-link" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" - -[[package]] -name = "windows-sys" -version = "0.61.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" -dependencies = [ - "windows-link", -] - -[[package]] -name = "zerocopy" -version = "0.8.33" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "668f5168d10b9ee831de31933dc111a459c97ec93225beb307aed970d1372dfd" -dependencies = [ - "zerocopy-derive", -] - -[[package]] -name = "zerocopy-derive" -version = "0.8.33" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c7962b26b0a8685668b671ee4b54d007a67d4eaf05fda79ac0ecf41e32270f1" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "zmij" -version = "1.0.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fc5a66a20078bf1251bde995aa2fdcc4b800c70b5d92dd2c62abc5c60f679f8" diff --git a/Cargo.toml b/Cargo.toml index 0dc9288..07a8655 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,10 +16,45 @@ keywords = [ ] categories = ["algorithms", "mathematics", "science"] +[features] +default = [] +serde = ["dep:serde"] +backend-wgpu = ["dep:pollster", "dep:wgpu"] + [dependencies] +serde = { version = "1.0", features = ["derive"], optional = true } +pollster = { version = "0.4", optional = true } +wgpu = { version = "25", optional = true } [dev-dependencies] -criterion = "0.8.1" +criterion = { version = "0.8.1", default-features = false } +serde_json = "1.0" + +[profile.release] +# Optimize for speed and binary size +opt-level = 3 +lto = "thin" +codegen-units = 1 +strip = true +panic = "abort" + +[profile.dev] +# Faster debug builds +codegen-units = 16 +lto = false +opt-level = 0 +incremental = true + +[profile.test] +# Fast test builds with some optimizations +codegen-units = 16 +opt-level = 1 +lto = false + +[profile.bench] +# Maximum performance for benchmarks +inherits = "release" +debug = true [[bench]] name = "compute_benchmark" diff --git a/README.md b/README.md index 3dcc068..d18018d 100644 --- a/README.md +++ b/README.md @@ -4,13 +4,19 @@ A pure Rust automatic differentiation library supporting both single-variable an ## Features -- **Single-variable autodiff** ([`MonoAD`](src/mono/mono_ad.rs)) - Chain operations like `sin`, `cos`, `exp` with automatic gradient computation +- **Single-variable autodiff** ([`MonoAD`](src/mono/mono_ad.rs)) - Chain operations like `sin`, `cos`, `tan`, `exp`, `ln`, `sqrt`, and `abs` with automatic gradient computation - **Multi-variable autodiff** ([`MultiAD`](src/multi/multi_ad.rs)) - Build computational graphs for functions with multiple inputs - **Box-wrapped by default** - Results use `Box` for flexibility; convert to `Arc` when needed for thread-safety - **Zero-copy backward pass** - Gradients computed efficiently through closure chains - **Convenient macros** - Use `mono_ops![]` for concise operation lists - **Builder API** - Fluent interface for constructing computation graphs -- **Comprehensive tests** - 39 unit tests + 10 doctests covering all operations and edge cases +- **Reusable graph/tape API** - Build graphs with node handles, select explicit single or multiple outputs, and evaluate repeatedly +- **Compiled IR, batching, and backend hooks** - Evaluate closure-free compiled graphs over scalar inputs, flat row-major batches, or prototype SIMD batch backends +- **Graph validation/export** - Validate reusable graphs and export Mermaid/DOT diagrams +- **Second-order derivatives** - Hessian computation for both single and multi-variable functions +- **Public forward-mode AD** - Single-variable derivatives and multivariate directional derivatives / JVPs +- **Opt-in checked mode** - Real-domain validation for scalar `Ln`/`Sqrt` and graph `Div`/`Ln`/`Sqrt`/`Pow` +- **Comprehensive tests** - Unit tests covering operations, edge cases, graph migration, and Hessian methods ## Installation @@ -18,7 +24,7 @@ Add to your `Cargo.toml`: ```toml [dependencies] -petite-ad = "0.1.0" +petite-ad = "0.1.2" ``` ## Quick Start @@ -29,7 +35,7 @@ petite-ad = "0.1.0" use petite_ad::{mono_ops, MonoAD}; let exprs = mono_ops![sin, cos, exp]; -let (value, backprop) = MonoAD::compute(&exprs, 2.0); +let (value, backprop) = MonoAD::compute_grad(&exprs, 2.0); let gradient = backprop(1.0); println!("f(2.0) = {}", value); // exp(cos(sin(2.0))) @@ -38,27 +44,216 @@ println!("f'(2.0) = {}", gradient); // derivative ### Multi-Variable Functions -#### Using the GraphBuilder API (Recommended) +#### Using the reusable Graph API (Recommended) ```rust -use petite_ad::{GraphBuilder, MultiAD}; +use petite_ad::Graph; // Build: f(x, y) = sin(x) * (x + y) -let graph = GraphBuilder::new(2) // 2 inputs - .add(0, 1) // x + y at index 2 - .sin(0) // sin(x) at index 3 - .mul(2, 3) // sin(x) * (x + y) at index 4 - .build(); +let mut graph = Graph::new(2); +let x = graph.input(0); +let y = graph.input(1); +let sum = graph.add(x, y); +let sin_x = graph.sin(x); +graph.mul(sum, sin_x); let inputs = &[0.6, 1.4]; -let (value, backprop_fn) = MultiAD::compute_grad(&graph, inputs).unwrap(); +let (value, backprop_fn) = graph.compute_grad(inputs).unwrap(); let gradients = backprop_fn(1.0); println!("f(0.6, 1.4) = {}", value); println!("∇f = {:?}", gradients); // [∂f/∂x, ∂f/∂y] + +// Reuse the same graph but select a different output node. +graph.set_output(sum).unwrap(); +assert!((graph.compute(inputs).unwrap() - 2.0).abs() < 1e-10); + +// Or expose multiple outputs directly. +graph.set_outputs(&[sum, sin_x]).unwrap(); +let values = graph.compute_many(inputs).unwrap(); +let jacobian = graph.jacobian(inputs).unwrap(); +assert_eq!(values.len(), 2); +assert_eq!(jacobian.len(), 2); +``` + +You can also compile a reusable tape for repeated evaluation: + +```rust +# use petite_ad::Graph; +# let mut graph = Graph::new(2); +# let x = graph.input(0); +# let y = graph.input(1); +# let sum = graph.add(x, y); +# let sin_x = graph.sin(x); +# graph.mul(sum, sin_x); +let tape = graph.compile(); +let value = tape.compute(&[0.6, 1.4]).unwrap(); +assert!(value.is_finite()); +``` + +For repeated hot-loop evaluation, reuse a `TapeWorkspace` to avoid reallocating buffers: + +```rust +# use petite_ad::Graph; +# let mut graph = Graph::new(2); +# let x = graph.input(0); +# let y = graph.input(1); +# let sum = graph.add(x, y); +# graph.mul(sum, y); +let tape = graph.compile(); +let mut workspace = tape.workspace(); +let (value, grad) = tape.gradient_with_workspace(&[2.0, 3.0], &mut workspace).unwrap(); +assert_eq!(grad.len(), 2); +assert!(value.is_finite()); +``` + +For closure-free execution and batch evaluation, compile to the instruction IR: + +```rust +# use petite_ad::{BatchGradientsBuffer, BatchInputs, BatchValuesBuffer, ExecutionBackend, Graph, SimdBackend}; +# let mut graph = Graph::new(2); +# let x = graph.input(0); +# let y = graph.input(1); +# graph.mul(x, y); +let compiled = graph.compile_ir().unwrap(); +let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); +let values = compiled.compute_batch(batch).unwrap(); +assert_eq!(values.data, vec![6.0, 20.0]); +let (backend, auto_values) = compiled.compute_batch_auto(batch).unwrap(); +assert_eq!(auto_values.data, values.data); +let simd_report = compiled.simd_support_report().unwrap(); +assert_eq!(simd_report.lane_width, simd_report.backend.lane_width()); +assert_eq!(backend.name(), if simd_report.can_compute_batch() { simd_report.backend.name() } else { "scalar" }); +let mut auto_buffer = BatchValuesBuffer::new(); +backend.compute_batch(&compiled, batch, &mut auto_buffer).unwrap(); +assert_eq!(auto_buffer.data, values.data); +let plan = compiled.device_batch_plan(backend, batch.batch_size); +assert_eq!(plan.layout, petite_ad::BatchLayout::RowMajor); + +// Mock device execution uses explicit allocated buffers and transfer plans. +let mock = petite_ad::MockDeviceBackend; +let mut device_buffers = mock.allocate_batch_buffers(&compiled, batch.batch_size); +let mut device_values = BatchValuesBuffer::new(); +let trace = mock + .compute_batch_with_buffers(&compiled, batch, &mut device_buffers, &mut device_values) + .unwrap(); +assert_eq!(trace.mode, petite_ad::DeviceExecutionMode::ComputeBatch); +assert_eq!(device_values.data, values.data); + +// With the optional `backend-wgpu` feature, the first real GPU backend can +// initialize a WGPU device, allocate real GPU buffers, and run a restricted +// exact-safe native compute kernel for batch value execution. Today the native +// path is intentionally conservative: it only accepts graphs composed from +// exact f32-roundtrippable constants plus neg/relu/abs, and it also requires +// the concrete batch inputs to roundtrip through f32 exactly. Other graphs or +// batches still use the host fallback path. Gradients also still use host +// fallback for now. +# #[cfg(feature = "backend-wgpu")] +# { +let boundary = petite_ad::GpuBackendBoundary::new( + petite_ad::AcceleratorDeviceContext::wgpu(0), + petite_ad::DeviceTransferPolicy::Explicit, +); +let wgpu = boundary.initialize_wgpu().unwrap(); +let mut gpu_buffers = compiled.allocate_wgpu_buffers(&wgpu, batch.batch_size).unwrap(); +let mut gpu_values = BatchValuesBuffer::new(); +let native_ok = compiled.supports_native_wgpu_batch_compute_for_batch(&wgpu, batch); +let trace = compiled + .compute_batch_wgpu_into(&wgpu, batch, &mut gpu_buffers, &mut gpu_values) + .unwrap(); +assert_eq!(trace.used_native_kernel, native_ok); +if native_ok { + assert_eq!(gpu_values.data, values.data); +} +assert_eq!( + petite_ad::WgpuBackend::native_batch_compute_supported_opcodes(), + petite_ad::WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES, +); +# } + +// Hot loops can reuse output buffers and inspect static metadata. +let metadata = compiled.metadata(); +assert_eq!(metadata.num_inputs, 2); +let flat = compiled.flat_instructions().unwrap(); +assert_eq!(flat.len(), metadata.num_instructions); +let mut buffer = BatchValuesBuffer::new(); +compiled.compute_batch_into(batch, &mut buffer).unwrap(); +assert_eq!(buffer.data, values.data); + +// SIMD backends use the same flat batch ABI. The prototypes support f64x2 +// and f64x4 batch compute/gradients for constants, +, -, *, /, pow, +// log1p_exp, logaddexp, negation, sin, cos, tan, exp, ln, sqrt, ReLU, +// abs, and tanh. Non-native math uses exact scalar-lane fallback semantics. +// The optional WGPU backend adds real device allocation/upload/download plus +// a first native batch-compute kernel. To preserve strict parity, the native +// path is currently restricted to an exact-safe subset instead of running all +// graphs through WGSL f32 math. Unsupported graphs and inexact batches fall +// back to the host path; use `supports_native_wgpu_batch_compute_for_batch` +// and `DeviceExecutionTrace::used_native_kernel` to see whether a concrete +// batch stayed on the native path. The currently exposed subset is available +// as `WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES`. +let simd = SimdBackend; +let simd_capabilities = simd.capabilities(); +if simd_capabilities.supports_batch_compute { + simd.compute_batch(&compiled, batch, &mut buffer).unwrap(); + assert_eq!(buffer.data, values.data); +} +if simd_capabilities.supports_batch_gradient { + let mut gradients = BatchGradientsBuffer::new(); + simd.gradient_batch(&compiled, batch, &mut gradients).unwrap(); + assert_eq!(gradients.batch_size, 2); +} ``` -#### Using Manual Graph Construction +Optimizers operate directly on parameter and gradient slices: + +```rust +# use petite_ad::Adam; +let mut optimizer = Adam::new(1, 0.1); +let mut params = vec![0.0]; +optimizer.step(&mut params, &[-1.0]).unwrap(); +assert!(params[0] > 0.0); +``` + +You can convert reusable graphs back to the legacy tuple form when they do not contain constants: + +```rust +# use petite_ad::Graph; +# let mut graph = Graph::new(2); +# let x = graph.input(0); +# let y = graph.input(1); +# let sum = graph.add(x, y); +# graph.sin(sum); +let legacy_ops = graph.to_operations().unwrap(); +assert!(!legacy_ops.is_empty()); +``` + +#### Using the GraphBuilder API + +```rust +use petite_ad::GraphBuilder; + +let mut builder = GraphBuilder::new(2); +let x = builder.input_node(0); +let y = builder.input_node(1); +let sum = builder.add_node(x, y); +let sin_x = builder.sin_node(x); +builder.mul_node(sum, sin_x); + +let graph = builder.build_graph_with_output(sum).unwrap(); +let inputs = &[0.6, 1.4]; +let (value, backprop_fn) = graph.compute_grad(inputs).unwrap(); +let gradients = backprop_fn(1.0); + +println!("f(0.6, 1.4) = {}", value); +println!("∇f = {:?}", gradients); + +let graph_multi = builder.build_graph_with_outputs(&[sum, sin_x]).unwrap(); +let values = graph_multi.compute_many(inputs).unwrap(); +assert_eq!(values.len(), 2); +``` + +#### Using Manual Tuple Graph Construction ```rust use petite_ad::MultiAD; @@ -86,9 +281,14 @@ println!("∇f = {:?}", gradients); // [∂f/∂x, ∂f/∂y] | Operation | Description | Derivative | | --------- | ----------- | ---------- | -| `Sin` | Sine | `x.cos()` | -| `Cos` | Cosine | `-x.sin()` | -| `Exp` | Exponential | `exp(x)` | +| `Sin` | Sine | `x.cos()` | +| `Cos` | Cosine | `-x.sin()` | +| `Tan` | Tangent | `1 / cos²(x)` | +| `Exp` | Exponential | `exp(x)` | +| `Neg` | Negation | `-1` | +| `Ln` | Natural log | `1 / x` | +| `Sqrt` | Square root | `1 / (2 * sqrt(x))` | +| `Abs` | Absolute | `sign(x)` (0 at `x=0`) | ### MultiAD (Multi-Variable) @@ -103,11 +303,233 @@ println!("∇f = {:?}", gradients); // [∂f/∂x, ∂f/∂y] | `Sin` | 1 | Sine: `sin(x)` | | `Cos` | 1 | Cosine: `cos(x)` | | `Tan` | 1 | Tangent: `tan(x)` | +| `Tanh` | 1 | Hyperbolic tangent | +| `Neg` | 1 | Negation: `-x` | | `Exp` | 1 | Exponential: `exp(x)` | | `Ln` | 1 | Natural log: `ln(x)` | | `Sqrt` | 1 | Square root: `sqrt(x)` | | `Abs` | 1 | Absolute value: `abs(x)` | +## Higher-Order Derivatives (Second-Order / Hessian) + +### Documentation + +For comprehensive mathematical background and algorithm details, see: +- **[Univariate Hessian Theory (mono_ad_hessian.md)](docs/mono_ad_hessian.md)** - Complete derivations, algorithms, and examples for single-variable functions +- **[Multivariate Hessian Theory (multi_ad_hessian.md)](docs/multi_ad_hessian.md)** - Complete theory for MultiAD Hessian methods (RR, FR, RF) with examples + +### MonoAD Second Derivatives + +The library provides **four methods** to compute second derivatives for single-variable functions: + +#### 1. Finite Differences (MonoAD::compute_hessian) + +Approximate second derivative using numerical differentiation: + +```rust +use petite_ad::{MonoAD, mono_ops}; + +// f(x) = sin(x), f''(x) = -sin(x) +let ops = mono_ops![sin]; +let x = 0.5; + +let second_deriv = MonoAD::compute_hessian(&ops, x); +println!("f''({}) = {}", x, second_deriv); // ≈ -0.4794 +``` + +- **Accuracy**: ~1e-5 to 1e-6 (using ε = 1e-5) +- **Use case**: Quick prototyping, most practical applications + +#### 2. Exact Methods (MonoAD2RR/FR/RF) + +Compute exact second derivatives using automatic differentiation: + +```rust +use petite_ad::{MonoAD2RR, mono_ops_rr}; + +// f(x) = exp(sin(x)) +let ops = mono_ops_rr![sin, exp]; +let x = 0.5; + +let hessian = MonoAD2RR::compute_hessian(&ops, x); +// Result is exact up to machine precision (< 1e-12) +``` + +**Three exact methods available**: +- **MonoAD2RR** (Reverse-over-Reverse): Most direct, propagates second derivatives backward +- **MonoAD2FR** (Forward-over-Reverse): Uses dual numbers to differentiate gradient +- **MonoAD2RF** (Reverse-over-Forward): Equivalent to FR for univariate functions + +The exact mono Hessian types support `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, and `Abs`. `Abs` follows the library's raw convention at zero: first derivative `0`, second derivative `0`. + +**Comparison** (MonoAD - single variable functions): + +| Method | Time | Space | Accuracy | Use Case | +|--------|------|-------|----------|----------| +| Finite Diff | O(n) | O(1) | ~1e-5 | Prototyping | +| RR (Exact) | O(n) | O(n) | < 1e-12 | Production, optimization | +| FR (Exact) | O(n) | O(n) | < 1e-12 | Education, research | +| RF (Exact) | O(n) | O(n) | < 1e-12 | Research | + +### MultiAD Hessian Computation + +For multi-variable functions, compute the Hessian matrix (second-order partial derivatives): + +```rust +use petite_ad::{MultiAD, multi_ops}; + +// f(x, y) = x² + y² +// Hessian: [[2, 0], [0, 2]] +let exprs = multi_ops![ + (inp, 0), (inp, 1), + (mul, 0, 0), (mul, 1, 1), + (add, 2, 3) +]; + +let hessian = MultiAD::compute_hessian(&exprs, &[2.0, 3.0]).unwrap(); +println!("H[0][0] = {}", hessian[0][0]); // ∂²f/∂x² = 2 +println!("H[0][1] = {}", hessian[0][1]); // ∂²f/∂x∂y = 0 +println!("H[1][0] = {}", hessian[1][0]); // ∂²f/∂y∂x = 0 +println!("H[1][1] = {}", hessian[1][1]); // ∂²f/∂y² = 2 +``` + +Currently uses central finite differences on gradients (ε = 1e-5; approximate, unlike the exact methods below). + +The exact multivariate Hessian types currently cover a larger smooth subset of operations than before, including `Neg`, `Pow`, `Sub`, `Div`, `Tan`, `Ln`, and `Sqrt` in addition to `Inp`, `Sin`, `Cos`, `Exp`, `Add`, and `Mul`. + +### MultiAD Second-Order Methods (Exact Hessian) + +The library provides **three exact methods** for computing Hessians of multivariate functions: + +```rust +use petite_ad::MultiAD2RR; + +// f(x, y) = sin(x) + exp(y) +// Hessian: [[-sin(x), 0], [0, exp(y)]] +let ops = &[ + MultiAD2RR::Inp(0), MultiAD2RR::Sin, + MultiAD2RR::Inp(1), MultiAD2RR::Exp, + MultiAD2RR::Add, +]; + +let x = vec![1.0, 2.0]; +let hessian = MultiAD2RR::compute_hessian(ops, &x).unwrap(); +// Result is exact up to machine precision (< 1e-12) +``` + +**Three exact methods available** (all machine-precision, no finite differences): +- **MultiAD2RR** (Reverse-over-Reverse): Per-node gradient vectors with outer-product Hessian accumulation in reverse pass +- **MultiAD2FR** (Forward-over-Reverse): Dual-number forward pass + dual-adjoint reverse pass for each seed direction +- **MultiAD2RF** (Reverse-over-Forward): Same dual-number algorithm as FR (equivalent for scalar f: ℝⁿ → ℝ) + +**Comparison** (MultiAD - multi-variable functions, n inputs, G = graph size): + +| Method | Time | Space | Accuracy | Approach | +|--------|------|-------|----------|----------| +| RR | O(G·n²) | O(G·n) | < 1e-12 | Scalar reverse pass with per-node grad vectors | +| FR | O(n·G) | O(G) | < 1e-12 | Dual forward + dual reverse per seed | +| RF | O(n·G) | O(G) | < 1e-12 | Same algorithm as FR for scalar functions | + +For detailed theory and algorithms, see [docs/multi_ad_hessian.md](docs/multi_ad_hessian.md). + +## Forward-Mode AD + +You can also compute forward tangents directly: + +```rust +use petite_ad::{ForwardAD, MonoAD, mono_ops}; + +let exprs = mono_ops![sin, exp]; +let result = ForwardAD::differentiate(&exprs, 0.5); +assert!(result.value.is_finite()); +assert!(result.tangent.is_finite()); +``` + +For multivariate functions, forward mode returns a directional derivative / JVP: + +```rust +use petite_ad::{ForwardAD, MultiAD}; + +let exprs = vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), +]; +let result = ForwardAD::directional_derivative(&exprs, &[2.0, 3.0], &[1.0, -1.0]).unwrap(); +assert_eq!(result.value, 6.0); +assert_eq!(result.tangent, 1.0); +``` + +You can also assemble Jacobians for vector outputs represented as multiple scalar graphs: + +```rust +use petite_ad::{ForwardAD, MultiAD}; + +let outputs = vec![ + vec![(MultiAD::Inp, vec![0]), (MultiAD::Inp, vec![1]), (MultiAD::Add, vec![0, 1])], + vec![(MultiAD::Inp, vec![0]), (MultiAD::Inp, vec![1]), (MultiAD::Mul, vec![0, 1])], +]; +let jacobian = ForwardAD::jacobian(&outputs, &[2.0, 3.0]).unwrap(); +assert_eq!(jacobian.len(), 2); +``` + +Reusable graphs can now be validated and exported for visualization: + +```rust +use petite_ad::Graph; + +let mut graph = Graph::new(1); +let x = graph.input(0); +let neg_x = graph.neg(x); +graph.exp(neg_x); + +graph.validate().unwrap(); +let mermaid = graph.to_mermaid(); +let dot = graph.to_dot(); +assert!(mermaid.contains("flowchart LR")); +assert!(dot.contains("digraph Graph")); +``` + +If you want stricter real-domain behavior, use checked evaluation APIs: + +```rust +use petite_ad::{Graph, MonoAD, MultiAD, multi_ops}; + +let mono_exprs = [MonoAD::Sqrt]; +assert!(MonoAD::compute_checked(&mono_exprs, 4.0).is_ok()); +assert!(MonoAD::compute_checked(&mono_exprs, -1.0).is_err()); + +let exprs = multi_ops![(ln, 0)]; +assert!(MultiAD::compute_checked(&exprs, &[2.0]).is_ok()); +assert!(MultiAD::compute_checked(&exprs, &[0.0]).is_err()); + +let mut graph = Graph::new(1); +let x = graph.input(0); +graph.ln(x); +assert!(graph.compute_checked(&[2.0]).is_ok()); +assert!(graph.compute_checked(&[0.0]).is_err()); + +let mut multi_graph = Graph::new(2); +let x = multi_graph.input(0); +let y = multi_graph.input(1); +let ratio = multi_graph.div(x, y); +let log_y = multi_graph.ln(y); +multi_graph.set_outputs(&[ratio, log_y]).unwrap(); +assert!(multi_graph.compute_many_checked(&[4.0, 2.0]).is_ok()); +assert!(multi_graph.jacobian_checked(&[4.0, 2.0]).is_ok()); +assert!(multi_graph.compute_many_checked(&[4.0, 0.0]).is_err()); +``` + +### Applications + +Second derivatives enable: +- **Newton's method** for optimization (requires accurate Hessians) +- **Convexity analysis** (f''(x) > 0 means convex) +- **Taylor series** approximations +- **Curvature** analysis of functions + +See [examples/hessian_demo.rs](examples/hessian_demo.rs) for accuracy comparisons. + ## License MIT @@ -116,8 +538,7 @@ MIT Contributions are welcome! Areas for improvement: -- Higher-order derivatives (Hessian computation) - Vector/matrix operations - Optimization algorithms (SGD, Adam, etc.) -- Constant/literal values in computation graphs (e.g., `x^2` without needing a separate input) - Additional mathematical operations +- Higher-order derivatives beyond second-order (third derivatives, etc.) diff --git a/benches/compute_benchmark.rs b/benches/compute_benchmark.rs index c717333..a420c4e 100644 --- a/benches/compute_benchmark.rs +++ b/benches/compute_benchmark.rs @@ -2,7 +2,11 @@ use std::sync::Arc; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; -use petite_ad::{mono_ops, types::MonoGradientFn, types::MultiGradientFn, MonoAD, MultiAD}; +use petite_ad::{ + mono_ops, types::MonoGradientFn, types::MultiGradientFn, BackendKind, BatchGradientsBuffer, + BatchInputs, BatchValuesBuffer, ExecutionBackend, ForwardAD, Graph, MonoAD, MultiAD, + SimdBackend, +}; fn bench_single_operations(c: &mut Criterion) { let mut group = c.benchmark_group("single_operation"); @@ -119,6 +123,44 @@ fn bench_macro_usage(c: &mut Criterion) { group.finish(); } +fn bench_mono_checked(c: &mut Criterion) { + let mut group = c.benchmark_group("mono_checked"); + let exprs = mono_ops![sqrt, ln, exp]; + let x = 4.0; + + group.bench_function("compute_checked", |b| { + b.iter(|| { + let value = + MonoAD::compute_checked(std::hint::black_box(&exprs), std::hint::black_box(x)) + .unwrap(); + std::hint::black_box(value); + }) + }); + + group.bench_function("compute_grad_checked", |b| { + b.iter(|| { + let (value, backprop) = + MonoAD::compute_grad_checked(std::hint::black_box(&exprs), std::hint::black_box(x)) + .unwrap(); + std::hint::black_box(value); + std::hint::black_box(backprop(1.0)); + }) + }); + + group.bench_function("compute_hessian_checked", |b| { + b.iter(|| { + let hessian = MonoAD::compute_hessian_checked( + std::hint::black_box(&exprs), + std::hint::black_box(x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.finish(); +} + fn bench_backprop_execution(c: &mut Criterion) { let mut group = c.benchmark_group("backprop_only"); @@ -312,15 +354,548 @@ fn bench_multi_graph_complexity(c: &mut Criterion) { group.finish(); } +fn make_reusable_graph() -> Graph { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + let sin_x = graph.sin(x); + graph.mul(sum, sin_x); + graph +} + +fn make_simd_basic_graph() -> Graph { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let product = graph.mul(x, y); + let ratio = graph.div(x, y); + let sum = graph.add(product, ratio); + let shifted = graph.add_const(sum, 0.25); + let root = graph.sqrt(shifted); + let exponent = graph.log1p_exp(y); + let powered = graph.pow(root, exponent); + let mixed = graph.log_add_exp(powered, y); + graph.tanh(mixed); + graph +} + +fn make_multi_output_graph() -> Graph { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + let product = graph.mul(x, y); + let sin_x = graph.sin(x); + let out = graph.mul(sum, sin_x); + graph.set_outputs(&[sum, product, out]).unwrap(); + graph +} + +fn make_checked_graph() -> Graph { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sqrt_x = graph.sqrt(x); + let ln_y = graph.ln(y); + let ratio = graph.div(sqrt_x, ln_y); + graph.set_outputs(&[ratio, sqrt_x, ln_y]).unwrap(); + graph +} + +fn bench_graph_tape_compute(c: &mut Criterion) { + let mut group = c.benchmark_group("graph_tape_compute"); + let inputs = [0.6, 1.4]; + let legacy_exprs = &[ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Add, vec![0, 1]), + (MultiAD::Sin, vec![0]), + (MultiAD::Mul, vec![2, 3]), + ]; + let seed = [1.0, 0.0]; + let graph = make_reusable_graph(); + let tape = graph.compile(); + let compiled = graph.compile_ir().unwrap(); + let mut workspace = tape.workspace(); + let mut compiled_workspace = compiled.workspace(); + + group.bench_function("legacy_tuple_compute", |b| { + b.iter(|| { + let value = MultiAD::compute( + std::hint::black_box(legacy_exprs), + std::hint::black_box(&inputs), + ) + .unwrap(); + std::hint::black_box(value); + }) + }); + + group.bench_function("forward_directional_derivative", |b| { + b.iter(|| { + let value = ForwardAD::directional_derivative( + std::hint::black_box(legacy_exprs), + std::hint::black_box(&inputs), + std::hint::black_box(&seed), + ) + .unwrap(); + std::hint::black_box(value); + }) + }); + + group.bench_function("graph_compute", |b| { + b.iter(|| { + let value = graph.compute(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(value); + }) + }); + + group.bench_function("tape_compute", |b| { + b.iter(|| { + let value = tape.compute(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(value); + }) + }); + + group.bench_function("tape_compute_workspace", |b| { + b.iter(|| { + let value = tape + .compute_with_workspace(std::hint::black_box(&inputs), &mut workspace) + .unwrap(); + std::hint::black_box(value); + }) + }); + + group.bench_function("compiled_compute_workspace", |b| { + b.iter(|| { + let value = compiled + .compute_with_workspace(std::hint::black_box(&inputs), &mut compiled_workspace) + .unwrap(); + std::hint::black_box(value); + }) + }); + + let batch_data = [0.6, 1.4, 0.7, 1.3, 0.8, 1.2, 0.9, 1.1]; + let batch = BatchInputs::new(&batch_data, 4, 2).unwrap(); + group.bench_function("compiled_compute_batch", |b| { + b.iter(|| { + let value = compiled.compute_batch(std::hint::black_box(batch)).unwrap(); + std::hint::black_box(value); + }) + }); + + let mut batch_values_buffer = BatchValuesBuffer::new(); + group.bench_function("compiled_compute_batch_into", |b| { + b.iter(|| { + compiled + .compute_batch_into(std::hint::black_box(batch), &mut batch_values_buffer) + .unwrap(); + std::hint::black_box(&batch_values_buffer); + }) + }); + + group.finish(); +} + +fn make_batch_data(batch_size: usize) -> Vec { + let mut data = Vec::with_capacity(batch_size * 2); + for row in 0..batch_size { + let row_f64 = row as f64; + data.push(2.0 + row_f64 * 0.01); + data.push(1.0 + row_f64 * 0.005); + } + data +} + +fn bench_simd_batch_compute(c: &mut Criterion) { + let mut group = c.benchmark_group("simd_batch_compute"); + let graph = make_simd_basic_graph(); + let compiled = graph.compile_ir().unwrap(); + let simd = SimdBackend; + + for batch_size in [3_usize, 4, 8, 31, 64] { + let batch_data = make_batch_data(batch_size); + let batch = BatchInputs::new(&batch_data, batch_size, 2).unwrap(); + + let mut scalar_buffer = BatchValuesBuffer::new(); + group.bench_with_input( + BenchmarkId::new("scalar_compute_into", batch_size), + &batch, + |b, batch| { + b.iter(|| { + compiled + .compute_batch_into(std::hint::black_box(*batch), &mut scalar_buffer) + .unwrap(); + std::hint::black_box(&scalar_buffer); + }) + }, + ); + + let mut scalar_gradient_buffer = BatchGradientsBuffer::new(); + group.bench_with_input( + BenchmarkId::new("scalar_gradient_into", batch_size), + &batch, + |b, batch| { + b.iter(|| { + compiled + .gradient_batch_into( + std::hint::black_box(*batch), + &mut scalar_gradient_buffer, + ) + .unwrap(); + std::hint::black_box(&scalar_gradient_buffer); + }) + }, + ); + + let mut auto_buffer = BatchValuesBuffer::new(); + group.bench_with_input( + BenchmarkId::new("auto_compute_into", batch_size), + &batch, + |b, batch| { + b.iter(|| { + compiled + .compute_batch_auto_into(std::hint::black_box(*batch), &mut auto_buffer) + .unwrap(); + std::hint::black_box(&auto_buffer); + }) + }, + ); + + let mut auto_gradient_buffer = BatchGradientsBuffer::new(); + group.bench_with_input( + BenchmarkId::new("auto_gradient_into", batch_size), + &batch, + |b, batch| { + b.iter(|| { + compiled + .gradient_batch_auto_into( + std::hint::black_box(*batch), + &mut auto_gradient_buffer, + ) + .unwrap(); + std::hint::black_box(&auto_gradient_buffer); + }) + }, + ); + + for backend in [BackendKind::SimdF64x2, BackendKind::SimdF64x4] { + let report = compiled.backend_support_report(backend).unwrap(); + if report.can_compute_batch() { + let mut values_buffer = BatchValuesBuffer::new(); + group.bench_with_input( + BenchmarkId::new(format!("{}_compute_into", backend.name()), batch_size), + &batch, + |b, batch| { + b.iter(|| { + backend + .compute_batch( + std::hint::black_box(&compiled), + std::hint::black_box(*batch), + &mut values_buffer, + ) + .unwrap(); + std::hint::black_box(&values_buffer); + }) + }, + ); + } + if report.can_gradient_batch() { + let mut gradients_buffer = BatchGradientsBuffer::new(); + group.bench_with_input( + BenchmarkId::new(format!("{}_gradient_into", backend.name()), batch_size), + &batch, + |b, batch| { + b.iter(|| { + backend + .gradient_batch( + std::hint::black_box(&compiled), + std::hint::black_box(*batch), + &mut gradients_buffer, + ) + .unwrap(); + std::hint::black_box(&gradients_buffer); + }) + }, + ); + } + } + } + + if simd.capabilities().supports_batch_compute { + let batch_data = make_batch_data(64); + let batch = BatchInputs::new(&batch_data, 64, 2).unwrap(); + let mut simd_buffer = BatchValuesBuffer::new(); + group.bench_function("simd_trait_compute_into/64", |b| { + b.iter(|| { + simd.compute_batch(std::hint::black_box(&compiled), batch, &mut simd_buffer) + .unwrap(); + std::hint::black_box(&simd_buffer); + }) + }); + } + + group.finish(); +} + +fn bench_graph_tape_gradient(c: &mut Criterion) { + let mut group = c.benchmark_group("graph_tape_gradient"); + let inputs = [0.6, 1.4]; + let graph = make_reusable_graph(); + let tape = graph.compile(); + let compiled = graph.compile_ir().unwrap(); + let mut workspace = tape.workspace(); + let mut compiled_workspace = compiled.workspace(); + + group.bench_function("graph_compute_grad", |b| { + b.iter(|| { + let (value, backprop_fn) = graph.compute_grad(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(value); + std::hint::black_box(backprop_fn(1.0)); + }) + }); + + group.bench_function("tape_compute_grad", |b| { + b.iter(|| { + let (value, backprop_fn) = tape.compute_grad(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(value); + std::hint::black_box(backprop_fn(1.0)); + }) + }); + + group.bench_function("tape_gradient_workspace", |b| { + b.iter(|| { + let (value, grad) = tape + .gradient_with_workspace(std::hint::black_box(&inputs), &mut workspace) + .unwrap(); + std::hint::black_box(value); + std::hint::black_box(grad); + }) + }); + + group.bench_function("compiled_gradient_workspace", |b| { + b.iter(|| { + let (value, grad) = compiled + .gradient_with_workspace(std::hint::black_box(&inputs), &mut compiled_workspace) + .unwrap(); + std::hint::black_box(value); + std::hint::black_box(grad); + }) + }); + + let batch_data = [0.6, 1.4, 0.7, 1.3, 0.8, 1.2, 0.9, 1.1]; + let batch = BatchInputs::new(&batch_data, 4, 2).unwrap(); + group.bench_function("compiled_gradient_batch", |b| { + b.iter(|| { + let value = compiled + .gradient_batch(std::hint::black_box(batch)) + .unwrap(); + std::hint::black_box(value); + }) + }); + + let mut batch_gradients_buffer = BatchGradientsBuffer::new(); + group.bench_function("compiled_gradient_batch_into", |b| { + b.iter(|| { + compiled + .gradient_batch_into(std::hint::black_box(batch), &mut batch_gradients_buffer) + .unwrap(); + std::hint::black_box(&batch_gradients_buffer); + }) + }); + + group.finish(); +} + +fn bench_graph_tape_jacobian(c: &mut Criterion) { + let mut group = c.benchmark_group("graph_tape_jacobian"); + let inputs = [0.6, 1.4]; + let graph = make_multi_output_graph(); + let tape = graph.compile(); + let mut workspace = tape.workspace(); + + group.bench_function("graph_jacobian", |b| { + b.iter(|| { + let jacobian = graph.jacobian(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(jacobian); + }) + }); + + group.bench_function("tape_jacobian", |b| { + b.iter(|| { + let jacobian = tape.jacobian(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(jacobian); + }) + }); + + group.bench_function("tape_jacobian_workspace", |b| { + b.iter(|| { + let jacobian = tape + .jacobian_with_workspace(std::hint::black_box(&inputs), &mut workspace) + .unwrap(); + std::hint::black_box(jacobian); + }) + }); + + group.finish(); +} + +fn bench_graph_tape_checked(c: &mut Criterion) { + let mut compute_group = c.benchmark_group("graph_tape_checked_compute"); + let inputs = [4.0, 2.5]; + let graph = make_checked_graph(); + let tape = graph.compile(); + let mut compute_workspace = tape.workspace(); + + compute_group.bench_function("graph_compute_checked", |b| { + b.iter(|| { + let value = graph + .compute_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(value); + }) + }); + + compute_group.bench_function("tape_compute_checked", |b| { + b.iter(|| { + let value = tape.compute_checked(std::hint::black_box(&inputs)).unwrap(); + std::hint::black_box(value); + }) + }); + + compute_group.bench_function("tape_compute_workspace_checked", |b| { + b.iter(|| { + let value = tape + .compute_with_workspace_checked( + std::hint::black_box(&inputs), + &mut compute_workspace, + ) + .unwrap(); + std::hint::black_box(value); + }) + }); + + compute_group.bench_function("graph_compute_many_checked", |b| { + b.iter(|| { + let values = graph + .compute_many_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(values); + }) + }); + + compute_group.bench_function("tape_compute_many_checked", |b| { + b.iter(|| { + let values = tape + .compute_many_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(values); + }) + }); + + compute_group.bench_function("tape_compute_many_workspace_checked", |b| { + b.iter(|| { + let values = tape + .compute_many_with_workspace_checked( + std::hint::black_box(&inputs), + &mut compute_workspace, + ) + .unwrap(); + std::hint::black_box(values); + }) + }); + + compute_group.finish(); + + let mut gradient_group = c.benchmark_group("graph_tape_checked_gradient"); + let mut gradient_workspace = tape.workspace(); + + gradient_group.bench_function("graph_gradient_checked", |b| { + b.iter(|| { + let (value, grad) = graph + .gradient_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(value); + std::hint::black_box(grad); + }) + }); + + gradient_group.bench_function("tape_gradient_checked", |b| { + b.iter(|| { + let (value, grad) = tape + .gradient_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(value); + std::hint::black_box(grad); + }) + }); + + gradient_group.bench_function("tape_gradient_workspace_checked", |b| { + b.iter(|| { + let (value, grad) = tape + .gradient_with_workspace_checked( + std::hint::black_box(&inputs), + &mut gradient_workspace, + ) + .unwrap(); + std::hint::black_box(value); + std::hint::black_box(grad); + }) + }); + + gradient_group.finish(); + + let mut jacobian_group = c.benchmark_group("graph_tape_checked_jacobian"); + let mut jacobian_workspace = tape.workspace(); + + jacobian_group.bench_function("graph_jacobian_checked", |b| { + b.iter(|| { + let jacobian = graph + .jacobian_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(jacobian); + }) + }); + + jacobian_group.bench_function("tape_jacobian_checked", |b| { + b.iter(|| { + let jacobian = tape + .jacobian_checked(std::hint::black_box(&inputs)) + .unwrap(); + std::hint::black_box(jacobian); + }) + }); + + jacobian_group.bench_function("tape_jacobian_workspace_checked", |b| { + b.iter(|| { + let jacobian = tape + .jacobian_with_workspace_checked( + std::hint::black_box(&inputs), + &mut jacobian_workspace, + ) + .unwrap(); + std::hint::black_box(jacobian); + }) + }); + + jacobian_group.finish(); +} + criterion_group!( benches, bench_backprop_execution, bench_single_operations, bench_chained_operations, bench_macro_usage, + bench_mono_checked, bench_multi_forward_only, bench_multi_forward_backward, bench_multi_backward_only, bench_multi_graph_complexity, + bench_graph_tape_compute, + bench_simd_batch_compute, + bench_graph_tape_gradient, + bench_graph_tape_jacobian, + bench_graph_tape_checked, ); criterion_main!(benches); diff --git a/benches/hessian_benchmark.rs b/benches/hessian_benchmark.rs new file mode 100644 index 0000000..240ee4e --- /dev/null +++ b/benches/hessian_benchmark.rs @@ -0,0 +1,293 @@ +// Criterion benchmarks for Hessian computation performance +// Compares RR, FR, and RF methods on different problem sizes + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use petite_ad::{Graph, MultiAD2FR, MultiAD2RF, MultiAD2RR}; + +fn bench_hessian_rr_vs_fr_vs_rf(c: &mut Criterion) { + let mut group = c.benchmark_group("hessian_methods"); + + // Test 1: Simple quadratic f(x,y) = x² + y² + let ops_rr = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(0), + MultiAD2RR::Mul, + MultiAD2RR::Inp(1), + MultiAD2RR::Inp(1), + MultiAD2RR::Mul, + MultiAD2RR::Add, + ]; + + let ops_fr = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Inp(0), + MultiAD2FR::Mul, + MultiAD2FR::Inp(1), + MultiAD2FR::Inp(1), + MultiAD2FR::Mul, + MultiAD2FR::Add, + ]; + + let ops_rf = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Inp(0), + MultiAD2RF::Mul, + MultiAD2RF::Inp(1), + MultiAD2RF::Inp(1), + MultiAD2RF::Mul, + MultiAD2RF::Add, + ]; + + let x = vec![1.0, 2.0]; + + group.bench_function("rr_quadratic_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2RR::compute_hessian( + std::hint::black_box(&ops_rr), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_function("fr_quadratic_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2FR::compute_hessian( + std::hint::black_box(&ops_fr), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_function("rf_quadratic_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2RF::compute_hessian( + std::hint::black_box(&ops_rf), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + let mut graph = Graph::new(2); + let x0 = graph.input(0); + let x1 = graph.input(1); + let x0_sq = graph.square(x0); + let x1_sq = graph.square(x1); + graph.add(x0_sq, x1_sq); + + group.bench_function("graph_exact_rr_quadratic_2vars", |b| { + b.iter(|| { + let hessian = graph.exact_hessian_rr(std::hint::black_box(&x)).unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.finish(); +} + +fn bench_hessian_scalability(c: &mut Criterion) { + let mut group = c.benchmark_group("hessian_scalability"); + + // Test different numbers of variables + // For n variables: f(x) = sum(x_i²) + for n in [2, 3, 4, 5].iter() { + // Build operations for f(x) = x_0² + x_1² + ... + x_{n-1}² + let mut ops_rr = vec![]; + let mut ops_fr = vec![]; + let mut ops_rf = vec![]; + + for i in 0..*n { + ops_rr.push(MultiAD2RR::Inp(i)); + ops_rr.push(MultiAD2RR::Inp(i)); + ops_rr.push(MultiAD2RR::Mul); + + ops_fr.push(MultiAD2FR::Inp(i)); + ops_fr.push(MultiAD2FR::Inp(i)); + ops_fr.push(MultiAD2FR::Mul); + + ops_rf.push(MultiAD2RF::Inp(i)); + ops_rf.push(MultiAD2RF::Inp(i)); + ops_rf.push(MultiAD2RF::Mul); + + if i > 0 { + ops_rr.push(MultiAD2RR::Add); + ops_fr.push(MultiAD2FR::Add); + ops_rf.push(MultiAD2RF::Add); + } + } + + let x: Vec = (0..*n).map(|i| (i + 1) as f64).collect(); + + group.bench_with_input(BenchmarkId::new("rr", n), &n, |b, _| { + b.iter(|| { + let hessian = MultiAD2RR::compute_hessian( + std::hint::black_box(&ops_rr), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_with_input(BenchmarkId::new("fr", n), &n, |b, _| { + b.iter(|| { + let hessian = MultiAD2FR::compute_hessian( + std::hint::black_box(&ops_fr), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_with_input(BenchmarkId::new("rf", n), &n, |b, _| { + b.iter(|| { + let hessian = MultiAD2RF::compute_hessian( + std::hint::black_box(&ops_rf), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + } + + group.finish(); +} + +fn bench_hessian_complex_operations(c: &mut Criterion) { + let mut group = c.benchmark_group("hessian_complex"); + + // Test 1: Trigonometric function f(x,y) = sin(x) * cos(y) + let ops_rr = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Sin, + MultiAD2RR::Inp(1), + MultiAD2RR::Cos, + MultiAD2RR::Mul, + ]; + + let ops_fr = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Inp(1), + MultiAD2FR::Cos, + MultiAD2FR::Mul, + ]; + + let ops_rf = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Inp(1), + MultiAD2RF::Cos, + MultiAD2RF::Mul, + ]; + + let x = vec![1.0, 2.0]; + + group.bench_function("rr_sin_cos_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2RR::compute_hessian( + std::hint::black_box(&ops_rr), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_function("fr_sin_cos_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2FR::compute_hessian( + std::hint::black_box(&ops_fr), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_function("rf_sin_cos_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2RF::compute_hessian( + std::hint::black_box(&ops_rf), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + // Test 2: Exponential function f(x,y) = exp(x) + exp(y) + let ops_rr_exp = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Exp, + MultiAD2RR::Inp(1), + MultiAD2RR::Exp, + MultiAD2RR::Add, + ]; + + let ops_fr_exp = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Exp, + MultiAD2FR::Inp(1), + MultiAD2FR::Exp, + MultiAD2FR::Add, + ]; + + let ops_rf_exp = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Exp, + MultiAD2RF::Inp(1), + MultiAD2RF::Exp, + MultiAD2RF::Add, + ]; + + group.bench_function("rr_exp_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2RR::compute_hessian( + std::hint::black_box(&ops_rr_exp), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_function("fr_exp_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2FR::compute_hessian( + std::hint::black_box(&ops_fr_exp), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.bench_function("rf_exp_2vars", |b| { + b.iter(|| { + let hessian = MultiAD2RF::compute_hessian( + std::hint::black_box(&ops_rf_exp), + std::hint::black_box(&x), + ) + .unwrap(); + std::hint::black_box(hessian); + }) + }); + + group.finish(); +} + +criterion_group!( + benches, + bench_hessian_rr_vs_fr_vs_rf, + bench_hessian_scalability, + bench_hessian_complex_operations +); +criterion_main!(benches); diff --git a/docs/mono_ad_hessian.md b/docs/mono_ad_hessian.md new file mode 100644 index 0000000..8da8ead --- /dev/null +++ b/docs/mono_ad_hessian.md @@ -0,0 +1,1158 @@ +# Second-Order Automatic Differentiation for Univariate Functions + +## Table of Contents +1. [Mathematical Foundations](#1-mathematical-foundations) +2. [Chain Rule for Second Derivatives](#2-chain-rule-for-second-derivatives) +3. [Reverse-over-Reverse (RR) Method](#3-reverse-over-reverse-rr-method) +4. [Forward-over-Reverse (FR) Method](#4-forward-over-reverse-fr-method) +5. [Reverse-over-Forward (RF) Method](#5-reverse-over-forward-rf-method) +6. [Implementation Details](#6-implementation-details) +7. [Numerical Considerations](#7-numerical-considerations) +8. [Examples and Applications](#8-examples-and-applications) +9. [Appendices](#9-appendices) + +--- + +## 1. Mathematical Foundations + +### 1.1 First-Order Differentiation (Recap) + +For a univariate function f: ℝ → ℝ, the first derivative f'(x) represents the rate of change at point x. + +For a composition h(x) = f(g(x)), the **chain rule** states: +``` +h'(x) = f'(g(x)) · g'(x) +``` + +**Example**: h(x) = sin(exp(x)) +``` +h'(x) = cos(exp(x)) · exp(x) +``` + +**Reverse-mode automatic differentiation** (backpropagation) efficiently computes derivatives by: +1. **Forward pass**: Compute and store all intermediate values +2. **Backward pass**: Propagate derivatives from output to input using the chain rule + +This is the foundation of modern deep learning frameworks. + +### 1.2 Second-Order Differentiation + +The **second derivative** f''(x) represents the rate of change of the first derivative - the "curvature" or "acceleration" of the function. + +**Physical interpretation**: +- f(x): position +- f'(x): velocity +- f''(x): acceleration + +**Mathematical interpretation**: +- f''(x) > 0: function is **convex** (curving upward) at x +- f''(x) < 0: function is **concave** (curving downward) at x +- f''(x) = 0: potential **inflection point** at x + +For a composition h(x) = f(g(x)), the chain rule for second derivatives is: + +``` +h''(x) = f''(g(x)) · [g'(x)]² + f'(g(x)) · g''(x) + \_________/ \______________/ + Quadratic term Linear term +``` + +**Derivation**: + +Starting from h'(x) = f'(g(x)) · g'(x), apply the **product rule**: + +``` +h''(x) = d/dx[h'(x)] + = d/dx[f'(g(x)) · g'(x)] + = [d/dx f'(g(x))] · g'(x) + f'(g(x)) · [d/dx g'(x)] +``` + +For the first term, apply chain rule to differentiate f'(g(x)): +``` +d/dx[f'(g(x))] = f''(g(x)) · g'(x) +``` + +For the second term: +``` +d/dx[g'(x)] = g''(x) +``` + +Therefore: +``` +h''(x) = f''(g(x)) · g'(x) · g'(x) + f'(g(x)) · g''(x) + = f''(g(x)) · [g'(x)]² + f'(g(x)) · g''(x) +``` + +**Key insight**: The quadratic term [g'(x)]² means second derivatives grow rapidly through compositions. This makes symbolic differentiation tedious for long chains, motivating automatic differentiation. + +### 1.3 Second Derivatives of Elementary Functions + +| Function f(x) | First Derivative f'(x) | Second Derivative f''(x) | Notes | +|---------------|------------------------|--------------------------|-------| +| sin(x) | cos(x) | -sin(x) | Cycles every 4 derivatives | +| cos(x) | -sin(x) | -cos(x) | Cycles every 4 derivatives | +| exp(x) | exp(x) | exp(x) | Unchanged through differentiation | +| ln(x) | 1/x | -1/x² | Only defined for x > 0 | +| x^n | n·x^(n-1) | n(n-1)·x^(n-2) | Polynomial degree decreases | +| sqrt(x) | 1/(2√x) | -1/(4x^(3/2)) | Only defined for x > 0 | +| abs(x) | sign(x) | 0 | Non-smooth at 0; raw convention uses 0 | +| -x | -1 | 0 | Linear function, no curvature | +| tan(x) | 1/cos²(x) | 2·sin(x)/cos³(x) | Asymptotes at π/2 + kπ | + +**Pattern observations**: +- Trigonometric functions (sin, cos) cycle through differentiation +- Exponential function remains unchanged +- Polynomial degrees decrease by 1 with each differentiation +- Rational functions have increasingly complex derivatives + +--- + +## 2. Chain Rule for Second Derivatives (Detailed) + +### 2.1 Two-Function Composition + +Given: h(x) = f(g(x)) + +**First derivative** (standard chain rule): +``` +h'(x) = f'(g(x)) · g'(x) +``` + +**Second derivative** (extended chain rule): + +We need to differentiate h'(x) = f'(g(x)) · g'(x). + +Using the **product rule**: (uv)' = u'v + uv' + +Let u = f'(g(x)) and v = g'(x) + +``` +h''(x) = [d/dx f'(g(x))] · g'(x) + f'(g(x)) · [d/dx g'(x)] +``` + +For the first term, using chain rule: +``` +d/dx[f'(g(x))] = f''(g(x)) · g'(x) +``` + +For the second term: +``` +d/dx[g'(x)] = g''(x) +``` + +Therefore: +``` +h''(x) = f''(g(x)) · g'(x) · g'(x) + f'(g(x)) · g''(x) + = f''(g(x)) · [g'(x)]² + f'(g(x)) · g''(x) +``` + +**Concrete example**: h(x) = sin(x²) + +Let f(u) = sin(u) and g(x) = x² + +``` +g(x) = x² → g'(x) = 2x, g''(x) = 2 +f(u) = sin(u) → f'(u) = cos(u), f''(u) = -sin(u) +``` + +At x = 2: +``` +g(2) = 4 +g'(2) = 4 +g''(2) = 2 + +h''(2) = f''(g(2)) · [g'(2)]² + f'(g(2)) · g''(2) + = f''(4) · 16 + f'(4) · 2 + = -sin(4) · 16 + cos(4) · 2 + = -16·sin(4) + 2·cos(4) + ≈ -16·(-0.7568) + 2·(-0.6536) + ≈ 12.109 - 1.307 + ≈ 10.802 +``` + +### 2.2 Three-Function Composition + +Given: k(x) = f(g(h(x))) + +Denote: u = h(x), v = g(u) = g(h(x)), w = f(v) = k(x) + +**First derivative**: + +Applying chain rule twice: +``` +k'(x) = f'(g(h(x))) · g'(h(x)) · h'(x) + = f'(v) · g'(u) · h'(x) +``` + +**Second derivative**: + +We differentiate k'(x) = f'(v) · g'(u) · h'(x). + +This is a product of three functions. Using the product rule: +``` +(fgh)' = f'gh + fg'h + fgh' +``` + +Let's work step by step: + +``` +k''(x) = d/dx[f'(v) · g'(u) · h'(x)] +``` + +First, treat f'(v)·g'(u) as one unit and apply product rule: +``` += [d/dx(f'(v)·g'(u))] · h'(x) + f'(v)·g'(u) · h''(x) +``` + +For the first term, apply product rule again: +``` +d/dx[f'(v)·g'(u)] = [d/dx f'(v)] · g'(u) + f'(v) · [d/dx g'(u)] +``` + +Now compute each piece using chain rule: +``` +d/dx[f'(v)] = f''(v) · dv/dx = f''(v) · g'(u) · h'(x) +d/dx[g'(u)] = g''(u) · du/dx = g''(u) · h'(x) +``` + +Putting it all together: +``` +k''(x) = f''(v) · g'(u) · h'(x) · g'(u) · h'(x) + + f'(v) · g''(u) · h'(x) · h'(x) + + f'(v) · g'(u) · h''(x) + + = f''(v) · [g'(u)]² · [h'(x)]² + + f'(v) · g''(u) · [h'(x)]² + + f'(v) · g'(u) · h''(x) +``` + +**Interpretation**: +- **Term 1**: f''(v)·[g'(u)]²·[h'(x)]² - Second derivative of f, first derivatives of g and h squared +- **Term 2**: f'(v)·g''(u)·[h'(x)]² - First derivative of f, second derivative of g, first derivative of h squared +- **Term 3**: f'(v)·g'(u)·h''(x) - First derivatives of f and g, second derivative of h + +**Concrete example**: k(x) = exp(sin(x²)) + +Let h(x) = x², g(u) = sin(u), f(v) = exp(v) + +``` +h(x) = x² → h'(x) = 2x, h''(x) = 2 +g(u) = sin(u) → g'(u) = cos(u), g''(u) = -sin(u) +f(v) = exp(v) → f'(v) = exp(v), f''(v) = exp(v) +``` + +At x = 1: +``` +h(1) = 1, h'(1) = 2, h''(1) = 2 +g(1) = sin(1), g'(1) = cos(1), g''(1) = -sin(1) +v = sin(1) ≈ 0.8414 + +k''(1) = exp(sin(1)) · [cos(1)]² · [2]² + + exp(sin(1)) · [-sin(1)] · [2]² + + exp(sin(1)) · cos(1) · 2 + + = exp(0.8414) · cos²(1) · 4 + + exp(0.8414) · (-sin(1)) · 4 + + exp(0.8414) · cos(1) · 2 +``` + +### 2.3 General n-Function Composition + +For h = f₁ ∘ f₂ ∘ ... ∘ fₙ, the second derivative involves: +- All combinations of two second derivatives or products of first derivatives +- Terms grow as O(n) for second derivatives +- This is manageable compared to symbolic expansion + +The recursive structure makes automatic differentiation natural: +- Store all intermediate first and second derivatives during forward pass +- Apply chain rule systematically during backward pass + +--- + +## 3. Reverse-over-Reverse (RR) Method + +### 3.1 High-Level Algorithm + +RR extends reverse-mode AD to compute second derivatives by: +1. **Forward pass**: Store values, first derivatives, AND second derivatives for each operation +2. **Reverse pass**: Propagate both gradient (first derivative) and Hessian (second derivative) backward + +**Key idea**: Just as reverse-mode propagates ∂L/∂x backward, RR propagates ∂²L/∂x² backward. + +### 3.2 Forward Pass Detail + +For each operation in the computation chain, compute three quantities: + +**Example**: Computing y = sin(x) at x = 0.5 + +``` +value: y = sin(0.5) ≈ 0.479425538604203 +first_deriv: dy/dx = cos(0.5) ≈ 0.877582561890373 +second_deriv: d²y/dx² = -sin(0.5) ≈ -0.479425538604203 +``` + +Store all three for use in the backward pass. + +**Example**: Computing y = exp(x) at x = 1.5 + +``` +value: y = exp(1.5) ≈ 4.481689070338065 +first_deriv: dy/dx = exp(1.5) ≈ 4.481689070338065 +second_deriv: d²y/dx² = exp(1.5) ≈ 4.481689070338065 +``` + +All three values are identical for exp! + +### 3.3 Backward Pass Detail + +**Initialization** at output node: +``` +grad = 1.0 // ∂Output/∂Output = 1 +hessian = 0.0 // ∂²Output/∂Output² = 0 (scalar has no curvature w.r.t. itself) +``` + +**Propagation rule** for each operation going backward: + +Given current operation computes y = op(x) with stored values: +- dy/dx (first derivative) +- d²y/dx² (second derivative) + +Update: +```rust +new_grad = grad · (dy/dx) +new_hessian = hessian · (dy/dx)² + grad · (d²y/dx²) + \_________________/ \________________/ + Propagate output Add local + Hessian (quadratic) second derivative (linear) +``` + +**Why this works**: + +This is exactly the chain rule formula we derived: +``` +h''(x) = f''(g(x)) · [g'(x)]² + f'(g(x)) · g''(x) +``` + +Where: +- `hessian` corresponds to f''(g(x)) +- `dy/dx` corresponds to g'(x) +- `grad` corresponds to f'(g(x)) +- `d²y/dx²` corresponds to g''(x) + +### 3.4 Concrete Example: exp(sin(x)) + +Let's compute the second derivative of h(x) = exp(sin(x)) at x = 0.5 in complete detail. + +**Forward pass**: + +``` +Step 1: Compute sin(0.5) + Input: x = 0.5 + + value₁ = sin(0.5) = 0.479425538604203 + dy₁/dx = cos(0.5) = 0.877582561890373 + d²y₁/dx² = -sin(0.5) = -0.479425538604203 + + Store: (0.479425538604203, 0.877582561890373, -0.479425538604203) + +Step 2: Compute exp(value₁) + Input: value₁ = 0.479425538604203 + + value₂ = exp(0.479425538604203) = 1.615129883784566 + dy₂/dy₁ = exp(0.479425538604203) = 1.615129883784566 + d²y₂/dy₁² = exp(0.479425538604203) = 1.615129883784566 + + Store: (1.615129883784566, 1.615129883784566, 1.615129883784566) +``` + +**Backward pass**: + +``` +Initialize at output (Step 2): + grad = 1.0 + hessian = 0.0 + +Process Step 2 (exp): + dy/dy₁ = 1.615129883784566 + d²y/dy₁² = 1.615129883784566 + + new_grad = 1.0 · 1.615129883784566 = 1.615129883784566 + new_hessian = 0.0 · (1.615129883784566)² + 1.0 · 1.615129883784566 + = 0.0 + 1.615129883784566 + = 1.615129883784566 + + Update: grad = 1.615129883784566, hessian = 1.615129883784566 + +Process Step 1 (sin): + dy/dx = 0.877582561890373 + d²y/dx² = -0.479425538604203 + + new_grad = 1.615129883784566 · 0.877582561890373 + = 1.417466313257094 + + new_hessian = 1.615129883784566 · (0.877582561890373)² + + 1.615129883784566 · (-0.479425538604203) + = 1.615129883784566 · 0.770151548091452 + - 0.774301385544626 + = 1.244229173328810 - 0.774301385544626 + = 0.469927787784184 + + Final result: hessian = 0.469927787784184 +``` + +**Verification** (analytical formula): + +For h(x) = exp(sin(x)): +``` +h'(x) = exp(sin(x)) · cos(x) +h''(x) = [exp(sin(x)) · cos(x)] · cos(x) + exp(sin(x)) · [-sin(x)] + = exp(sin(x)) · cos²(x) - exp(sin(x)) · sin(x) + = exp(sin(x)) · [cos²(x) - sin(x)] +``` + +At x = 0.5: +``` +h''(0.5) = exp(sin(0.5)) · [cos²(0.5) - sin(0.5)] + = 1.615129883784566 · [0.770151548091452 - 0.479425538604203] + = 1.615129883784566 · 0.290726009487249 + = 0.469927787784184 ✓ +``` + +Perfect match! + +### 3.5 Implementation in Rust + +```rust +pub fn compute_hessian(exprs: &[MonoAD2RR], x: f64) -> f64 { + if exprs.is_empty() { + return 0.0; + } + + let n = exprs.len(); + + // Forward pass: compute and store all derivatives + let mut values: Vec = Vec::with_capacity(n + 1); + let mut first_derivs: Vec = Vec::with_capacity(n); + let mut second_derivs: Vec = Vec::with_capacity(n); + + values.push(x); // Initial input value + + for &op in exprs { + let (y, dy, ddy) = op.forward_d2(*values.last().unwrap()); + values.push(y); + first_derivs.push(dy); + second_derivs.push(ddy); + } + + // Reverse pass: propagate gradient and Hessian backward + let mut grad: f64 = 1.0; // ∂Output/∂Output = 1 + let mut hessian: f64 = 0.0; // ∂²Output/∂Output² = 0 + + for i in (0..n).rev() { + let dy = first_derivs[i]; + let ddy = second_derivs[i]; + + // Apply extended chain rule: + // h''(x) = f''(g(x))·[g'(x)]² + f'(g(x))·g''(x) + let new_grad = grad * dy; + let new_hessian = hessian * dy * dy + grad * ddy; + // \_________/ \________/ + // Quadratic term Linear term + + grad = new_grad; + hessian = new_hessian; + } + + hessian +} +``` + +### 3.6 Complexity Analysis + +**Time complexity**: +- Forward pass: O(n) - process each operation once +- Backward pass: O(n) - process each operation once +- Total: **O(n)** where n = number of operations + +**Space complexity**: +- Values: n+1 entries +- First derivatives: n entries +- Second derivatives: n entries +- Total: **O(n)** space + +**Comparison with finite differences**: +- Finite differences need 3 function evaluations: f(x), f'(x+ε), f'(x) +- Each f' evaluation already costs O(n) +- Total for finite differences: **O(n)** but with numerical approximation error + +RR is **exact** (up to floating-point precision) while finite differences have truncation error. + +--- + +## 4. Forward-over-Reverse (FR) Method + +### 4.1 Conceptual Approach + +FR computes second derivatives by: +1. Use **reverse-mode AD** to obtain the gradient function g(x) = f'(x) +2. Use **forward-mode AD** to differentiate g to get g'(x) = f''(x) + +**Key insight**: Forward-mode AD can differentiate any function, including a gradient function computed by reverse-mode! + +### 4.2 Dual Numbers + +Forward-mode AD uses **dual numbers** to automatically track derivatives. + +A dual number has two components: +``` +d = (value, tangent) +``` + +Where: +- `value`: the function value f(x) +- `tangent`: the derivative f'(x) + +**Arithmetic rules** for dual numbers: + +``` +Addition: (a, a') + (b, b') = (a + b, a' + b') +Subtraction: (a, a') - (b, b') = (a - b, a' - b') +Multiplication: (a, a') · (b, b') = (a · b, a'·b + a·b') +Division: (a, a') / (b, b') = (a / b, (a'·b - a·b') / b²) + +sin: sin(a, a') = (sin(a), cos(a) · a') +cos: cos(a, a') = (cos(a), -sin(a) · a') +exp: exp(a, a') = (exp(a), exp(a) · a') +``` + +**Example**: Compute derivative of f(x) = x · sin(x) at x = 2 + +Using dual numbers with tangent = 1 (differentiating w.r.t. x): +``` +x_dual = (2, 1) + +sin(x_dual) = sin((2, 1)) + = (sin(2), cos(2) · 1) + = (0.909, -0.416) + +x_dual · sin(x_dual) = (2, 1) · (0.909, -0.416) + = (2 · 0.909, 1 · 0.909 + 2 · (-0.416)) + = (1.818, 0.909 - 0.832) + = (1.818, 0.077) +``` + +The tangent part 0.077 is the derivative! + +Verify: f'(x) = sin(x) + x·cos(x) = sin(2) + 2·cos(2) ≈ 0.909 - 0.832 = 0.077 ✓ + +### 4.3 FR Algorithm for Second Derivatives + +**Step 1**: Compute gradient function using reverse-mode +```rust +let (value, grad_fn) = MonoAD::compute_grad(exprs, x); +// grad_fn(1.0) gives us f'(x) +``` + +**Step 2**: Differentiate gradient function using dual numbers + +For simple operations (Sin, Cos, Tan, Exp, Neg, Ln, Sqrt, Abs): +```rust +// Create dual number with tangent = 1 +let x_dual = Dual::variable(x, 1.0); + +// Evaluate operations using dual arithmetic +let result_dual = evaluate_with_dual(exprs, x_dual); + +// Extract second derivative from tangent +let hessian = result_dual.tangent; +``` + +For composed operations: May delegate to RR for correctness. + +### 4.4 Comparison: FR vs RR + +| Aspect | RR | FR | +|--------|----|----| +| **Conceptual model** | Propagate Hessian backward | Differentiate gradient forward | +| **Implementation** | Single backward pass with Hessian | Two passes (reverse then forward) | +| **Natural for** | Deep learning practitioners | Forward-mode enthusiasts | +| **Pedagogical value** | Shows extension of backprop | Shows composition of AD modes | +| **Performance** | O(n) | O(n) | +| **Code complexity** | Moderate | Moderate | + +For **univariate functions**, both are equally efficient: O(n) time and space. + +### 4.5 When FR is Preferable + +- **Directional derivatives**: Computing f''(x)·v for a direction v +- **Educational purposes**: Forward-mode is often easier to understand initially +- **Hardware considerations**: Some architectures favor forward-mode +- **Composition exploration**: Demonstrates how different AD modes can be combined + +--- + +## 5. Reverse-over-Forward (RF) Method + +### 5.1 Conceptual Approach + +RF is conceptually the "opposite" of FR: +1. Use **forward-mode AD** to compute value and first derivative simultaneously +2. Use **reverse-mode AD** to differentiate the derivative computation + +### 5.2 Relationship to FR + +For **univariate functions**, RF and FR are **mathematically equivalent**: +- Both compute f''(x) using one forward and one reverse pass +- They differ in the conceptual order of operations +- Implementation may differ slightly but results are identical (up to floating-point rounding) + +For **multivariate functions** (discussed in multi_ad_hessian.md): +- FR computes Hessian column-by-column +- RF computes Hessian row-by-row +- This distinction matters for efficiency in high dimensions + +### 5.3 Implementation Note + +In the `petite-ad` implementation: +- `MonoAD2RF` has a similar interface to `MonoAD2FR` +- For simple operations, may use direct computation or delegate to RR +- Results are numerically identical (within floating-point precision) + +### 5.4 When to Choose RF + +- **Research purposes**: Exploring different AD mode combinations +- **Consistency testing**: Verifying that multiple methods produce same result +- **Multivariate preparation**: Understanding RF helps with multivariate Hessians + +For most practical purposes with univariate functions, **use RR** (most direct) or **FR** (most pedagogical). + +--- + +## 6. Implementation Details + +### 6.1 Operations and Their Second Derivatives + +Each operation in `MonoAD2RR` implements `forward_d2` to compute value, first derivative, and second derivative: + +```rust +impl MonoAD2RR { + fn forward_d2(&self, x: f64) -> (f64, f64, f64) { + match self { + MonoAD2RR::Sin => { + // f(x) = sin(x) + // f'(x) = cos(x) + // f''(x) = d/dx[cos(x)] = -sin(x) + let val = x.sin(); + let d1 = x.cos(); + let d2 = -x.sin(); // Note: -sin, not +sin + (val, d1, d2) + } + + MonoAD2RR::Cos => { + // f(x) = cos(x) + // f'(x) = -sin(x) + // f''(x) = d/dx[-sin(x)] = -cos(x) + let val = x.cos(); + let d1 = -x.sin(); + let d2 = -x.cos(); // Note: both derivatives negative + (val, d1, d2) + } + + MonoAD2RR::Exp => { + // f(x) = exp(x) + // f'(x) = exp(x) + // f''(x) = exp(x) + // All three are identical! + let val = x.exp(); + let d1 = val; + let d2 = val; + (val, d1, d2) + } + + MonoAD2RR::Neg => { + // f(x) = -x + // f'(x) = -1 (constant) + // f''(x) = 0 (no curvature) + let val = -x; + let d1 = -1.0; + let d2 = 0.0; + (val, d1, d2) + } + } + } +} +``` + +### 6.2 Memory Layout + +For a chain of n operations computing f(x): + +**Storage during forward pass**: +``` +values: [x, y₁, y₂, ..., yₙ] length: n+1 +first_derivs: [dy₁/dx, dy₂/dy₁, ..., dyₙ/dyₙ₋₁] length: n +second_derivs: [d²y₁/dx², d²y₂/dy₁², ...] length: n +``` + +**Memory usage**: +- Each f64 value: 8 bytes +- Total: (n+1 + n + n) × 8 = (3n+1) × 8 bytes +- For n=1000: ~24 KB (negligible) + +**Cache considerations**: +- Sequential access pattern is cache-friendly +- All arrays fit in L1/L2 cache for typical computation graphs +- No pointer chasing (unlike tape-based AD systems) + +### 6.3 Time Complexity Proof + +**Theorem**: RR computes f''(x) in O(n) time for n operations. + +**Proof**: + +Forward pass processes each operation exactly once: +- For operation i: constant time to compute (yᵢ, dyᵢ, ddyᵢ) +- Total: O(n) + +Backward pass processes each operation exactly once in reverse: +- For operation i: constant time arithmetic for chain rule +- Total: O(n) + +Overall: O(n) + O(n) = O(n) ∎ + +**Comparison with naive symbolic differentiation**: +- Symbolic: O(n²) or worse due to expression growth +- RR: O(n) by sharing computations + +### 6.4 Numerical Stability + +**Question**: Is RR numerically stable? + +**Answer**: Yes, with caveats: + +**Stable aspects**: +1. No repeated finite-difference approximations (source of cancellation error) +2. Each arithmetic operation uses exact formulas +3. Chain rule application is mathematically exact + +**Potential instability sources**: +1. **Large derivatives**: If g'(x) is large, [g'(x)]² grows rapidly +2. **Cancellation**: When quadratic and linear terms are large but opposite sign +3. **Overflow**: exp(large) can overflow to infinity + +**Mitigation**: +- Use `f64` (not `f32`) for better precision +- Document when operations may overflow +- Let caller handle edge cases (don't panic) + +--- + +## 7. Numerical Considerations + +### 7.1 Finite Differences vs Exact Methods + +**Finite differences** (existing `MonoAD::compute_hessian`): + +```rust +f''(x) ≈ [f'(x + ε) - f'(x)] / ε +``` + +With ε = 1e-5: + +**Error sources**: +1. **Truncation error**: O(ε) from Taylor series remainder + - At ε = 1e-5: error ≈ 1e-5 × |f'''(x)| + +2. **Roundoff error**: O(ε_machine/ε) from floating-point arithmetic + - With ε_machine ≈ 1e-16 and ε = 1e-5: error ≈ 1e-11 + +**Total error**: Typically 1e-4 to 1e-6 in practice + +**Optimal ε**: Balance truncation vs roundoff, often ε ≈ ∛(ε_machine) ≈ 1e-5 + +**Exact autodiff** (MonoAD2RR/FR/RF): + +Computes exact derivatives up to floating-point roundoff. + +**Error sources**: +1. **Roundoff only**: Each operation adds ~ε_machine relative error +2. For n operations: accumulated error ≈ n × ε_machine ≈ 1e-15 × n + +**Typical accuracy**: 1e-15 to 1e-12 (machine precision) + +**Accuracy improvement**: ~10,000× better than finite differences! + +### 7.2 Overflow and Special Values + +| Situation | Behavior | f''(x) Value | Design Choice | +|-----------|----------|--------------|---------------| +| exp(1000) | Overflow | `inf` | Follow f64::exp semantics | +| exp(-1000) | Underflow | 0.0 | f''(x) also underflows to 0 | +| sin(x), cos(x) | Always finite | Always finite | No overflow possible | +| -x | Linear | 0.0 | No curvature by definition | + +**Philosophy**: Don't panic on overflow/underflow +- Return `inf` or `NaN` as appropriate +- Let caller decide how to handle edge cases +- Matches Rust's f64 behavior (no exceptions) + +**Example**: Computing exp(1000) +```rust +let ops = mono_ops_rr![exp]; +let hessian = MonoAD2RR::compute_hessian(&ops, 1000.0); +assert!(hessian.is_infinite()); +assert!(hessian > 0.0); // Positive infinity +``` + +### 7.3 Catastrophic Cancellation + +**Problem**: When computing a - b where a ≈ b and both are large, lose precision. + +**In second derivative computation**: +``` +h''(x) = f''(g(x))·[g'(x)]² + f'(g(x))·g''(x) + \_________________/ \_____________/ + Term 1 Term 2 +``` + +If Term 1 and Term 2 are large and opposite sign, cancellation occurs. + +**Example**: h(x) = exp(x) - exp(x) (pathological case) +``` +h(x) = 0 (exactly) +h'(x) = 0 +h''(x) = 0 + +But numerically: +Term 1 = large positive number +Term 2 = large negative number +Sum ≈ 0 but with error from cancellation +``` + +**Mitigation in RR**: +- Compute each term exactly (no finite differences) +- Minimize intermediate steps +- Still limited by f64 precision (15-16 significant digits) + +**Comparison**: Finite differences suffer from **both** cancellation (in f'(x+ε) - f'(x)) **and** truncation error. RR eliminates truncation error. + +### 7.4 When Finite Differences Are Acceptable + +Use finite differences when: +1. **Prototyping**: Quick implementation, don't need high accuracy +2. **Accuracy sufficient**: 1e-6 is good enough for your application +3. **Simple codebase**: Don't want to add RR complexity +4. **Black-box functions**: Can't compute exact derivatives + +Use exact methods (RR/FR/RF) when: +1. **Production code**: Need reliability and accuracy +2. **Optimization**: Newton's method requires accurate Hessians +3. **Scientific computing**: Every bit of precision matters +4. **Composition of many operations**: Finite-diff errors accumulate + +--- + +## 8. Examples and Applications + +### 8.1 Simple Function Examples + +#### Example 1: Quadratic Function + +```rust +use petite_ad::{MonoAD2RR, mono_ops_rr}; + +// f(x) = x² (composed as x · x, or using mul twice) +// We can't directly express x² with current ops, so use sin(x) instead + +// f(x) = sin(x) +// f'(x) = cos(x) +// f''(x) = -sin(x) + +let ops = mono_ops_rr![sin]; +let x = 0.5; + +let hessian = MonoAD2RR::compute_hessian(&ops, x); +let expected = -x.sin(); + +assert!((hessian - expected).abs() < 1e-12); +println!("f''({}) = {} (exact)", x, hessian); +``` + +#### Example 2: Exponential Function + +```rust +// f(x) = exp(x) +// f'(x) = exp(x) +// f''(x) = exp(x) (all derivatives are identical!) + +let ops = mono_ops_rr![exp]; +let x = 2.0; + +let hessian = MonoAD2RR::compute_hessian(&ops, x); +let expected = x.exp(); + +assert!((hessian - expected).abs() < 1e-12); +println!("f''({}) = {:.10}", x, hessian); +// Output: f''(2) = 7.3890560989 +``` + +### 8.2 Composed Function Examples + +#### Example 3: exp(sin(x)) + +```rust +// f(x) = exp(sin(x)) +// f''(x) = exp(sin(x))·cos²(x) - exp(sin(x))·sin(x) + +let ops = mono_ops_rr![sin, exp]; +let x = 0.5; + +let hessian = MonoAD2RR::compute_hessian(&ops, x); +let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + +assert!((hessian - expected).abs() < 1e-10); +println!("f''({}) = {:.10}", x, hessian); +// Output: f''(0.5) = 0.4699277878 +``` + +#### Example 4: Long Composition + +```rust +// f(x) = exp(sin(sin(x))) +// Very complex to compute symbolically, but RR handles it easily! + +let ops = mono_ops_rr![sin, sin, exp]; +let x = 2.0; + +let hessian = MonoAD2RR::compute_hessian(&ops, x); +println!("f''({}) = {:.10}", x, hessian); + +// Verify it's finite and reasonable +assert!(hessian.is_finite()); +``` + +### 8.3 Newton's Method for Optimization + +Use second derivatives to find local minima/maxima of f(x). + +**Newton's method update**: +``` +x_{n+1} = x_n - f'(x_n) / f''(x_n) +``` + +```rust +use petite_ad::{MonoAD, MonoAD2RR, mono_ops, mono_ops_rr}; + +// Find minimum of f(x) = x² - 4x + 5 (expressed as operations) +// Minimum is at x = 2 with f(2) = 1 + +// For this example, we'll use sin(x) since we have those operations +// f(x) = sin(x), find critical point near x = π/2 + +let ops = mono_ops![sin]; +let ops_rr = mono_ops_rr![sin]; + +let mut x = 1.0; // Initial guess + +for iteration in 0..10 { + let (value, grad_fn) = MonoAD::compute_grad(&ops, x); + let gradient = grad_fn(1.0); + let hessian = MonoAD2RR::compute_hessian(&ops_rr, x); + + if hessian.abs() < 1e-10 { + println!("Hessian too small, stopping"); + break; + } + + let delta = gradient / hessian; + x = x - delta; + + println!("Iteration {}: x = {:.10}, f(x) = {:.10}", iteration, x, value); + + if delta.abs() < 1e-10 { + println!("Converged!"); + break; + } +} + +// Should converge to x ≈ π/2 where sin(x) has maximum +println!("Final x = {:.10} (π/2 = {:.10})", x, std::f64::consts::FRAC_PI_2); +``` + +### 8.4 Taylor Series Approximation + +Use derivatives to build polynomial approximations. + +**Taylor series around x₀**: +``` +f(x₀ + h) ≈ f(x₀) + f'(x₀)·h + (1/2)·f''(x₀)·h² +``` + +```rust +use petite_ad::{MonoAD, MonoAD2RR, mono_ops, mono_ops_rr}; + +// Approximate sin(x) near x₀ = 0 +let ops = mono_ops![sin]; +let ops_rr = mono_ops_rr![sin]; + +let x0 = 0.0; +let h = 0.1; // Small displacement + +// Compute derivatives at x₀ +let value = MonoAD::compute(&ops, x0); +let (_, grad_fn) = MonoAD::compute_grad(&ops, x0); +let gradient = grad_fn(1.0); +let hessian = MonoAD2RR::compute_hessian(&ops_rr, x0); + +// Taylor approximation +let approx = value + gradient * h + 0.5 * hessian * h * h; +let exact = (x0 + h).sin(); + +println!("f({}) exact: {:.10}", x0 + h, exact); +println!("f({}) approximation: {:.10}", x0 + h, approx); +println!("Error: {:.2e}", (approx - exact).abs()); + +// Output: +// f(0.1) exact: 0.0998334166 +// f(0.1) approximation: 0.0998334166 +// Error: 6.66e-11 +``` + +The second-order Taylor approximation is very accurate for small h! + +### 8.5 Convexity Analysis + +Second derivatives determine whether a function is convex or concave. + +```rust +// Analyze convexity of f(x) = exp(x) on interval [0, 2] + +let ops_rr = mono_ops_rr![exp]; + +for i in 0..21 { + let x = 0.1 * (i as f64); + let hessian = MonoAD2RR::compute_hessian(&ops_rr, x); + + let curvature = if hessian > 0.0 { + "convex (curving upward)" + } else if hessian < 0.0 { + "concave (curving downward)" + } else { + "linear (no curvature)" + }; + + println!("x = {:.1}: f''(x) = {:.4}, {}", x, hessian, curvature); +} + +// Output: f''(x) > 0 for all x, so exp(x) is convex everywhere +``` + +--- + +## 9. Appendices + +### Appendix A: Comparison Table + +| Aspect | Finite Diff | RR (Exact) | FR (Exact) | RF (Exact) | +|--------|-------------|------------|------------|------------| +| **Accuracy** | ~1e-4 to 1e-6 | ~1e-15 | ~1e-15 | ~1e-15 | +| **Time Complexity** | O(n) | O(n) | O(n) | O(n) | +| **Space Complexity** | O(n) | O(n) | O(n) | O(n) | +| **Implementation Complexity** | Simple | Moderate | Moderate | Moderate | +| **Lines of Code** | ~30 | ~200 | ~250 | ~250 | +| **Numerically Stable** | No (cancellation) | Yes | Yes | Yes | +| **Conceptual Model** | Approximation | Exact backward pass | Differentiate gradient | Reverse-over-forward | +| **Use Case** | Prototyping | Production | Education | Research | +| **Works on Black-box** | Yes | No | No | No | +| **Requires Source Code** | No | Yes | Yes | Yes | + +### Appendix B: Second Derivative Formulas + +Quick reference for common functions: + +| f(x) | f'(x) | f''(x) | Domain Notes | +|------|-------|--------|--------------| +| c (constant) | 0 | 0 | All x | +| x | 1 | 0 | All x | +| x² | 2x | 2 | All x | +| xⁿ | nxⁿ⁻¹ | n(n-1)xⁿ⁻² | All x (for n ≥ 2) | +| sin(x) | cos(x) | -sin(x) | All x | +| cos(x) | -sin(x) | -cos(x) | All x | +| tan(x) | 1/cos²(x) | 2sin(x)/cos³(x) | x ≠ π/2 + kπ | +| exp(x) | exp(x) | exp(x) | All x | +| ln(x) | 1/x | -1/x² | x > 0 | +| √x | 1/(2√x) | -1/(4x^(3/2)) | x > 0 for Hessian; checked forward allows x = 0 | +| abs(x) | sign(x) | 0 | Non-smooth at 0; raw convention uses 0 | +| 1/x | -1/x² | 2/x³ | x ≠ 0 | +| aˣ | aˣ ln(a) | aˣ (ln a)² | x ∈ ℝ, a > 0 | + +### Appendix C: Common Pitfalls + +**Pitfall 1**: Forgetting the quadratic term in chain rule +``` +✗ WRONG: h''(x) = f'(g(x))·g''(x) +✓ RIGHT: h''(x) = f''(g(x))·[g'(x)]² + f'(g(x))·g''(x) +``` + +**Pitfall 2**: Sign error in trigonometric derivatives +``` +✗ WRONG: (sin x)'' = sin x +✓ RIGHT: (sin x)'' = -sin x (note the negative sign!) +``` + +**Pitfall 3**: Using float32 instead of float64 +``` +// With f32: accuracy ~1e-7 (barely better than finite diff) +// With f64: accuracy ~1e-15 (true machine precision) +``` + +**Pitfall 4**: Not checking for overflow +```rust +// May panic or give wrong results: +let hessian = MonoAD2RR::compute_hessian(&ops, 1e100); + +// Better: +let hessian = MonoAD2RR::compute_hessian(&ops, 1e100); +if hessian.is_finite() { + // Use hessian +} else { + // Handle overflow +} +``` + +### Appendix D: References + +1. **Griewank, A., & Walther, A. (2008).** *Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation* (2nd ed.). SIAM. + - The definitive reference on automatic differentiation theory + +2. **Baydin, A. G., Pearlmutter, B. A., Radul, A. A., & Siskind, J. M. (2018).** "Automatic Differentiation in Machine Learning: a Survey." *Journal of Machine Learning Research*, 18(153), 1-43. + - Comprehensive survey of AD techniques in ML context + +3. **Nocedal, J., & Wright, S. J. (2006).** *Numerical Optimization* (2nd ed.). Springer. + - Applications of second derivatives in optimization algorithms + +4. **Margossian, C. C. (2019).** "A Review of automatic differentiation and its efficient implementation." *Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery*, 9(4), e1305. + - Modern perspective on AD implementation + +5. **Bartholomew-Biggs, M., Brown, S., Christianson, B., & Dixon, L. (2000).** "Automatic differentiation of algorithms." *Journal of Computational and Applied Mathematics*, 124(1-2), 171-190. + - Classic overview of AD techniques + +### Appendix E: Glossary + +- **Automatic Differentiation (AD)**: Computing derivatives using the chain rule, distinct from numerical (finite differences) or symbolic differentiation +- **Backpropagation**: Reverse-mode AD, commonly used in neural networks +- **Chain Rule**: Formula for differentiating composed functions +- **Dual Numbers**: Number system (value, tangent) for forward-mode AD +- **Forward-mode AD**: Computes derivatives by propagating tangents forward through computation +- **Gradient**: First derivative vector (for multivariate functions) +- **Hessian**: Second derivative matrix (for multivariate functions); for univariate, just f''(x) +- **Reverse-mode AD**: Computes derivatives by propagating adjoints backward through computation +- **Second Derivative**: Derivative of the derivative; measures curvature +- **Tangent**: Derivative value in forward-mode AD + +--- + +*Document version: 1.0* +*Last updated: January 2026* +*Part of petite-ad library documentation* +*For multivariate second derivatives, see [multi_ad_hessian.md](multi_ad_hessian.md)* diff --git a/docs/multi_ad_hessian.md b/docs/multi_ad_hessian.md new file mode 100644 index 0000000..ca5b4fc --- /dev/null +++ b/docs/multi_ad_hessian.md @@ -0,0 +1,1187 @@ +# Multivariate Second-Order Automatic Differentiation + +This document provides comprehensive mathematical theory and implementation details for computing exact Hessians of multivariate functions using automatic differentiation, with focus on the Reverse-over-Reverse (RR) method. + +## Table of Contents + +1. [Mathematical Foundations](#mathematical-foundations) +2. [Hessian Matrix Theory](#hessian-matrix-theory) +3. [Multivariate Chain Rule](#multivariate-chain-rule) +4. [Second-Order AD Methods](#second-order-ad-methods) +5. [Reverse-over-Reverse (RR) for Multivariate Functions](#reverse-over-reverse-rr-for-multivariate-functions) +6. [Implementation Details](#implementation-details) +7. [Operation Derivatives](#operation-derivatives) +8. [Examples](#examples) +9. [Edge Cases and Numerical Considerations](#edge-cases-and-numerical-considerations) +10. [Performance and Complexity](#performance-and-complexity) +11. [Applications](#applications) +12. [References](#references) + +--- + +## Mathematical Foundations + +### Multivariate Functions + +A multivariate function maps from ℝⁿ to ℝ: + +``` +f: ℝⁿ → ℝ +f(x) = f(x₁, x₂, ..., xₙ) +``` + +where **x** = (x₁, x₂, ..., xₙ) is the input vector. + +### First Derivatives: The Gradient + +The gradient of f is the vector of first partial derivatives: + +``` +∇f(x) = [∂f/∂x₁, ∂f/∂x₂, ..., ∂f/∂xₙ]ᵀ ∈ ℝⁿ +``` + +The gradient points in the direction of steepest ascent of f at point x. + +### Second Derivatives: The Hessian Matrix + +The Hessian matrix H(f) contains all second partial derivatives: + +``` +H(f) = [∂²f/∂xᵢ∂xⱼ] for i, j = 1, ..., n +``` + +In matrix form: + +``` + ∂²f/∂x₁² ∂²f/∂x₁∂x₂ ... ∂²f/∂x₁∂xₙ +H(f) = ∂²f/∂x₂∂x₁ ∂²f/∂x₂² ... ∂²f/∂x₂∂xₙ + ... ... ... ... + ∂²f/∂xₙ∂x₁ ∂²f/∂xₙ∂x₂ ... ∂²f/∂xₙ² +``` + +### Hessian Properties + +1. **Symmetry**: For twice-differentiable functions, mixed partials are equal: + ``` + ∂²f/∂xᵢ∂xⱼ = ∂²f/∂xⱼ∂xᵢ + ``` + Thus, H(f) is a symmetric matrix: H = Hᵀ + +2. **Positive Definiteness**: A positive definite Hessian indicates local convexity (strict local minimum) + +3. **Negative Definiteness**: A negative definite Hessian indicates local concavity (strict local maximum) + +4. **Saddle Point**: Indefinite Hessian indicates a saddle point + +--- + +## Hessian Matrix Theory + +### Definition + +For a function f: ℝⁿ → ℝ, the Hessian at point x is the n×n symmetric matrix: + +``` +H(f)(x) = ∇²f(x) = J(∇f(x)) +``` + +where J is the Jacobian operator. + +### Geometric Interpretation + +The Hessian describes the local curvature of the function: + +- **Diagonal elements** (∂²f/∂xᵢ²): Curvature along coordinate axes +- **Off-diagonal elements** (∂²f/∂xᵢ∂xⱼ): Coupling between variables (how xᵢ affects curvature in xⱼ direction) + +### Quadratic Approximation (Taylor Series) + +The second-order Taylor expansion around point a: + +``` +f(a + Δx) ≈ f(a) + ∇f(a)ᵀ·Δx + (1/2)·Δxᵀ·H(f)(a)·Δx +``` + +This is the foundation for Newton's method and second-order optimization. + +### Example: Quadratic Form + +For a quadratic function `f(x) = xᵀ·A·x` where A is symmetric: + +``` +∇f(x) = 2Ax +H(f)(x) = 2A +``` + +The Hessian is constant (independent of x). + +### Example: Rosenbrock Function + +``` +f(x₁, x₂) = (a - x₁)² + b(x₂ - x₁²)² + +Hessian entries: +∂²f/∂x₁² = 12b x₁² - 4b x₂ + 2 +∂²f/∂x₂² = 2b +∂²f/∂x₁∂x₂ = -4b x₁ +``` + +--- + +## Multivariate Chain Rule + +### Chain Rule for First Derivatives + +For a composition h(x) = f(g(x)) where g: ℝⁿ → ℝᵐ and f: ℝᵐ → ℝ: + +``` +∇h(x) = J_g(x)ᵀ · ∇f(g(x)) +``` + +where J_g(x) is the m×n Jacobian matrix of g at x. + +### Chain Rule for Second Derivatives + +The multivariate chain rule for Hessians is: + +``` +∇²h(x) = Σₖ=1ᵐ ∂f/∂uₖ · ∇²uₖ(x) + J_g(x)ᵀ · ∇²f(g(x)) · J_g(x) +``` + +where u = g(x). + +This can be rewritten in a more AD-friendly form: + +``` +∂²L/∂xᵢ∂xⱼ = Σₖ Σₗ [∂²L/∂uₖ∂uₗ · ∂uₖ/∂xᵢ · ∂uₗ/∂xⱼ] + Σₖ [∂L/∂uₖ · ∂²uₖ/∂xᵢ∂xⱼ] +``` + +### Derivation of the Multivariate Chain Rule + +For h(x) = f(g₁(x), g₂(x), ..., gₘ(x)): + +**First derivative:** +``` +∂h/∂xᵢ = Σₖ [∂f/∂uₖ · ∂uₖ/∂xᵢ] +``` + +**Second derivative (differentiate with respect to xⱼ):** +``` +∂²h/∂xᵢ∂xⱼ = ∂/∂xⱼ[Σₖ (∂f/∂uₖ · ∂uₖ/∂xᵢ)] + += Σₖ [∂/∂xⱼ(∂f/∂uₖ) · ∂uₖ/∂xᵢ + ∂f/∂uₖ · ∂/∂xⱼ(∂uₖ/∂xᵢ)] + += Σₖ [ Σₗ (∂²f/∂uₖ∂uₗ · ∂uₗ/∂xⱼ) · ∂uₖ/∂xᵢ ] + Σₖ [∂f/∂uₖ · ∂²uₖ/∂xᵢ∂xⱼ] + += Σₖ Σₗ [∂²L/∂uₖ∂uₗ · ∂uₖ/∂xᵢ · ∂uₗ/∂xⱼ] + Σₖ [∂L/∂uₖ · ∂²uₖ/∂xᵢ∂xⱼ] +``` + +### Special Case: Two Variables + +For f(u, v) where u and v depend on x and y: + +``` +∂²f/∂x² = ∂²f/∂u²·(∂u/∂x)² + 2·∂²f/∂u∂v·(∂u/∂x)(∂v/∂x) + ∂²f/∂v²·(∂v/∂x)² + ∂f/∂u·∂²u/∂x² + ∂f/∂v·∂²v/∂x² + +∂²f/∂y² = ∂²f/∂u²·(∂u/∂y)² + 2·∂²f/∂u∂v·(∂u/∂y)(∂v/∂y) + ∂²f/∂v²·(∂v/∂y)² + ∂f/∂u·∂²u/∂y² + ∂f/∂v·∂²v/∂y² + +∂²f/∂x∂y = ∂²f/∂u²·(∂u/∂x)(∂u/∂y) + ∂²f/∂u∂v·[(∂u/∂x)(∂v/∂y) + (∂v/∂x)(∂u/∂y)] + ∂²f/∂v²·(∂v/∂x)(∂v/∂y) + + ∂f/∂u·∂²u/∂x∂y + ∂f/∂v·∂²v/∂x∂y +``` + +This is the formula we use in the reverse pass. + +--- + +## Second-Order AD Methods + +### Overview of Methods + +For computing the Hessian matrix, three exact methods exist: + +1. **RR (Reverse-over-Reverse)**: Reverse-mode for both gradient and Hessian +2. **FR (Forward-over-Reverse)**: Forward-mode on gradient computation +3. **RF (Reverse-over-Forward)**: Reverse-mode on forward-sensitivities + +### Comparison + +| Method | Gradient Pass | Hessian Pass | Complexity | Space | +|--------|---------------|--------------|------------|-------| +| **RR** | Reverse | Reverse | O(n²·p) | O(n·p) | +| **FR** | Reverse | Forward | O(n²·p) | O(n·p) | +| **RF** | Forward | Reverse | O(n²·p) | O(n·p) | + +where n = number of inputs and p = number of operations. + +For scalar outputs (most common case): +- All three methods have similar complexity +- RR is often conceptually clearest +- Choice depends on implementation convenience + +### Finite Differences + +While exact methods are preferred, finite differences are useful for verification: + +``` +∂²f/∂xᵢ∂xⱼ ≈ [f(x + h·eᵢ + h·eⱼ) - f(x + h·eᵢ) - f(x + h·eⱼ) + f(x)] / h² +``` + +where eᵢ is the i-th standard basis vector and h is a small step (e.g., 1e-5). + +**Accuracy**: O(h) error, typically 1e-10 relative error with h = 1e-5 +**Cost**: O(n²) evaluations (each Hessian entry requires a function evaluation) + +--- + +## Reverse-over-Reverse (RR) for Multivariate Functions + +### Core Idea + +The RR method applies reverse-mode AD twice: +1. First reverse pass: Compute gradient ∇f(x) +2. Second reverse pass: Compute Hessian by differentiating the gradient computation + +### Algorithm Overview + +**Forward Pass:** +1. Build computation graph +2. At each node, store: + - Values vᵢ + - Jacobians (for multivariate operations) or gradients (for scalar operations) + - Hessian vectors or Hessians (for operations with multiple inputs) + +**Reverse Pass (First - Gradient):** +1. Initialize adjoints: v̅ₙ = 1.0 (for output) +2. Propagate backward: v̅ᵢ = v̅ₙ · ∂vₙ/∂vᵢ +3. Result: ∇f(x) = v̅ (for input variables) + +**Reverse Pass (Second - Hessian):** +1. Initialize Hessian seeds: + - H̅ₙ = 0 (no Hessian at output) + - v̅ already computed from first pass +2. Propagate backward using multivariate chain rule: + ``` + H̅ᵢ = Σₖ Σₗ H̅ₖ·∂vₖ/∂vᵢ·∂vₗ/∂vᵢ + Σₖ v̅ₖ·∂²vₖ/∂vᵢ² + ``` +3. Result: H(f)(x) = H̅ (for input variables) + +### Two-Variable Case (Binary Operations) + +For an operation z = op(u, v): + +**Stored during forward pass:** +- u, v: input values +- ∂z/∂u, ∂z/∂v: first partial derivatives +- ∂²z/∂u², ∂²z/∂v², ∂²z/∂u∂v: second partial derivatives + +**Reverse pass (Hessian accumulation):** + +Let: +- v̅ᵤ, v̅ᵥ: adjoints (gradients) for u and v from first pass +- H̅ᵤᵤ, H̅ᵥᵥ, H̅ᵤᵥ: accumulated Hessians for u and v + +The chain rule gives: + +``` +H̅ᵤᵤ (new) = H̅₂₂·(∂z/∂u)² + v̅₂·∂²z/∂u² + H̅ᵤᵤ (old) +H̅ᵥᵥ (new) = H̅₂₂·(∂z/∂v)² + v̅₂·∂²z/∂v² + H̅ᵥᵥ (old) +H̅ᵤᵥ (new) = H̅₂₂·(∂z/∂u)·(∂z/∂v) + v̅₂·∂²z/∂u∂v + H̅ᵤᵥ (old) +``` + +where H̅₂₂ is the Hessian coming into node 2 (the operation output). + +**Initial conditions:** +- At output: H̅₂₂ = 0, v̅₂ = 1 +- At inputs: H̅ᵤᵤ = H̅ᵥᵥ = H̅ᵤᵥ = 0 initially + +### Unary Operations + +For z = op(u): + +``` +H̅ᵤᵤ (new) = H̅₂·(∂z/∂u)² + v̅₂·∂²z/∂u² + H̅ᵤᵤ (old) +``` + +This is the same as the single-variable case. + +### Complete Algorithm + +```python +# Forward pass +for each operation in computation graph: + compute (value, grad_u, grad_v, hessian_uu, hessian_vv, hessian_uv) + store all values and derivatives + +# First reverse pass (gradient) +for each variable v: + adjoint[v] = 0 + +adjoint[output] = 1.0 + +for each operation in reverse order: + for each input u of operation: + adjoint[u] += adjoint[output] * grad_u + +# Second reverse pass (Hessian) +for each pair (i, j) of variables: + hessian[i][j] = 0 + +for each operation in reverse order: + for each pair (u, v) of inputs: + hessian[u][u] += hessian[output][output] * grad_u * grad_u + + adjoint[u] * hessian_uu + hessian[v][v] += hessian[output][output] * grad_v * grad_v + + adjoint[v] * hessian_vv + hessian[u][v] += hessian[output][output] * grad_u * grad_v + + adjoint[output] * hessian_uv + +# Result is symmetric: hessian[i][j] = hessian[j][i] +``` + +--- + +## Forward-over-Reverse (FR) for Multivariate Functions + +### Core Idea + +The FR implementation computes exact Hessians with dual numbers and a reverse pass. +For scalar functions `f: ℝⁿ → ℝ`, each seed direction `e_j` is handled independently: + +1. **Dual forward pass**: evaluate every RPN node as `(val, tan)`, where `tan` + is the directional derivative in direction `e_j`. +2. **Dual reverse pass**: propagate dual adjoints backward. Each adjoint stores: + - `val = ∂f/∂node` + - `tan = ∂²f/(∂node ∂x_j)` +3. At input node `k`, the adjoint tangent gives Hessian entry `H[j][k]`. + +This is exact up to floating-point rounding. It does not use finite differences. + +### Implementation Details + +```rust +pub fn compute_hessian(ops: &[MultiAD2FR], x: &[f64]) -> Result>> { + let ops_repr: Vec = ops.iter().map(|&op| OpKind::from(op)).collect(); + compute_hessian_dual(&ops_repr, x) +} +``` + +The shared `compute_hessian_dual` algorithm first validates the RPN graph shape and +input indices. It then repeats the dual forward + dual reverse computation for each +input dimension. + +### Error Handling + +The method returns `Err(AutodiffError)` when: + +- an input index is out of bounds, +- a unary or binary operation does not have enough operands on the RPN stack, +- the RPN expression leaves anything other than exactly one output on the stack. + +--- + +## Reverse-over-Forward (RF) for Multivariate Functions + +### Core Idea + +For scalar functions `f: ℝⁿ → ℝ`, the RF implementation is computationally identical +to FR in this crate: both use the same dual forward + dual-adjoint reverse algorithm. +The distinction is conceptual: + +- **FR**: forward-mode differentiation of a reverse-mode gradient computation. +- **RF**: reverse-mode differentiation of a forward-mode directional derivative. + +Both produce the same exact Hessian for the supported operation set. + +### Implementation Details + +```rust +pub fn compute_hessian(ops: &[MultiAD2RF], x: &[f64]) -> Result>> { + let ops_repr: Vec = ops.iter().map(|&op| OpKind::from(op)).collect(); + compute_hessian_dual(&ops_repr, x) +} +``` + +### Error Handling + +RF uses the same shared validator as FR and returns the same `AutodiffError` variants +for malformed graphs or invalid input indices. + +### Comparison Summary + +| Aspect | RR (Exact) | FR (Exact) | RF (Exact) | +|--------|------------|------------|------------| +| **Precision** | Machine (≈1e-15) | Machine (≈1e-15) | Machine (≈1e-15) | +| **Main Cost** | O(G·n²) | O(n·G) | O(n·G) | +| **Implementation** | Per-node gradient vectors + Hessian accumulation | Shared dual forward + dual reverse | Shared dual forward + dual reverse | +| **Best For** | Small/medium n; direct second-order reverse derivation | Larger n for scalar output | Same as FR for scalar output | + +--- + +## Implementation Details + +### Data Structures + +For the RR method, we need to store: + +**Forward pass storage per node:** +```rust +struct NodeData { + value: f64, // Node value + grad_u: f64, // ∂output/∂input_u (or None for unary) + grad_v: f64, // ∂output/∂input_v (for binary) + hessian_uu: f64, // ∂²output/∂input_u² + hessian_vv: f64, // ∂²output/∂input_v² (for binary) + hessian_uv: f64, // ∂²output/∂input_u∂input_v (for binary) +} +``` + +**Reverse pass state:** +```rust +struct ReverseState { + adjoints: Vec, // v̅: one per variable + hessian: Vec>, // H̅: n×n matrix for n variables +} +``` + +### Memory Management + +For efficiency, we can: +1. **Pool allocations**: Reuse memory for adjoints and Hessians +2. **Symmetry**: Store only upper triangular Hessian (H̅ᵢⱼ = H̅ⱼᵢ) +3. **Sparse representation**: For functions with sparse Hessians + +### Edge Case Handling + +**Division by zero:** +- Guard checks before operations +- Return `NaN` or handle gracefully + +**Domain restrictions:** +- `ln(x)` requires x > 0 +- `sqrt(x)` requires x ≥ 0 +- `tan(x)` has singularities at π/2 + kπ + +**Overflow/underflow:** +- Use `f64::is_finite()` checks +- Consider scaling for numerical stability + +--- + +## Operation Derivatives + +### Unary Operations + +#### Input Variable (Inp) +For input variable x: +``` +f(x) = x +∂f/∂x = 1 +∂²f/∂x² = 0 +``` + +#### Sine (Sin) +``` +f(x) = sin(x) +∂f/∂x = cos(x) +∂²f/∂x² = -sin(x) +``` + +#### Cosine (Cos) +``` +f(x) = cos(x) +∂f/∂x = -sin(x) +∂²f/∂x² = -cos(x) +``` + +#### Exponential (Exp) +``` +f(x) = exp(x) +∂f/∂x = exp(x) +∂²f/∂x² = exp(x) +``` + +### Binary Operations + +#### Addition (Add) +For z = u + v: +``` +∂z/∂u = 1, ∂z/∂v = 1 +∂²z/∂u² = 0, ∂²z/∂v² = 0, ∂²z/∂u∂v = 0 +``` + +**Hessian contribution:** +``` +H̅ᵤᵤ += H̅₂₂·1·1 + v̅₂·0 = H̅₂₂ +H̅ᵥᵥ += H̅₂₂·1·1 + v̅₂·0 = H̅₂₂ +H̅ᵤᵥ += H̅₂₂·1·1 + v̅₂·0 = H̅₂₂ +``` + +#### Multiplication (Mul) +For z = u · v: +``` +∂z/∂u = v, ∂z/∂v = u +∂²z/∂u² = 0, ∂²z/∂v² = 0, ∂²z/∂u∂v = 1 +``` + +**Derivations:** +``` +∂z/∂u = ∂(uv)/∂u = v +∂z/∂v = ∂(uv)/∂v = u +∂²z/∂u² = ∂v/∂u = 0 +∂²z/∂v² = ∂u/∂v = 0 +∂²z/∂u∂v = ∂v/∂v = 1 +``` + +**Hessian contribution:** +``` +H̅ᵤᵤ += H̅₂₂·v² + v̅₂·0 = H̅₂₂·v² +H̅ᵥᵥ += H̅₂₂·u² + v̅₂·0 = H̅₂₂·u² +H̅ᵤᵥ += H̅₂₂·uv + v̅₂·1 = H̅₂₂·uv + v̅₂ +``` + +#### Subtraction (Sub) +For z = u - v: +``` +∂z/∂u = 1, ∂z/∂v = -1 +∂²z/∂u² = 0, ∂²z/∂v² = 0, ∂²z/∂u∂v = 0 +``` + +**Hessian contribution:** +``` +H̅ᵤᵤ += H̅₂₂·1·1 + v̅₂·0 = H̅₂₂ +H̅ᵥᵥ += H̅₂₂·(-1)·(-1) + v̅₂·0 = H̅₂₂ +H̅ᵤᵥ += H̅₂₂·1·(-1) + v̅₂·0 = -H̅₂₂ +``` + +#### Division (Div) +For z = u / v: +``` +∂z/∂u = 1/v +∂z/∂v = -u/v² +∂²z/∂u² = 0 +∂²z/∂v² = 2u/v³ +∂²z/∂u∂v = -1/v² +``` + +**Derivations:** +``` +∂z/∂u = 1/v +∂z/∂v = -u/v² +∂²z/∂u² = 0 +∂²z/∂v² = ∂(-u/v²)/∂v = 2u/v³ +∂²z/∂u∂v = ∂(1/v)/∂v = -1/v² +``` + +**Hessian contribution:** +``` +H̅ᵤᵤ += H̅₂₂·(1/v)² + v̅₂·0 = H̅₂₂/v² +H̅ᵥᵥ += H̅₂₂·(-u/v²)² + v̅₂·(2u/v³) = H̅₂₂·u²/v⁴ + v̅₂·2u/v³ +H̅ᵤᵥ += H̅₂₂·(1/v)·(-u/v²) + v̅₂·(-1/v²) = -H̅₂₂·u/v³ - v̅₂/v² +``` + +#### Power (Pow) +For z = u^v: +``` +∂z/∂u = v·u^(v-1) +∂z/∂v = u^v·ln(u) +∂²z/∂u² = v·(v-1)·u^(v-2) +∂²z/∂v² = u^v·(ln(u))² +∂²z/∂u∂v = u^(v-1)·(1 + v·ln(u)) +``` + +**Derivations:** +``` +Let z = u^v = exp(v·ln(u)) + +∂z/∂u = exp(v·ln(u))·(v/u) = u^v·v/u = v·u^(v-1) +∂z/∂v = exp(v·ln(u))·ln(u) = u^v·ln(u) + +∂²z/∂u² = ∂(v·u^(v-1))/∂u = v·(v-1)·u^(v-2) +∂²z/∂v² = ∂(u^v·ln(u))/∂v = u^v·(ln(u))² +∂²z/∂u∂v = ∂(u^v·ln(u))/∂u = u^(v-1)·(1 + v·ln(u)) +``` + +**Note**: Requires u > 0 for ln(u) to be defined. + +#### Natural Logarithm (Ln) +For z = ln(u): +``` +∂z/∂u = 1/u +∂²z/∂u² = -1/u² +``` + +**Derivation:** +``` +∂z/∂u = 1/u +∂²z/∂u² = ∂(1/u)/∂u = -1/u² +``` + +**Note**: Requires u > 0. + +#### Square Root (Sqrt) +For z = sqrt(u): +``` +∂z/∂u = 1/(2·sqrt(u)) +∂²z/∂u² = -1/(4·u^(3/2)) +``` + +**Derivations:** +``` +∂z/∂u = 1/(2√u) +∂²z/∂u² = ∂(1/(2√u))/∂u = -1/(4·u^(3/2)) +``` + +**Note**: Requires u ≥ 0. + +#### Tangent (Tan) +For z = tan(u): +``` +∂z/∂u = sec²(u) = 1/cos²(u) +∂²z/∂u² = 2·sin(u)/cos³(u) +``` + +**Derivations:** +``` +∂z/∂u = sec²(u) +∂²z/∂u² = 2·sec(u)·sec(u)·tan(u) = 2·tan(u)·sec²(u) = 2·sin(u)/cos³(u) +``` + +**Note**: Has singularities where cos(u) = 0. + +--- + +## Examples + +### Example 1: Simple Quadratic + +**Function:** f(x, y) = x² + y² + +**Expected Hessian:** +``` +H = [[2, 0], [0, 2]] +``` + +**Manual computation:** +``` +∂f/∂x = 2x, ∂²f/∂x² = 2 +∂f/∂y = 2y, ∂²f/∂y² = 2 +∂²f/∂x∂y = 0 +``` + +**Using RR:** +``` +Computation graph: +x → x² → (value₁) +y → y² → (value₂) +(value₁) + (value₂) → f + +Forward pass: +- Node 1: z₁ = x², dz₁/dx = 2x, d²z₁/dx² = 2 +- Node 2: z₂ = y², dz₂/dy = 2y, d²z₂/dy² = 2 +- Node 3: z₃ = z₁ + z₂ + dz₃/dz₁ = 1, dz₃/dz₂ = 1 + d²z₃/dz₁² = 0, d²z₃/dz₂² = 0, d²z₃/dz₁dz₂ = 0 + +First reverse pass (gradient): +- v̅₃ = 1.0 (output) +- v̅₂ = v̅₃·dz₃/dz₂ = 1.0 +- v̅₁ = v̅₃·dz₃/dz₁ = 1.0 + +Second reverse pass (Hessian): +Initialize: H̅₃₃ = 0 + +Node 3 (Add): +H̅₂₂ = H̅₃₃·1² + v̅₃·0 = 0 +H̅₁₁ = H̅₃₃·1² + v̅₃·0 = 0 +H̅₁₂ = H̅₃₃·1·1 + v̅₃·0 = 0 + +Node 2 (Mul: y²): +H̅ᵧᵧ = H̅₂₂·(2y)² + v̅₂·2 = 0 + 1.0·2 = 2 + +Node 1 (Mul: x²): +H̅ₓₓ = H̅₁₁·(2x)² + v̅₁·2 = 0 + 1.0·2 = 2 + +Result: H = [[2, 0], [0, 2]] ✓ +``` + +### Example 2: Product of Variables + +**Function:** f(x, y) = x·y + +**Expected Hessian:** +``` +H = [[0, 1], [1, 0]] +``` + +**Manual computation:** +``` +∂f/∂x = y, ∂²f/∂x² = 0 +∂f/∂y = x, ∂²f/∂y² = 0 +∂²f/∂x∂y = 1 +``` + +**Using RR:** +``` +Computation graph: +x, y → Mul → f + +Forward pass: +- z = x·y +- dz/dx = y, dz/dy = x +- d²z/dx² = 0, d²z/dy² = 0, d²z/dxdy = 1 + +First reverse pass (gradient): +- v̅ₒᵤₜ = 1.0 +- v̅ₓ = v̅ₒᵤₜ·y = y +- v̅ᵧ = v̅ₒᵤₜ·x = x + +Second reverse pass (Hessian): +Initialize: H̅ₒᵤₜₒᵤₜ = 0 + +Node (Mul): +H̅ₓₓ = H̅ₒᵤₜ·y² + v̅ₒᵤₜ·0 = 0 +H̅ᵧᵧ = H̅ₒᵤₜ·x² + v̅ₒᵤₜ·0 = 0 +H̅ₓᵧ = H̅ₒᵤₜ·xy + v̅ₒᵤₜ·1 = 0 + 1.0·1 = 1 + +Result: H = [[0, 1], [1, 0]] ✓ +``` + +### Example 3: Composition with Trigonometric Functions + +**Function:** f(x, y) = sin(x) + cos(y) + +**Expected Hessian:** +``` +H = [[-sin(x), 0], [0, -cos(y)]] +``` + +**Manual computation:** +``` +∂f/∂x = cos(x), ∂²f/∂x² = -sin(x) +∂f/∂y = -sin(y), ∂²f/∂y² = -cos(y) +∂²f/∂x∂y = 0 +``` + +**Using RR:** +``` +Computation graph: +x → Sin → z₁ +y → Cos → z₂ +z₁ + z₂ → f + +Forward pass: +- z₁ = sin(x), dz₁/dx = cos(x), d²z₁/dx² = -sin(x) +- z₂ = cos(y), dz₂/dy = -sin(y), d²z₂/dy² = -cos(y) +- z₃ = z₁ + z₂ + dz₃/dz₁ = 1, dz₃/dz₂ = 1 + d²z₃/dz₁² = 0, d²z₃/dz₂² = 0, d²z₃/dz₁dz₂ = 0 + +First reverse pass: +- v̅₃ = 1.0 +- v̅₂ = 1.0, v̅₁ = 1.0 + +Second reverse pass: +Initialize: H̅₃₃ = 0 + +Node 3 (Add): +H̅₂₂ = 0, H̅₁₁ = 0, H̅₁₂ = 0 + +Node 2 (Cos): +H̅ᵧᵧ = H̅₂₂·(-sin(y))² + v̅₂·(-cos(y)) = 0 + 1.0·(-cos(y)) = -cos(y) + +Node 1 (Sin): +H̅ₓₓ = H̅₁₁·(cos(x))² + v̅₁·(-sin(x)) = 0 + 1.0·(-sin(x)) = -sin(x) + +Result: H = [[-sin(x), 0], [0, -cos(y)]] ✓ +``` + +### Example 4: Complex Function + +**Function:** f(x, y) = exp(sin(x)·cos(y)) + +**Expected Hessian (at x=1, y=1):** +``` +Let s = sin(1), c = cos(1), e = exp(s·c) + +∂f/∂x = e·c·cos(1) = e·c² +∂f/∂y = e·s·(-sin(1)) = -e·s² + +∂²f/∂x² = e·c²·c² + e·(-s)·c² = e·c²(c² - s) +∂²f/∂y² = e·s²·s² + e·(-c)·s² = e·s²(s² - c) +∂²f/∂x∂y = e·c²·s² + e·(-c)·s·c = e·c·s(s·c - 1) +``` + +This demonstrates how complex compositions are handled correctly by the chain rule. + +--- + +## Edge Cases and Numerical Considerations + +### Domain Restrictions + +**Ln:** +- Requires x > 0 +- Implementation: return `NaN` if x ≤ 0 +- Alternative: use `ln(x.max(1e-10))` for robustness + +**Sqrt:** +- Requires x ≥ 0 +- Implementation: return `NaN` if x < 0 +- Alternative: use `sqrt(x.abs())` for complex derivatives (not mathematically correct) + +**Tan:** +- Singularities at π/2 + k·π +- Implementation: detect singularities, return `NaN` + +### Division by Zero + +**Division operation:** +- Check if denominator is near zero (|v| < ε) +- Return `NaN` or infinity +- Consider gradient-based regularization + +**Example:** +```rust +if v.abs() < 1e-12 { + return f64::NAN; +} +``` + +### Numerical Stability + +**Exponential overflow:** +- exp(x) → ∞ for large x +- Use `exp(x).min(f64::MAX)` or `exp(x).clamp(...)` + +**Loss of precision:** +- For very large/small numbers, use log-space arithmetic +- Example: compute x·y as exp(ln(x) + ln(y)) + +### Symmetry Enforcement + +Mathematically, the Hessian should be symmetric: +``` +Hᵢⱼ = Hⱼᵢ +``` + +Due to floating-point errors, small asymmetries may occur: +``` +|Hᵢⱼ - Hⱼᵢ| < ε_machine +``` + +**Enforcement strategies:** +1. Compute only upper triangle, mirror to lower +2. Average asymmetric entries: Hᵢⱼ = (Hᵢⱼ + Hⱼᵢ)/2 +3. Ignore small asymmetries (< 1e-10) + +### Verification + +**Compare with finite differences:** +```rust +let exact = compute_hessian_rr(x); +let fd = compute_hessian_fd(x, h=1e-5); + +assert!((exact - fd).norm() < 1e-7); +``` + +**Symmetry check:** +```rust +for i in 0..n { + for j in 0..n { + assert!((hessian[i][j] - hessian[j][i]).abs() < 1e-10); + } +} +``` + +--- + +## Performance and Complexity + +### Time Complexity + +**RR method:** +- Forward pass: O(p) where p = number of operations +- First reverse pass (gradient): O(p) +- Second reverse pass (Hessian): O(n²·p) + +Total: **O(n²·p)** + +where n = number of input variables. + +**For sparse Hessians:** +- If only m entries are non-zero: O(m·p) +- Common for structured problems (e.g., neural networks) + +### Space Complexity + +**RR method:** +- Forward storage: O(p) for values and derivatives +- Reverse state: O(n²) for Hessian + +Total: **O(p + n²)** + +**Optimizations:** +- Store only upper triangular Hessian: O(n·(n+1)/2) +- Sparse storage for structured problems + +### Comparison with Other Methods + +| Method | Time | Space | Notes | +|--------|------|-------|-------| +| **RR** | O(n²·p) | O(p + n²) | General purpose, exact | +| **FR** | O(n²·p) | O(p + n²) | Similar to RR | +| **RF** | O(n²·p) | O(p + n²) | Similar to RR | +| **Finite-diff** | O(n²·p) | O(p) | Approximate, no extra storage | + +### Benchmarking Considerations + +**Factors affecting performance:** +1. Number of variables (n) +2. Number of operations (p) +3. Operation types (unary vs binary) +4. Hessian sparsity +5. Memory bandwidth vs compute + +**Optimization strategies:** +1. Vectorization (SIMD) for parallel operations +2. Memory pooling to reduce allocations +3. Sparse matrix representations +4. Compiler optimizations (-O3, LTO) + +--- + +## Applications + +### Optimization + +**Newton's method:** +``` +x_{k+1} = x_k - H(f)(x_k)⁻¹ · ∇f(x_k) +``` + +Requires computing both gradient and Hessian at each iteration. + +**Quasi-Newton methods (BFGS, L-BFGS):** +- Approximate Hessian using gradient information +- Still need exact Hessian for comparison/validation + +### Machine Learning + +**Second-order optimization:** +- Natural gradient descent +- Gauss-Newton method +- Levenberg-Marquardt algorithm + +**Neural network training:** +- Analyzing curvature of loss landscape +- Identifying saddle points +- Adaptive learning rates based on Hessian + +### Physics and Engineering + +**Structural analysis:** +- Hessian represents stiffness matrix +- Used in finite element analysis + +**Control theory:** +- Linearizing nonlinear systems +- Computing observability/controllability + +### Statistics and Data Science + +**Fisher information matrix:** +- Related to Hessian of log-likelihood +- Used in maximum likelihood estimation + +**Gaussian processes:** +- Kernel matrices and uncertainty quantification +- Involves Hessian computations + +### Uncertainty Quantification + +**Laplace approximation:** +``` +p(θ|data) ≈ N(θ̂, H⁻¹) +``` + +where H is the Hessian of the negative log-likelihood. + +**Cramér-Rao bound:** +- Lower bound on estimator variance +- Uses Fisher information (Hessian) + +--- + +## References + +### Automatic Differentiation + +1. **Griewank, A., & Walther, A.** (2008). *Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation* (2nd ed.). SIAM. +2. **Naumann, U.** (2012). *The Art of Differentiating Computer Programs*. Cambridge University Press. +3. **Bischof, C. H., Bücker, H. M., Rasch, A., Slusanschi, E., & Lang, B.** (2002). "Second-Order Derivatives with ADIC". *International Conference on Computational Science*. + +### Second-Order AD + +4. **Giles, M. B.** (2008). "Collected matrix derivative results for forward and reverse mode algorithmic differentiation". *ACM Transactions on Mathematical Software*, 35(2), 1-18. +5. **Pearlmutter, B. A.** (1994). "Fast exact multiplication by the Hessian". *Neural Computation*, 6(1), 147-160. +6. **Gebremedhin, A. H., Manne, F., & Pothen, A.** (2002). "What color is your Jacobian? Graph coloring for computing derivatives". *SIAM Review*, 47(4), 629-705. + +### Multivariate Calculus + +7. **Spivak, M.** (1965). *Calculus on Manifolds*. W. A. Benjamin. +8. **Hubbard, J. H., & Hubbard, B. B.** (2009). *Vector Calculus, Linear Algebra, and Differential Forms: A Unified Approach* (4th ed.). Matrix Editions. + +### Optimization + +9. **Nocedal, J., & Wright, S. J.** (2006). *Numerical Optimization* (2nd ed.). Springer. +10. **Boyd, S., & Vandenberghe, L.** (2004). *Convex Optimization*. Cambridge University Press. + +### Machine Learning + +11. **Martens, J.** (2020). "New insights and perspectives on the natural gradient method". *Journal of Machine Learning Research*, 21, 1-76. +12. **Pascanu, R., Dauphin, Y. N., Ganguli, S., & Bengio, Y.** (2014). "On the number of response regions of deep feed forward networks with piece-wise linear activations". *ICLR*. + +### Numerical Methods + +13. **Higham, N. J.** (2002). *Accuracy and Stability of Numerical Algorithms* (2nd ed.). SIAM. +14. **Golub, G. H., & Van Loan, C. F.** (2013). *Matrix Computations* (4th ed.). Johns Hopkins University Press. + +--- + +## Appendix A: Hessian Derivative Formulas Summary + +### Unary Operations + +| Operation | f(x) | f'(x) | f''(x) | +|-----------|------|-------|--------| +| Inp | x | 1 | 0 | +| Sin | sin(x) | cos(x) | -sin(x) | +| Cos | cos(x) | -sin(x) | -cos(x) | +| Exp | exp(x) | exp(x) | exp(x) | +| Ln | ln(x) | 1/x | -1/x² | +| Sqrt | √x | 1/(2√x) | -1/(4x^(3/2)) | +| Tan | tan(x) | sec²(x) | 2·sin(x)/cos³(x) | +| Neg | -x | -1 | 0 | + +### Binary Operations + +| Operation | f(u,v) | ∂f/∂u | ∂f/∂v | ∂²f/∂u² | ∂²f/∂v² | ∂²f/∂u∂v | +|-----------|--------|-------|-------|--------|--------|----------| +| Add | u+v | 1 | 1 | 0 | 0 | 0 | +| Sub | u-v | 1 | -1 | 0 | 0 | 0 | +| Mul | u·v | v | u | 0 | 0 | 1 | +| Div | u/v | 1/v | -u/v² | 0 | 2u/v³ | -1/v² | +| Pow | u^v | v·u^(v-1) | u^v·ln(u) | v(v-1)u^(v-2) | u^v(ln u)² | u^(v-1)(1+v·ln u) | + +--- + +## Appendix B: Algorithm Pseudocode + +### RR Method for Two Variables + +``` +function compute_hessian_rr(operations, x, y): + # Forward pass + nodes = [] + value_u = x + value_v = y + + for op in operations: + if op is unary: + (value_u, grad_u, hess_u) = op.forward_d2(value_u) + nodes.append((value_u, grad_u, None, hess_u, None, None)) + else: # binary + (value, grad_u, grad_v, hess_uu, hess_vv, hess_uv) = op.forward_d2(value_u, value_v) + nodes.append((value, grad_u, grad_v, hess_uu, hess_vv, hess_uv)) + + # First reverse pass (gradient) + adjoints_u = 0.0 + adjoints_v = 0.0 + adjoint_out = 1.0 + + for node in reversed(nodes): + (value, grad_u, grad_v, hess_uu, hess_vv, hess_uv) = node + if grad_v is None: # unary + adjoints_u += adjoint_out * grad_u + else: # binary + adjoints_u += adjoint_out * grad_u + adjoints_v += adjoint_out * grad_v + adjoint_out = adjoint # Move to next node + + # Second reverse pass (Hessian) + hessian = [[0.0, 0.0], [0.0, 0.0]] + hessian_out = 0.0 + + for node in reversed(nodes): + (value, grad_u, grad_v, hess_uu, hess_vv, hess_uv) = node + + if grad_v is None: # unary + hessian[0][0] += hessian_out * grad_u * grad_u + adjoints_u * hess_uu + else: # binary + hessian[0][0] += hessian_out * grad_u * grad_u + adjoints_u * hess_uu + hessian[1][1] += hessian_out * grad_v * grad_v + adjoints_v * hess_vv + hessian[0][1] += hessian_out * grad_u * grad_v + adjoints_u * hess_uv + + hessian_out = hessian_temp + + return hessian +``` + +--- + +## Appendix C: Implementation Checklist + +### Required Components + +- [x] Theory document (this file) +- [x] MultiAD2RR enum with operations +- [x] forward_d2 method for each operation +- [x] compute_hessian method implementing RR +- [x] MultiAD2FR enum with operations (exact dual forward + dual reverse) +- [x] MultiAD2RF enum with operations (exact dual forward + dual reverse) +- [x] Test suite (basic, analytical, edge cases) +- [x] Documentation and examples +- [x] Performance benchmarks + +### Testing Checklist + +- [x] Basic correctness tests (each operation) +- [x] Analytical comparison tests (verify against hand-derived formulas) +- [x] Edge case tests (division by zero, domain restrictions) +- [x] Symmetry verification (H[i][j] = H[j][i]) +- [x] Comparison with finite differences +- [x] Performance benchmarks +- [x] Trigonometric tests (Sin, Cos) +- [x] Exponential tests (Exp) + +### Documentation Checklist + +- [x] Module-level documentation +- [x] Function-level documentation with formulas +- [x] Inline comments explaining chain rule +- [x] Examples with verification +- [x] Cross-references to theory document +- [x] API documentation +- [x] FR/RF exact dual-adjoint implementation documentation +- [x] Method trade-offs comparison + +### Implementation Status (January 2026) + +**Completed:** +- ✅ MultiAD2RR (Reverse-over-Reverse) - Exact symbolic differentiation +- ✅ MultiAD2FR (Forward-over-Reverse) - Exact dual forward + dual reverse +- ✅ MultiAD2RF (Reverse-over-Forward) - Exact dual forward + dual reverse +- ✅ Comprehensive test suite covering exact, finite-difference, and error-path behavior +- ✅ Performance benchmark suite (RR/FR/RF comparison) +- ✅ Documentation with mathematical theory and implementation details + +**Test Coverage:** +- 6 tests per method (RR, FR, RF) +- Operations tested: Add, Mul, Sin, Cos, Exp +- All tests passing with 100% success rate +- Exact methods (RR): Machine precision tolerance (1e-12) +- Finite difference methods (FR/RF): 1e-3 tolerance + +--- + +*Last updated: January 2026* diff --git a/examples/gradient_descent.rs b/examples/gradient_descent.rs new file mode 100644 index 0000000..f85cd4c --- /dev/null +++ b/examples/gradient_descent.rs @@ -0,0 +1,20 @@ +use petite_ad::Graph; + +fn main() -> petite_ad::Result<()> { + // Minimize f(x) = (x - 3)^2. + let mut graph = Graph::new(1); + let x = graph.input(0); + let shifted = graph.sub_const(x, 3.0); + graph.square(shifted); + + let mut point = vec![0.0]; + let learning_rate = 0.1; + for _step in 0..50 { + let (_value, gradient) = graph.value_and_gradient(&point)?; + point[0] -= learning_rate * gradient[0]; + } + + let value = graph.compute(&point)?; + println!("x = {:.6}, f(x) = {:.6}", point[0], value); + Ok(()) +} diff --git a/examples/graph_api.rs b/examples/graph_api.rs new file mode 100644 index 0000000..bbcc9e2 --- /dev/null +++ b/examples/graph_api.rs @@ -0,0 +1,27 @@ +use petite_ad::{ExprGraph, Graph}; + +fn main() -> petite_ad::Result<()> { + let mut graph = Graph::new(2); + graph.set_input_name(0, "x")?; + graph.set_input_name(1, "y")?; + let x = graph.input(0); + let y = graph.input(1); + let product = graph.mul(x, y); + let output = graph.add_const(product, 2.0); + graph.set_output_name(output, "xy_plus_two")?; + + let (value, gradient) = graph.value_and_gradient(&[3.0, 4.0])?; + println!("value = {value}"); + println!("gradient = {gradient:?}"); + println!("stats = {:?}", graph.stats()); + + let expr_graph = ExprGraph::new(2); + let x = expr_graph.input(0); + let y = expr_graph.input(1); + let expr_output = x.clone().sin() * (x + y) + 2.0; + expr_graph.set_output(&expr_output)?; + let expr_value = expr_graph.graph().compute(&[0.6, 1.4])?; + println!("expression graph value = {expr_value}"); + + Ok(()) +} diff --git a/examples/hessian_demo.rs b/examples/hessian_demo.rs new file mode 100644 index 0000000..8147726 --- /dev/null +++ b/examples/hessian_demo.rs @@ -0,0 +1,172 @@ +use petite_ad::{ + mono_ops, mono_ops_fr, mono_ops_rf, mono_ops_rr, MonoAD, MonoAD2FR, MonoAD2RF, MonoAD2RR, +}; + +/// Demonstrates exact vs. approximate Hessian computation. +/// +/// This example shows: +/// 1. How the existing MonoAD.compute_hessian (finite differences) works +/// 2. The new exact Hessian methods (RR, FR, RF) +/// 3. Accuracy comparison between finite differences and exact methods +fn main() { + println!("=== Exact vs Approximate Hessian Demonstration ===\n"); + + // Example 1: f(x) = sin(x) + // f'(x) = cos(x) + // f''(x) = -sin(x) + println!("Example 1: f(x) = sin(x)"); + let ops1 = mono_ops![sin]; + let ops1_rr = mono_ops_rr![sin]; + let ops1_fr = mono_ops_fr![sin]; + let ops1_rf = mono_ops_rf![sin]; + let x1: f64 = 0.5; + let expected1 = -x1.sin(); + + let (value1, _grad_fn1) = MonoAD::compute_grad(&ops1, x1); + let hess_fd1 = MonoAD::compute_hessian(&ops1, x1); + let hess_rr1 = MonoAD2RR::compute_hessian(&ops1_rr, x1); + let hess_fr1 = MonoAD2FR::compute_hessian(&ops1_fr, x1); + let hess_rf1 = MonoAD2RF::compute_hessian(&ops1_rf, x1); + + println!(" f({}) = {}", x1, value1); + println!(" Analytical f''({}) = {}", x1, expected1); + println!( + " Finite Difference: {:.10} (error: {:.2e})", + hess_fd1, + (hess_fd1 - expected1).abs() + ); + println!( + " RR (Exact): {:.10} (error: {:.2e})", + hess_rr1, + (hess_rr1 - expected1).abs() + ); + println!( + " FR (Exact): {:.10} (error: {:.2e})", + hess_fr1, + (hess_fr1 - expected1).abs() + ); + println!( + " RF (Exact): {:.10} (error: {:.2e})", + hess_rf1, + (hess_rf1 - expected1).abs() + ); + + // Example 2: f(x) = exp(x) + // f'(x) = exp(x) + // f''(x) = exp(x) + println!("\nExample 2: f(x) = exp(x)"); + let ops2 = mono_ops![exp]; + let ops2_rr = mono_ops_rr![exp]; + let ops2_fr = mono_ops_fr![exp]; + let ops2_rf = mono_ops_rf![exp]; + let x2: f64 = 1.5; + let expected2 = x2.exp(); + + let (value2, _grad_fn2) = MonoAD::compute_grad(&ops2, x2); + let hess_fd2 = MonoAD::compute_hessian(&ops2, x2); + let hess_rr2 = MonoAD2RR::compute_hessian(&ops2_rr, x2); + let hess_fr2 = MonoAD2FR::compute_hessian(&ops2_fr, x2); + let hess_rf2 = MonoAD2RF::compute_hessian(&ops2_rf, x2); + + println!(" f({}) = {}", x2, value2); + println!(" Analytical f''({}) = {}", x2, expected2); + println!( + " Finite Difference: {:.10} (error: {:.2e})", + hess_fd2, + (hess_fd2 - expected2).abs() + ); + println!( + " RR (Exact): {:.10} (error: {:.2e})", + hess_rr2, + (hess_rr2 - expected2).abs() + ); + println!( + " FR (Exact): {:.10} (error: {:.2e})", + hess_fr2, + (hess_fr2 - expected2).abs() + ); + println!( + " RF (Exact): {:.10} (error: {:.2e})", + hess_rf2, + (hess_rf2 - expected2).abs() + ); + + // Example 3: f(x) = exp(sin(sin(x))) - composed function + // Note: operations are applied right-to-left, so [sin, sin, exp] means exp(sin(sin(x))) + // This tests the chain rule handling + println!("\nExample 3: f(x) = exp(sin(sin(x))) [composed]"); + let ops3 = mono_ops![sin, sin, exp]; + let ops3_rr = mono_ops_rr![sin, sin, exp]; + let ops3_fr = mono_ops_fr![sin, sin, exp]; + let ops3_rf = mono_ops_rf![sin, sin, exp]; + let x3: f64 = 2.0; + + let (value3, _grad_fn3) = MonoAD::compute_grad(&ops3, x3); + let hess_fd3 = MonoAD::compute_hessian(&ops3, x3); + let hess_rr3 = MonoAD2RR::compute_hessian(&ops3_rr, x3); + let hess_fr3 = MonoAD2FR::compute_hessian(&ops3_fr, x3); + let hess_rf3 = MonoAD2RF::compute_hessian(&ops3_rf, x3); + + // Compute exact value manually for verification: f(x) = exp(sin(sin(x))) + let t1 = x3.sin(); + let t2 = t1.sin(); + let dt1 = x3.cos(); + let dt2 = t1.cos() * dt1; + let ddt1 = -x3.sin(); + let ddt2 = -t1.sin() * dt1 * dt1 + t1.cos() * ddt1; + let expected3 = t2.exp() * dt2 * dt2 + t2.exp() * ddt2; + + println!(" f({}) = {}", x3, value3); + println!(" Analytical f''({}) = {:.10}", x3, expected3); + println!( + " Finite Difference: {:.10} (error: {:.2e})", + hess_fd3, + (hess_fd3 - expected3).abs() + ); + println!( + " RR (Exact): {:.10} (error: {:.2e})", + hess_rr3, + (hess_rr3 - expected3).abs() + ); + println!( + " FR (Exact): {:.10} (error: {:.2e})", + hess_fr3, + (hess_fr3 - expected3).abs() + ); + println!( + " RF (Exact): {:.10} (error: {:.2e})", + hess_rf3, + (hess_rf3 - expected3).abs() + ); + + println!("\n=== Summary ===\n"); + + println!("1. Finite Difference Method (MonoAD::compute_hessian):"); + println!(" - Uses ε = 1e-5"); + println!(" - Accuracy: ~1e-4 to 1e-6"); + println!(" - Pros: Simple, works for all operations"); + println!(" - Cons: Numerical approximation, not machine precision"); + println!(); + + println!("2. Exact Methods (MonoAD2RR/FR/RF::compute_hessian):"); + println!(" - RR: Reverse-over-Reverse (tracks 2nd derivatives during backward pass)"); + println!(" - FR: Forward-over-Reverse (forward-mode on gradient function)"); + println!(" - RF: Reverse-over-Forward (reverse-mode on forward-mode)"); + println!(" - Accuracy: Machine precision (~1e-15)"); + println!(" - Pros: Exact differentiation, no approximation error"); + println!(" - Cons: More complex implementation, slightly more computation"); + println!(); + + println!("3. When to Use Which:"); + println!(" - Finite differences: Quick prototyping, most practical applications"); + println!(" - Exact methods: When you need machine precision for:"); + println!(" * Numerical optimization requiring exact Hessians"); + println!(" * Scientific computing with strict accuracy requirements"); + println!(" * Research applications in automatic differentiation"); + println!(); + + println!("4. Performance Note:"); + println!(" - All three exact methods have similar performance"); + println!(" - For univariate functions, FR/RF delegate to RR for composed operations"); + println!(" - Choice between FR/RF/RR is more relevant for multivariate functions"); +} diff --git a/examples/hessian_performance.rs b/examples/hessian_performance.rs new file mode 100644 index 0000000..a5bb333 --- /dev/null +++ b/examples/hessian_performance.rs @@ -0,0 +1,124 @@ +use petite_ad::{MultiAD2FR, MultiAD2RF, MultiAD2RR}; +use std::time::Instant; + +fn main() { + println!("Hessian Performance Comparison (all exact methods)"); + println!("==================================================\n"); + + // Test 1: Simple quadratic f(x,y) = x² + y² + let ops_rr = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(0), + MultiAD2RR::Mul, + MultiAD2RR::Inp(1), + MultiAD2RR::Inp(1), + MultiAD2RR::Mul, + MultiAD2RR::Add, + ]; + + let ops_fr = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Inp(0), + MultiAD2FR::Mul, + MultiAD2FR::Inp(1), + MultiAD2FR::Inp(1), + MultiAD2FR::Mul, + MultiAD2FR::Add, + ]; + + let ops_rf = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Inp(0), + MultiAD2RF::Mul, + MultiAD2RF::Inp(1), + MultiAD2RF::Inp(1), + MultiAD2RF::Mul, + MultiAD2RF::Add, + ]; + + let x = vec![1.0, 2.0]; + + println!("Test 1: Simple quadratic f(x,y) = x² + y²"); + println!("Expected Hessian: [[2, 0], [0, 2]]\n"); + + let iterations = 10000; + + let start = Instant::now(); + for _ in 0..iterations { + let _ = MultiAD2RR::compute_hessian(&ops_rr, &x).unwrap(); + } + println!( + "RR (exact, reverse-reverse): {:?}", + start.elapsed() / iterations + ); + + let start = Instant::now(); + for _ in 0..iterations { + let _ = MultiAD2FR::compute_hessian(&ops_fr, &x).unwrap(); + } + println!( + "FR (exact, forward-reverse): {:?}", + start.elapsed() / iterations + ); + + let start = Instant::now(); + for _ in 0..iterations { + let _ = MultiAD2RF::compute_hessian(&ops_rf, &x).unwrap(); + } + println!( + "RF (exact, reverse-forward): {:?}", + start.elapsed() / iterations + ); + + // Test 2: Trigonometric function f(x,y) = sin(x) * cos(y) + let ops_trig_rr = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Sin, + MultiAD2RR::Inp(1), + MultiAD2RR::Cos, + MultiAD2RR::Mul, + ]; + + let ops_trig_fr = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Inp(1), + MultiAD2FR::Cos, + MultiAD2FR::Mul, + ]; + + let ops_trig_rf = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Inp(1), + MultiAD2RF::Cos, + MultiAD2RF::Mul, + ]; + + println!("\nTest 2: Trigonometric f(x,y) = sin(x) * cos(y)"); + println!("x = {}, y = {}\n", x[0], x[1]); + + let start = Instant::now(); + for _ in 0..iterations { + let _ = MultiAD2RR::compute_hessian(&ops_trig_rr, &x).unwrap(); + } + println!("RR (exact): {:?}", start.elapsed() / iterations); + + let start = Instant::now(); + for _ in 0..iterations { + let _ = MultiAD2FR::compute_hessian(&ops_trig_fr, &x).unwrap(); + } + println!("FR (exact): {:?}", start.elapsed() / iterations); + + let start = Instant::now(); + for _ in 0..iterations { + let _ = MultiAD2RF::compute_hessian(&ops_trig_rf, &x).unwrap(); + } + println!("RF (exact): {:?}", start.elapsed() / iterations); + + println!("\nSummary:"); + println!("- All three methods produce machine-precision exact results"); + println!("- RR uses per-node gradient vectors with outer products in reverse pass"); + println!("- FR uses dual-number forward + dual-adjoint reverse passes"); + println!("- RF uses the same dual algorithm as FR (equivalent for scalar fns)"); +} diff --git a/examples/newton_method.rs b/examples/newton_method.rs new file mode 100644 index 0000000..fdb2029 --- /dev/null +++ b/examples/newton_method.rs @@ -0,0 +1,20 @@ +use petite_ad::Graph; + +fn main() -> petite_ad::Result<()> { + // Minimize f(x) = (x - 3)^2 using Newton's method. + let mut graph = Graph::new(1); + let x = graph.input(0); + let shifted = graph.sub_const(x, 3.0); + graph.square(shifted); + + let mut point = vec![0.0]; + for _step in 0..5 { + let (_value, gradient) = graph.value_and_gradient(&point)?; + let hessian = graph.exact_hessian_rr(&point)?; + point[0] -= gradient[0] / hessian[0][0]; + } + + let value = graph.compute(&point)?; + println!("x = {:.6}, f(x) = {:.6}", point[0], value); + Ok(()) +} diff --git a/examples/optimizer_adam.rs b/examples/optimizer_adam.rs new file mode 100644 index 0000000..bfddadf --- /dev/null +++ b/examples/optimizer_adam.rs @@ -0,0 +1,22 @@ +use petite_ad::{Adam, Graph}; + +fn main() -> petite_ad::Result<()> { + let mut graph = Graph::new(1); + let x = graph.input(0); + let shifted = graph.sub_const(x, 3.0); + graph.square(shifted); + + let mut params = vec![0.0]; + let mut optimizer = Adam::new(1, 0.1); + for _ in 0..100 { + let (_value, gradient) = graph.value_and_gradient(¶ms)?; + optimizer.step(&mut params, &gradient)?; + } + + println!( + "x = {:.6}, f(x) = {:.6}", + params[0], + graph.compute(¶ms)? + ); + Ok(()) +} diff --git a/examples/wasm_demo.html b/examples/wasm_demo.html new file mode 100644 index 0000000..47fb2d6 --- /dev/null +++ b/examples/wasm_demo.html @@ -0,0 +1,21 @@ + + + + + petite-ad WASM demo sketch + + +

petite-ad WASM demo sketch

+

+ This static file documents the intended browser demo shape: compile a small + wrapper crate with wasm-bindgen, expose graph construction and + to_mermaid(), then render the Mermaid text in the browser. +

+ 
flowchart LR
+    n0["x: Input 0"]
+    n1["y: Input 1"]
+    n0 --> n2["Mul"]
+    n1 --> n2
+
+ + diff --git a/src/error.rs b/src/error.rs index ecbadb3..087d898 100644 --- a/src/error.rs +++ b/src/error.rs @@ -16,6 +16,11 @@ pub enum AutodiffError { }, /// The computation graph is empty or invalid. EmptyGraph, + /// The computation graph is malformed. + InvalidGraph { + /// Human-readable reason for the graph validation failure. + reason: &'static str, + }, /// An index references a non-existent value in the computation. IndexOutOfBounds { /// The invalid index @@ -23,6 +28,18 @@ pub enum AutodiffError { /// The maximum valid index max_index: usize, }, + /// A real-domain restriction was violated in checked evaluation mode. + DomainError { + /// Name of the operation. + operation: &'static str, + /// Human-readable reason for the domain failure. + reason: &'static str, + }, + /// Invalid arguments were passed to a function. + InvalidArguments { + /// Human-readable reason for the argument validation failure. + reason: &'static str, + }, } impl fmt::Display for AutodiffError { @@ -38,9 +55,18 @@ impl fmt::Display for AutodiffError { operation, expected, actual ), AutodiffError::EmptyGraph => write!(f, "Computation graph is empty"), + AutodiffError::InvalidGraph { reason } => { + write!(f, "Computation graph is invalid: {}", reason) + } AutodiffError::IndexOutOfBounds { index, max_index } => { write!(f, "Index {} is out of bounds (max: {})", index, max_index) } + AutodiffError::DomainError { operation, reason } => { + write!(f, "Domain error in {}: {}", operation, reason) + } + AutodiffError::InvalidArguments { reason } => { + write!(f, "Invalid arguments: {}", reason) + } } } } @@ -69,7 +95,114 @@ impl AutodiffError { Err(AutodiffError::arity(operation, expected, actual)) } } + + /// Create a domain error for checked evaluation mode. + pub fn domain(operation: &'static str, reason: &'static str) -> Self { + AutodiffError::DomainError { operation, reason } + } + + /// Create an invalid-arguments error. + pub fn invalid_arguments(reason: &'static str) -> Self { + AutodiffError::InvalidArguments { reason } + } } /// Result type for automatic differentiation operations. pub type Result = std::result::Result; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_display_empty_graph() { + let err = AutodiffError::EmptyGraph; + assert_eq!(format!("{}", err), "Computation graph is empty"); + } + + #[test] + fn test_display_index_out_of_bounds() { + let err = AutodiffError::IndexOutOfBounds { + index: 5, + max_index: 3, + }; + assert_eq!(format!("{}", err), "Index 5 is out of bounds (max: 3)"); + } + + #[test] + fn test_display_domain_error() { + let err = AutodiffError::DomainError { + operation: "Sqrt", + reason: "input must be non-negative", + }; + assert_eq!( + format!("{}", err), + "Domain error in Sqrt: input must be non-negative" + ); + } + + #[test] + fn test_display_arity_error() { + let err = AutodiffError::ArityError { + operation: "Mul", + expected: 2, + actual: 1, + }; + assert_eq!(format!("{}", err), "Arity error in Mul: expected 2, got 1"); + } + + #[test] + fn test_display_invalid_graph() { + let err = AutodiffError::InvalidGraph { reason: "bad node" }; + assert_eq!(format!("{}", err), "Computation graph is invalid: bad node"); + } + + #[test] + fn test_arity_constructor() { + let err = AutodiffError::arity("Mul", 2, 1); + assert_eq!( + err, + AutodiffError::ArityError { + operation: "Mul", + expected: 2, + actual: 1, + } + ); + } + + #[test] + fn test_check_arity_success() { + assert!(AutodiffError::check_arity("Sin", 1, 1).is_ok()); + } + + #[test] + fn test_check_arity_failure() { + let err = AutodiffError::check_arity("Sin", 1, 2).unwrap_err(); + assert_eq!( + err, + AutodiffError::ArityError { + operation: "Sin", + expected: 1, + actual: 2, + } + ); + } + + #[test] + fn test_domain_constructor() { + let err = AutodiffError::domain("Ln", "input must be positive"); + assert_eq!( + err, + AutodiffError::DomainError { + operation: "Ln", + reason: "input must be positive", + } + ); + } + + #[test] + fn test_error_is_send_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + } +} diff --git a/src/forward.rs b/src/forward.rs new file mode 100644 index 0000000..9efc6d5 --- /dev/null +++ b/src/forward.rs @@ -0,0 +1,840 @@ +//! Public forward-mode automatic differentiation APIs. +//! +//! This module exposes a small forward-mode surface on top of the library's +//! existing operation sets: +//! +//! - [`ForwardAD::differentiate`] for single-variable `MonoAD` expressions +//! - [`ForwardAD::directional_derivative`] for multivariate `MultiAD` tuple graphs +//! - [`ForwardAD::directional_derivative_graph`] for reusable [`crate::Graph`] values + +use crate::multi::op_rules; +use crate::{AutodiffError, Graph, GraphNode, MonoAD, MultiAD, Result}; + +/// Result of a forward-mode evaluation. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ForwardValue { + /// Primal function value. + pub value: f64, + /// Forward tangent / directional derivative. + pub tangent: f64, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +struct Dual { + value: f64, + tangent: f64, +} + +impl Dual { + #[inline(always)] + fn new(value: f64, tangent: f64) -> Self { + Self { value, tangent } + } + + #[inline(always)] + fn constant(value: f64) -> Self { + Self { + value, + tangent: 0.0, + } + } +} + +/// Public forward-mode utilities. +#[derive(Debug, Clone, Copy, Default)] +pub struct ForwardAD; + +impl ForwardAD { + /// Fill a reusable seed buffer with the `active`-th basis vector. + fn basis_seed_into(dim: usize, active: usize, seed: &mut Vec) { + seed.clear(); + seed.resize(dim, 0.0); + seed[active] = 1.0; + } + + /// Create a fresh basis seed vector (convenience wrapper). + #[allow(dead_code)] + fn basis_seed(dim: usize, active: usize) -> Vec { + let mut seed = Vec::with_capacity(dim); + Self::basis_seed_into(dim, active, &mut seed); + seed + } + + /// Evaluate a mono expression and its derivative at `x`. + #[must_use] + pub fn differentiate(exprs: &[MonoAD], x: f64) -> ForwardValue { + let mut dual = Dual::new(x, 1.0); + for expr in exprs { + dual = match expr { + MonoAD::Sin => Dual::new(dual.value.sin(), dual.value.cos() * dual.tangent), + MonoAD::Cos => Dual::new(dual.value.cos(), -dual.value.sin() * dual.tangent), + MonoAD::Tan => { + let cos_val = dual.value.cos(); + Dual::new(dual.value.tan(), dual.tangent / (cos_val * cos_val)) + } + MonoAD::Exp => { + let exp_val = dual.value.exp(); + Dual::new(exp_val, exp_val * dual.tangent) + } + MonoAD::Neg => Dual::new(-dual.value, -dual.tangent), + MonoAD::Ln => Dual::new(dual.value.ln(), dual.tangent / dual.value), + MonoAD::Sqrt => { + let sqrt_val = dual.value.sqrt(); + Dual::new(sqrt_val, dual.tangent / (2.0 * sqrt_val)) + } + MonoAD::Abs => { + let sign = if dual.value > 0.0 { + 1.0 + } else if dual.value < 0.0 { + -1.0 + } else { + 0.0 + }; + Dual::new(dual.value.abs(), sign * dual.tangent) + } + }; + } + ForwardValue { + value: dual.value, + tangent: dual.tangent, + } + } + + /// Evaluate a mono expression and derivative with checked real-domain validation. + pub fn differentiate_checked(exprs: &[MonoAD], x: f64) -> Result { + let mut dual = Dual::new(x, 1.0); + for expr in exprs { + match expr { + MonoAD::Ln if dual.value <= 0.0 => { + return Err(AutodiffError::domain("Ln", "input must be positive")); + } + MonoAD::Sqrt if dual.value < 0.0 => { + return Err(AutodiffError::domain("Sqrt", "input must be non-negative")); + } + _ => {} + } + let step = Self::differentiate(&[*expr], dual.value); + dual = Dual::new(step.value, step.tangent * dual.tangent); + } + Ok(ForwardValue { + value: dual.value, + tangent: dual.tangent, + }) + } + + /// Evaluate only the primal mono expression value, skipping tangent computation. + #[must_use] + pub fn compute(exprs: &[MonoAD], x: f64) -> f64 { + let mut value = x; + for expr in exprs { + value = match expr { + MonoAD::Sin => value.sin(), + MonoAD::Cos => value.cos(), + MonoAD::Tan => value.tan(), + MonoAD::Exp => value.exp(), + MonoAD::Neg => -value, + MonoAD::Ln => value.ln(), + MonoAD::Sqrt => value.sqrt(), + MonoAD::Abs => value.abs(), + }; + } + value + } + + /// Evaluate only the primal mono value with checked real-domain validation. + pub fn compute_checked(exprs: &[MonoAD], x: f64) -> Result { + let mut value = x; + for expr in exprs { + match expr { + MonoAD::Ln if value <= 0.0 => { + return Err(AutodiffError::domain("Ln", "input must be positive")); + } + MonoAD::Sqrt if value < 0.0 => { + return Err(AutodiffError::domain("Sqrt", "input must be non-negative")); + } + _ => {} + } + value = match expr { + MonoAD::Sin => value.sin(), + MonoAD::Cos => value.cos(), + MonoAD::Tan => value.tan(), + MonoAD::Exp => value.exp(), + MonoAD::Neg => -value, + MonoAD::Ln => value.ln(), + MonoAD::Sqrt => value.sqrt(), + MonoAD::Abs => value.abs(), + }; + } + Ok(value) + } + + /// Evaluate a multivariate tuple graph and its directional derivative. + /// + /// `seed` must have the same length as `inputs`; the returned tangent is the + /// Jacobian-vector product `J(x) · seed` for scalar-valued graphs. + pub fn directional_derivative( + exprs: &[(MultiAD, Vec)], + inputs: &[f64], + seed: &[f64], + ) -> Result { + if seed.len() != inputs.len() { + return Err(AutodiffError::InvalidGraph { + reason: "seed vector length must match input length", + }); + } + + let mut values: Vec = Vec::with_capacity(inputs.len() + exprs.len()); + for (&value, &tangent) in inputs.iter().zip(seed.iter()) { + values.push(Dual::new(value, tangent)); + } + + let mut final_output_index = inputs.len().checked_sub(1); + + for (op, arg_indices) in exprs { + if *op == MultiAD::Inp { + AutodiffError::check_arity("Inp", 1, arg_indices.len())?; + MultiAD::check_value_index(arg_indices[0], inputs.len())?; + final_output_index = Some(arg_indices[0]); + continue; + } + + let mut arg_values = [0.0f64; 2]; + let mut arg_duals = [Dual::constant(0.0); 2]; + AutodiffError::check_arity( + op_rules::op_name(*op), + op_rules::expected_arity(*op), + arg_indices.len(), + )?; + for (i, &index) in arg_indices.iter().enumerate() { + MultiAD::check_value_index(index, values.len())?; + arg_values[i] = values[index].value; + arg_duals[i] = values[index]; + } + + let value = op.forward(&arg_values[..arg_indices.len()])?; + let arg_tangents = [arg_duals[0].tangent, arg_duals[1].tangent]; + let tangent = match op { + MultiAD::Inp => unreachable!("input markers are handled earlier"), + _ => op_rules::directional_tangent( + *op, + &arg_values[..arg_indices.len()], + &arg_tangents[..arg_indices.len()], + value, + )?, + }; + + values.push(Dual::new(value, tangent)); + final_output_index = Some(values.len() - 1); + } + + let Some(final_output_index) = final_output_index else { + return Ok(ForwardValue { + value: 0.0, + tangent: 0.0, + }); + }; + + Ok(ForwardValue { + value: values[final_output_index].value, + tangent: values[final_output_index].tangent, + }) + } + + /// Evaluate a multivariate tuple graph and directional derivative with checked domains. + pub fn directional_derivative_checked( + exprs: &[(MultiAD, Vec)], + inputs: &[f64], + seed: &[f64], + ) -> Result { + if seed.len() != inputs.len() { + return Err(AutodiffError::InvalidGraph { + reason: "seed vector length must match input length", + }); + } + + let mut values: Vec = Vec::with_capacity(inputs.len() + exprs.len()); + for (&value, &tangent) in inputs.iter().zip(seed.iter()) { + values.push(Dual::new(value, tangent)); + } + + let mut final_output_index = inputs.len().checked_sub(1); + + for (op, arg_indices) in exprs { + if *op == MultiAD::Inp { + AutodiffError::check_arity("Inp", 1, arg_indices.len())?; + MultiAD::check_value_index(arg_indices[0], inputs.len())?; + final_output_index = Some(arg_indices[0]); + continue; + } + + let mut arg_values = [0.0f64; 2]; + let mut arg_duals = [Dual::constant(0.0); 2]; + AutodiffError::check_arity( + op_rules::op_name(*op), + op_rules::expected_arity(*op), + arg_indices.len(), + )?; + for (i, &index) in arg_indices.iter().enumerate() { + MultiAD::check_value_index(index, values.len())?; + arg_values[i] = values[index].value; + arg_duals[i] = values[index]; + } + + let value = op.forward_checked(&arg_values[..arg_indices.len()])?; + let arg_tangents = [arg_duals[0].tangent, arg_duals[1].tangent]; + let tangent = op_rules::directional_tangent( + *op, + &arg_values[..arg_indices.len()], + &arg_tangents[..arg_indices.len()], + value, + )?; + + values.push(Dual::new(value, tangent)); + final_output_index = Some(values.len() - 1); + } + + let Some(final_output_index) = final_output_index else { + return Ok(ForwardValue { + value: 0.0, + tangent: 0.0, + }); + }; + + Ok(ForwardValue { + value: values[final_output_index].value, + tangent: values[final_output_index].tangent, + }) + } + + /// Compute the full gradient of a scalar multivariate tuple graph. + pub fn gradient(exprs: &[(MultiAD, Vec)], inputs: &[f64]) -> Result> { + let mut gradient = Vec::with_capacity(inputs.len()); + let mut seed = Vec::with_capacity(inputs.len()); + for active in 0..inputs.len() { + Self::basis_seed_into(inputs.len(), active, &mut seed); + let value = Self::directional_derivative(exprs, inputs, &seed)?; + gradient.push(value.tangent); + } + Ok(gradient) + } + + /// Compute the full gradient of a scalar reusable graph. + pub fn gradient_graph(graph: &Graph, inputs: &[f64]) -> Result> { + let mut gradient = Vec::with_capacity(inputs.len()); + let mut seed = Vec::with_capacity(inputs.len()); + for active in 0..inputs.len() { + Self::basis_seed_into(inputs.len(), active, &mut seed); + let value = Self::directional_derivative_graph(graph, inputs, &seed)?; + gradient.push(value.tangent); + } + Ok(gradient) + } + + /// Compute a Jacobian for multiple scalar outputs represented as tuple graphs. + /// + /// Each entry in `outputs` is treated as one scalar component of a vector-valued + /// function `f: R^n -> R^m`, and the returned matrix has shape `m x n`. + pub fn jacobian( + outputs: &[Vec<(MultiAD, Vec)>], + inputs: &[f64], + ) -> Result>> { + let mut jacobian = Vec::with_capacity(outputs.len()); + for exprs in outputs { + jacobian.push(Self::gradient(exprs, inputs)?); + } + Ok(jacobian) + } + + /// Compute a Jacobian for multiple scalar outputs represented as reusable graphs. + pub fn jacobian_graphs(outputs: &[Graph], inputs: &[f64]) -> Result>> { + let mut jacobian = Vec::with_capacity(outputs.len()); + for graph in outputs { + jacobian.push(Self::gradient_graph(graph, inputs)?); + } + Ok(jacobian) + } + + /// Evaluate a reusable graph and its directional derivative. + pub fn directional_derivative_graph( + graph: &Graph, + inputs: &[f64], + seed: &[f64], + ) -> Result { + Self::directional_derivative_graph_inner(graph, inputs, seed, false) + } + + /// Evaluate a reusable graph and directional derivative with checked domains. + pub fn directional_derivative_graph_checked( + graph: &Graph, + inputs: &[f64], + seed: &[f64], + ) -> Result { + Self::directional_derivative_graph_inner(graph, inputs, seed, true) + } + + fn directional_derivative_graph_inner( + graph: &Graph, + inputs: &[f64], + seed: &[f64], + checked: bool, + ) -> Result { + if inputs.len() != graph.num_inputs() { + return Err(AutodiffError::InvalidGraph { + reason: "graph input length must match graph.num_inputs()", + }); + } + if seed.len() != inputs.len() { + return Err(AutodiffError::InvalidGraph { + reason: "seed vector length must match input length", + }); + } + + let output_index = graph.effective_output_node(); + let mut values: Vec = Vec::with_capacity(graph.num_inputs() + graph.len()); + for (&value, &tangent) in inputs.iter().zip(seed.iter()) { + values.push(Dual::new(value, tangent)); + } + + for node in graph.nodes() { + match node { + GraphNode::Constant(value) => { + values.push(Dual::constant(*value)); + } + GraphNode::Operation { op, inputs } => { + let mut arg_values: Vec = Vec::with_capacity(inputs.len()); + let mut arg_duals: Vec = Vec::with_capacity(inputs.len()); + for &index in inputs { + MultiAD::check_value_index(index, values.len())?; + arg_values.push(values[index].value); + arg_duals.push(values[index]); + } + let value = if checked { + op.forward_checked(&arg_values)? + } else { + op.forward(&arg_values)? + }; + let arg_tangents: Vec = + arg_duals.iter().map(|dual| dual.tangent).collect(); + let tangent = match op { + MultiAD::Inp => unreachable!("graph nodes do not store input markers"), + _ => op_rules::directional_tangent(*op, &arg_values, &arg_tangents, value)?, + }; + values.push(Dual::new(value, tangent)); + } + } + } + + let Some(output_index) = output_index else { + return Ok(ForwardValue { + value: 0.0, + tangent: 0.0, + }); + }; + MultiAD::check_value_index(output_index, values.len())?; + + Ok(ForwardValue { + value: values[output_index].value, + tangent: values[output_index].tangent, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mono_ops; + use crate::test_utils::approx_eq_eps as approx_eq; + + #[test] + fn test_forward_mono_derivative() { + let exprs = mono_ops![sin, exp]; + let result = ForwardAD::differentiate(&exprs, 0.5); + let expected_value = 0.5_f64.sin().exp(); + let expected_tangent = expected_value * 0.5_f64.cos(); + assert!(approx_eq(result.value, expected_value, 1e-10)); + assert!(approx_eq(result.tangent, expected_tangent, 1e-10)); + } + + #[test] + fn test_forward_multi_directional_derivative() { + let exprs = vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), + (MultiAD::Sin, vec![2]), + ]; + let result = ForwardAD::directional_derivative(&exprs, &[2.0, 3.0], &[1.0, -1.0]).unwrap(); + let xy = 2.0_f64 * 3.0_f64; + let expected_value = xy.sin(); + let expected_tangent = xy.cos() * (3.0 - 2.0); + assert!(approx_eq(result.value, expected_value, 1e-10)); + assert!(approx_eq(result.tangent, expected_tangent, 1e-10)); + } + + #[test] + fn test_forward_graph_directional_derivative_with_constant() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let c = graph.constant(2.0); + let x_sq = graph.mul(x, x); + graph.add(x_sq, c); + + let result = ForwardAD::directional_derivative_graph(&graph, &[3.0], &[1.0]).unwrap(); + assert!(approx_eq(result.value, 11.0, 1e-10)); + assert!(approx_eq(result.tangent, 6.0, 1e-10)); + } + + #[test] + fn test_forward_graph_respects_explicit_output() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let x_sq = graph.mul(x, x); + graph.mul(x_sq, x); + graph.set_output(x_sq).unwrap(); + + let result = ForwardAD::directional_derivative_graph(&graph, &[2.0], &[1.0]).unwrap(); + assert!(approx_eq(result.value, 4.0, 1e-10)); + assert!(approx_eq(result.tangent, 4.0, 1e-10)); + } + + #[test] + fn test_forward_gradient_matches_scalar_graph() { + let exprs = vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), + (MultiAD::Sin, vec![2]), + ]; + let gradient = ForwardAD::gradient(&exprs, &[2.0, 3.0]).unwrap(); + let xy = 2.0_f64 * 3.0_f64; + assert!(approx_eq(gradient[0], xy.cos() * 3.0, 1e-10)); + assert!(approx_eq(gradient[1], xy.cos() * 2.0, 1e-10)); + } + + #[test] + fn test_forward_checked_domain_errors() { + let mono_error = ForwardAD::differentiate_checked(&[MonoAD::Ln], -1.0).unwrap_err(); + assert_eq!( + mono_error, + AutodiffError::DomainError { + operation: "Ln", + reason: "input must be positive", + } + ); + + let exprs = vec![(MultiAD::Ln, vec![0])]; + assert!(ForwardAD::directional_derivative_checked(&exprs, &[0.0], &[1.0]).is_err()); + + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sqrt(x); + assert!(ForwardAD::directional_derivative_graph_checked(&graph, &[-1.0], &[1.0]).is_err()); + } + + #[test] + fn test_forward_reverse_gradient_table() { + type MultiExpr = Vec<(MultiAD, Vec)>; + let cases: Vec<(MultiExpr, Vec)> = vec![ + (vec![(MultiAD::Tan, vec![0])], vec![0.25]), + (vec![(MultiAD::Ln, vec![0])], vec![2.0]), + (vec![(MultiAD::Sqrt, vec![0])], vec![4.0]), + ( + vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Div, vec![0, 1]), + ], + vec![4.0, 2.0], + ), + ( + vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Pow, vec![0, 1]), + ], + vec![2.0, 3.0], + ), + ]; + + for (exprs, inputs) in cases { + let reverse = MultiAD::compute_grad(&exprs, &inputs).unwrap().1(1.0); + let forward = ForwardAD::gradient(&exprs, &inputs).unwrap(); + for (left, right) in reverse.iter().zip(forward.iter()) { + assert!(approx_eq(*left, *right, 1e-8)); + } + } + } + + #[test] + fn test_forward_jacobian_for_two_outputs() { + let outputs = vec![ + vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Add, vec![0, 1]), + ], + vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), + ], + ]; + let jacobian = ForwardAD::jacobian(&outputs, &[2.0, 3.0]).unwrap(); + assert_eq!(jacobian.len(), 2); + assert!(approx_eq(jacobian[0][0], 1.0, 1e-10)); + assert!(approx_eq(jacobian[0][1], 1.0, 1e-10)); + assert!(approx_eq(jacobian[1][0], 3.0, 1e-10)); + assert!(approx_eq(jacobian[1][1], 2.0, 1e-10)); + } + + // --- Additional tests for uncovered lines --- + + #[test] + fn test_differentiate_cos() { + let result = ForwardAD::differentiate(&[MonoAD::Cos], 0.5); + assert!(approx_eq(result.value, 0.5_f64.cos(), 1e-10)); + assert!(approx_eq(result.tangent, -0.5_f64.sin(), 1e-10)); + } + + #[test] + fn test_differentiate_tan() { + let result = ForwardAD::differentiate(&[MonoAD::Tan], 0.5); + let cos_val = 0.5_f64.cos(); + assert!(approx_eq(result.value, 0.5_f64.tan(), 1e-10)); + assert!(approx_eq(result.tangent, 1.0 / (cos_val * cos_val), 1e-10)); + } + + #[test] + fn test_differentiate_neg() { + let result = ForwardAD::differentiate(&[MonoAD::Neg], 2.0); + assert!(approx_eq(result.value, -2.0, 1e-10)); + assert!(approx_eq(result.tangent, -1.0, 1e-10)); + } + + #[test] + fn test_differentiate_ln() { + let result = ForwardAD::differentiate(&[MonoAD::Ln], 2.0); + assert!(approx_eq(result.value, 2.0_f64.ln(), 1e-10)); + assert!(approx_eq(result.tangent, 0.5, 1e-10)); + } + + #[test] + fn test_differentiate_sqrt() { + let result = ForwardAD::differentiate(&[MonoAD::Sqrt], 4.0); + assert!(approx_eq(result.value, 2.0, 1e-10)); + assert!(approx_eq(result.tangent, 0.25, 1e-10)); + } + + #[test] + fn test_differentiate_abs_positive() { + let result = ForwardAD::differentiate(&[MonoAD::Abs], 3.0); + assert!(approx_eq(result.value, 3.0, 1e-10)); + assert!(approx_eq(result.tangent, 1.0, 1e-10)); + } + + #[test] + fn test_differentiate_abs_negative() { + let result = ForwardAD::differentiate(&[MonoAD::Abs], -3.0); + assert!(approx_eq(result.value, 3.0, 1e-10)); + assert!(approx_eq(result.tangent, -1.0, 1e-10)); + } + + #[test] + fn test_differentiate_abs_zero() { + let result = ForwardAD::differentiate(&[MonoAD::Abs], 0.0); + assert!(approx_eq(result.value, 0.0, 1e-10)); + assert!(approx_eq(result.tangent, 0.0, 1e-10)); + } + + #[test] + fn test_compute_value_only() { + let value = ForwardAD::compute(&mono_ops![sin, exp], 0.5); + assert!(approx_eq(value, 0.5_f64.sin().exp(), 1e-10)); + } + + #[test] + fn test_compute_all_ops() { + // Exercise all MonoAD variants in compute() + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Sin], 1.0), + 1.0_f64.sin(), + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Cos], 1.0), + 1.0_f64.cos(), + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Tan], 1.0), + 1.0_f64.tan(), + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Exp], 1.0), + 1.0_f64.exp(), + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Neg], 1.0), + -1.0, + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Ln], 2.0), + 2.0_f64.ln(), + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Sqrt], 4.0), + 2.0, + 1e-10 + )); + assert!(approx_eq( + ForwardAD::compute(&[MonoAD::Abs], -3.0), + 3.0, + 1e-10 + )); + } + + #[test] + fn test_compute_checked_success() { + let value = ForwardAD::compute_checked(&mono_ops![sin, exp], 0.5).unwrap(); + assert!(approx_eq(value, 0.5_f64.sin().exp(), 1e-10)); + } + + #[test] + fn test_compute_checked_ln_error() { + assert!(ForwardAD::compute_checked(&[MonoAD::Ln], -1.0).is_err()); + } + + #[test] + fn test_compute_checked_sqrt_error() { + assert!(ForwardAD::compute_checked(&[MonoAD::Sqrt], -1.0).is_err()); + } + + #[test] + fn test_differentiate_checked_success() { + let result = ForwardAD::differentiate_checked(&mono_ops![sin, exp], 0.5).unwrap(); + let expected_value = 0.5_f64.sin().exp(); + let expected_tangent = expected_value * 0.5_f64.cos(); + assert!(approx_eq(result.value, expected_value, 1e-10)); + assert!(approx_eq(result.tangent, expected_tangent, 1e-10)); + } + + #[test] + fn test_differentiate_checked_sqrt_error() { + assert!(ForwardAD::differentiate_checked(&[MonoAD::Sqrt], -1.0).is_err()); + } + + #[test] + fn test_directional_derivative_seed_mismatch() { + let result = ForwardAD::directional_derivative(&[], &[1.0], &[1.0, 2.0]); + assert!(result.is_err()); + } + + #[test] + fn test_directional_derivative_no_inputs_no_exprs() { + let result = ForwardAD::directional_derivative(&[], &[], &[]).unwrap(); + assert_eq!(result.value, 0.0); + assert_eq!(result.tangent, 0.0); + } + + #[test] + fn test_directional_derivative_checked_success() { + let exprs = vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), + ]; + let result = + ForwardAD::directional_derivative_checked(&exprs, &[2.0, 3.0], &[1.0, 0.0]).unwrap(); + assert!(approx_eq(result.value, 6.0, 1e-10)); + assert!(approx_eq(result.tangent, 3.0, 1e-10)); + } + + #[test] + fn test_directional_derivative_checked_seed_mismatch() { + let result = ForwardAD::directional_derivative_checked(&[], &[1.0], &[1.0, 2.0]); + assert!(result.is_err()); + } + + #[test] + fn test_directional_derivative_checked_no_inputs_no_exprs() { + let result = ForwardAD::directional_derivative_checked(&[], &[], &[]).unwrap(); + assert_eq!(result.value, 0.0); + assert_eq!(result.tangent, 0.0); + } + + #[test] + fn test_directional_derivative_checked_domain_error() { + let exprs = vec![(MultiAD::Ln, vec![0])]; + assert!(ForwardAD::directional_derivative_checked(&exprs, &[0.0], &[1.0]).is_err()); + } + + #[test] + fn test_gradient_graph() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let grad = ForwardAD::gradient_graph(&graph, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(grad[0], 3.0, 1e-10)); + assert!(approx_eq(grad[1], 2.0, 1e-10)); + } + + #[test] + fn test_jacobian_graphs() { + let mut g1 = Graph::new(2); + let x = g1.input(0); + let y = g1.input(1); + g1.add(x, y); + + let mut g2 = Graph::new(2); + let x2 = g2.input(0); + let y2 = g2.input(1); + g2.mul(x2, y2); + + let jacobian = ForwardAD::jacobian_graphs(&[g1, g2], &[2.0, 3.0]).unwrap(); + assert_eq!(jacobian.len(), 2); + assert!(approx_eq(jacobian[0][0], 1.0, 1e-10)); + assert!(approx_eq(jacobian[0][1], 1.0, 1e-10)); + assert!(approx_eq(jacobian[1][0], 3.0, 1e-10)); + assert!(approx_eq(jacobian[1][1], 2.0, 1e-10)); + } + + #[test] + fn test_directional_derivative_graph_checked_success() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let result = + ForwardAD::directional_derivative_graph_checked(&graph, &[2.0, 3.0], &[1.0, 0.0]) + .unwrap(); + assert!(approx_eq(result.value, 6.0, 1e-10)); + assert!(approx_eq(result.tangent, 3.0, 1e-10)); + } + + #[test] + fn test_directional_derivative_graph_input_mismatch() { + let graph = Graph::new(2); + assert!(ForwardAD::directional_derivative_graph(&graph, &[1.0], &[1.0]).is_err()); + } + + #[test] + fn test_directional_derivative_graph_seed_mismatch() { + let graph = Graph::new(2); + assert!(ForwardAD::directional_derivative_graph(&graph, &[1.0, 2.0], &[1.0]).is_err()); + } + + #[test] + fn test_directional_derivative_graph_empty() { + let graph = Graph::new(0); + let result = ForwardAD::directional_derivative_graph(&graph, &[], &[]).unwrap(); + assert_eq!(result.value, 0.0); + assert_eq!(result.tangent, 0.0); + } +} diff --git a/src/lib.rs b/src/lib.rs index 1730732..fd98968 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,8 +5,12 @@ //! //! ## Features //! -//! - **Single-variable autodiff** - Chain operations like `sin`, `cos`, `exp` +//! - **Single-variable autodiff** - Chain operations like `sin`, `cos`, `tan`, `exp`, `ln`, `sqrt`, and `abs` //! - **Multi-variable autodiff** - Build computational graphs for multiple inputs +//! - **Reusable graph/tape API** - Build once with node handles, select explicit outputs, and evaluate repeatedly +//! - **Graph validation/export** - Check reusable graphs and export Mermaid/DOT views +//! - **Public forward-mode AD** - Compute derivatives, gradients, and directional derivatives directly +//! - **Opt-in checked mode** - Validate real-domain restrictions for sensitive operations //! - **Zero-copy backward pass** - Efficient gradient computation through closure chains //! - **Convenient macros** - Use `mono_ops![]` and `multi_ops![]` for concise notation //! @@ -24,38 +28,148 @@ //! //! ### Multi-variable function //! ``` -//! use petite_ad::{MultiAD, multi_ops}; +//! use petite_ad::Graph; //! -//! let exprs = multi_ops![ -//! (inp, 0), // x₁ -//! (inp, 1), // x₂ -//! (add, 0, 1), // x₁ + x₂ -//! (sin, 0), // sin(x₁) -//! (mul, 2, 3), // sin(x₁) * (x₁ + x₂) -//! ]; +//! let mut graph = Graph::new(2); +//! let x = graph.input(0); +//! let y = graph.input(1); +//! let sum = graph.add(x, y); +//! let sin_x = graph.sin(x); +//! graph.mul(sum, sin_x); //! -//! let (value, grad_fn) = MultiAD::compute_grad(&exprs, &[0.6, 1.4]).unwrap(); +//! let (value, grad_fn) = graph.compute_grad(&[0.6, 1.4]).unwrap(); //! let gradients = grad_fn(1.0); //! println!("f(0.6, 1.4) = {}", value); //! println!("∇f = {:?}", gradients); +//! +//! // Reuse the same graph but choose a different output node. +//! graph.set_output(sum).unwrap(); +//! assert!((graph.compute(&[0.6, 1.4]).unwrap() - 2.0).abs() < 1e-10); +//! +//! // Or expose multiple outputs and get a Jacobian directly. +//! graph.set_outputs(&[sum, sin_x]).unwrap(); +//! let values = graph.compute_many(&[0.6, 1.4]).unwrap(); +//! assert_eq!(values.len(), 2); +//! let jacobian = graph.jacobian(&[0.6, 1.4]).unwrap(); +//! assert_eq!(jacobian.len(), 2); +//! ``` +//! +//! ### Forward-mode derivative +//! ``` +//! use petite_ad::{ForwardAD, MonoAD, mono_ops}; +//! +//! let exprs = mono_ops![sin, exp]; +//! let result = ForwardAD::differentiate(&exprs, 0.5); +//! assert!(result.value.is_finite()); +//! assert!(result.tangent.is_finite()); +//! ``` +//! +//! ### Graph validation/export +//! ``` +//! use petite_ad::Graph; +//! +//! let mut graph = Graph::new(1); +//! let x = graph.input(0); +//! let neg_x = graph.neg(x); +//! graph.exp(neg_x); +//! +//! graph.validate().unwrap(); +//! assert!(graph.to_mermaid().contains("flowchart LR")); +//! assert!(graph.to_dot().contains("digraph Graph")); +//! ``` +//! +//! ### Reusable tape workspace +//! ``` +//! use petite_ad::Graph; +//! +//! let mut graph = Graph::new(2); +//! let x = graph.input(0); +//! let y = graph.input(1); +//! let sum = graph.add(x, y); +//! graph.mul(sum, y); +//! +//! let tape = graph.compile(); +//! let mut workspace = tape.workspace(); +//! let (value, grad) = tape.gradient_with_workspace(&[2.0, 3.0], &mut workspace).unwrap(); +//! assert!(value.is_finite()); +//! assert_eq!(grad.len(), 2); +//! ``` +//! +//! ### Checked-domain evaluation +//! ``` +//! use petite_ad::{Graph, MonoAD, MultiAD, multi_ops}; +//! +//! let mono_exprs = [MonoAD::Sqrt]; +//! assert!(MonoAD::compute_checked(&mono_exprs, 4.0).is_ok()); +//! assert!(MonoAD::compute_checked(&mono_exprs, -1.0).is_err()); +//! +//! let exprs = multi_ops![(ln, 0)]; +//! assert!(MultiAD::compute_checked(&exprs, &[2.0]).is_ok()); +//! assert!(MultiAD::compute_checked(&exprs, &[0.0]).is_err()); +//! +//! let mut graph = Graph::new(1); +//! let x = graph.input(0); +//! graph.ln(x); +//! assert!(graph.compute_checked(&[2.0]).is_ok()); +//! assert!(graph.compute_checked(&[0.0]).is_err()); +//! +//! let mut multi_graph = Graph::new(2); +//! let x = multi_graph.input(0); +//! let y = multi_graph.input(1); +//! let ratio = multi_graph.div(x, y); +//! let log_y = multi_graph.ln(y); +//! multi_graph.set_outputs(&[ratio, log_y]).unwrap(); +//! assert!(multi_graph.compute_many_checked(&[4.0, 2.0]).is_ok()); +//! assert!(multi_graph.jacobian_checked(&[4.0, 2.0]).is_ok()); +//! assert!(multi_graph.compute_many_checked(&[4.0, 0.0]).is_err()); //! ``` mod error; +pub mod forward; mod macros; +pub mod optim; #[cfg(test)] mod test_utils; +#[cfg(test)] +mod tests_comprehensive; + mod mono; mod multi; // Core types pub use mono::MonoAD; + +// Higher-order autodiff methods (exact Hessian computation) +pub use mono::{MonoAD2FR, MonoAD2RF, MonoAD2RR}; + pub use multi::builder::GraphBuilder; +pub use multi::multi_ad_fr::MultiAD2FR; +pub use multi::multi_ad_rf::MultiAD2RF; +pub use multi::multi_ad_rr::MultiAD2RR; pub use multi::MultiAD; +pub use multi::{ + AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind, + BackendRejectionReason, BackendSupportReport, BatchGradients, BatchGradientsBuffer, + BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, + CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle, + DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, + DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, DeviceTransferPlan, + DeviceTransferPolicy, DomainPolicy, ExecutionBackend, ExprGraph, ExprNode, FlatInstruction, + GpuBackendBoundary, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, GraphStats, + Instruction, MockDeviceBackend, NodeId, OpCode, ScalarBackend, SimdBackend, Tape, + TapeWorkspace, UNUSED_NODE_ID, +}; +#[cfg(feature = "backend-wgpu")] +pub use multi::{ + WgpuBackend, WgpuBuffer, WgpuBufferSet, WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES, +}; // Error handling pub use error::{AutodiffError, Result}; +pub use forward::{ForwardAD, ForwardValue}; +pub use optim::{Adam, GradientDescent}; /// Type definitions for autodiff results and gradient functions. /// diff --git a/src/macros.rs b/src/macros.rs index 609f731..012c0ae 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -5,7 +5,7 @@ /// ``` /// use petite_ad::{mono_ops, MonoAD}; /// -/// let (value, backprop) = MonoAD::compute_grad(&mono_ops![sin, sin, exp], 2.0); +/// let (value, backprop) = MonoAD::compute_grad(&mono_ops![sin, tan, exp], 2.0); /// println!("backprop: {} {}", value, backprop(1.0)); /// ``` /// @@ -13,16 +13,77 @@ macro_rules! mono_ops { (@one sin) => { $crate::MonoAD::Sin }; (@one cos) => { $crate::MonoAD::Cos }; + (@one tan) => { $crate::MonoAD::Tan }; (@one exp) => { $crate::MonoAD::Exp }; (@one neg) => { $crate::MonoAD::Neg }; + (@one ln) => { $crate::MonoAD::Ln }; + (@one sqrt) => { $crate::MonoAD::Sqrt }; + (@one abs) => { $crate::MonoAD::Abs }; (@one $x:ident) => { - compile_error!(concat!("Unsupported math operation: ", stringify!($x), ". Use: sin, cos, or exp")) + compile_error!(concat!("Unsupported math operation: ", stringify!($x), ". Use: sin, cos, tan, exp, neg, ln, sqrt, or abs")) }; ($($x:ident),* $(,)?) => { [$($crate::mono_ops!(@one $x)),*] }; } +/// Macro for MonoAD2RR (Reverse-over-Reverse) operations. +#[macro_export] +macro_rules! mono_ops_rr { + (@one sin) => { $crate::MonoAD2RR::Sin }; + (@one cos) => { $crate::MonoAD2RR::Cos }; + (@one tan) => { $crate::MonoAD2RR::Tan }; + (@one exp) => { $crate::MonoAD2RR::Exp }; + (@one neg) => { $crate::MonoAD2RR::Neg }; + (@one ln) => { $crate::MonoAD2RR::Ln }; + (@one sqrt) => { $crate::MonoAD2RR::Sqrt }; + (@one abs) => { $crate::MonoAD2RR::Abs }; + (@one $x:ident) => { + compile_error!(concat!("Unsupported math operation: ", stringify!($x), ". Use: sin, cos, tan, exp, neg, ln, sqrt, or abs")) + }; + ($($x:ident),* $(,)?) => { + [$($crate::mono_ops_rr!(@one $x)),*] + }; +} + +/// Macro for MonoAD2FR (Forward-over-Reverse) operations. +#[macro_export] +macro_rules! mono_ops_fr { + (@one sin) => { $crate::MonoAD2FR::Sin }; + (@one cos) => { $crate::MonoAD2FR::Cos }; + (@one tan) => { $crate::MonoAD2FR::Tan }; + (@one exp) => { $crate::MonoAD2FR::Exp }; + (@one neg) => { $crate::MonoAD2FR::Neg }; + (@one ln) => { $crate::MonoAD2FR::Ln }; + (@one sqrt) => { $crate::MonoAD2FR::Sqrt }; + (@one abs) => { $crate::MonoAD2FR::Abs }; + (@one $x:ident) => { + compile_error!(concat!("Unsupported math operation: ", stringify!($x), ". Use: sin, cos, tan, exp, neg, ln, sqrt, or abs")) + }; + ($($x:ident),* $(,)?) => { + [$($crate::mono_ops_fr!(@one $x)),*] + }; +} + +/// Macro for MonoAD2RF (Reverse-over-Forward) operations. +#[macro_export] +macro_rules! mono_ops_rf { + (@one sin) => { $crate::MonoAD2RF::Sin }; + (@one cos) => { $crate::MonoAD2RF::Cos }; + (@one tan) => { $crate::MonoAD2RF::Tan }; + (@one exp) => { $crate::MonoAD2RF::Exp }; + (@one neg) => { $crate::MonoAD2RF::Neg }; + (@one ln) => { $crate::MonoAD2RF::Ln }; + (@one sqrt) => { $crate::MonoAD2RF::Sqrt }; + (@one abs) => { $crate::MonoAD2RF::Abs }; + (@one $x:ident) => { + compile_error!(concat!("Unsupported math operation: ", stringify!($x), ". Use: sin, cos, tan, exp, neg, ln, sqrt, or abs")) + }; + ($($x:ident),* $(,)?) => { + [$($crate::mono_ops_rf!(@one $x)),*] + }; +} + /// Macro to build multi-variable computation graphs with lowercase operation names. /// Converts lowercase identifiers to MultiAD enum variants. /// @@ -36,7 +97,8 @@ macro_rules! mono_ops { /// - `add`, `sub`, `mul`, `div` - Binary operations (takes two indices) /// - `pow` - Power operation (takes two indices: base, exponent) /// - `sin`, `cos`, `tan`, `exp`, `ln` - Unary operations (takes single index) -/// - `sqrt`, `abs` - Unary operations (takes single index) +/// - `sqrt`, `abs`, `log1p_exp` - Unary operations (takes single index) +/// - `log_add_exp` - Binary stable log-sum-exp (takes two indices) /// /// # Example /// ``` @@ -60,6 +122,10 @@ macro_rules! multi_ops { (@op sin) => { $crate::MultiAD::Sin }; (@op cos) => { $crate::MultiAD::Cos }; (@op tan) => { $crate::MultiAD::Tan }; + (@op tanh) => { $crate::MultiAD::Tanh }; + (@op relu) => { $crate::MultiAD::Relu }; + (@op log1p_exp) => { $crate::MultiAD::Log1pExp }; + (@op neg) => { $crate::MultiAD::Neg }; (@op exp) => { $crate::MultiAD::Exp }; (@op ln) => { $crate::MultiAD::Ln }; (@op sqrt) => { $crate::MultiAD::Sqrt }; @@ -70,6 +136,7 @@ macro_rules! multi_ops { (@op mul) => { $crate::MultiAD::Mul }; (@op div) => { $crate::MultiAD::Div }; (@op pow) => { $crate::MultiAD::Pow }; + (@op log_add_exp) => { $crate::MultiAD::LogAddExp }; // Input (@op inp) => { $crate::MultiAD::Inp }; // Error for unknown operations @@ -78,7 +145,7 @@ macro_rules! multi_ops { concat!( "Unsupported operation: ", stringify!($x), - ". Use: inp, add, sub, mul, div, pow, sin, cos, tan, exp, ln, sqrt, or abs" + ". Use: inp, add, sub, mul, div, pow, log_add_exp, sin, cos, tan, tanh, relu, log1p_exp, neg, exp, ln, sqrt, or abs" ) ) }; diff --git a/src/main.rs b/src/main.rs index 4c38a44..74900c4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,8 +7,6 @@ use petite_ad::{ fn main() { println!("=== Autodiff Library Demo ===\n"); - println!("=== Autodiff Library Demo ===\n"); - // Mono-variable automatic differentiation println!("--- Mono-variable automatic differentiation ---"); // Example 1. @@ -26,6 +24,23 @@ fn main() { println!(" Arc gradient: {:.4}", grad_fn_arc(1.0)); println!(" Cloned gradient: {:.4}", grad_fn_clone(1.0)); + // Example 3: Second derivative (Hessian for single-variable) + println!("\n3. Second derivative computation:"); + let sin_exprs = mono_ops![sin]; + let x = 0.5; + let second_deriv = MonoAD::compute_hessian(&sin_exprs, x); + println!("f(x) = sin(x), f''(x) = -sin(x)"); + println!("At x = {:.1}, f''({:.1}) = {:.4}", x, x, second_deriv); + println!("Expected: f''(0.5) = -sin(0.5) = {:.4}", -x.sin()); + + let exp_sin_exprs = mono_ops![sin, exp]; + let x2 = 0.5; + let second_deriv2 = MonoAD::compute_hessian(&exp_sin_exprs, x2); + println!("\nf(x) = exp(sin(x))"); + println!("At x = {:.1}, f''({:.1}) = {:.4}", x2, x2, second_deriv2); + let expected2 = x2.sin().exp() * x2.cos().powi(2) - x2.sin().exp() * x2.sin(); + println!("Expected: {:.4}", expected2); + // Multi-variable automatic differentiation println!("\n--- Multi-variable automatic differentiation ---"); @@ -97,4 +112,51 @@ fn main() { "Analytical Gradients: ∂f/∂x = {:.4}, ∂f/∂y = {:.4}, ∂f/∂z = {:.4}", analytical_dx, analytical_dy, analytical_dz ); + + // Example 4: Hessian computation (second-order derivatives) + println!("\n=== Hessian Computation (Second-Order Derivatives) ===\n"); + + // Example 1: f(x, y) = x² + y² + // Hessian: [[2, 0], [0, 2]] + let quadratic_exprs = multi_ops![(inp, 0), (inp, 1), (mul, 0, 0), (mul, 1, 1), (add, 2, 3)]; + let inputs = [2.0, 3.0]; + let hessian = MultiAD::compute_hessian(&quadratic_exprs, &inputs).unwrap(); + + println!("Function: f(x, y) = x² + y²"); + println!("Point: (x, y) = ({}, {})", inputs[0], inputs[1]); + println!("Hessian matrix:"); + println!(" [[{:.1}, {:.1}],", hessian[0][0], hessian[0][1]); + println!(" [{:.1}, {:.1}]]", hessian[1][0], hessian[1][1]); + println!("Expected: [[2.0, 0.0], [0.0, 2.0]]"); + println!( + "Interpretation: ∂²f/∂x² = {:.1}, ∂²f/∂x∂y = {:.1}, ∂²f/∂y∂x = {:.1}, ∂²f/∂y² = {:.1}", + hessian[0][0], hessian[0][1], hessian[1][0], hessian[1][1] + ); + + // Example 2: f(x, y) = x * y + // Hessian: [[0, 1], [1, 0]] + println!("\nFunction: f(x, y) = x * y"); + let cross_exprs = multi_ops![(inp, 0), (inp, 1), (mul, 0, 1)]; + let cross_hessian = MultiAD::compute_hessian(&cross_exprs, &[2.0, 3.0]).unwrap(); + println!("Point: (x, y) = (2.0, 3.0)"); + println!("Hessian matrix:"); + println!( + " [[{:.1}, {:.1}],", + cross_hessian[0][0], cross_hessian[0][1] + ); + println!( + " [{:.1}, {:.1}]]", + cross_hessian[1][0], cross_hessian[1][1] + ); + println!("Expected: [[0.0, 1.0], [1.0, 0.0]]"); + + // Example 3: Single-variable second derivative + // f(x) = x³, f''(x) = 6x + println!("\nFunction: f(x) = x³"); + let cubic_exprs = multi_ops![(inp, 0), (mul, 0, 0), (mul, 1, 0)]; + let x = 3.0; + let cubic_hessian = MultiAD::compute_hessian(&cubic_exprs, &[x]).unwrap(); + println!("Point: x = {}", x); + println!("Hessian (1x1 matrix): [[{:.4}]]", cubic_hessian[0][0]); + println!("Expected: [[18.0]] (since f''(x) = 6x = 6*3 = 18)"); } diff --git a/src/mono.rs b/src/mono.rs index 8d1197e..08dd72f 100644 --- a/src/mono.rs +++ b/src/mono.rs @@ -3,9 +3,22 @@ pub mod types; #[cfg(test)] mod tests; -mod mono_ad; +#[cfg(test)] +mod tests_ho; + +pub mod mono_ad; pub use mono_ad::MonoAD; +pub mod mono_ad_rr; +pub use mono_ad_rr::MonoAD2RR; + +pub mod mono_ad_fr; +mod mono_hessian_common; +pub use mono_ad_fr::MonoAD2FR; + +pub mod mono_ad_rf; +pub use mono_ad_rf::MonoAD2RF; + mod mono_fn; // Re-export trait for library extension - users can implement custom mono functions #[allow(unused_imports)] // May not be used internally, but part of public API diff --git a/src/mono/mono_ad.rs b/src/mono/mono_ad.rs index 69bb454..2c12f24 100644 --- a/src/mono/mono_ad.rs +++ b/src/mono/mono_ad.rs @@ -1,3 +1,5 @@ +use crate::{AutodiffError, Result}; + use super::types::*; /// Single-variable automatic differentiation operations. @@ -35,6 +37,14 @@ pub enum MonoAD { /// - Delegates to `f64::cos()`, which operates in radians /// - Returns values in the range `[-1.0, 1.0]` Cos, + /// Tangent function: tan(x) + /// + /// Derivative: sec²(x) = 1 / cos²(x) + /// + /// # Notes + /// - Delegates to `f64::tan()`, which operates in radians + /// - Returns very large values near `π/2 + kπ` (asymptotes) + Tan, /// Exponential function: exp(x) /// /// Derivative: exp(x) @@ -44,7 +54,20 @@ pub enum MonoAD { /// - Returns `inf` for very large inputs (> ~709 for f64) /// - Returns `0.0` for very large negative inputs (< ~-745 for f64) Exp, + /// Negation: -x Neg, + /// Natural logarithm: ln(x) + /// + /// Derivative: 1 / x + Ln, + /// Square root: sqrt(x) + /// + /// Derivative: 1 / (2 sqrt(x)) + Sqrt, + /// Absolute value: abs(x) + /// + /// Derivative: sign(x), with subgradient sign(0) = 0 + Abs, } impl MonoAD { @@ -52,15 +75,38 @@ impl MonoAD { /// /// This is an internal helper that computes just the forward value /// without building gradient closures. + #[inline(always)] fn forward(&self, x: f64) -> f64 { match self { MonoAD::Sin => x.sin(), MonoAD::Cos => x.cos(), + MonoAD::Tan => x.tan(), MonoAD::Exp => x.exp(), MonoAD::Neg => -x, + MonoAD::Ln => x.ln(), + MonoAD::Sqrt => x.sqrt(), + MonoAD::Abs => x.abs(), } } + /// Validate real-domain restrictions for checked mono evaluation. + #[inline(always)] + fn check_domain(self, x: f64) -> Result<()> { + match self { + MonoAD::Ln if x <= 0.0 => Err(AutodiffError::domain("Ln", "input must be positive")), + MonoAD::Sqrt if x < 0.0 => { + Err(AutodiffError::domain("Sqrt", "input must be non-negative")) + } + _ => Ok(()), + } + } + + #[inline(always)] + fn forward_checked(&self, x: f64) -> Result { + self.check_domain(x)?; + Ok(self.forward(x)) + } + /// Compute the forward pass only (no gradient computation). /// /// Evaluates the composed function by applying operations sequentially. @@ -79,6 +125,7 @@ impl MonoAD { /// let result = MonoAD::compute(&ops, 2.0); /// assert!((result - 2.0_f64.sin().exp()).abs() < 1e-10); /// ``` + #[inline] pub fn compute(exprs: &[MonoAD], x: f64) -> f64 { let mut value = x; for expr in exprs { @@ -87,8 +134,27 @@ impl MonoAD { value } + /// Compute the forward pass with opt-in real-domain validation. + /// + /// Existing unchecked methods preserve raw `f64` behavior. This checked variant + /// returns a domain error for `Ln` inputs `<= 0` and `Sqrt` inputs `< 0`. + /// + /// # Errors + /// + /// Returns [`AutodiffError::DomainError`] if an intermediate value violates a + /// real-domain restriction. + #[inline] + pub fn compute_checked(exprs: &[MonoAD], x: f64) -> Result { + let mut value = x; + for expr in exprs { + value = expr.forward_checked(value)?; + } + Ok(value) + } + // Helper that works with Box wrapper type // Box is the common type that all arms return + #[inline(always)] fn backward_generic(self, x: f64) -> (f64, W) where W: From>, @@ -96,12 +162,20 @@ impl MonoAD { let (y, grad_fn): (f64, Box f64>) = match self { MonoAD::Sin => { let y = x.sin(); - let grad = Box::new(move |dy: f64| -> f64 { dy * x.cos() }); + let x_cos = x.cos(); + let grad = Box::new(move |dy: f64| -> f64 { dy * x_cos }); (y, grad) } MonoAD::Cos => { let y = x.cos(); - let grad = Box::new(move |dy: f64| -> f64 { dy * -x.sin() }); + let x_sin = x.sin(); + let grad = Box::new(move |dy: f64| -> f64 { -dy * x_sin }); + (y, grad) + } + MonoAD::Tan => { + let y = x.tan(); + let sec_sq = 1.0 / x.cos().powi(2); + let grad = Box::new(move |dy: f64| -> f64 { dy * sec_sq }); (y, grad) } MonoAD::Exp => { @@ -111,7 +185,29 @@ impl MonoAD { } MonoAD::Neg => { let y = -x; - let grad = Box::new(move |dy: f64| -> f64 { -dy * 1.0 }); + let grad = Box::new(move |dy: f64| -> f64 { -dy }); + (y, grad) + } + MonoAD::Ln => { + let y = x.ln(); + let grad = Box::new(move |dy: f64| -> f64 { dy / x }); + (y, grad) + } + MonoAD::Sqrt => { + let y = x.sqrt(); + let grad = Box::new(move |dy: f64| -> f64 { dy / (2.0 * y) }); + (y, grad) + } + MonoAD::Abs => { + let y = x.abs(); + let sign = if x > 0.0 { + 1.0 + } else if x < 0.0 { + -1.0 + } else { + 0.0 + }; + let grad = Box::new(move |dy: f64| -> f64 { dy * sign }); (y, grad) } }; @@ -152,12 +248,14 @@ impl MonoAD { /// let arc_grad_fn: Arc f64> = Arc::from(grad_fn); /// ``` #[must_use = "gradient computation is expensive; discarding the result is likely a bug"] + #[inline] pub fn compute_grad_generic(exprs: &[MonoAD], x: f64) -> (f64, W) where W: From> + std::ops::Deref + 'static, { let mut value = x; - let mut backprops: Vec = Vec::new(); + // Pre-allocate with capacity to avoid reallocations + let mut backprops: Vec = Vec::with_capacity(exprs.len()); // Compute backward pass for each operation for &op in exprs { @@ -182,4 +280,105 @@ impl MonoAD { pub fn compute_grad(exprs: &[MonoAD], x: f64) -> BackwardResultBox { Self::compute_grad_generic::>(exprs, x) } + + /// Compute forward pass and gradient function with checked-domain validation. + /// + /// # Errors + /// + /// Returns [`AutodiffError::DomainError`] if an intermediate value violates a + /// real-domain restriction before its operation is evaluated. + #[must_use = "gradient computation is expensive; discarding the result is likely a bug"] + pub fn compute_grad_checked(exprs: &[MonoAD], x: f64) -> Result { + let mut value = x; + let mut backprops: Vec> = Vec::with_capacity(exprs.len()); + + for &op in exprs { + op.check_domain(value)?; + let (new_value, backprop) = op.backward_generic(value); + value = new_value; + backprops.push(backprop); + } + + let backward_fn = Box::new(move |cotangent: f64| -> f64 { + let mut grad = cotangent; + for backprop in backprops.iter().rev() { + grad = backprop(grad); + } + grad + }); + + Ok((value, backward_fn)) + } + + /// Compute the second derivative (Hessian for single-variable functions). + /// + /// For a single-variable function f(x), the Hessian is a 1x1 matrix + /// containing f''(x), which we return as a scalar value. + /// + /// This uses finite differences on the gradient function: + /// f''(x) ≈ (f'(x + ε) - f'(x - ε)) / (2ε) + /// + /// # Arguments + /// + /// * `exprs` - Slice of operations to apply in sequence + /// * `x` - Input value to evaluate at + /// + /// # Returns + /// + /// The second derivative f''(x) at the input point + /// + /// # Accuracy + /// + /// Uses central difference with ε = 1e-5, giving approximately 1e-4 accuracy. + /// For higher precision, use analytical derivatives or smaller ε values. + /// + /// # Examples + /// + /// ``` + /// use petite_ad::{MonoAD, mono_ops}; + /// + /// // f(x) = sin(x), f''(x) = -sin(x) + /// let ops = mono_ops![sin]; + /// let second_deriv = MonoAD::compute_hessian(&ops, 0.5); + /// // At x = 0.5: -sin(0.5) ≈ -0.4794 + /// assert!((second_deriv - (-0.5_f64.sin())).abs() < 1e-4); + /// ``` + #[must_use = "second derivative computation is expensive; discarding the result is likely a bug"] + pub fn compute_hessian(exprs: &[MonoAD], x: f64) -> f64 { + let epsilon = 1e-5; + + // Compute gradient at x + ε + let (_value_plus, grad_fn_plus) = Self::compute_grad(exprs, x + epsilon); + let grad_plus = grad_fn_plus(1.0); + + // Compute gradient at x - ε + let (_value_minus, grad_fn_minus) = Self::compute_grad(exprs, x - epsilon); + let grad_minus = grad_fn_minus(1.0); + + // Central difference formula for second derivative + (grad_plus - grad_minus) / (2.0 * epsilon) + } + + /// Compute the finite-difference Hessian with checked-domain validation. + /// + /// This validates the two perturbed gradient evaluations used by the central + /// difference. Points near a domain boundary may therefore return an error even + /// when the unperturbed forward value is defined. + /// + /// # Errors + /// + /// Returns [`AutodiffError::DomainError`] if `x + ε` or `x - ε` violates a + /// real-domain restriction at any intermediate operation. + #[must_use = "second derivative computation is expensive; discarding the result is likely a bug"] + pub fn compute_hessian_checked(exprs: &[MonoAD], x: f64) -> Result { + let epsilon = 1e-5; + + let (_value_plus, grad_fn_plus) = Self::compute_grad_checked(exprs, x + epsilon)?; + let grad_plus = grad_fn_plus(1.0); + + let (_value_minus, grad_fn_minus) = Self::compute_grad_checked(exprs, x - epsilon)?; + let grad_minus = grad_fn_minus(1.0); + + Ok((grad_plus - grad_minus) / (2.0 * epsilon)) + } } diff --git a/src/mono/mono_ad_fr.rs b/src/mono/mono_ad_fr.rs new file mode 100644 index 0000000..e65d3e2 --- /dev/null +++ b/src/mono/mono_ad_fr.rs @@ -0,0 +1,203 @@ +//! Exact second-order autodiff using Forward-over-Reverse (FR) mode. +//! +//! This module implements the **Forward-over-Reverse (FR)** method for computing exact +//! Hessians (second derivatives) of single-variable functions. It combines reverse-mode +//! gradient computation with forward-mode differentiation on the gradient function. +//! +//! # Comprehensive Documentation +//! +//! For complete mathematical theory, detailed derivations, and comparison with other methods, see: +//! **[`/docs/mono_ad_hessian.md`](../../docs/mono_ad_hessian.md)** +//! +//! # Supported Operations +//! +//! This type supports a subset of operations compared to [`crate::MonoAD`]: +//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. See [`super::mono_hessian_common`] for details. +//! +//! # Mathematical Foundation +//! +//! The Forward-over-Reverse method computes the Hessian by treating the gradient function +//! itself as a differentiable function and applying forward-mode AD to it. +//! +//! For a function f(x): +//! 1. **Reverse pass**: Compute g(x) = f'(x) using reverse-mode AD +//! 2. **Forward pass**: Compute g'(x) = d/dx[f'(x)] = f''(x) using forward-mode AD with dual numbers +//! +//! # Algorithm +//! +//! ## For Single Operations +//! +//! Direct dual-number evaluation of the second derivative. +//! +//! ## For Multiple Operations (Compositions) +//! +//! Delegates to RR mode ([`crate::MonoAD2RR`]), which handles general compositions +//! correctly. +//! +//! # Accuracy +//! +//! - **Method**: Exact symbolic differentiation via dual numbers +//! - **Error source**: Only floating-point rounding (machine epsilon) +//! - **Typical relative error**: < 1e-14 +//! - **No truncation error**: Mathematically exact (unlike finite differences) +//! +//! # Example Usage +//! +//! ```rust +//! use petite_ad::MonoAD2FR; +//! +//! // f(x) = sin(x), f''(x) = -sin(x) +//! let ops = [MonoAD2FR::Sin]; +//! let x = 1.0; +//! +//! let value = MonoAD2FR::compute(&ops, x); +//! let hessian = MonoAD2FR::compute_hessian(&ops, x); +//! +//! println!("f({}) = {}", x, value); // sin(1) ≈ 0.8414 +//! println!("f''({}) = {}", x, hessian); // -sin(1) ≈ -0.8414 +//! ``` +//! +//! # References +//! +//! - Griewank & Walther (2008): *Evaluating Derivatives* +//! - Naumann (2012): *The Art of Differentiating Computer Programs* +//! - Pearlmutter (1994): "Fast exact multiplication by the Hessian" + +use crate::Result; + +use super::mono_ad_rr::MonoAD2RR; +use super::mono_hessian_common::{self, MonoHessianOpKind}; +use super::types::*; + +/// Single-variable automatic differentiation operations for Forward-over-Reverse Hessian computation. +/// +/// Supports: `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MonoAD2FR { + Sin, + Cos, + Tan, + Exp, + Neg, + Ln, + Sqrt, + Abs, +} + +impl From for MonoHessianOpKind { + fn from(op: MonoAD2FR) -> MonoHessianOpKind { + match op { + MonoAD2FR::Sin => MonoHessianOpKind::Sin, + MonoAD2FR::Cos => MonoHessianOpKind::Cos, + MonoAD2FR::Tan => MonoHessianOpKind::Tan, + MonoAD2FR::Exp => MonoHessianOpKind::Exp, + MonoAD2FR::Neg => MonoHessianOpKind::Neg, + MonoAD2FR::Ln => MonoHessianOpKind::Ln, + MonoAD2FR::Sqrt => MonoHessianOpKind::Sqrt, + MonoAD2FR::Abs => MonoHessianOpKind::Abs, + } + } +} + +impl MonoAD2FR { + /// Compute forward pass only. + pub fn compute(exprs: &[MonoAD2FR], x: f64) -> f64 { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_forward(&ops, x) + } + + /// Compute forward pass with opt-in checked-domain validation. + pub fn compute_checked(exprs: &[MonoAD2FR], x: f64) -> Result { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_forward_checked(&ops, x) + } + + /// Compute forward pass and return gradient function using reverse-mode. + pub fn compute_grad(exprs: &[MonoAD2FR], x: f64) -> BackwardResultBox { + Self::compute_grad_generic::>(exprs, x) + } + + /// Compute forward pass and gradient function with checked-domain validation. + pub fn compute_grad_checked(exprs: &[MonoAD2FR], x: f64) -> Result { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_grad_generic_checked::>(&ops, x) + } + + fn compute_grad_generic(exprs: &[MonoAD2FR], x: f64) -> (f64, W) + where + W: From> + std::ops::Deref + 'static, + { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_grad_generic::(&ops, x) + } + + /// Compute exact Hessian using Forward-over-Reverse mode. + /// + /// For single-operation expressions, uses direct dual-number differentiation. + /// For multi-operation compositions, delegates to [`crate::MonoAD2RR`] (which + /// handles compositions correctly via explicit second-order chain rule). + /// + /// # Arguments + /// + /// * `exprs` - Slice of operations to apply in sequence + /// * `x` - Input value to evaluate at + /// + /// # Returns + /// + /// The exact second derivative f''(x) at the given point + /// + /// # Examples + /// + /// ``` + /// use petite_ad::MonoAD2FR; + /// + /// // f(x) = sin(x), f''(x) = -sin(x) + /// let ops = [MonoAD2FR::Sin]; + /// let x = 0.5; + /// let hessian = MonoAD2FR::compute_hessian(&ops, x); + /// assert!((hessian - (-0.5_f64.sin())).abs() < 1e-12); + /// + /// // Composition f(x) = exp(sin(x)) — uses RR fallback + /// let ops = [MonoAD2FR::Sin, MonoAD2FR::Exp]; + /// let x = 1.0; + /// let hessian = MonoAD2FR::compute_hessian(&ops, x); + /// let expected = x.sin().exp() * (x.cos().powi(2) - x.sin()); + /// assert!((hessian - expected).abs() < 1e-12); + /// ``` + pub fn compute_hessian(exprs: &[MonoAD2FR], x: f64) -> f64 { + if exprs.is_empty() { + return 0.0; + } + + // Single operation: direct dual-number evaluation + if let [op] = exprs { + return mono_hessian_common::compute_single_op_hessian((*op).into(), x) + .expect("all single ops supported"); + } + + let rr_ops = Self::to_rr_ops(exprs); + MonoAD2RR::compute_hessian(&rr_ops, x) + } + + /// Compute exact Hessian with checked-domain validation. + pub fn compute_hessian_checked(exprs: &[MonoAD2FR], x: f64) -> Result { + let rr_ops = Self::to_rr_ops(exprs); + MonoAD2RR::compute_hessian_checked(&rr_ops, x) + } + + fn to_rr_ops(exprs: &[MonoAD2FR]) -> Vec { + exprs + .iter() + .map(|&op| match op { + MonoAD2FR::Sin => MonoAD2RR::Sin, + MonoAD2FR::Cos => MonoAD2RR::Cos, + MonoAD2FR::Tan => MonoAD2RR::Tan, + MonoAD2FR::Exp => MonoAD2RR::Exp, + MonoAD2FR::Neg => MonoAD2RR::Neg, + MonoAD2FR::Ln => MonoAD2RR::Ln, + MonoAD2FR::Sqrt => MonoAD2RR::Sqrt, + MonoAD2FR::Abs => MonoAD2RR::Abs, + }) + .collect() + } +} diff --git a/src/mono/mono_ad_rf.rs b/src/mono/mono_ad_rf.rs new file mode 100644 index 0000000..f119b99 --- /dev/null +++ b/src/mono/mono_ad_rf.rs @@ -0,0 +1,259 @@ +//! Exact second-order autodiff using Reverse-over-Forward (RF) mode. +//! +//! This module implements the **Reverse-over-Forward (RF)** method for computing exact +//! Hessians (second derivatives) of single-variable functions. +//! +//! # Comprehensive Documentation +//! +//! For complete mathematical theory, detailed derivations, and comparison with other methods, see: +//! **[`/docs/mono_ad_hessian.md`](../../docs/mono_ad_hessian.md)** +//! +//! # Supported Operations +//! +//! This type supports a subset of operations compared to [`crate::MonoAD`]: +//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. See [`super::mono_hessian_common`] for details. +//! +//! For single-variable functions, RF and FR are mathematically equivalent (both use dual +//! numbers), but they differ in conceptual organization: +//! +//! - **FR**: Reverse-mode for gradient, then forward-mode on the gradient function +//! - **RF**: Forward-mode embedded within reverse-mode computation +//! +//! For a function f(x): +//! 1. Use dual numbers to track both value and derivative +//! 2. The tangent component propagates derivatives through operations +//! 3. Result: f''(x) from the final tangent +//! +//! ## Dual Number Arithmetic +//! +//! A dual number is a pair (val, tan) where: +//! - `val`: the function value +//! - `tan`: the tangent (derivative) value +//! +//! For an input x with seed tangent 1.0: Dual::variable(x) = (x, 1.0) +//! +//! Each operation propagates derivatives using the chain rule: +//! +//! ```text +//! Operation | Value | Tangent (derivative) +//! ----------|---------------|---------------------- +//! sin(x) | sin(x.val) | cos(x.val) · x.tan +//! cos(x) | cos(x.val) | -sin(x.val) · x.tan +//! tan(x) | tan(x.val) | sec²(x.val) · x.tan +//! exp(x) | exp(x.val) | exp(x.val) · x.tan +//! -x | -x.val | -x.tan +//! ln(x) | ln(x.val) | x.tan / x.val +//! sqrt(x) | sqrt(x.val) | x.tan / (2 sqrt(x.val)) +//! abs(x) | abs(x.val) | sign_or_zero(x.val) · x.tan +//! ``` +//! +//! ## Gradient Functions +//! +//! For elementary operations, the gradient functions are: +//! +//! | Original f(x) | Gradient g(x) = f'(x) | Second derivative g'(x) = f''(x) | +//! |---------------|----------------------|----------------------------------| +//! | sin(x) | cos(x) | -sin(x) | +//! | cos(x) | -sin(x) | -cos(x) | +//! | exp(x) | exp(x) | exp(x) | +//! | -x | -1 | 0 | +//! | ln(x) | 1/x | -1/x² | +//! | sqrt(x) | 1/(2√x) | -1/(4x√x) | +//! | abs(x) | sign(x) | 0 (raw convention at x = 0) | +//! +//! ## RF vs FR: What's the Difference? +//! +//! For **single-variable** functions, RF and FR are computationally identical: +//! - Both use dual number arithmetic +//! - Both compute the same second derivatives +//! - Both have the same complexity +//! +//! The distinction becomes meaningful for **multi-variable** functions: +//! - **FR**: Compute full Jacobian (reverse), then differentiate each row (forward) +//! - **RF**: Differentiate function (forward) for each output (reverse) +//! +//! For MonoAD (single variable), this implementation is shared with FR via +//! [`super::mono_hessian_common`]. +//! +//! # Algorithm +//! +//! ## For Single Operations +//! +//! Direct dual-number evaluation of the second derivative. +//! +//! ## For Multiple Operations (Compositions) +//! +//! Delegates to RR mode ([`crate::MonoAD2RR`]), which handles general compositions +//! correctly. +//! +//! # Accuracy +//! +//! - **Method**: Exact symbolic differentiation via dual numbers +//! - **Error source**: Only floating-point rounding (machine epsilon) +//! - **Typical relative error**: < 1e-14 +//! - **No truncation error**: Mathematically exact (unlike finite differences) +//! +//! # Example Usage +//! +//! ```rust +//! use petite_ad::MonoAD2RF; +//! +//! // f(x) = cos(x), f''(x) = -cos(x) +//! let ops = [MonoAD2RF::Cos]; +//! let x = 1.0; +//! +//! let value = MonoAD2RF::compute(&ops, x); +//! let hessian = MonoAD2RF::compute_hessian(&ops, x); +//! +//! println!("f({}) = {}", x, value); // cos(1) ≈ 0.5403 +//! println!("f''({}) = {}", x, hessian); // -cos(1) ≈ -0.5403 +//! ``` +//! +//! # References +//! +//! - Griewank & Walther (2008): *Evaluating Derivatives* +//! - Naumann (2012): *The Art of Differentiating Computer Programs* +//! - Giles (2008): "Collected matrix derivative results for forward and reverse mode AD" + +use crate::Result; + +use super::mono_ad_rr::MonoAD2RR; +use super::mono_hessian_common::{self, MonoHessianOpKind}; +use super::types::*; + +/// Single-variable automatic differentiation operations for Reverse-over-Forward Hessian computation. +/// +/// Supports: `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MonoAD2RF { + Sin, + Cos, + Tan, + Exp, + Neg, + Ln, + Sqrt, + Abs, +} + +impl From for MonoHessianOpKind { + fn from(op: MonoAD2RF) -> MonoHessianOpKind { + match op { + MonoAD2RF::Sin => MonoHessianOpKind::Sin, + MonoAD2RF::Cos => MonoHessianOpKind::Cos, + MonoAD2RF::Tan => MonoHessianOpKind::Tan, + MonoAD2RF::Exp => MonoHessianOpKind::Exp, + MonoAD2RF::Neg => MonoHessianOpKind::Neg, + MonoAD2RF::Ln => MonoHessianOpKind::Ln, + MonoAD2RF::Sqrt => MonoHessianOpKind::Sqrt, + MonoAD2RF::Abs => MonoHessianOpKind::Abs, + } + } +} + +impl MonoAD2RF { + /// Compute forward pass only. + pub fn compute(exprs: &[MonoAD2RF], x: f64) -> f64 { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_forward(&ops, x) + } + + /// Compute forward pass with opt-in checked-domain validation. + pub fn compute_checked(exprs: &[MonoAD2RF], x: f64) -> Result { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_forward_checked(&ops, x) + } + + /// Compute forward pass and return gradient function using reverse-mode. + pub fn compute_grad(exprs: &[MonoAD2RF], x: f64) -> BackwardResultBox { + Self::compute_grad_generic::>(exprs, x) + } + + /// Compute forward pass and gradient function with checked-domain validation. + pub fn compute_grad_checked(exprs: &[MonoAD2RF], x: f64) -> Result { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_grad_generic_checked::>(&ops, x) + } + + fn compute_grad_generic(exprs: &[MonoAD2RF], x: f64) -> (f64, W) + where + W: From> + std::ops::Deref + 'static, + { + let ops: Vec = exprs.iter().map(|&op| op.into()).collect(); + mono_hessian_common::compute_grad_generic::(&ops, x) + } + + /// Compute exact Hessian using Reverse-over-Forward mode. + /// + /// For single-variable functions, this is computationally identical to + /// [`crate::MonoAD2FR`] but represents a different conceptual organization. + /// + /// For single-operation expressions, uses direct dual-number differentiation. + /// For multi-operation compositions, delegates to [`crate::MonoAD2RR`]. + /// + /// # Arguments + /// + /// * `exprs` - Slice of operations to apply in sequence + /// * `x` - Input value to evaluate at + /// + /// # Returns + /// + /// The exact second derivative f''(x) at the given point + /// + /// # Examples + /// + /// ``` + /// use petite_ad::MonoAD2RF; + /// + /// // f(x) = cos(x), f''(x) = -cos(x) + /// let ops = [MonoAD2RF::Cos]; + /// let x = 0.5; + /// let hessian = MonoAD2RF::compute_hessian(&ops, x); + /// assert!((hessian - (-0.5_f64.cos())).abs() < 1e-12); + /// + /// // Composition f(x) = sin(sin(x)) — uses RR fallback + /// let ops = [MonoAD2RF::Sin, MonoAD2RF::Sin]; + /// let x = 0.5; + /// let hessian = MonoAD2RF::compute_hessian(&ops, x); + /// let sin_x = x.sin(); + /// let cos_x = x.cos(); + /// let expected = -sin_x.sin() * cos_x.powi(2) - sin_x.cos() * x.sin(); + /// assert!((hessian - expected).abs() < 1e-12); + /// ``` + pub fn compute_hessian(exprs: &[MonoAD2RF], x: f64) -> f64 { + if exprs.is_empty() { + return 0.0; + } + + // Single operation: direct dual-number evaluation + if let [op] = exprs { + return mono_hessian_common::compute_single_op_hessian((*op).into(), x) + .expect("all single ops supported"); + } + + let rr_ops = Self::to_rr_ops(exprs); + MonoAD2RR::compute_hessian(&rr_ops, x) + } + + /// Compute exact Hessian with checked-domain validation. + pub fn compute_hessian_checked(exprs: &[MonoAD2RF], x: f64) -> Result { + let rr_ops = Self::to_rr_ops(exprs); + MonoAD2RR::compute_hessian_checked(&rr_ops, x) + } + + fn to_rr_ops(exprs: &[MonoAD2RF]) -> Vec { + exprs + .iter() + .map(|&op| match op { + MonoAD2RF::Sin => MonoAD2RR::Sin, + MonoAD2RF::Cos => MonoAD2RR::Cos, + MonoAD2RF::Tan => MonoAD2RR::Tan, + MonoAD2RF::Exp => MonoAD2RR::Exp, + MonoAD2RF::Neg => MonoAD2RR::Neg, + MonoAD2RF::Ln => MonoAD2RR::Ln, + MonoAD2RF::Sqrt => MonoAD2RR::Sqrt, + MonoAD2RF::Abs => MonoAD2RR::Abs, + }) + .collect() + } +} diff --git a/src/mono/mono_ad_rr.rs b/src/mono/mono_ad_rr.rs new file mode 100644 index 0000000..34dff7a --- /dev/null +++ b/src/mono/mono_ad_rr.rs @@ -0,0 +1,661 @@ +//! Exact second-order autodiff using Reverse-over-Reverse (RR) mode. +//! +//! This module implements the **Reverse-over-Reverse (RR)** method for computing exact +//! Hessians (second derivatives) of single-variable functions. It tracks second-order +//! derivatives during the reverse pass using the chain rule. +//! +//! # Supported Operations +//! +//! `MonoAD2RR` supports a subset of the operations available in [`crate::MonoAD`]: +//! `Sin`, `Cos`, `Tan`, `Exp`, `Neg`, `Ln`, `Sqrt`, `Abs`. The `Abs` operation is +//! non-smooth at zero and follows the raw `f64` convention used by [`crate::MonoAD`]: +//! derivative `0` and curvature `0` at `x = 0`. For first-order differentiation with +//! the full operation set, use +//! [`crate::MonoAD`]. For Hessian approximation with all operations, use +//! [`MultiAD::compute_hessian`](crate::MultiAD::compute_hessian) (finite-difference based). +//! +//! # Comprehensive Documentation +//! +//! For complete mathematical theory, detailed derivations, complexity analysis, and +//! comparison with other methods, see: +//! **[`/docs/mono_ad_hessian.md`](../../docs/mono_ad_hessian.md)** +//! +//! # Mathematical Foundation +//! +//! ## Chain Rule for Second Derivatives +//! +//! For a composition h(x) = f(g(x)), the chain rule for derivatives is: +//! +//! ```text +//! First derivative: h'(x) = f'(g(x)) · g'(x) +//! Second derivative: h''(x) = f''(g(x)) · [g'(x)]² + f'(g(x)) · g''(x) +//! ``` +//! +//! This formula is fundamental to the RR method. The second derivative has two terms: +//! - **Product term**: f''(g(x)) · [g'(x)]² — contribution from second derivative of outer function +//! - **Chain term**: f'(g(x)) · g''(x) — contribution from first derivative of outer times second of inner +//! +//! ## Backward Propagation +//! +//! During the reverse pass, we propagate two values backward through the computation graph: +//! - `grad`: the accumulated gradient (first derivative) ∂L/∂u +//! - `hessian`: the accumulated second derivative ∂²L/∂u² +//! +//! For each operation y = op(u) with: +//! - First derivative: dy = op'(u) = ∂y/∂u +//! - Second derivative: ddy = op''(u) = ∂²y/∂u² +//! +//! The backward accumulation follows the chain rule: +//! +//! ```text +//! new_grad = grad · dy (standard reverse-mode AD) +//! new_hessian = hessian · dy² + grad · ddy (reverse-over-reverse formula) +//! ``` +//! +//! This is derived from applying the chain rule to the gradient computation itself. +//! See [docs/mono_ad_hessian.md](../../docs/mono_ad_hessian.md#reverse-over-reverse-rr) for full derivation. +//! +//! # Algorithm Overview +//! +//! The RR method operates in two phases: +//! +//! ## Phase 1: Forward Pass +//! +//! Traverse the computation graph forward, storing at each node i: +//! - `values[i]`: the value vᵢ +//! - `first_derivs[i]`: the first derivative dvᵢ/du (where u is the input) +//! - `second_derivs[i]`: the second derivative d²vᵢ/du² +//! +//! ## Phase 2: Reverse Pass +//! +//! Traverse backward from output to input, accumulating: +//! - Start: `grad = 1.0` (∂F/∂F = 1), `hessian = 0.0` (∂²F/∂F² = 0) +//! - For each operation with derivatives (dy, ddy): +//! - Apply chain rule: `new_hessian = hessian · dy² + grad · ddy` +//! - Update: `grad = grad · dy` +//! - Result: `hessian` contains ∂²F/∂x² +//! +//! # Computational Complexity +//! +//! For a computation graph with n operations: +//! - **Time**: O(n) — single forward pass + single reverse pass +//! - **Space**: O(n) — must store all intermediate values and derivatives +//! - **Overhead**: ~3x compared to first-order reverse-mode AD +//! +//! # Accuracy +//! +//! This provides **exact** second derivatives up to floating-point precision: +//! - No finite difference approximations +//! - No truncation error +//! - Error bounded by machine epsilon (~2.2e-16 for f64) +//! - Typical relative error: < 1e-14 +//! +//! # Supported Operations +//! +//! Currently supports eight elementary operations. Each operation's second derivative: +//! +//! | Operation | Function | First Derivative f'(x) | Second Derivative f''(x) | +//! |-----------|----------|----------------------|------------------------| +//! | `Sin` | sin(x) | cos(x) | -sin(x) | +//! | `Cos` | cos(x) | -sin(x) | -cos(x) | +//! | `Tan` | tan(x) | sec²(x) | 2 sec²(x) tan(x) | +//! | `Exp` | exp(x) | exp(x) | exp(x) | +//! | `Neg` | -x | -1 | 0 | +//! | `Ln` | ln(x) | 1/x | -1/x² | +//! | `Sqrt` | sqrt(x) | 1/(2sqrt(x)) | -1/(4x sqrt(x)) | +//! | `Abs` | abs(x) | sign(x) | 0 | +//! +//! # Example Usage +//! +//! ```rust +//! use petite_ad::MonoAD2RR; +//! +//! // Compute f(x) = exp(sin(x)) +//! // f'(x) = cos(x) · exp(sin(x)) +//! // f''(x) = -sin(x) · exp(sin(x)) + cos²(x) · exp(sin(x)) +//! let ops = [MonoAD2RR::Sin, MonoAD2RR::Exp]; +//! let x = 1.0; +//! +//! let value = MonoAD2RR::compute(&ops, x); +//! let hessian = MonoAD2RR::compute_hessian(&ops, x); +//! +//! // Expected: f''(1.0) ≈ -0.9318... +//! println!("f''({}) = {}", x, hessian); +//! ``` +//! +//! # Comparison with Other Methods +//! +//! | Method | Accuracy | Time | Space | Best For | +//! |--------|----------|------|-------|----------| +//! | **RR (this)** | Exact | O(n) | O(n) | General purpose, balanced | +//! | FR | Exact | O(n) | O(n) | Slightly faster than RR | +//! | RF | Exact | O(n) | O(n) | Slightly slower than RR | +//! | Finite-diff | ~1e-5 | O(n) | O(1) | Quick approximations | +//! +//! See [docs/mono_ad_hessian.md](../../docs/mono_ad_hessian.md) for detailed comparisons. +//! +//! # References +//! +//! - Griewank & Walther (2008): *Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation* +//! - Gebremedhin et al. (2002): "Efficient computation of Hessian matrices" +//! - Pearlmutter (1994): "Fast exact multiplication by the Hessian" + +use crate::Result; + +use super::mono_hessian_common::{self, MonoHessianOpKind}; +use super::types::*; + +/// Single-variable automatic differentiation operations for Reverse-over-Reverse Hessian computation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MonoAD2RR { + Sin, + Cos, + Tan, + Exp, + Neg, + Ln, + Sqrt, + Abs, +} + +impl MonoAD2RR { + #[inline(always)] + fn as_hessian_op(self) -> MonoHessianOpKind { + match self { + MonoAD2RR::Sin => MonoHessianOpKind::Sin, + MonoAD2RR::Cos => MonoHessianOpKind::Cos, + MonoAD2RR::Tan => MonoHessianOpKind::Tan, + MonoAD2RR::Exp => MonoHessianOpKind::Exp, + MonoAD2RR::Neg => MonoHessianOpKind::Neg, + MonoAD2RR::Ln => MonoHessianOpKind::Ln, + MonoAD2RR::Sqrt => MonoHessianOpKind::Sqrt, + MonoAD2RR::Abs => MonoHessianOpKind::Abs, + } + } + + #[inline(always)] + fn check_domain(self, x: f64) -> Result<()> { + mono_hessian_common::check_domain(self.as_hessian_op(), x) + } + + /// Forward pass computing value, first derivative, and second derivative. + /// + /// Evaluates an elementary operation at a given point, returning the function value + /// and its first two derivatives. This is the core primitive for the RR method. + /// + /// # Mathematical Formulas + /// + /// For each operation, we compute: + /// + /// ## Sin: y = sin(x) + /// - Value: y = sin(x) + /// - First derivative: dy/dx = cos(x) + /// - Second derivative: d²y/dx² = -sin(x) + /// + /// Derivation: d/dx[cos(x)] = -sin(x) + /// + /// ## Cos: y = cos(x) + /// - Value: y = cos(x) + /// - First derivative: dy/dx = -sin(x) + /// - Second derivative: d²y/dx² = -cos(x) + /// + /// Derivation: d/dx[-sin(x)] = -cos(x) + /// + /// ## Exp: y = exp(x) + /// - Value: y = exp(x) + /// - First derivative: dy/dx = exp(x) + /// - Second derivative: d²y/dx² = exp(x) + /// + /// Derivation: The exponential function is its own derivative at all orders + /// + /// ## Neg: y = -x + /// - Value: y = -x + /// - First derivative: dy/dx = -1 + /// - Second derivative: d²y/dx² = 0 + /// + /// Derivation: Linear functions have zero second derivative + /// + /// # Arguments + /// + /// * `self` - The operation to evaluate (Sin, Cos, Tan, Exp, Neg, Ln, Sqrt, or Abs) + /// * `x` - The point at which to evaluate the operation + /// + /// # Returns + /// + /// A tuple `(value, first_derivative, second_derivative)` containing: + /// - `value`: f(x) + /// - `first_derivative`: f'(x) + /// - `second_derivative`: f''(x) + /// + /// # Examples + /// + /// ```rust,ignore + /// // This is a private method used internally by compute_hessian + /// use petite_ad::MonoAD2RR; + /// + /// let op = MonoAD2RR::Sin; + /// // forward_d2 is private and used internally + /// // See compute_hessian for public API examples + /// ``` + /// + /// For public usage examples, see [`MonoAD2RR::compute_hessian`]. + /// + /// # Complexity + /// + /// - **Time**: O(1) — constant time for each operation + /// - **Space**: O(1) — returns three f64 values + fn forward_d2(&self, x: f64) -> (f64, f64, f64) { + match self { + MonoAD2RR::Sin => { + // f(x) = sin(x) + // f'(x) = cos(x) + // f''(x) = -sin(x) + let y = x.sin(); + let dy = x.cos(); // First derivative + let ddy = -x.sin(); // Second derivative: d/dx[cos(x)] = -sin(x) + (y, dy, ddy) + } + MonoAD2RR::Cos => { + // f(x) = cos(x) + // f'(x) = -sin(x) + // f''(x) = -cos(x) + let y = x.cos(); + let dy = -x.sin(); // First derivative + let ddy = -x.cos(); // Second derivative: d/dx[-sin(x)] = -cos(x) + (y, dy, ddy) + } + MonoAD2RR::Tan => { + // f(x) = tan(x) + // f'(x) = sec^2(x) + // f''(x) = 2 sec^2(x) tan(x) + let y = x.tan(); + let sec_sq = 1.0 / x.cos().powi(2); + let dy = sec_sq; + let ddy = 2.0 * sec_sq * y; + (y, dy, ddy) + } + MonoAD2RR::Exp => { + // f(x) = exp(x) + // f'(x) = exp(x) + // f''(x) = exp(x) + // The exponential function is its own derivative at all orders + let y = x.exp(); + let dy = y; // First derivative: exp(x) + let ddy = y; // Second derivative: exp(x) + (y, dy, ddy) + } + MonoAD2RR::Neg => { + // f(x) = -x + // f'(x) = -1 + // f''(x) = 0 + // Linear functions have zero second derivative + let y = -x; + let dy = -1.0; // First derivative: constant + let ddy = 0.0; // Second derivative: zero (linear function) + (y, dy, ddy) + } + MonoAD2RR::Ln => { + // f(x) = ln(x) + // f'(x) = 1/x + // f''(x) = -1/x^2 + let y = x.ln(); + let dy = 1.0 / x; + let ddy = -1.0 / x.powi(2); + (y, dy, ddy) + } + MonoAD2RR::Sqrt => { + // f(x) = sqrt(x) + // f'(x) = 1/(2 sqrt(x)) + // f''(x) = -1/(4 x sqrt(x)) + let y = x.sqrt(); + let dy = 1.0 / (2.0 * y); + let ddy = -1.0 / (4.0 * x * y); + (y, dy, ddy) + } + MonoAD2RR::Abs => { + // f(x) = abs(x) + // f'(x) = sign(x), with raw convention f'(0) = 0 + // f''(x) = 0 away from zero; raw convention f''(0) = 0 + let y = x.abs(); + let dy = mono_hessian_common::sign_or_zero(x); + let ddy = 0.0; + (y, dy, ddy) + } + } + } + + /// Compute forward pass only. + pub fn compute(exprs: &[MonoAD2RR], x: f64) -> f64 { + let mut value = x; + for expr in exprs { + value = expr.forward_d2(value).0; + } + value + } + + /// Compute forward pass with opt-in checked-domain validation. + pub fn compute_checked(exprs: &[MonoAD2RR], x: f64) -> Result { + let mut value = x; + for &op in exprs { + op.check_domain(value)?; + value = op.forward_d2(value).0; + } + Ok(value) + } + + /// Compute forward pass and return gradient function using exact reverse-mode. + pub fn compute_grad(exprs: &[MonoAD2RR], x: f64) -> BackwardResultBox { + Self::compute_grad_generic::>(exprs, x) + } + + /// Compute forward pass and gradient function with checked-domain validation. + pub fn compute_grad_checked(exprs: &[MonoAD2RR], x: f64) -> Result { + let mut value = x; + let mut backprops: Vec> = Vec::new(); + + for &op in exprs { + op.check_domain(value)?; + let (new_value, _dy, _ddy) = op.forward_d2(value); + let backprop = Self::make_backward_fn(op, value); + value = new_value; + backprops.push(backprop); + } + + let backward_fn = Box::new(move |cotangent: f64| -> f64 { + let mut grad = cotangent; + for backprop in backprops.iter().rev() { + grad = backprop(grad); + } + grad + }); + + Ok((value, backward_fn)) + } + + /// Generic gradient computation. + fn compute_grad_generic(exprs: &[MonoAD2RR], x: f64) -> (f64, W) + where + W: From> + std::ops::Deref + 'static, + { + let mut value = x; + let mut backprops: Vec = Vec::new(); + + for &op in exprs { + let (new_value, _dy, _ddy) = op.forward_d2(value); + let backprop = Self::make_backward_fn(op, value); + value = new_value; + backprops.push(backprop); + } + + let backward_fn = Box::new(move |cotangent: f64| -> f64 { + let mut grad = cotangent; + for backprop in backprops.iter().rev() { + grad = backprop(grad); + } + grad + }); + + (value, W::from(backward_fn)) + } + + /// Create backward function for an operation (first-order only). + fn make_backward_fn(op: MonoAD2RR, x: f64) -> W + where + W: From>, + { + let grad_fn: Box = match op { + MonoAD2RR::Sin => Box::new(move |dy: f64| -> f64 { dy * x.cos() }), + MonoAD2RR::Cos => Box::new(move |dy: f64| -> f64 { dy * -x.sin() }), + MonoAD2RR::Tan => { + let sec_sq = 1.0 / x.cos().powi(2); + Box::new(move |dy: f64| -> f64 { dy * sec_sq }) + } + MonoAD2RR::Exp => { + let exp_val = x.exp(); + Box::new(move |dy: f64| -> f64 { dy * exp_val }) + } + MonoAD2RR::Neg => Box::new(move |dy: f64| -> f64 { -dy * 1.0 }), + MonoAD2RR::Ln => Box::new(move |dy: f64| -> f64 { dy / x }), + MonoAD2RR::Sqrt => { + let sqrt_x = x.sqrt(); + Box::new(move |dy: f64| -> f64 { dy / (2.0 * sqrt_x) }) + } + MonoAD2RR::Abs => { + let sign = mono_hessian_common::sign_or_zero(x); + Box::new(move |dy: f64| -> f64 { dy * sign }) + } + }; + W::from(grad_fn) + } + + /// Compute exact Hessian using Reverse-over-Reverse mode. + /// + /// This is the main entry point for computing exact second derivatives. + /// It implements the RR algorithm by performing a forward pass to collect + /// derivative information, then a reverse pass to accumulate the Hessian + /// using the chain rule. + /// + /// # Algorithm Walkthrough + /// + /// Given a sequence of operations [op₁, op₂, ..., opₙ] applied to input x: + /// + /// ## Forward Pass + /// + /// For each operation opᵢ in sequence: + /// 1. Compute (yᵢ, dyᵢ, ddyᵢ) = opᵢ.forward_d2(yᵢ₋₁) + /// 2. Store yᵢ, dyᵢ, ddyᵢ + /// + /// where y₀ = x (the input) + /// + /// This gives us the computation graph with all intermediate values and local derivatives. + /// + /// ## Reverse Pass + /// + /// Initialize: + /// - `grad = 1.0` (seed: ∂F/∂F = 1, where F is the final output) + /// - `hessian = 0.0` (seed: ∂²F/∂F² = 0) + /// + /// For each operation in reverse order (opₙ, ..., op₂, op₁): + /// 1. Retrieve stored derivatives: (dyᵢ, ddyᵢ) + /// 2. Apply chain rule: + /// ```text + /// new_hessian = hessian · dyᵢ² + grad · ddyᵢ + /// new_grad = grad · dyᵢ + /// ``` + /// 3. Update: grad ← new_grad, hessian ← new_hessian + /// + /// The final `hessian` value is ∂²F/∂x². + /// + /// ## Chain Rule Derivation + /// + /// At each step, we're computing derivatives with respect to earlier variables. + /// If we have L = f(u) and u = g(x), then by the chain rule: + /// + /// ```text + /// ∂L/∂x = ∂L/∂u · ∂u/∂x = grad · dy + /// + /// ∂²L/∂x² = ∂/∂x[∂L/∂x] + /// = ∂/∂x[∂L/∂u · ∂u/∂x] + /// = ∂²L/∂u² · (∂u/∂x)² + ∂L/∂u · ∂²u/∂x² + /// = hessian · dy² + grad · ddy + /// ``` + /// + /// This is the fundamental formula used in the reverse pass. + /// + /// # Example Trace + /// + /// For f(x) = exp(sin(x)) at x = 1.0: + /// + /// ## Forward Pass: + /// ```text + /// Input: x = 1.0 + /// + /// Op 1: Sin + /// y₁ = sin(1.0) ≈ 0.8414709848 + /// dy₁ = cos(1.0) ≈ 0.5403023059 + /// ddy₁ = -sin(1.0) ≈ -0.8414709848 + /// + /// Op 2: Exp + /// y₂ = exp(0.8414709848) ≈ 2.3198323620 + /// dy₂ = exp(0.8414709848) ≈ 2.3198323620 + /// ddy₂ = exp(0.8414709848) ≈ 2.3198323620 + /// ``` + /// + /// ## Reverse Pass: + /// ```text + /// Initialize: grad = 1.0, hessian = 0.0 + /// + /// Step 1 (Op 2: Exp): + /// new_hessian = 0.0 · (2.3198)² + 1.0 · 2.3198 ≈ 2.3198 + /// new_grad = 1.0 · 2.3198 ≈ 2.3198 + /// + /// Step 2 (Op 1: Sin): + /// new_hessian = 2.3198 · (0.5403)² + 2.3198 · (-0.8414) + /// ≈ 0.6766 - 1.9523 ≈ -1.2757 + /// new_grad = 2.3198 · 0.5403 ≈ 1.2533 + /// + /// Result: f''(1.0) ≈ -1.2757 + /// ``` + /// + /// # Arguments + /// + /// * `exprs` - Slice of operations to apply in sequence + /// * `x` - Input value to evaluate at + /// + /// # Returns + /// + /// The exact second derivative f''(x) at the given point + /// + /// # Edge Cases + /// + /// * **Empty expression**: Returns 0.0 (constant function has zero second derivative) + /// * **Single operation**: Returns the operation's second derivative at x + /// * **Linear composition**: Correctly returns 0.0 for linear functions (e.g., Neg) + /// + /// # Accuracy + /// + /// - **Method**: Exact symbolic differentiation evaluated numerically + /// - **Error source**: Only floating-point rounding (machine epsilon) + /// - **Typical relative error**: < 1e-12 + /// - **No truncation error**: Unlike finite differences, this is mathematically exact + /// + /// # Complexity + /// + /// For n operations: + /// - **Time**: O(n) — one forward pass + one reverse pass, each visiting n nodes + /// - **Space**: O(n) — must store n values and 2n derivatives + /// - **Overhead vs first-order AD**: ~3x (stores second derivatives + extra arithmetic) + /// + /// # Examples + /// + /// ``` + /// use petite_ad::MonoAD2RR; + /// + /// // Example 1: f(x) = sin(x), f''(x) = -sin(x) + /// let ops = [MonoAD2RR::Sin]; + /// let x = 0.5; + /// let hessian = MonoAD2RR::compute_hessian(&ops, x); + /// assert!((hessian - (-0.5_f64.sin())).abs() < 1e-12); + /// + /// // Example 2: f(x) = exp(sin(x)) + /// // f'(x) = cos(x) · exp(sin(x)) + /// // f''(x) = -sin(x) · exp(sin(x)) + cos²(x) · exp(sin(x)) + /// let ops = [MonoAD2RR::Sin, MonoAD2RR::Exp]; + /// let x = 1.0; + /// let hessian = MonoAD2RR::compute_hessian(&ops, x); + /// + /// // Manual calculation: + /// let sin_x = x.sin(); + /// let cos_x = x.cos(); + /// let exp_sin_x = sin_x.exp(); + /// let expected = -sin_x * exp_sin_x + cos_x * cos_x * exp_sin_x; + /// assert!((hessian - expected).abs() < 1e-12); + /// + /// // Example 3: f(x) = -x (linear, so f''(x) = 0) + /// let ops = [MonoAD2RR::Neg]; + /// let hessian = MonoAD2RR::compute_hessian(&ops, 1.0); + /// assert_eq!(hessian, 0.0); + /// ``` + /// + /// # See Also + /// + /// - [`MonoAD2RR::compute`]: For computing just the function value + /// - [`MonoAD2RR::compute_grad`]: For computing first derivatives + /// - [`crate::MonoAD2FR`]: Forward-over-Reverse method (alternative exact method) + /// - [`crate::MonoAD2RF`]: Reverse-over-Forward method (alternative exact method) + /// - [docs/mono_ad_hessian.md](../../docs/mono_ad_hessian.md): Complete mathematical theory + pub fn compute_hessian(exprs: &[MonoAD2RR], x: f64) -> f64 { + Self::compute_hessian_impl(exprs, x, false).expect("unchecked Hessian cannot fail") + } + + /// Compute exact Hessian with checked-domain validation. + pub fn compute_hessian_checked(exprs: &[MonoAD2RR], x: f64) -> Result { + Self::compute_hessian_impl(exprs, x, true) + } + + fn compute_hessian_impl(exprs: &[MonoAD2RR], x: f64, checked: bool) -> Result { + // Edge case: empty expression represents constant function + if exprs.is_empty() { + return Ok(0.0); + } + + let n = exprs.len(); + + // ======================================== + // FORWARD PASS: Collect derivative information + // ======================================== + // Store all intermediate values (needed for computing derivatives) + let mut values: Vec = Vec::with_capacity(n + 1); + values.push(x); // y₀ = x (input) + + // Store first and second derivatives at each operation + let mut first_derivs: Vec = Vec::with_capacity(n); + let mut second_derivs: Vec = Vec::with_capacity(n); + + // Traverse forward through operations, computing and storing derivatives + for &op in exprs { + let input_val = *values.last().unwrap(); // yᵢ₋₁ + if checked { + op.check_domain(input_val)?; + } + let (y, dy, ddy) = op.forward_d2(input_val); + + values.push(y); // Store yᵢ + first_derivs.push(dy); // Store ∂yᵢ/∂yᵢ₋₁ + second_derivs.push(ddy); // Store ∂²yᵢ/∂yᵢ₋₁² + } + + // ======================================== + // REVERSE PASS: Accumulate Hessian using chain rule + // ======================================== + // Initialize reverse pass seeds + // grad = ∂F/∂F = 1.0 (derivative of output with respect to itself) + // hessian = ∂²F/∂F² = 0.0 (second derivative of output with respect to itself) + let mut grad: f64 = 1.0; + let mut hessian: f64 = 0.0; + + // Traverse backward through operations, applying chain rule + for i in (0..n).rev() { + let dy = first_derivs[i]; // ∂yᵢ/∂yᵢ₋₁ + let ddy = second_derivs[i]; // ∂²yᵢ/∂yᵢ₋₁² + + // Apply second-order chain rule: + // If L depends on yᵢ and yᵢ depends on yᵢ₋₁, then: + // ∂²L/∂yᵢ₋₁² = ∂²L/∂yᵢ² · (∂yᵢ/∂yᵢ₋₁)² + ∂L/∂yᵢ · ∂²yᵢ/∂yᵢ₋₁² + // = hessian · dy² + grad · ddy + // + // See docs/mono_ad_hessian.md for detailed derivation + let new_hessian = hessian * dy * dy + grad * ddy; + + // Standard reverse-mode gradient update + let new_grad = grad * dy; + + // Update for next iteration (moving backward) + grad = new_grad; + hessian = new_hessian; + } + + // After processing all operations, hessian contains ∂²F/∂x² + Ok(hessian) + } +} diff --git a/src/mono/mono_fn.rs b/src/mono/mono_fn.rs index e1773b3..97cf0eb 100644 --- a/src/mono/mono_fn.rs +++ b/src/mono/mono_fn.rs @@ -67,7 +67,7 @@ pub trait MonoFn { assert!((result - self.expected_value()).abs() < 1e-10); } println!("\nForward pass only:"); - println!("f({:?}) = {}", &self.input(), result); + println!("f({:?}) = {}", self.input(), result); // Forward + backward (automatic differentiation) let (value, backprop_fn) = self.compute_with_gradient(); @@ -76,13 +76,13 @@ pub trait MonoFn { } let grad = backprop_fn(1.0); println!("\nForward + backward (automatic differentiation):"); - println!("f({:?}) = {}", &self.input(), value); + println!("f({:?}) = {}", self.input(), value); println!("∂f/∂x = {}", grad); // Verify against analytical solution let expected_grad = self.expected_gradient(); println!("\nAnalytical gradients:"); - println!("∂f/∂x = {:?}", &expected_grad); + println!("∂f/∂x = {:?}", expected_grad); println!("\nGradient differences:"); println!( @@ -94,3 +94,68 @@ pub trait MonoFn { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::approx_eq_eps as approx_eq; + + /// f(x) = exp(sin(x)), f'(x) = exp(sin(x)) * cos(x) + struct SinExpFn; + + impl MonoFn for SinExpFn { + fn input(&self) -> f64 { + 0.5 + } + fn graph(&self) -> &'static GraphType { + static G: [MonoAD; 2] = [MonoAD::Sin, MonoAD::Exp]; + &G + } + fn expected_value(&self) -> f64 { + 0.5_f64.sin().exp() + } + fn expected_gradient(&self) -> f64 { + 0.5_f64.sin().exp() * 0.5_f64.cos() + } + } + + #[test] + fn test_mono_fn_compute() { + let f = SinExpFn; + let value = f.compute(); + assert!( + approx_eq(value, f.expected_value(), 1e-10), + "compute value mismatch" + ); + } + + #[test] + fn test_mono_fn_compute_with_gradient() { + let f = SinExpFn; + let (value, backprop) = f.compute_with_gradient(); + assert!( + approx_eq(value, f.expected_value(), 1e-10), + "value mismatch" + ); + let grad = backprop(1.0); + assert!( + approx_eq(grad, f.expected_gradient(), 1e-10), + "gradient mismatch" + ); + } + + #[test] + fn test_mono_fn_test_mono_ad() { + SinExpFn.test_mono_ad(); + } + + #[test] + fn test_mono_fn_demonstrate_with_assert() { + SinExpFn.demonstrate(true); + } + + #[test] + fn test_mono_fn_demonstrate_without_assert() { + SinExpFn.demonstrate(false); + } +} diff --git a/src/mono/mono_hessian_common.rs b/src/mono/mono_hessian_common.rs new file mode 100644 index 0000000..34b7178 --- /dev/null +++ b/src/mono/mono_hessian_common.rs @@ -0,0 +1,264 @@ +//! Shared utilities for mono-variable FR/RF Hessian computation. +//! +//! For single-variable functions, Forward-over-Reverse (FR) and Reverse-over-Forward (RF) +//! are mathematically equivalent: both use dual-number arithmetic to differentiate the +//! gradient function. This module provides the shared implementation to avoid code +//! duplication between `MonoAD2FR` and `MonoAD2RF`. +//! +//! # Supported Operations +//! +//! The Hessian types (`MonoAD2FR`, `MonoAD2RF`, `MonoAD2RR`) support a subset of the +//! operations available in [`crate::MonoAD`]: +//! +//! | Operation | Description | First derivative | Second derivative | +//! |-----------|-------------------|-------------------|-------------------| +//! | `Sin` | sin(x) | cos(x) | -sin(x) | +//! | `Cos` | cos(x) | -sin(x) | -cos(x) | +//! | `Tan` | tan(x) | sec²(x) | 2 sec²(x) tan(x) | +//! | `Exp` | exp(x) | exp(x) | exp(x) | +//! | `Neg` | -x | -1 | 0 | +//! | `Ln` | ln(x) | 1/x | -1/x² | +//! | `Sqrt` | sqrt(x) | 1/(2 sqrt(x)) | -1/(4x sqrt(x)) | +//! | `Abs` | abs(x) | sign(x) | 0 | +//! +//! `Abs` is non-smooth at zero; exact Hessian types follow the same raw `f64` +//! convention as [`crate::MonoAD`] by using derivative `0` and curvature `0` at `x = 0`. +//! Operations like `Pow` are not yet supported in the +//! Hessian types. Use [`crate::MonoAD`] for first-order differentiation with the full +//! operation set, or `MultiAD::compute_hessian` for finite-difference Hessian +//! approximation with all operations. + +use crate::{AutodiffError, Result}; + +use super::types::*; + +/// Operation kind for mono-variable second-order AD (shared by FR and RF). +/// +/// This enum abstracts over `MonoAD2FR` and `MonoAD2RF`, which have identical +/// variants. It is used by the shared [`compute_hessian_dual`] implementation. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum MonoHessianOpKind { + Sin, + Cos, + Tan, + Exp, + Neg, + Ln, + Sqrt, + Abs, +} + +/// Validate real-domain restrictions for checked mono Hessian evaluation. +#[inline(always)] +pub(crate) fn check_domain(op: MonoHessianOpKind, x: f64) -> Result<()> { + match op { + MonoHessianOpKind::Ln if x <= 0.0 => { + Err(AutodiffError::domain("Ln", "input must be positive")) + } + MonoHessianOpKind::Sqrt if x < 0.0 => { + Err(AutodiffError::domain("Sqrt", "input must be non-negative")) + } + _ => Ok(()), + } +} + +/// Evaluate a single operation on a dual number. +#[inline(always)] +pub(crate) fn forward_dual(op: MonoHessianOpKind, x: Dual) -> Dual { + match op { + MonoHessianOpKind::Sin => Dual { + val: x.val.sin(), + tan: x.val.cos() * x.tan, + }, + MonoHessianOpKind::Cos => Dual { + val: x.val.cos(), + tan: -x.val.sin() * x.tan, + }, + MonoHessianOpKind::Tan => { + let cos_x = x.val.cos(); + let sec_sq = 1.0 / cos_x.powi(2); + Dual { + val: x.val.tan(), + tan: sec_sq * x.tan, + } + } + MonoHessianOpKind::Exp => Dual { + val: x.val.exp(), + tan: x.val.exp() * x.tan, + }, + MonoHessianOpKind::Neg => Dual { + val: -x.val, + tan: -x.tan, + }, + MonoHessianOpKind::Ln => Dual { + val: x.val.ln(), + tan: x.tan / x.val, + }, + MonoHessianOpKind::Sqrt => { + let sqrt_x = x.val.sqrt(); + Dual { + val: sqrt_x, + tan: x.tan / (2.0 * sqrt_x), + } + } + MonoHessianOpKind::Abs => Dual { + val: x.val.abs(), + tan: x.tan * sign_or_zero(x.val), + }, + } +} + +/// Evaluate a single operation's scalar value. +#[inline(always)] +pub(crate) fn eval_scalar(op: MonoHessianOpKind, x: f64) -> f64 { + match op { + MonoHessianOpKind::Sin => x.sin(), + MonoHessianOpKind::Cos => x.cos(), + MonoHessianOpKind::Tan => x.tan(), + MonoHessianOpKind::Exp => x.exp(), + MonoHessianOpKind::Neg => -x, + MonoHessianOpKind::Ln => x.ln(), + MonoHessianOpKind::Sqrt => x.sqrt(), + MonoHessianOpKind::Abs => x.abs(), + } +} + +/// Evaluate a single operation's scalar value with checked-domain validation. +#[inline(always)] +pub(crate) fn eval_scalar_checked(op: MonoHessianOpKind, x: f64) -> Result { + check_domain(op, x)?; + Ok(eval_scalar(op, x)) +} + +/// Create a first-order backward closure for an operation at a given input value. +pub(crate) fn make_backward_fn(op: MonoHessianOpKind, x: f64) -> Box { + match op { + MonoHessianOpKind::Sin => Box::new(move |dy: f64| -> f64 { dy * x.cos() }), + MonoHessianOpKind::Cos => Box::new(move |dy: f64| -> f64 { dy * -x.sin() }), + MonoHessianOpKind::Tan => { + let sec_sq = 1.0 / x.cos().powi(2); + Box::new(move |dy: f64| -> f64 { dy * sec_sq }) + } + MonoHessianOpKind::Exp => { + let exp_val = x.exp(); + Box::new(move |dy: f64| -> f64 { dy * exp_val }) + } + MonoHessianOpKind::Neg => Box::new(move |dy: f64| -> f64 { -dy }), + MonoHessianOpKind::Ln => Box::new(move |dy: f64| -> f64 { dy / x }), + MonoHessianOpKind::Sqrt => { + let sqrt_x = x.sqrt(); + Box::new(move |dy: f64| -> f64 { dy / (2.0 * sqrt_x) }) + } + MonoHessianOpKind::Abs => { + let sign = sign_or_zero(x); + Box::new(move |dy: f64| -> f64 { dy * sign }) + } + } +} + +/// Return the sign of `x`, using `0` at `x = 0` to match [`crate::MonoAD`]. +#[inline] +pub(crate) fn sign_or_zero(x: f64) -> f64 { + if x > 0.0 { + 1.0 + } else if x < 0.0 { + -1.0 + } else { + 0.0 + } +} + +/// Compute the second derivative for a single-operation expression using dual numbers. +/// +/// For multi-operation compositions, returns `None` (caller should delegate to RR). +pub(crate) fn compute_single_op_hessian(op: MonoHessianOpKind, x: f64) -> Option { + match op { + MonoHessianOpKind::Sin => Some(-x.sin()), + MonoHessianOpKind::Cos => Some(-x.cos()), + MonoHessianOpKind::Tan => { + let sec_sq = 1.0 / x.cos().powi(2); + Some(2.0 * sec_sq * x.tan()) + } + MonoHessianOpKind::Exp => Some(x.exp()), + MonoHessianOpKind::Neg => Some(0.0), + MonoHessianOpKind::Ln => Some(-1.0 / x.powi(2)), + MonoHessianOpKind::Sqrt => { + let sqrt_x = x.sqrt(); + Some(-1.0 / (4.0 * x * sqrt_x)) + } + MonoHessianOpKind::Abs => Some(0.0), + } +} + +/// Compute forward pass with checked-domain validation. +pub(crate) fn compute_forward_checked(ops: &[MonoHessianOpKind], x: f64) -> Result { + let mut value = x; + for &op in ops { + value = eval_scalar_checked(op, value)?; + } + Ok(value) +} + +/// Generic gradient computation using reverse-mode (first-order only). +/// +/// Shared by `MonoAD2FR` and `MonoAD2RF` to avoid duplicating the backward pass logic. +pub(crate) fn compute_grad_generic(ops: &[MonoHessianOpKind], x: f64) -> (f64, W) +where + W: From> + std::ops::Deref + 'static, +{ + let mut value = x; + let mut backprops: Vec = Vec::new(); + + for &op in ops { + let new_value = eval_scalar(op, value); + let backprop: W = W::from(make_backward_fn(op, value)); + value = new_value; + backprops.push(backprop); + } + + let backward_fn = Box::new(move |cotangent: f64| -> f64 { + let mut grad = cotangent; + for backprop in backprops.iter().rev() { + grad = backprop(grad); + } + grad + }); + + (value, W::from(backward_fn)) +} + +/// Generic checked gradient computation using reverse-mode (first-order only). +pub(crate) fn compute_grad_generic_checked(ops: &[MonoHessianOpKind], x: f64) -> Result<(f64, W)> +where + W: From> + std::ops::Deref + 'static, +{ + let mut value = x; + let mut backprops: Vec = Vec::new(); + + for &op in ops { + check_domain(op, value)?; + let new_value = eval_scalar(op, value); + let backprop = W::from(make_backward_fn(op, value)); + value = new_value; + backprops.push(backprop); + } + + let backward_fn = Box::new(move |cotangent: f64| -> f64 { + let mut grad = cotangent; + for backprop in backprops.iter().rev() { + grad = backprop(grad); + } + grad + }); + + Ok((value, W::from(backward_fn))) +} + +/// Compute forward pass only using the shared operation kind. +pub(crate) fn compute_forward(ops: &[MonoHessianOpKind], x: f64) -> f64 { + let mut dual = Dual::variable(x); + for &op in ops { + dual = forward_dual(op, dual); + } + dual.val +} diff --git a/src/mono/tests.rs b/src/mono/tests.rs index 499f188..aeb621f 100644 --- a/src/mono/tests.rs +++ b/src/mono/tests.rs @@ -3,6 +3,7 @@ use super::mono_fn::MonoFn; use super::{mf1::MF1, mf2::MF2, mf3::MF3, mf4::MF4}; use crate::mono_ops; use crate::test_utils::approx_eq_eps as approx_eq; +use crate::AutodiffError; #[test] fn test_single_sin_compute() { @@ -52,7 +53,7 @@ fn test_compute_arc_same_result() { let (value_box, backprop_box) = MonoAD::compute_grad(ops, 2.0); let grad_box = backprop_box(1.0); - // Verify the computation is correct + // Verify that computation is correct assert!(approx_eq(value_box, 2.2013533791690376, 1e-10)); assert!(approx_eq(grad_box, -0.562752038662712, 1e-10)); } @@ -116,7 +117,7 @@ fn test_compute_arc_consistency() { let (v1, b1) = MonoAD::compute_grad(&ops, 1.5); let g1 = b1(1.0); - // Verify computation succeeds + // Verify that computation succeeds assert!(v1.is_finite(), "value should be finite for ops: {:?}", ops); assert!( g1.is_finite(), @@ -143,3 +144,449 @@ fn test_different_cotangents() { assert!(approx_eq(grad, expected, 1e-10), "cotangent {}", cotangent); } } + +#[test] +fn test_hessian_sin() { + // f(x) = sin(x), f''(x) = -sin(x) + let ops = &[MonoAD::Sin]; + let x = 0.5; + let second_deriv = MonoAD::compute_hessian(ops, x); + + let expected = -x.sin(); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_cos() { + // f(x) = cos(x), f''(x) = -cos(x) + let ops = &[MonoAD::Cos]; + let x = 0.5; + let second_deriv = MonoAD::compute_hessian(ops, x); + + let expected = -x.cos(); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_exp() { + // f(x) = exp(x), f''(x) = exp(x) + let ops = &[MonoAD::Exp]; + let x = 2.0; + let second_deriv = MonoAD::compute_hessian(ops, x); + + let expected = x.exp(); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_composed() { + // f(x) = exp(sin(x)), f''(x) = exp(sin(x)) * cos²(x) - exp(sin(x)) * sin(x) + let ops = &[MonoAD::Sin, MonoAD::Exp]; + let x = 0.5; + let second_deriv = MonoAD::compute_hessian(ops, x); + + let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_chain() { + // f(x) = sin(sin(sin(x))) + // f'(x) = cos(sin(sin(x))) * cos(sin(x)) * cos(x) + // This is complex, so we'll just verify that the computation doesn't panic and returns a value + let ops = &[MonoAD::Sin, MonoAD::Sin, MonoAD::Sin]; + let x = 1.0; + let second_deriv = MonoAD::compute_hessian(ops, x); + + // Just verify that it's a finite number + assert!(second_deriv.is_finite()); +} + +#[test] +fn test_single_tan_compute() { + let ops = &[MonoAD::Tan]; + let x: f64 = 0.5; + let (value, backprop) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(value, x.tan(), 1e-10)); + assert!(approx_eq(backprop(1.0), 1.0 / x.cos().powi(2), 1e-10)); +} + +#[test] +fn test_single_ln_compute() { + let ops = &[MonoAD::Ln]; + let x: f64 = 2.0; + let (value, backprop) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(value, x.ln(), 1e-10)); + assert!(approx_eq(backprop(1.0), 1.0 / x, 1e-10)); +} + +#[test] +fn test_single_sqrt_compute() { + let ops = &[MonoAD::Sqrt]; + let x: f64 = 4.0; + let (value, backprop) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(value, x.sqrt(), 1e-10)); + assert!(approx_eq(backprop(1.0), 1.0 / (2.0 * x.sqrt()), 1e-10)); +} + +#[test] +fn test_single_abs_compute() { + let ops = &[MonoAD::Abs]; + let (value_pos, backprop_pos) = MonoAD::compute_grad(ops, 3.0); + assert!(approx_eq(value_pos, 3.0, 1e-10)); + assert!(approx_eq(backprop_pos(1.0), 1.0, 1e-10)); + + let (value_neg, backprop_neg) = MonoAD::compute_grad(ops, -3.0); + assert!(approx_eq(value_neg, 3.0, 1e-10)); + assert!(approx_eq(backprop_neg(1.0), -1.0, 1e-10)); + + let (value_zero, backprop_zero) = MonoAD::compute_grad(ops, 0.0); + assert!(approx_eq(value_zero, 0.0, 1e-10)); + assert!(approx_eq(backprop_zero(1.0), 0.0, 1e-10)); +} + +// ============================================================================ +// MonoAD::compute tests (forward pass only) +// ============================================================================ + +#[test] +fn test_compute_sin() { + let ops = mono_ops![sin]; + let result = MonoAD::compute(&ops, 2.0); + assert!(approx_eq(result, 2.0_f64.sin(), 1e-10)); +} + +#[test] +fn test_compute_cos() { + let ops = mono_ops![cos]; + let result = MonoAD::compute(&ops, 2.0); + assert!(approx_eq(result, 2.0_f64.cos(), 1e-10)); +} + +#[test] +fn test_compute_tan() { + let ops = mono_ops![tan]; + let result = MonoAD::compute(&ops, 0.5); + assert!(approx_eq(result, 0.5_f64.tan(), 1e-10)); +} + +#[test] +fn test_compute_exp() { + let ops = mono_ops![exp]; + let result = MonoAD::compute(&ops, 2.0); + assert!(approx_eq(result, 2.0_f64.exp(), 1e-10)); +} + +#[test] +fn test_compute_neg() { + let ops = mono_ops![neg]; + let result = MonoAD::compute(&ops, 3.0); + assert!(approx_eq(result, -3.0, 1e-10)); +} + +#[test] +fn test_compute_ln() { + let ops = mono_ops![ln]; + let result = MonoAD::compute(&ops, 2.0); + assert!(approx_eq(result, 2.0_f64.ln(), 1e-10)); +} + +#[test] +fn test_compute_sqrt() { + let ops = mono_ops![sqrt]; + let result = MonoAD::compute(&ops, 4.0); + assert!(approx_eq(result, 4.0_f64.sqrt(), 1e-10)); +} + +#[test] +fn test_compute_abs() { + let ops = mono_ops![abs]; + let result = MonoAD::compute(&ops, -3.0); + assert!(approx_eq(result, 3.0, 1e-10)); +} + +#[test] +fn test_compute_chained() { + let ops = mono_ops![sin, exp, neg]; + let result = MonoAD::compute(&ops, 1.0); + let expected = -1.0_f64.sin().exp(); + assert!(approx_eq(result, expected, 1e-10)); +} + +// ============================================================================ +// MonoAD::compute_checked tests +// ============================================================================ + +#[test] +fn test_compute_checked_sin_success() { + let ops = mono_ops![sin]; + let result = MonoAD::compute_checked(&ops, 2.0).unwrap(); + assert!(approx_eq(result, 2.0_f64.sin(), 1e-10)); +} + +#[test] +fn test_compute_checked_cos_success() { + let ops = mono_ops![cos]; + let result = MonoAD::compute_checked(&ops, 2.0).unwrap(); + assert!(approx_eq(result, 2.0_f64.cos(), 1e-10)); +} + +#[test] +fn test_compute_checked_tan_success() { + let ops = mono_ops![tan]; + let result = MonoAD::compute_checked(&ops, 0.5).unwrap(); + assert!(approx_eq(result, 0.5_f64.tan(), 1e-10)); +} + +#[test] +fn test_compute_checked_exp_success() { + let ops = mono_ops![exp]; + let result = MonoAD::compute_checked(&ops, 2.0).unwrap(); + assert!(approx_eq(result, 2.0_f64.exp(), 1e-10)); +} + +#[test] +fn test_compute_checked_neg_success() { + let ops = mono_ops![neg]; + let result = MonoAD::compute_checked(&ops, 3.0).unwrap(); + assert!(approx_eq(result, -3.0, 1e-10)); +} + +#[test] +fn test_compute_checked_ln_success() { + let ops = mono_ops![ln]; + let result = MonoAD::compute_checked(&ops, 2.0).unwrap(); + assert!(approx_eq(result, 2.0_f64.ln(), 1e-10)); +} + +#[test] +fn test_compute_checked_ln_error_negative() { + let ops = mono_ops![ln]; + let result = MonoAD::compute_checked(&ops, -1.0); + assert!(result.is_err()); + match result.unwrap_err() { + AutodiffError::DomainError { operation, reason } => { + assert_eq!(operation, "Ln"); + assert!(reason.contains("positive")); + } + _ => panic!("Expected DomainError"), + } +} + +#[test] +fn test_compute_checked_ln_error_zero() { + let ops = mono_ops![ln]; + let result = MonoAD::compute_checked(&ops, 0.0); + assert!(result.is_err()); +} + +#[test] +fn test_compute_checked_sqrt_success() { + let ops = mono_ops![sqrt]; + let result = MonoAD::compute_checked(&ops, 4.0).unwrap(); + assert!(approx_eq(result, 2.0, 1e-10)); +} + +#[test] +fn test_compute_checked_sqrt_error_negative() { + let ops = mono_ops![sqrt]; + let result = MonoAD::compute_checked(&ops, -1.0); + assert!(result.is_err()); + match result.unwrap_err() { + AutodiffError::DomainError { operation, reason } => { + assert_eq!(operation, "Sqrt"); + assert!(reason.contains("non-negative")); + } + _ => panic!("Expected DomainError"), + } +} + +#[test] +fn test_compute_checked_abs_success() { + let ops = mono_ops![abs]; + let result = MonoAD::compute_checked(&ops, -3.0).unwrap(); + assert!(approx_eq(result, 3.0, 1e-10)); +} + +#[test] +fn test_compute_checked_chained() { + let ops = mono_ops![sin, exp]; + let result = MonoAD::compute_checked(&ops, 1.0).unwrap(); + assert!(approx_eq(result, 1.0_f64.sin().exp(), 1e-10)); +} + +#[test] +fn test_compute_checked_chained_error() { + // exp(sqrt(x)): forward pass goes through sqrt(-4) → error + let ops = mono_ops![sqrt, exp]; + let result = MonoAD::compute_checked(&ops, -4.0); + assert!(result.is_err()); +} + +// ============================================================================ +// MonoAD::compute_grad_checked tests +// ============================================================================ + +#[test] +fn test_compute_grad_checked_sin_success() { + let ops = mono_ops![sin]; + let (value, grad_fn) = MonoAD::compute_grad_checked(&ops, 2.0).unwrap(); + assert!(approx_eq(value, 2.0_f64.sin(), 1e-10)); + assert!(approx_eq(grad_fn(1.0), 2.0_f64.cos(), 1e-10)); +} + +#[test] +fn test_compute_grad_checked_exp_success() { + let ops = mono_ops![exp]; + let (value, grad_fn) = MonoAD::compute_grad_checked(&ops, 2.0).unwrap(); + assert!(approx_eq(value, 2.0_f64.exp(), 1e-10)); + assert!(approx_eq(grad_fn(1.0), 2.0_f64.exp(), 1e-10)); +} + +#[test] +fn test_compute_grad_checked_ln_error() { + let ops = mono_ops![ln]; + let result = MonoAD::compute_grad_checked(&ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_compute_grad_checked_chained() { + let ops = mono_ops![sin, cos]; + let (value, grad_fn) = MonoAD::compute_grad_checked(&ops, 1.0).unwrap(); + assert!(approx_eq(value, 1.0_f64.sin().cos(), 1e-10)); + assert!(grad_fn(1.0).is_finite()); +} + +// ============================================================================ +// MonoAD::compute_hessian tests (all variants) +// ============================================================================ + +#[test] +fn test_hessian_tan() { + let ops = &[MonoAD::Tan]; + let x = 0.4; + let second_deriv = MonoAD::compute_hessian(ops, x); + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_ln() { + let ops = &[MonoAD::Ln]; + let x = 1.7; + let second_deriv = MonoAD::compute_hessian(ops, x); + let expected = -1.0 / x.powi(2); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_sqrt() { + let ops = &[MonoAD::Sqrt]; + let x = 2.5; + let second_deriv = MonoAD::compute_hessian(ops, x); + let expected = -1.0 / (4.0 * x * x.sqrt()); + assert!(approx_eq(second_deriv, expected, 1e-4)); +} + +#[test] +fn test_hessian_abs() { + // abs(x) has second derivative 0 for x ≠ 0 (finite diff unstable at 0) + let ops = &[MonoAD::Abs]; + for x in [-2.5, 2.5] { + let second_deriv = MonoAD::compute_hessian(ops, x); + assert!(approx_eq(second_deriv, 0.0, 1e-4)); + } +} + +#[test] +fn test_hessian_neg() { + let ops = &[MonoAD::Neg]; + let second_deriv = MonoAD::compute_hessian(ops, 3.0); + assert!(approx_eq(second_deriv, 0.0, 1e-4)); +} + +// ============================================================================ +// MonoAD::compute_hessian_checked tests +// ============================================================================ + +#[test] +fn test_hessian_checked_sin_success() { + let ops = &[MonoAD::Sin]; + let result = MonoAD::compute_hessian_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, -1.0_f64.sin(), 1e-4)); +} + +#[test] +fn test_hessian_checked_ln_error() { + let ops = &[MonoAD::Ln]; + // x = 0.000001 → x - ε ≈ -0.000009 (negative after perturbation step), + let result = MonoAD::compute_hessian_checked(ops, 0.000001); + assert!(result.is_err()); +} + +#[test] +fn test_hessian_checked_chained() { + let ops = &[MonoAD::Sin, MonoAD::Exp]; + let x = 0.5; + let result = MonoAD::compute_hessian_checked(ops, x).unwrap(); + let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + assert!(approx_eq(result, expected, 1e-4)); +} + +// ============================================================================ +// MonoAD::compute_grad with non-1.0 cotangent for all variants +// ============================================================================ + +#[test] +fn test_grad_non_unit_cotangent_cos() { + let ops = &[MonoAD::Cos]; + let x = 0.5; + let (_value, grad_fn) = MonoAD::compute_grad(ops, x); + let grad_2 = grad_fn(2.0); + assert!(approx_eq(grad_2, 2.0 * -x.sin(), 1e-10)); +} + +#[test] +fn test_grad_non_unit_cotangent_tan() { + let ops = &[MonoAD::Tan]; + let x = 0.4; + let (_value, grad_fn) = MonoAD::compute_grad(ops, x); + let grad_3 = grad_fn(3.0); + let sec_sq = 1.0 / x.cos().powi(2); + assert!(approx_eq(grad_3, 3.0 * sec_sq, 1e-10)); +} + +#[test] +fn test_grad_non_unit_cotangent_neg() { + let ops = &[MonoAD::Neg]; + let (_value, grad_fn) = MonoAD::compute_grad(ops, 5.0); + assert!(approx_eq(grad_fn(2.5), -2.5, 1e-10)); + assert!(approx_eq(grad_fn(0.0), 0.0, 1e-10)); +} + +#[test] +fn test_grad_non_unit_cotangent_ln() { + let ops = &[MonoAD::Ln]; + let x = 3.0; + let (_value, grad_fn) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(grad_fn(5.0), 5.0 / x, 1e-10)); +} + +#[test] +fn test_grad_non_unit_cotangent_sqrt() { + let ops = &[MonoAD::Sqrt]; + let x = 9.0; + let (_value, grad_fn) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(grad_fn(4.0), 4.0 / (2.0 * x.sqrt()), 1e-10)); +} + +#[test] +fn test_grad_non_unit_cotangent_abs() { + let ops = &[MonoAD::Abs]; + let (_value, grad_fn) = MonoAD::compute_grad(ops, 5.0); + assert!(approx_eq(grad_fn(3.0), 3.0, 1e-10)); + let (_value2, grad_fn2) = MonoAD::compute_grad(ops, -5.0); + assert!(approx_eq(grad_fn2(3.0), -3.0, 1e-10)); +} diff --git a/src/mono/tests_ho.rs b/src/mono/tests_ho.rs new file mode 100644 index 0000000..9f08908 --- /dev/null +++ b/src/mono/tests_ho.rs @@ -0,0 +1,657 @@ +//! Tests for higher-order autodiff methods (RR, FR, RF). + +use super::mono_ad_fr::MonoAD2FR; +use super::mono_ad_rf::MonoAD2RF; +use super::mono_ad_rr::MonoAD2RR; +use crate::test_utils::approx_eq_eps as approx_eq; + +// Common tolerance for exact autodiff (machine precision) +const EXACT_TOL: f64 = 1e-12; +const COMPARISON_TOL: f64 = 1e-10; + +// ============================================================================ +// MonoAD2RR (Reverse-over-Reverse) Tests +// ============================================================================ + +#[test] +fn test_rr_sin_hessian() { + // f(x) = sin(x), f''(x) = -sin(x) + let ops = &[MonoAD2RR::Sin]; + let x = 0.5; + let hessian = MonoAD2RR::compute_hessian(ops, x); + let expected = -x.sin(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rr_cos_hessian() { + // f(x) = cos(x), f''(x) = -cos(x) + let ops = &[MonoAD2RR::Cos]; + let x = 0.5; + let hessian = MonoAD2RR::compute_hessian(ops, x); + let expected = -x.cos(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rr_exp_hessian() { + // f(x) = exp(x), f''(x) = exp(x) + let ops = &[MonoAD2RR::Exp]; + let x = 2.0; + let hessian = MonoAD2RR::compute_hessian(ops, x); + let expected = x.exp(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rr_tan_hessian() { + let ops = &[MonoAD2RR::Tan]; + let x: f64 = 0.4; + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rr_ln_hessian() { + let ops = &[MonoAD2RR::Ln]; + let x: f64 = 1.7; + let expected = -1.0 / x.powi(2); + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rr_sqrt_hessian() { + let ops = &[MonoAD2RR::Sqrt]; + let x: f64 = 2.5; + let expected = -1.0 / (4.0 * x * x.sqrt()); + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rr_abs_hessian() { + let ops = &[MonoAD2RR::Abs]; + for x in [-2.5, 0.0, 2.5] { + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); + } +} + +#[test] +fn test_rr_neg_hessian() { + // f(x) = -x, f''(x) = 0 + let ops = &[MonoAD2RR::Neg]; + let x = 3.0; + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); +} + +#[test] +fn test_rr_composed_hessian() { + // f(x) = exp(sin(x)), f''(x) = exp(sin(x)) * cos²(x) - exp(sin(x)) * sin(x) + let ops = &[MonoAD2RR::Sin, MonoAD2RR::Exp]; + let x = 0.5; + let hessian = MonoAD2RR::compute_hessian(ops, x); + let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + assert!(approx_eq(hessian, expected, COMPARISON_TOL)); +} + +#[test] +fn test_rr_complex_chain() { + // f(x) = -exp(sin(x)) + let ops = &[MonoAD2RR::Sin, MonoAD2RR::Exp, MonoAD2RR::Neg]; + let x = 1.0; + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(hessian.is_finite()); +} + +// ============================================================================ +// MonoAD2FR compute / compute_checked / compute_grad_checked tests +// ============================================================================ + +#[test] +fn test_fr_compute_sin() { + let ops = &[MonoAD2FR::Sin]; + let result = MonoAD2FR::compute(ops, 1.0); + assert!(approx_eq(result, 1.0_f64.sin(), COMPARISON_TOL)); +} + +#[test] +fn test_fr_compute_neg() { + let ops = &[MonoAD2FR::Neg]; + let result = MonoAD2FR::compute(ops, 3.0); + assert!(approx_eq(result, -3.0, COMPARISON_TOL)); +} + +#[test] +fn test_fr_compute_checked_success() { + let ops = &[MonoAD2FR::Sin, MonoAD2FR::Exp]; + let result = MonoAD2FR::compute_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, 1.0_f64.sin().exp(), COMPARISON_TOL)); +} + +#[test] +fn test_fr_compute_checked_ln_error() { + let ops = &[MonoAD2FR::Ln]; + let result = MonoAD2FR::compute_checked(ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_fr_compute_checked_sqrt_error() { + let ops = &[MonoAD2FR::Sqrt]; + let result = MonoAD2FR::compute_checked(ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_fr_compute_grad_checked_success() { + let ops = &[MonoAD2FR::Sin]; + let (val, grad_fn) = MonoAD2FR::compute_grad_checked(ops, 1.0).unwrap(); + assert!(approx_eq(val, 1.0_f64.sin(), COMPARISON_TOL)); + assert!(approx_eq(grad_fn(1.0), 1.0_f64.cos(), COMPARISON_TOL)); +} + +#[test] +fn test_fr_compute_grad_checked_error() { + let ops = &[MonoAD2FR::Ln]; + let result = MonoAD2FR::compute_grad_checked(ops, -2.0); + assert!(result.is_err()); +} + +#[test] +fn test_fr_hessian_empty() { + let ops: &[MonoAD2FR] = &[]; + let hessian = MonoAD2FR::compute_hessian(ops, 1.0); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); +} + +#[test] +fn test_fr_hessian_checked_success() { + let ops = &[MonoAD2FR::Sin]; + let result = MonoAD2FR::compute_hessian_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, -1.0_f64.sin(), EXACT_TOL)); +} + +// ============================================================================ +// MonoAD2FR (Forward-over-Reverse) Hessian Tests +// ============================================================================ + +#[test] +fn test_fr_sin_hessian() { + let ops = &[MonoAD2FR::Sin]; + let x = 0.5; + let hessian = MonoAD2FR::compute_hessian(ops, x); + let expected = -x.sin(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_fr_cos_hessian() { + let ops = &[MonoAD2FR::Cos]; + let x = 0.5; + let hessian = MonoAD2FR::compute_hessian(ops, x); + let expected = -x.cos(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_fr_exp_hessian() { + let ops = &[MonoAD2FR::Exp]; + let x = 2.0; + let hessian = MonoAD2FR::compute_hessian(ops, x); + let expected = x.exp(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_fr_tan_hessian() { + let ops = &[MonoAD2FR::Tan]; + let x: f64 = 0.4; + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_fr_ln_hessian() { + let ops = &[MonoAD2FR::Ln]; + let x: f64 = 1.7; + let expected = -1.0 / x.powi(2); + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_fr_sqrt_hessian() { + let ops = &[MonoAD2FR::Sqrt]; + let x: f64 = 2.5; + let expected = -1.0 / (4.0 * x * x.sqrt()); + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_fr_abs_hessian() { + let ops = &[MonoAD2FR::Abs]; + for x in [-2.5, 0.0, 2.5] { + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); + } +} + +#[test] +fn test_fr_neg_hessian() { + let ops = &[MonoAD2FR::Neg]; + let x = 3.0; + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); +} + +#[test] +fn test_fr_composed_hessian() { + let ops = &[MonoAD2FR::Sin, MonoAD2FR::Exp]; + let x = 0.5; + let hessian = MonoAD2FR::compute_hessian(ops, x); + let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + assert!(approx_eq(hessian, expected, COMPARISON_TOL)); +} + +// ============================================================================ +// MonoAD2RF compute / compute_checked / compute_grad_checked tests +// ============================================================================ + +#[test] +fn test_rf_compute_sin() { + let ops = &[MonoAD2RF::Sin]; + let result = MonoAD2RF::compute(ops, 1.0); + assert!(approx_eq(result, 1.0_f64.sin(), COMPARISON_TOL)); +} + +#[test] +fn test_rf_compute_neg() { + let ops = &[MonoAD2RF::Neg]; + let result = MonoAD2RF::compute(ops, 3.0); + assert!(approx_eq(result, -3.0, COMPARISON_TOL)); +} + +#[test] +fn test_rf_compute_checked_success() { + let ops = &[MonoAD2RF::Sin, MonoAD2RF::Exp]; + let result = MonoAD2RF::compute_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, 1.0_f64.sin().exp(), COMPARISON_TOL)); +} + +#[test] +fn test_rf_compute_checked_ln_error() { + let ops = &[MonoAD2RF::Ln]; + let result = MonoAD2RF::compute_checked(ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_rf_compute_checked_sqrt_error() { + let ops = &[MonoAD2RF::Sqrt]; + let result = MonoAD2RF::compute_checked(ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_rf_compute_grad_checked_success() { + let ops = &[MonoAD2RF::Sin]; + let (val, grad_fn) = MonoAD2RF::compute_grad_checked(ops, 1.0).unwrap(); + assert!(approx_eq(val, 1.0_f64.sin(), COMPARISON_TOL)); + assert!(approx_eq(grad_fn(1.0), 1.0_f64.cos(), COMPARISON_TOL)); +} + +#[test] +fn test_rf_compute_grad_checked_error() { + let ops = &[MonoAD2RF::Ln]; + let result = MonoAD2RF::compute_grad_checked(ops, -2.0); + assert!(result.is_err()); +} + +#[test] +fn test_rf_hessian_empty() { + let ops: &[MonoAD2RF] = &[]; + let hessian = MonoAD2RF::compute_hessian(ops, 1.0); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); +} + +#[test] +fn test_rf_hessian_checked_success() { + let ops = &[MonoAD2RF::Sin]; + let result = MonoAD2RF::compute_hessian_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, -1.0_f64.sin(), EXACT_TOL)); +} + +// ============================================================================ +// MonoAD2RF (Reverse-over-Forward) Hessian Tests +// ============================================================================ + +#[test] +fn test_rf_sin_hessian() { + let ops = &[MonoAD2RF::Sin]; + let x = 0.5; + let hessian = MonoAD2RF::compute_hessian(ops, x); + let expected = -x.sin(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rf_cos_hessian() { + let ops = &[MonoAD2RF::Cos]; + let x = 0.5; + let hessian = MonoAD2RF::compute_hessian(ops, x); + let expected = -x.cos(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rf_exp_hessian() { + let ops = &[MonoAD2RF::Exp]; + let x = 2.0; + let hessian = MonoAD2RF::compute_hessian(ops, x); + let expected = x.exp(); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rf_tan_hessian() { + let ops = &[MonoAD2RF::Tan]; + let x: f64 = 0.4; + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rf_ln_hessian() { + let ops = &[MonoAD2RF::Ln]; + let x: f64 = 1.7; + let expected = -1.0 / x.powi(2); + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rf_sqrt_hessian() { + let ops = &[MonoAD2RF::Sqrt]; + let x: f64 = 2.5; + let expected = -1.0 / (4.0 * x * x.sqrt()); + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, expected, EXACT_TOL)); +} + +#[test] +fn test_rf_abs_hessian() { + let ops = &[MonoAD2RF::Abs]; + for x in [-2.5, 0.0, 2.5] { + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); + } +} + +#[test] +fn test_rf_neg_hessian() { + let ops = &[MonoAD2RF::Neg]; + let x = 3.0; + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, EXACT_TOL)); +} + +#[test] +fn test_rf_composed_hessian() { + let ops = &[MonoAD2RF::Sin, MonoAD2RF::Exp]; + let x = 0.5; + let hessian = MonoAD2RF::compute_hessian(ops, x); + let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + assert!(approx_eq(hessian, expected, COMPARISON_TOL)); +} + +// ============================================================================ +// MonoAD2RR compute / compute_checked / compute_grad_checked tests +// ============================================================================ + +#[test] +fn test_rr_compute_sin() { + let ops = &[MonoAD2RR::Sin]; + let result = MonoAD2RR::compute(ops, 1.0); + assert!(approx_eq(result, 1.0_f64.sin(), COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_neg() { + let ops = &[MonoAD2RR::Neg]; + let result = MonoAD2RR::compute(ops, 3.0); + assert!(approx_eq(result, -3.0, COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_ln() { + let ops = &[MonoAD2RR::Ln]; + let result = MonoAD2RR::compute(ops, 2.0); + assert!(approx_eq(result, 2.0_f64.ln(), COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_sqrt() { + let ops = &[MonoAD2RR::Sqrt]; + let result = MonoAD2RR::compute(ops, 4.0); + assert!(approx_eq(result, 2.0, COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_abs() { + let ops = &[MonoAD2RR::Abs]; + let result = MonoAD2RR::compute(ops, -5.0); + assert!(approx_eq(result, 5.0, COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_checked_success() { + let ops = &[MonoAD2RR::Sin, MonoAD2RR::Exp]; + let result = MonoAD2RR::compute_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, 1.0_f64.sin().exp(), COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_checked_ln_error() { + let ops = &[MonoAD2RR::Ln]; + let result = MonoAD2RR::compute_checked(ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_rr_compute_checked_sqrt_error() { + let ops = &[MonoAD2RR::Sqrt]; + let result = MonoAD2RR::compute_checked(ops, -1.0); + assert!(result.is_err()); +} + +#[test] +fn test_rr_compute_grad_non_unit_cotangent() { + let ops = &[MonoAD2RR::Sin]; + let (_val, grad_fn) = MonoAD2RR::compute_grad(ops, 1.0); + assert!(approx_eq(grad_fn(3.0), 3.0 * 1.0_f64.cos(), COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_grad_checked_success() { + let ops = &[MonoAD2RR::Sin]; + let (val, grad_fn) = MonoAD2RR::compute_grad_checked(ops, 1.0).unwrap(); + assert!(approx_eq(val, 1.0_f64.sin(), COMPARISON_TOL)); + assert!(approx_eq(grad_fn(1.0), 1.0_f64.cos(), COMPARISON_TOL)); +} + +#[test] +fn test_rr_compute_grad_checked_error() { + let ops = &[MonoAD2RR::Ln]; + let result = MonoAD2RR::compute_grad_checked(ops, -2.0); + assert!(result.is_err()); +} + +#[test] +fn test_rr_hessian_checked_success() { + let ops = &[MonoAD2RR::Sin]; + let result = MonoAD2RR::compute_hessian_checked(ops, 1.0).unwrap(); + assert!(approx_eq(result, -1.0_f64.sin(), EXACT_TOL)); +} + +#[test] +fn test_rr_hessian_checked_error() { + // ln(-1.0) fails domain check + let ops = &[MonoAD2RR::Ln]; + let result = MonoAD2RR::compute_hessian_checked(ops, -1.0); + assert!(result.is_err()); +} + +// ============================================================================ +// Consistency Tests: All three methods should produce identical results +// ============================================================================ + +#[test] +fn test_consistency_all_methods_sin() { + let ops_rr = &[MonoAD2RR::Sin]; + let ops_fr = &[MonoAD2FR::Sin]; + let ops_rf = &[MonoAD2RF::Sin]; + let x = 0.5; + + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + let expected = -x.sin(); + assert!(approx_eq(h_rr, expected, EXACT_TOL)); + assert!(approx_eq(h_fr, expected, EXACT_TOL)); + assert!(approx_eq(h_rf, expected, EXACT_TOL)); + + // All methods should agree + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); +} + +#[test] +fn test_consistency_all_methods_exp() { + let ops_rr = &[MonoAD2RR::Exp]; + let ops_fr = &[MonoAD2FR::Exp]; + let ops_rf = &[MonoAD2RF::Exp]; + let x = 2.0; + + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + let expected = x.exp(); + assert!(approx_eq(h_rr, expected, EXACT_TOL)); + assert!(approx_eq(h_fr, expected, EXACT_TOL)); + assert!(approx_eq(h_rf, expected, EXACT_TOL)); + + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); +} + +#[test] +fn test_consistency_all_methods_tan() { + let ops_rr = &[MonoAD2RR::Tan]; + let ops_fr = &[MonoAD2FR::Tan]; + let ops_rf = &[MonoAD2RF::Tan]; + let x: f64 = 0.4; + + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + assert!(approx_eq(h_rr, expected, EXACT_TOL)); + assert!(approx_eq(h_fr, expected, EXACT_TOL)); + assert!(approx_eq(h_rf, expected, EXACT_TOL)); + + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); +} + +#[test] +fn test_consistency_all_methods_ln() { + let ops_rr = &[MonoAD2RR::Ln]; + let ops_fr = &[MonoAD2FR::Ln]; + let ops_rf = &[MonoAD2RF::Ln]; + let x: f64 = 1.7; + + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + let expected = -1.0 / x.powi(2); + assert!(approx_eq(h_rr, expected, EXACT_TOL)); + assert!(approx_eq(h_fr, expected, EXACT_TOL)); + assert!(approx_eq(h_rf, expected, EXACT_TOL)); + + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); +} + +#[test] +fn test_consistency_all_methods_sqrt() { + let ops_rr = &[MonoAD2RR::Sqrt]; + let ops_fr = &[MonoAD2FR::Sqrt]; + let ops_rf = &[MonoAD2RF::Sqrt]; + let x: f64 = 2.5; + + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + let expected = -1.0 / (4.0 * x * x.sqrt()); + assert!(approx_eq(h_rr, expected, EXACT_TOL)); + assert!(approx_eq(h_fr, expected, EXACT_TOL)); + assert!(approx_eq(h_rf, expected, EXACT_TOL)); + + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); +} + +#[test] +fn test_consistency_all_methods_abs() { + let ops_rr = &[MonoAD2RR::Abs]; + let ops_fr = &[MonoAD2FR::Abs]; + let ops_rf = &[MonoAD2RF::Abs]; + + for x in [-2.5, 0.0, 2.5] { + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + assert!(approx_eq(h_rr, 0.0, EXACT_TOL)); + assert!(approx_eq(h_fr, 0.0, EXACT_TOL)); + assert!(approx_eq(h_rf, 0.0, EXACT_TOL)); + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); + } +} + +#[test] +fn test_consistency_all_methods_composed() { + let ops_rr = &[MonoAD2RR::Sin, MonoAD2RR::Exp]; + let ops_fr = &[MonoAD2FR::Sin, MonoAD2FR::Exp]; + let ops_rf = &[MonoAD2RF::Sin, MonoAD2RF::Exp]; + let x = 0.5; + + let h_rr = MonoAD2RR::compute_hessian(ops_rr, x); + let h_fr = MonoAD2FR::compute_hessian(ops_fr, x); + let h_rf = MonoAD2RF::compute_hessian(ops_rf, x); + + let expected = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + assert!(approx_eq(h_rr, expected, COMPARISON_TOL)); + assert!(approx_eq(h_fr, expected, COMPARISON_TOL)); + assert!(approx_eq(h_rf, expected, COMPARISON_TOL)); + + assert!(approx_eq(h_rr, h_fr, COMPARISON_TOL)); + assert!(approx_eq(h_rr, h_rf, COMPARISON_TOL)); +} diff --git a/src/mono/types.rs b/src/mono/types.rs index f86690a..7010480 100644 --- a/src/mono/types.rs +++ b/src/mono/types.rs @@ -10,3 +10,92 @@ pub type BackwardResultBox = (f64, Box); /// Result type containing value and gradient function (Arc-wrapped for sharing) pub type BackwardResultArc = (f64, Arc); + +/// Dual number for forward-mode automatic differentiation. +/// +/// A dual number represents a value and its derivative (tangent) in a specified direction. +/// Used in Forward-over-Reverse (FR) and Reverse-over-Forward (RF) modes to compute +/// second-order derivatives by differentiating first-order gradients. +/// +/// # Structure +/// +/// - `val`: The primal (function) value +/// - `tan`: The tangent (derivative in the forward direction) +/// +/// # Arithmetic Rules +/// +/// For an operation y = op(x) where x is a dual number: +/// +/// ```text +/// Sin: (sin(x.val), cos(x.val) * x.tan) +/// Cos: (cos(x.val), -sin(x.val) * x.tan) +/// Exp: (exp(x.val), exp(x.val) * x.tan) +/// Neg: (-x.val, -x.tan) +/// ``` +/// +/// # Example +/// +/// ```rust,ignore +/// // Note: Dual is an internal type accessible from crate::mono::types. +/// +/// // Starting from a variable x with derivative 1.0 +/// let x = Dual::variable(2.0); +/// +/// // Forward-mode computes both value and derivative +/// let y = Dual { val: x.val.sin(), tan: x.val.cos() * x.tan }; +/// // y.val ≈ sin(2.0), y.tan ≈ cos(2.0) +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct Dual { + pub val: f64, + pub tan: f64, +} + +impl Dual { + /// Create a new dual number with specified value and tangent. + #[allow(dead_code)] + pub fn new(val: f64, tan: f64) -> Self { + Self { val, tan } + } + + /// Create a constant dual number (zero derivative). + /// + /// Used for constants in the computation graph that don't depend on inputs. + #[allow(dead_code)] + pub fn constant(val: f64) -> Self { + Self { val, tan: 0.0 } + } + + /// Create a variable dual number (unit derivative). + /// + /// Used for input variable when computing derivatives with respect to it. + pub fn variable(val: f64) -> Self { + Self { val, tan: 1.0 } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_dual_variable() { + let d = Dual::variable(2.0); + assert_eq!(d.val, 2.0); + assert_eq!(d.tan, 1.0); + } + + #[test] + fn test_dual_new() { + let d = Dual::new(3.0, 0.5); + assert_eq!(d.val, 3.0); + assert_eq!(d.tan, 0.5); + } + + #[test] + fn test_dual_constant() { + let d = Dual::constant(5.0); + assert_eq!(d.val, 5.0); + assert_eq!(d.tan, 0.0); + } +} diff --git a/src/multi.rs b/src/multi.rs index f1fb3e2..12bed78 100644 --- a/src/multi.rs +++ b/src/multi.rs @@ -9,13 +9,48 @@ mod f2; mod f3; pub mod builder; +pub mod compiled; +pub mod graph; mod multi_ad; +pub mod multi_ad_fr; +pub mod multi_ad_rf; +pub mod multi_ad_rr; +mod parser; + +// Shared internal modules for multivariate derivative rules and Hessian computation. +mod multi_hessian_common; +pub(crate) mod op_rules; + mod multi_fn; #[cfg(test)] mod tests; pub mod types; +pub use compiled::{ + AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind, + BackendRejectionReason, BackendSupportReport, BatchGradients, BatchGradientsBuffer, + BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, + CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle, + DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode, + DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, DeviceTransferPlan, + DeviceTransferPolicy, ExecutionBackend, FlatInstruction, GpuBackendBoundary, Instruction, + MockDeviceBackend, OpCode, ScalarBackend, SimdBackend, UNUSED_NODE_ID, +}; +#[cfg(feature = "backend-wgpu")] +pub use compiled::{ + WgpuBackend, WgpuBuffer, WgpuBufferSet, WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES, +}; +pub use graph::{ + DomainPolicy, ExprGraph, ExprNode, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, + GraphStats, NodeId, Tape, TapeWorkspace, +}; pub use multi_ad::MultiAD; -// Re-export trait for library extension - users can implement custom multi-variable functions -#[allow(unused_imports)] // May not be used internally, but part of public API pub use multi_fn::MultiFn; + +// Re-exported at crate root via lib.rs — suppress unused-import warnings in this module +#[allow(unused_imports)] +pub use multi_ad_fr::MultiAD2FR; +#[allow(unused_imports)] +pub use multi_ad_rf::MultiAD2RF; +#[allow(unused_imports)] +pub use multi_ad_rr::MultiAD2RR; diff --git a/src/multi/builder.rs b/src/multi/builder.rs index da6e622..66d7454 100644 --- a/src/multi/builder.rs +++ b/src/multi/builder.rs @@ -3,6 +3,7 @@ //! This module provides a fluent, type-safe interface for building computational //! graphs without manually managing indices and vectors. +use super::graph::{Graph, GraphNode, NodeId}; use super::multi_ad::MultiAD; /// Builder for constructing multi-variable computation graphs. @@ -14,30 +15,58 @@ use super::multi_ad::MultiAD; /// # Examples /// /// ```rust -/// use petite_ad::{GraphBuilder, MultiAD}; +/// use petite_ad::GraphBuilder; /// -/// // Build: f(x, y) = sin(x) * (x + y) -/// let graph = GraphBuilder::new(2) // 2 inputs -/// .add(0, 1) // x + y at index 2 -/// .sin(0) // sin(x) at index 3 -/// .mul(2, 3) // sin(x) * (x + y) at index 4 +/// // Legacy tuple graph API. +/// let graph = GraphBuilder::new(2) +/// .add(0, 1) +/// .sin(0) +/// .mul(2, 3) /// .build(); /// /// let inputs = &[0.6, 1.4]; -/// let (value, grad_fn) = MultiAD::compute_grad(&graph, inputs).unwrap(); +/// let (value, grad_fn) = petite_ad::MultiAD::compute_grad(&graph, inputs).unwrap(); +/// +/// // Node-handle API for reusable graphs. +/// let mut builder = GraphBuilder::new(2); +/// let x = builder.input_node(0); +/// let y = builder.input_node(1); +/// let sum = builder.add_node(x, y); +/// let sin_x = builder.sin_node(x); +/// let reusable = builder.mul_node(sum, sin_x); +/// assert_eq!(reusable, 4); +/// let graph = builder.build_graph(); +/// let value_from_graph = graph.compute(inputs).unwrap(); +/// assert!((value - value_from_graph).abs() < 1e-10); /// ``` #[derive(Debug, Clone)] pub struct GraphBuilder { /// Number of input variables #[allow(dead_code)] num_inputs: usize, - /// Operations in the computation graph + /// Operations in the legacy tuple computation graph operations: Vec<(MultiAD, Vec)>, + /// Nodes for the reusable graph representation + graph_nodes: Vec, /// Next available index for new operations next_index: usize, } impl GraphBuilder { + #[inline] + fn push_node(&mut self, op: MultiAD, indices: Vec) -> NodeId { + let node_id = self.next_index; + self.operations.push((op, indices.clone())); + if op != MultiAD::Inp { + self.graph_nodes.push(GraphNode::Operation { + op, + inputs: indices, + }); + self.next_index += 1; + } + node_id + } + /// Creates a new graph builder with the specified number of inputs. /// /// # Arguments @@ -55,24 +84,44 @@ impl GraphBuilder { Self { num_inputs, operations: Vec::new(), + graph_nodes: Vec::new(), next_index: num_inputs, } } /// Adds an input placeholder operation. /// - /// This is rarely needed directly as inputs are automatically available, - /// but can be useful for explicit graph construction. + /// Inputs are already available at indices `0..num_inputs`, so this marker + /// does not allocate a new graph value and does not change `next_index()`. + /// It is mainly useful for keeping builder output visually aligned with + /// graphs built using the `multi_ops!` macro. /// /// # Arguments /// /// * `input_index` - Which input variable to reference (0 to num_inputs-1) pub fn input(&mut self, input_index: usize) -> &mut Self { self.operations.push((MultiAD::Inp, vec![input_index])); - self.next_index += 1; self } + /// Returns the node id for an input. + #[must_use] + pub fn input_node(&self, input_index: usize) -> NodeId { + input_index + } + + /// Adds a graph-only constant node and returns its output node id. + /// + /// Constants cannot be represented in the legacy tuple graph returned by + /// [`GraphBuilder::build`]. Use [`GraphBuilder::build_graph`] after calling + /// this method. + pub fn constant_node(&mut self, value: f64) -> NodeId { + let node_id = self.next_index; + self.graph_nodes.push(GraphNode::Constant(value)); + self.next_index += 1; + node_id + } + /// Adds a sine operation. /// /// # Arguments @@ -83,77 +132,120 @@ impl GraphBuilder { /// /// The index where this operation's result will be stored pub fn sin(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Sin, vec![arg_index])); - self.next_index += 1; + let _ = self.sin_node(arg_index); self } + /// Adds a sine operation and returns its output node id. + pub fn sin_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Sin, vec![arg_index]) + } + /// Adds a cosine operation. /// /// # Arguments /// /// * `arg_index` - Index of the input value pub fn cos(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Cos, vec![arg_index])); - self.next_index += 1; + let _ = self.cos_node(arg_index); self } + /// Adds a cosine operation and returns its output node id. + pub fn cos_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Cos, vec![arg_index]) + } + /// Adds a tangent operation. /// /// # Arguments /// /// * `arg_index` - Index of the input value pub fn tan(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Tan, vec![arg_index])); - self.next_index += 1; + let _ = self.tan_node(arg_index); + self + } + + /// Adds a tangent operation and returns its output node id. + pub fn tan_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Tan, vec![arg_index]) + } + + /// Adds a negation operation. + /// + /// # Arguments + /// + /// * `arg_index` - Index of the input value + pub fn neg(&mut self, arg_index: usize) -> &mut Self { + let _ = self.neg_node(arg_index); self } + /// Adds a negation operation and returns its output node id. + pub fn neg_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Neg, vec![arg_index]) + } + /// Adds an exponential operation. /// /// # Arguments /// /// * `arg_index` - Index of the input value pub fn exp(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Exp, vec![arg_index])); - self.next_index += 1; + let _ = self.exp_node(arg_index); self } + /// Adds an exponential operation and returns its output node id. + pub fn exp_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Exp, vec![arg_index]) + } + /// Adds a natural logarithm operation. /// /// # Arguments /// /// * `arg_index` - Index of the input value pub fn ln(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Ln, vec![arg_index])); - self.next_index += 1; + let _ = self.ln_node(arg_index); self } + /// Adds a natural logarithm operation and returns its output node id. + pub fn ln_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Ln, vec![arg_index]) + } + /// Adds a square root operation. /// /// # Arguments /// /// * `arg_index` - Index of the input value pub fn sqrt(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Sqrt, vec![arg_index])); - self.next_index += 1; + let _ = self.sqrt_node(arg_index); self } + /// Adds a square root operation and returns its output node id. + pub fn sqrt_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Sqrt, vec![arg_index]) + } + /// Adds an absolute value operation. /// /// # Arguments /// /// * `arg_index` - Index of the input value pub fn abs(&mut self, arg_index: usize) -> &mut Self { - self.operations.push((MultiAD::Abs, vec![arg_index])); - self.next_index += 1; + let _ = self.abs_node(arg_index); self } + /// Adds an absolute-value operation and returns its output node id. + pub fn abs_node(&mut self, arg_index: usize) -> NodeId { + self.push_node(MultiAD::Abs, vec![arg_index]) + } + /// Adds an addition operation. /// /// # Arguments @@ -161,12 +253,15 @@ impl GraphBuilder { /// * `left_index` - Index of the left operand /// * `right_index` - Index of the right operand pub fn add(&mut self, left_index: usize, right_index: usize) -> &mut Self { - self.operations - .push((MultiAD::Add, vec![left_index, right_index])); - self.next_index += 1; + let _ = self.add_node(left_index, right_index); self } + /// Adds an addition operation and returns its output node id. + pub fn add_node(&mut self, left_index: usize, right_index: usize) -> NodeId { + self.push_node(MultiAD::Add, vec![left_index, right_index]) + } + /// Adds a subtraction operation. /// /// # Arguments @@ -174,12 +269,15 @@ impl GraphBuilder { /// * `left_index` - Index of the left operand /// * `right_index` - Index of the right operand pub fn sub(&mut self, left_index: usize, right_index: usize) -> &mut Self { - self.operations - .push((MultiAD::Sub, vec![left_index, right_index])); - self.next_index += 1; + let _ = self.sub_node(left_index, right_index); self } + /// Adds a subtraction operation and returns its output node id. + pub fn sub_node(&mut self, left_index: usize, right_index: usize) -> NodeId { + self.push_node(MultiAD::Sub, vec![left_index, right_index]) + } + /// Adds a multiplication operation. /// /// # Arguments @@ -187,12 +285,15 @@ impl GraphBuilder { /// * `left_index` - Index of the left operand /// * `right_index` - Index of the right operand pub fn mul(&mut self, left_index: usize, right_index: usize) -> &mut Self { - self.operations - .push((MultiAD::Mul, vec![left_index, right_index])); - self.next_index += 1; + let _ = self.mul_node(left_index, right_index); self } + /// Adds a multiplication operation and returns its output node id. + pub fn mul_node(&mut self, left_index: usize, right_index: usize) -> NodeId { + self.push_node(MultiAD::Mul, vec![left_index, right_index]) + } + /// Adds a division operation. /// /// # Arguments @@ -200,12 +301,15 @@ impl GraphBuilder { /// * `left_index` - Index of the numerator /// * `right_index` - Index of the denominator pub fn div(&mut self, left_index: usize, right_index: usize) -> &mut Self { - self.operations - .push((MultiAD::Div, vec![left_index, right_index])); - self.next_index += 1; + let _ = self.div_node(left_index, right_index); self } + /// Adds a division operation and returns its output node id. + pub fn div_node(&mut self, left_index: usize, right_index: usize) -> NodeId { + self.push_node(MultiAD::Div, vec![left_index, right_index]) + } + /// Adds a power operation. /// /// # Arguments @@ -213,13 +317,16 @@ impl GraphBuilder { /// * `base_index` - Index of the base /// * `exp_index` - Index of the exponent pub fn pow(&mut self, base_index: usize, exp_index: usize) -> &mut Self { - self.operations - .push((MultiAD::Pow, vec![base_index, exp_index])); - self.next_index += 1; + let _ = self.pow_node(base_index, exp_index); self } - /// Builds the final computation graph. + /// Adds a power operation and returns its output node id. + pub fn pow_node(&mut self, base_index: usize, exp_index: usize) -> NodeId { + self.push_node(MultiAD::Pow, vec![base_index, exp_index]) + } + + /// Builds the final legacy computation graph. /// /// Returns a vector of `(operation, indices)` pairs that can be used /// with `MultiAD::compute()` and `MultiAD::compute_grad()`. @@ -239,9 +346,47 @@ impl GraphBuilder { /// let (value, grad_fn) = MultiAD::compute_grad(&graph, inputs).unwrap(); /// ``` pub fn build(&self) -> Vec<(MultiAD, Vec)> { + debug_assert!( + !self + .graph_nodes + .iter() + .any(|n| matches!(n, GraphNode::Constant(_))), + "build() does not include constant nodes; use build_graph() instead" + ); self.operations.clone() } + /// Builds a reusable [`Graph`] with node-handle semantics. + #[must_use] + pub fn build_graph(&self) -> Graph { + let mut graph = Graph::new(self.num_inputs); + for node in &self.graph_nodes { + match node { + GraphNode::Constant(value) => { + graph.constant(*value); + } + GraphNode::Operation { op, inputs } => { + graph.push_operation(*op, inputs.clone()); + } + } + } + graph + } + + /// Builds a reusable graph and selects an explicit output node. + pub fn build_graph_with_output(&self, output: NodeId) -> crate::Result { + let mut graph = self.build_graph(); + graph.set_output(output)?; + Ok(graph) + } + + /// Builds a reusable graph and selects multiple explicit outputs. + pub fn build_graph_with_outputs(&self, outputs: &[NodeId]) -> crate::Result { + let mut graph = self.build_graph(); + graph.set_outputs(outputs)?; + Ok(graph) + } + /// Returns the current number of operations in the graph. pub fn len(&self) -> usize { self.operations.len() @@ -265,15 +410,37 @@ impl GraphBuilder { /// This allows extending the builder with operations not directly /// supported by the fluent API. /// + /// Input marker operations do not allocate a new graph value, matching + /// [`GraphBuilder::input`]. All other operations increment `next_index()`. + /// /// # Arguments /// /// * `op` - The operation to add /// * `indices` - Argument indices for the operation pub fn custom(&mut self, op: MultiAD, indices: Vec) -> &mut Self { - self.operations.push((op, indices)); - self.next_index += 1; + let _ = self.custom_node(op, indices); self } + + /// Adds a custom operation and returns its output node id. + /// + /// For `Inp`, the referenced input index is returned because input markers + /// do not allocate new graph values. + pub fn custom_node(&mut self, op: MultiAD, indices: Vec) -> NodeId { + if op == MultiAD::Inp { + assert_eq!( + indices.len(), + 1, + "Inp requires exactly one index, got {}", + indices.len() + ); + let input_index = indices[0]; + self.operations.push((op, indices)); + input_index + } else { + self.push_node(op, indices) + } + } } #[cfg(test)] @@ -347,6 +514,9 @@ mod tests { let mut builder = GraphBuilder::new(2); assert_eq!(builder.next_index(), 2); // Start at 2 (after inputs) + builder.input(0); + assert_eq!(builder.next_index(), 2); // Input markers do not allocate values + builder.add(0, 1); assert_eq!(builder.next_index(), 3); // After add operation @@ -383,4 +553,296 @@ mod tests { let result = MultiAD::compute(&graph, inputs).unwrap(); assert!(approx_eq(result, 5.0, 1e-10)); } + + #[test] + fn test_builder_custom_input_does_not_advance_index() { + let mut builder = GraphBuilder::new(2); + builder.custom(MultiAD::Inp, vec![0]); + + assert_eq!(builder.next_index(), 2); + } + + #[test] + fn test_builder_node_handle_api() { + let mut builder = GraphBuilder::new(2); + let x = builder.input_node(0); + let y = builder.input_node(1); + let sum = builder.add_node(x, y); + let sin_x = builder.sin_node(x); + let product = builder.mul_node(sum, sin_x); + + assert_eq!(sum, 2); + assert_eq!(sin_x, 3); + assert_eq!(product, 4); + + let graph = builder.build_graph(); + let value = graph.compute(&[0.6, 1.4]).unwrap(); + let expected = (0.6_f64 + 1.4_f64) * 0.6_f64.sin(); + assert!(approx_eq(value, expected, 1e-10)); + } + + #[test] + fn test_builder_build_graph_with_output() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let x_sq = builder.mul_node(x, x); + builder.mul_node(x_sq, x); + + let graph = builder.build_graph_with_output(x_sq).unwrap(); + let value = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(value, 4.0, 1e-10)); + } + + #[test] + fn test_builder_build_graph_with_outputs() { + let mut builder = GraphBuilder::new(2); + let x = builder.input_node(0); + let y = builder.input_node(1); + let sum = builder.add_node(x, y); + let product = builder.mul_node(x, y); + + let graph = builder.build_graph_with_outputs(&[sum, product]).unwrap(); + let values = graph.compute_many(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(values[0], 5.0, 1e-10)); + assert!(approx_eq(values[1], 6.0, 1e-10)); + } + + // ---- Cover all builder operation methods ---- + + #[test] + fn test_builder_cos() { + let graph = GraphBuilder::new(1).cos(0).build(); + let result = MultiAD::compute(&graph, &[0.0]).unwrap(); + assert!(approx_eq(result, 1.0, 1e-10)); + } + + #[test] + fn test_builder_cos_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let cos_x = builder.cos_node(x); + assert_eq!(cos_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(val, 1.0, 1e-10)); + } + + #[test] + fn test_builder_tan() { + let graph = GraphBuilder::new(1).tan(0).build(); + let result = MultiAD::compute(&graph, &[0.0]).unwrap(); + assert!(approx_eq(result, 0.0, 1e-10)); + } + + #[test] + fn test_builder_tan_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let tan_x = builder.tan_node(x); + assert_eq!(tan_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(val, 0.0, 1e-10)); + } + + #[test] + fn test_builder_neg() { + let graph = GraphBuilder::new(1).neg(0).build(); + let result = MultiAD::compute(&graph, &[3.0]).unwrap(); + assert!(approx_eq(result, -3.0, 1e-10)); + } + + #[test] + fn test_builder_neg_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let neg_x = builder.neg_node(x); + assert_eq!(neg_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(val, -3.0, 1e-10)); + } + + #[test] + fn test_builder_exp() { + let graph = GraphBuilder::new(1).exp(0).build(); + let result = MultiAD::compute(&graph, &[1.0]).unwrap(); + assert!(approx_eq(result, std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_builder_exp_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let exp_x = builder.exp_node(x); + assert_eq!(exp_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[1.0]).unwrap(); + assert!(approx_eq(val, std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_builder_ln() { + let graph = GraphBuilder::new(1).ln(0).build(); + let result = MultiAD::compute(&graph, &[std::f64::consts::E]).unwrap(); + assert!(approx_eq(result, 1.0, 1e-10)); + } + + #[test] + fn test_builder_ln_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let ln_x = builder.ln_node(x); + assert_eq!(ln_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[std::f64::consts::E]).unwrap(); + assert!(approx_eq(val, 1.0, 1e-10)); + } + + #[test] + fn test_builder_sqrt() { + let graph = GraphBuilder::new(1).sqrt(0).build(); + let result = MultiAD::compute(&graph, &[4.0]).unwrap(); + assert!(approx_eq(result, 2.0, 1e-10)); + } + + #[test] + fn test_builder_sqrt_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let sqrt_x = builder.sqrt_node(x); + assert_eq!(sqrt_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[4.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-10)); + } + + #[test] + fn test_builder_abs() { + let graph = GraphBuilder::new(1).abs(0).build(); + let result = MultiAD::compute(&graph, &[-5.0]).unwrap(); + assert!(approx_eq(result, 5.0, 1e-10)); + } + + #[test] + fn test_builder_abs_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let abs_x = builder.abs_node(x); + assert_eq!(abs_x, 1); + let graph = builder.build_graph(); + let val = graph.compute(&[-5.0]).unwrap(); + assert!(approx_eq(val, 5.0, 1e-10)); + } + + #[test] + fn test_builder_sub() { + let graph = GraphBuilder::new(2).sub(0, 1).build(); + let result = MultiAD::compute(&graph, &[5.0, 3.0]).unwrap(); + assert!(approx_eq(result, 2.0, 1e-10)); + } + + #[test] + fn test_builder_sub_node() { + let mut builder = GraphBuilder::new(2); + let x = builder.input_node(0); + let y = builder.input_node(1); + let diff = builder.sub_node(x, y); + assert_eq!(diff, 2); + let graph = builder.build_graph(); + let val = graph.compute(&[5.0, 3.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-10)); + } + + #[test] + fn test_builder_div() { + let graph = GraphBuilder::new(2).div(0, 1).build(); + let result = MultiAD::compute(&graph, &[6.0, 3.0]).unwrap(); + assert!(approx_eq(result, 2.0, 1e-10)); + } + + #[test] + fn test_builder_div_node() { + let mut builder = GraphBuilder::new(2); + let x = builder.input_node(0); + let y = builder.input_node(1); + let quotient = builder.div_node(x, y); + assert_eq!(quotient, 2); + let graph = builder.build_graph(); + let val = graph.compute(&[6.0, 3.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-10)); + } + + #[test] + fn test_builder_pow_node() { + let mut builder = GraphBuilder::new(2); + let x = builder.input_node(0); + let y = builder.input_node(1); + let result = builder.pow_node(x, y); + assert_eq!(result, 2); + let graph = builder.build_graph(); + let val = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(val, 8.0, 1e-10)); + } + + #[test] + fn test_builder_constant_node() { + let mut builder = GraphBuilder::new(1); + let x = builder.input_node(0); + let c = builder.constant_node(3.0); + let _sum = builder.add_node(x, c); + + let graph = builder.build_graph(); + let val = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(val, 5.0, 1e-10)); + } + + #[test] + fn test_builder_custom_node_non_inp() { + let mut builder = GraphBuilder::new(2); + let x = builder.input_node(0); + let y = builder.input_node(1); + let node = builder.custom_node(MultiAD::Mul, vec![x, y]); + assert_eq!(node, 2); + let graph = builder.build_graph(); + let val = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(val, 6.0, 1e-10)); + } + + #[test] + #[should_panic(expected = "Inp requires exactly one index")] + fn test_builder_custom_node_inp_wrong_arity() { + let mut builder = GraphBuilder::new(2); + builder.custom_node(MultiAD::Inp, vec![0, 1]); + } + + #[test] + fn test_builder_fluent_all_unary_ops() { + // Test sin via fluent API already tested above; test chaining multiple ops + let graph = GraphBuilder::new(1) + .exp(0) // exp(x) at index 1 + .ln(1) // ln(exp(x)) = x at index 2 + .build(); + let result = MultiAD::compute(&graph, &[2.0]).unwrap(); + assert!(approx_eq(result, 2.0, 1e-10)); + } + + #[test] + fn test_builder_input_marker() { + let graph = GraphBuilder::new(2).input(0).input(1).add(0, 1).build(); + let result = MultiAD::compute(&graph, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(result, 5.0, 1e-10)); + } + + #[test] + fn test_builder_chain_multiple_binary_ops() { + // f(x, y) = (x + y) - (x * y) => 5 - 6 = -1 + let graph = GraphBuilder::new(2) + .add(0, 1) // index 2: x+y + .mul(0, 1) // index 3: x*y + .sub(2, 3) // index 4: (x+y)-(x*y) + .build(); + let result = MultiAD::compute(&graph, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(result, -1.0, 1e-10)); + } } diff --git a/src/multi/compiled.rs b/src/multi/compiled.rs new file mode 100644 index 0000000..6b8640c --- /dev/null +++ b/src/multi/compiled.rs @@ -0,0 +1,4615 @@ +//! Closure-free compiled instruction IR for acceleration-ready graph execution. + +use super::multi_ad::MultiAD; +use super::op_rules; +use crate::{AutodiffError, NodeId, Result}; + +#[cfg(feature = "backend-wgpu")] +use std::sync::mpsc; + +#[cfg(target_arch = "x86_64")] +use std::arch::x86_64::{__m128d, __m256d}; + +#[cfg(feature = "backend-wgpu")] +use pollster::block_on; +#[cfg(feature = "backend-wgpu")] +use wgpu::{self, BufferUsages}; + +/// One closure-free instruction in a compiled scalar graph. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum Instruction { + /// Literal constant node. + Constant(f64), + /// Unary operation with one input node id. + Unary { op: MultiAD, arg: NodeId }, + /// Binary operation with two input node ids. + Binary { + op: MultiAD, + left: NodeId, + right: NodeId, + }, +} + +/// Sentinel node id used by flat backend instructions when an argument is absent. +pub const UNUSED_NODE_ID: NodeId = usize::MAX; + +/// Compact operation code used by flat backend instructions. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum OpCode { + Constant, + Add, + Sub, + Mul, + Div, + Pow, + Sin, + Cos, + Tan, + Tanh, + Relu, + Log1pExp, + LogAddExp, + Neg, + Exp, + Ln, + Sqrt, + Abs, +} + +impl OpCode { + fn from_multi_ad(op: MultiAD) -> Result { + Ok(match op { + MultiAD::Add => OpCode::Add, + MultiAD::Sub => OpCode::Sub, + MultiAD::Mul => OpCode::Mul, + MultiAD::Div => OpCode::Div, + MultiAD::Pow => OpCode::Pow, + MultiAD::Sin => OpCode::Sin, + MultiAD::Cos => OpCode::Cos, + MultiAD::Tan => OpCode::Tan, + MultiAD::Tanh => OpCode::Tanh, + MultiAD::Relu => OpCode::Relu, + MultiAD::Log1pExp => OpCode::Log1pExp, + MultiAD::LogAddExp => OpCode::LogAddExp, + MultiAD::Neg => OpCode::Neg, + MultiAD::Exp => OpCode::Exp, + MultiAD::Ln => OpCode::Ln, + MultiAD::Sqrt => OpCode::Sqrt, + MultiAD::Abs => OpCode::Abs, + MultiAD::Inp => { + return Err(AutodiffError::InvalidGraph { + reason: "input markers are not backend instructions", + }); + } + }) + } + + fn to_multi_ad(self) -> Option { + Some(match self { + OpCode::Constant => return None, + OpCode::Add => MultiAD::Add, + OpCode::Sub => MultiAD::Sub, + OpCode::Mul => MultiAD::Mul, + OpCode::Div => MultiAD::Div, + OpCode::Pow => MultiAD::Pow, + OpCode::Sin => MultiAD::Sin, + OpCode::Cos => MultiAD::Cos, + OpCode::Tan => MultiAD::Tan, + OpCode::Tanh => MultiAD::Tanh, + OpCode::Relu => MultiAD::Relu, + OpCode::Log1pExp => MultiAD::Log1pExp, + OpCode::LogAddExp => MultiAD::LogAddExp, + OpCode::Neg => MultiAD::Neg, + OpCode::Exp => MultiAD::Exp, + OpCode::Ln => MultiAD::Ln, + OpCode::Sqrt => MultiAD::Sqrt, + OpCode::Abs => MultiAD::Abs, + }) + } + + /// Return the number of input slots consumed by this opcode. + #[must_use] + pub fn arity(&self) -> usize { + match self { + OpCode::Constant => 0, + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Relu + | OpCode::Log1pExp + | OpCode::Neg + | OpCode::Exp + | OpCode::Ln + | OpCode::Sqrt + | OpCode::Abs => 1, + OpCode::Add + | OpCode::Sub + | OpCode::Mul + | OpCode::Div + | OpCode::Pow + | OpCode::LogAddExp => 2, + } + } +} + +/// Fully explicit instruction shape for SIMD/GPU backend lowering. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct FlatInstruction { + pub opcode: OpCode, + pub output: NodeId, + pub left: NodeId, + pub right: NodeId, + pub value: f64, +} + +#[cfg(target_arch = "x86_64")] +#[inline] +fn supports_simd_f64x2_runtime() -> bool { + true +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn supports_simd_f64x2_runtime() -> bool { + false +} + +#[cfg(target_arch = "x86_64")] +#[inline] +fn supports_simd_f64x4_runtime() -> bool { + std::is_x86_feature_detected!("avx") +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn supports_simd_f64x4_runtime() -> bool { + false +} + +#[inline] +fn supports_wgpu_runtime() -> bool { + cfg!(feature = "backend-wgpu") +} + +/// Backend feature declaration used before dispatching compiled graphs. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BackendCapabilities { + pub supports_f64: bool, + pub supports_f32: bool, + pub supports_constants: bool, + pub supports_unary: bool, + pub supports_binary: bool, + pub supports_multi_output: bool, + pub supports_reverse_gradient: bool, + pub supports_batch_compute: bool, + pub supports_batch_gradient: bool, + pub supported_opcodes: Vec, +} + +impl BackendCapabilities { + /// Capabilities for the reference scalar f64 backend. + #[must_use] + pub fn scalar_f64() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: true, + supports_batch_compute: true, + supports_batch_gradient: true, + supported_opcodes: vec![ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Relu, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Abs, + ], + } + } + + /// Capabilities for the prototype f64x2 SIMD batch backend. + #[must_use] + pub fn simd_f64x2() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: false, + supports_batch_compute: supports_simd_f64x2_runtime(), + supports_batch_gradient: supports_simd_f64x2_runtime(), + supported_opcodes: vec![ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Relu, + OpCode::Abs, + ], + } + } + + /// Capabilities for the prototype f64x4 SIMD batch backend. + #[must_use] + pub fn simd_f64x4() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: false, + supports_batch_compute: supports_simd_f64x4_runtime(), + supports_batch_gradient: supports_simd_f64x4_runtime(), + supported_opcodes: vec![ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Relu, + OpCode::Abs, + ], + } + } + + /// Capabilities for the feature-gated WGPU batch backend skeleton. + #[must_use] + pub fn wgpu_f64() -> Self { + Self { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: true, + supports_batch_compute: supports_wgpu_runtime(), + supports_batch_gradient: supports_wgpu_runtime(), + supported_opcodes: BackendCapabilities::scalar_f64().supported_opcodes, + } + } + + /// Return whether the backend declares support for an opcode. + #[must_use] + pub fn supports_opcode(&self, opcode: OpCode) -> bool { + self.supported_opcodes.contains(&opcode) + } +} + +/// Reference scalar backend for the execution-backend abstraction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ScalarBackend; + +/// Mock device-style backend that executes on CPU while using device-oriented plans. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct MockDeviceBackend; + +/// Prototype f64x2 SIMD backend for batch compute and batch gradients. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct SimdBackend; + +/// Backend selected by automatic compiled-graph dispatch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackendKind { + Scalar, + MockDeviceCpu, + Wgpu, + SimdF64x4, + SimdF64x2, +} + +impl BackendKind { + /// Return the number of f64 lanes processed by this backend. + #[must_use] + pub fn lane_width(&self) -> usize { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => 1, + BackendKind::SimdF64x4 => 4, + BackendKind::SimdF64x2 => 2, + } + } + + /// Return runtime CPU features required by this backend. + #[must_use] + pub fn required_runtime_features(&self) -> &'static [&'static str] { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu => &[], + BackendKind::Wgpu => &["backend-wgpu"], + BackendKind::SimdF64x4 => &["x86_64-avx"], + BackendKind::SimdF64x2 => &["x86_64-sse2"], + } + } + + /// Return whether the backend is available on this runtime target. + #[must_use] + pub fn runtime_available(&self) -> bool { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu => true, + BackendKind::Wgpu => supports_wgpu_runtime(), + BackendKind::SimdF64x4 => supports_simd_f64x4_runtime(), + BackendKind::SimdF64x2 => supports_simd_f64x2_runtime(), + } + } + + /// Return runtime CPU features that are required but unavailable. + #[must_use] + pub fn unavailable_runtime_features(&self) -> Vec<&'static str> { + if self.runtime_available() { + Vec::new() + } else { + self.required_runtime_features().to_vec() + } + } + + /// Return a stable backend name. + #[must_use] + pub fn name(&self) -> &'static str { + match self { + BackendKind::Scalar => "scalar", + BackendKind::MockDeviceCpu => "mock-device-cpu", + BackendKind::Wgpu => "wgpu", + BackendKind::SimdF64x4 => "simd-f64x4", + BackendKind::SimdF64x2 => "simd-f64x2", + } + } + + /// Return the logical memory location used by this backend's batch plan. + #[must_use] + pub fn memory_location(&self) -> DeviceMemoryLocation { + match self { + BackendKind::MockDeviceCpu | BackendKind::Wgpu => DeviceMemoryLocation::Device, + BackendKind::Scalar | BackendKind::SimdF64x4 | BackendKind::SimdF64x2 => { + DeviceMemoryLocation::Host + } + } + } + + /// Return declared capabilities for this backend. + #[must_use] + pub fn capabilities(&self) -> BackendCapabilities { + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu => BackendCapabilities::scalar_f64(), + BackendKind::Wgpu => BackendCapabilities::wgpu_f64(), + BackendKind::SimdF64x4 => BackendCapabilities::simd_f64x4(), + BackendKind::SimdF64x2 => BackendCapabilities::simd_f64x2(), + } + } + + /// Execute batch value computation with this backend. + pub fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + if matches!(self, BackendKind::Scalar | BackendKind::MockDeviceCpu) { + return ScalarBackend.compute_batch(graph, batch, buffer); + } + if matches!(self, BackendKind::Wgpu) { + #[cfg(feature = "backend-wgpu")] + { + return WgpuBackend::new_default()?.compute_batch(graph, batch, buffer); + } + #[cfg(not(feature = "backend-wgpu"))] + { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu backend requires the backend-wgpu cargo feature", + }); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_compute { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch compute on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => unreachable!(), + BackendKind::SimdF64x4 => compute_batch_simd_f64x4(graph, batch, buffer), + BackendKind::SimdF64x2 => compute_batch_simd_f64x2(graph, batch, buffer), + } + } + + /// Execute batch gradient computation with this backend. + pub fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + if matches!(self, BackendKind::Scalar | BackendKind::MockDeviceCpu) { + return ScalarBackend.gradient_batch(graph, batch, buffer); + } + if matches!(self, BackendKind::Wgpu) { + #[cfg(feature = "backend-wgpu")] + { + return WgpuBackend::new_default()?.gradient_batch(graph, batch, buffer); + } + #[cfg(not(feature = "backend-wgpu"))] + { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu backend requires the backend-wgpu cargo feature", + }); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch gradients on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + match self { + BackendKind::Scalar | BackendKind::MockDeviceCpu | BackendKind::Wgpu => unreachable!(), + BackendKind::SimdF64x4 => gradient_batch_simd_f64x4(graph, batch, buffer), + BackendKind::SimdF64x2 => gradient_batch_simd_f64x2(graph, batch, buffer), + } + } +} + +/// Batch memory layout exposed for backend/device planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BatchLayout { + RowMajor, +} + +/// Logical device buffer role for batch execution planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceBufferKind { + Inputs, + Values, + Outputs, + PrimaryValues, + Gradients, +} + +/// One logical device buffer description. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceBufferLayout { + pub kind: DeviceBufferKind, + pub len: usize, + pub element_size: usize, +} + +/// Logical memory location for a planned backend buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceMemoryLocation { + Host, + Device, +} + +/// Handle-like description for one planned backend buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceBufferHandle { + pub kind: DeviceBufferKind, + pub location: DeviceMemoryLocation, + pub offset: usize, + pub len: usize, +} + +/// Logical transfer direction for backend execution planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceTransferKind { + HostToDevice, + DeviceToHost, +} + +/// One logical host/device transfer needed by a batch plan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceTransferPlan { + pub kind: DeviceTransferKind, + pub buffer: DeviceBufferKind, + pub len: usize, +} + +/// Device buffer owned by a planned backend execution. +#[derive(Debug, Clone, PartialEq)] +pub struct DeviceBuffer { + handle: DeviceBufferHandle, + data: Vec, +} + +impl DeviceBuffer { + fn new(handle: DeviceBufferHandle) -> Self { + Self { + handle, + data: vec![0.0; handle.len], + } + } + + /// Return the immutable planned buffer handle. + #[must_use] + pub fn handle(&self) -> DeviceBufferHandle { + self.handle + } + + /// Return the owned buffer data. + #[must_use] + pub fn data(&self) -> &[f64] { + &self.data + } +} + +/// Owned buffer set allocated from a [`DeviceBatchPlan`]. +#[derive(Debug, Clone, PartialEq)] +pub struct DeviceBufferSet { + plan: DeviceBatchPlan, + buffers: Vec, +} + +impl DeviceBufferSet { + /// Allocate zero-initialized buffers for a batch plan. + #[must_use] + pub fn new(plan: DeviceBatchPlan) -> Self { + let buffers = plan + .buffer_handles + .iter() + .copied() + .map(DeviceBuffer::new) + .collect(); + Self { plan, buffers } + } + + /// Return the plan used to allocate this buffer set. + #[must_use] + pub fn plan(&self) -> &DeviceBatchPlan { + &self.plan + } + + /// Return all allocated buffers in plan order. + #[must_use] + pub fn buffers(&self) -> &[DeviceBuffer] { + &self.buffers + } + + /// Return an immutable buffer by logical kind. + pub fn buffer(&self, kind: DeviceBufferKind) -> Result<&DeviceBuffer> { + self.buffers + .iter() + .find(|buffer| buffer.handle.kind == kind) + .ok_or(AutodiffError::InvalidGraph { + reason: "device buffer kind is not in the plan", + }) + } + + /// Return a mutable buffer by logical kind. + fn buffer_mut(&mut self, kind: DeviceBufferKind) -> Result<&mut DeviceBuffer> { + self.buffers + .iter_mut() + .find(|buffer| buffer.handle.kind == kind) + .ok_or(AutodiffError::InvalidGraph { + reason: "device buffer kind is not in the plan", + }) + } + + /// Upload host data into a planned buffer. + pub fn upload(&mut self, kind: DeviceBufferKind, data: &[f64]) -> Result<()> { + let buffer = self.buffer_mut(kind)?; + if buffer.data.len() != data.len() { + return Err(AutodiffError::InvalidGraph { + reason: "upload length must match planned buffer length", + }); + } + buffer.data.copy_from_slice(data); + Ok(()) + } + + /// Download a planned buffer into an owned vector. + pub fn download(&self, kind: DeviceBufferKind) -> Result> { + Ok(self.buffer(kind)?.data.clone()) + } +} + +/// Device-oriented batch execution plan for a backend. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeviceBatchPlan { + pub backend: BackendKind, + pub layout: BatchLayout, + pub batch_size: usize, + pub input_dim: usize, + pub output_dim: usize, + pub gradient_dim: usize, + pub value_count: usize, + pub buffers: Vec, + pub buffer_handles: Vec, + pub compute_transfer_plan: Vec, + pub gradient_transfer_plan: Vec, +} + +/// Batch execution mode for explicit device transfer planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceExecutionMode { + ComputeBatch, + GradientBatch, +} + +/// Trace returned by explicit device-style execution. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DeviceExecutionTrace { + pub backend: BackendKind, + pub mode: DeviceExecutionMode, + pub transfers: Vec, + /// Whether this execution used a native accelerator kernel instead of host fallback logic. + pub used_native_kernel: bool, +} + +/// Feature-neutral accelerator device family. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AcceleratorDeviceKind { + MockCpu, + Cuda, + Wgpu, +} + +/// Feature-neutral accelerator context descriptor for future GPU backends. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AcceleratorDeviceContext { + pub kind: AcceleratorDeviceKind, + pub device_id: usize, + pub name: String, +} + +impl AcceleratorDeviceContext { + /// Return a context descriptor for the CPU mock-device backend. + #[must_use] + pub fn mock_cpu() -> Self { + Self { + kind: AcceleratorDeviceKind::MockCpu, + device_id: 0, + name: "mock-device-cpu".to_string(), + } + } + + /// Return a context descriptor for a future CUDA backend. + #[must_use] + pub fn cuda(device_id: usize) -> Self { + Self { + kind: AcceleratorDeviceKind::Cuda, + device_id, + name: format!("cuda:{device_id}"), + } + } + + /// Return a context descriptor for a future WGPU backend. + #[must_use] + pub fn wgpu(device_id: usize) -> Self { + Self { + kind: AcceleratorDeviceKind::Wgpu, + device_id, + name: format!("wgpu:{device_id}"), + } + } +} + +/// Host/device transfer policy for future accelerator backends. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DeviceTransferPolicy { + Explicit, + Automatic, +} + +/// Boundary object for future GPU backends. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GpuBackendBoundary { + pub context: AcceleratorDeviceContext, + pub transfer_policy: DeviceTransferPolicy, +} + +impl GpuBackendBoundary { + /// Create a boundary descriptor for a future accelerator backend. + #[must_use] + pub fn new(context: AcceleratorDeviceContext, transfer_policy: DeviceTransferPolicy) -> Self { + Self { + context, + transfer_policy, + } + } + + #[cfg(feature = "backend-wgpu")] + /// Initialize the real WGPU backend skeleton from this boundary. + pub fn initialize_wgpu(&self) -> Result { + WgpuBackend::from_boundary(self.clone()) + } + + /// Return an explicit error because generic GPU execution still requires a concrete backend. + pub fn unsupported_execution_error(&self) -> Result { + Err(AutodiffError::InvalidGraph { + reason: + "real GPU execution requires a concrete backend instance; initialize WgpuBackend", + }) + } +} + +#[cfg(feature = "backend-wgpu")] +/// Initialized WGPU backend skeleton with real device allocation and transfer plumbing. +#[derive(Debug, Clone)] +pub struct WgpuBackend { + boundary: GpuBackendBoundary, + device: wgpu::Device, + queue: wgpu::Queue, + adapter_info: wgpu::AdapterInfo, +} + +#[cfg(feature = "backend-wgpu")] +/// One real WGPU buffer allocated for a logical batch-plan role. +#[derive(Debug, Clone)] +pub struct WgpuBuffer { + handle: DeviceBufferHandle, + buffer: wgpu::Buffer, +} + +#[cfg(feature = "backend-wgpu")] +impl WgpuBuffer { + /// Return the immutable planned buffer handle. + #[must_use] + pub fn handle(&self) -> DeviceBufferHandle { + self.handle + } +} + +#[cfg(feature = "backend-wgpu")] +/// Owned WGPU buffers allocated from a [`DeviceBatchPlan`]. +#[derive(Debug, Clone)] +pub struct WgpuBufferSet { + plan: DeviceBatchPlan, + buffers: Vec, +} + +#[cfg(feature = "backend-wgpu")] +impl WgpuBufferSet { + /// Return the plan used to allocate this buffer set. + #[must_use] + pub fn plan(&self) -> &DeviceBatchPlan { + &self.plan + } + + /// Return all allocated GPU buffers in plan order. + #[must_use] + pub fn buffers(&self) -> &[WgpuBuffer] { + &self.buffers + } + + /// Return an immutable GPU buffer by logical kind. + pub fn buffer(&self, kind: DeviceBufferKind) -> Result<&WgpuBuffer> { + self.buffers + .iter() + .find(|buffer| buffer.handle.kind == kind) + .ok_or(AutodiffError::InvalidGraph { + reason: "wgpu buffer kind is not in the plan", + }) + } + + /// Upload host data into a planned GPU buffer. + pub fn upload( + &self, + backend: &WgpuBackend, + kind: DeviceBufferKind, + data: &[f64], + ) -> Result<()> { + backend.upload_buffer(self, kind, data) + } + + /// Download a planned GPU buffer into an owned vector. + pub fn download(&self, backend: &WgpuBackend, kind: DeviceBufferKind) -> Result> { + backend.download_buffer(self, kind) + } +} + +/// Device-oriented batch backend planning interface. +pub trait DeviceBackend { + fn backend_kind(&self, graph: &CompiledGraph) -> BackendKind; + + fn batch_layout(&self) -> BatchLayout { + BatchLayout::RowMajor + } + + fn batch_plan(&self, graph: &CompiledGraph, batch_size: usize) -> DeviceBatchPlan { + let metadata = graph.metadata(); + let value_count = metadata.value_count.saturating_mul(batch_size); + let input_count = metadata.num_inputs.saturating_mul(batch_size); + let output_count = metadata.num_outputs.saturating_mul(batch_size); + let gradient_count = metadata.num_inputs.saturating_mul(batch_size); + let backend = self.backend_kind(graph); + let buffer_location = backend.memory_location(); + let buffers = vec![ + DeviceBufferLayout { + kind: DeviceBufferKind::Inputs, + len: input_count, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::Values, + len: value_count, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::Outputs, + len: output_count, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::PrimaryValues, + len: batch_size, + element_size: std::mem::size_of::(), + }, + DeviceBufferLayout { + kind: DeviceBufferKind::Gradients, + len: gradient_count, + element_size: std::mem::size_of::(), + }, + ]; + let mut offset = 0; + let buffer_handles = buffers + .iter() + .map(|buffer| { + let handle = DeviceBufferHandle { + kind: buffer.kind, + location: buffer_location, + offset, + len: buffer.len, + }; + offset = offset.saturating_add(buffer.len); + handle + }) + .collect(); + let (compute_transfer_plan, gradient_transfer_plan) = + if buffer_location == DeviceMemoryLocation::Device { + ( + vec![ + DeviceTransferPlan { + kind: DeviceTransferKind::HostToDevice, + buffer: DeviceBufferKind::Inputs, + len: input_count, + }, + DeviceTransferPlan { + kind: DeviceTransferKind::DeviceToHost, + buffer: DeviceBufferKind::Outputs, + len: output_count, + }, + ], + vec![ + DeviceTransferPlan { + kind: DeviceTransferKind::HostToDevice, + buffer: DeviceBufferKind::Inputs, + len: input_count, + }, + DeviceTransferPlan { + kind: DeviceTransferKind::DeviceToHost, + buffer: DeviceBufferKind::PrimaryValues, + len: batch_size, + }, + DeviceTransferPlan { + kind: DeviceTransferKind::DeviceToHost, + buffer: DeviceBufferKind::Gradients, + len: gradient_count, + }, + ], + ) + } else { + (Vec::new(), Vec::new()) + }; + DeviceBatchPlan { + backend, + layout: self.batch_layout(), + batch_size, + input_dim: metadata.num_inputs, + output_dim: metadata.num_outputs, + gradient_dim: metadata.num_inputs, + value_count: metadata.value_count, + buffers, + buffer_handles, + compute_transfer_plan, + gradient_transfer_plan, + } + } +} + +/// Reason a backend cannot execute a graph for a requested batch mode. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BackendRejectionReason { + MissingF64, + UnsupportedOutputs, + UnsupportedOpcodes, + UnsupportedArities, + UnavailableRuntime, + NoBatchCompute, + NoBatchGradient, +} + +/// Static backend compatibility report for a compiled graph. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BackendSupportReport { + pub backend: BackendKind, + pub supports_f64: bool, + pub supports_required_outputs: bool, + pub supports_required_opcodes: bool, + pub supports_required_arities: bool, + pub supports_batch_compute: bool, + pub supports_batch_gradient: bool, + pub lane_width: usize, + pub runtime_available: bool, + pub required_runtime_features: Vec<&'static str>, + pub unavailable_runtime_features: Vec<&'static str>, + pub missing_opcodes: Vec, + pub batch_compute_rejection_reasons: Vec, + pub batch_gradient_rejection_reasons: Vec, +} + +impl BackendSupportReport { + /// Return whether this backend can compute batch values for the graph. + #[must_use] + pub fn can_compute_batch(&self) -> bool { + self.batch_compute_rejection_reasons.is_empty() + } + + /// Return whether this backend can compute batch gradients for the graph. + #[must_use] + pub fn can_gradient_batch(&self) -> bool { + self.batch_gradient_rejection_reasons.is_empty() + } +} + +/// Common execution-backend interface for future scalar, SIMD, and GPU backends. +pub trait ExecutionBackend { + fn name(&self) -> &'static str; + fn capabilities(&self) -> BackendCapabilities; + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result; + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()>; + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)>; + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()>; +} + +/// Flat row-major batch input view. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct BatchInputs<'a> { + /// Flat row-major input data. + pub data: &'a [f64], + /// Number of rows in the batch. + pub batch_size: usize, + /// Number of inputs per row. + pub input_dim: usize, +} + +impl<'a> BatchInputs<'a> { + /// Create a validated row-major input view. + pub fn new(data: &'a [f64], batch_size: usize, input_dim: usize) -> Result { + if data.len() != batch_size.saturating_mul(input_dim) { + return Err(AutodiffError::InvalidGraph { + reason: "batch data length must equal batch_size * input_dim", + }); + } + Ok(Self { + data, + batch_size, + input_dim, + }) + } + + /// Return one input row. + /// + /// Prefer [`BatchInputs::try_row`] when invalid row indices should return an error. + #[must_use] + pub fn row(&self, index: usize) -> &[f64] { + self.try_row(index).expect("batch row index out of bounds") + } + + /// Return one input row, or an error when `index` is out of range. + pub fn try_row(&self, index: usize) -> Result<&[f64]> { + if index >= self.batch_size { + return Err(AutodiffError::IndexOutOfBounds { + index, + max_index: self.batch_size.saturating_sub(1), + }); + } + let start = index * self.input_dim; + Ok(&self.data[start..start + self.input_dim]) + } +} + +/// Flat row-major batch values. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchValues { + pub data: Vec, + pub batch_size: usize, + pub output_dim: usize, +} + +/// Reusable flat row-major batch value buffer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchValuesBuffer { + pub data: Vec, + pub batch_size: usize, + pub output_dim: usize, +} + +impl BatchValuesBuffer { + /// Create an empty reusable output buffer. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + fn reset(&mut self, batch_size: usize, output_dim: usize) { + self.data.clear(); + self.data.reserve(batch_size.saturating_mul(output_dim)); + self.batch_size = batch_size; + self.output_dim = output_dim; + } + + /// Clone the current buffer contents into an owned result value. + #[must_use] + pub fn to_values(&self) -> BatchValues { + BatchValues { + data: self.data.clone(), + batch_size: self.batch_size, + output_dim: self.output_dim, + } + } +} + +/// Flat row-major batch scalar-output gradients. +#[derive(Debug, Clone, PartialEq)] +pub struct BatchGradients { + pub values: Vec, + pub gradients: Vec, + pub batch_size: usize, + pub input_dim: usize, +} + +/// Reusable flat row-major batch scalar-output gradient buffer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct BatchGradientsBuffer { + pub values: Vec, + pub gradients: Vec, + pub batch_size: usize, + pub input_dim: usize, +} + +impl BatchGradientsBuffer { + /// Create an empty reusable gradient buffer. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + fn reset(&mut self, batch_size: usize, input_dim: usize) { + self.values.clear(); + self.gradients.clear(); + self.values.reserve(batch_size); + self.gradients.reserve(batch_size.saturating_mul(input_dim)); + self.batch_size = batch_size; + self.input_dim = input_dim; + } + + /// Clone the current buffer contents into an owned result value. + #[must_use] + pub fn to_gradients(&self) -> BatchGradients { + BatchGradients { + values: self.values.clone(), + gradients: self.gradients.clone(), + batch_size: self.batch_size, + input_dim: self.input_dim, + } + } +} + +/// Static metadata for backend selection and workspace planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CompiledGraphMetadata { + pub num_inputs: usize, + pub num_outputs: usize, + pub num_instructions: usize, + pub num_constants: usize, + pub num_unary: usize, + pub num_binary: usize, + pub value_count: usize, + pub is_scalar_output: bool, +} + +/// Reusable buffers for [`CompiledGraph`] execution. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct CompiledWorkspace { + values: Vec, + cotangents: Vec, + gradients: Vec, + outputs: Vec, +} + +/// Closure-free compiled graph representation. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq)] +pub struct CompiledGraph { + pub(crate) num_inputs: usize, + pub(crate) instructions: Vec, + pub(crate) output_nodes: Vec, + flat_instructions: Vec, +} + +impl CompiledGraph { + pub(crate) fn new( + num_inputs: usize, + instructions: Vec, + output_nodes: Vec, + ) -> Result { + let flat_instructions = Self::lower_flat_instructions(num_inputs, &instructions)?; + Ok(Self { + num_inputs, + instructions, + output_nodes, + flat_instructions, + }) + } + + /// Return number of graph inputs. + #[must_use] + pub fn num_inputs(&self) -> usize { + self.num_inputs + } + + /// Return compiled instructions. + #[must_use] + pub fn instructions(&self) -> &[Instruction] { + &self.instructions + } + + /// Return selected output nodes. + #[must_use] + pub fn output_nodes(&self) -> &[NodeId] { + &self.output_nodes + } + + fn lower_flat_instructions( + num_inputs: usize, + instructions: &[Instruction], + ) -> Result> { + let mut flat = Vec::with_capacity(instructions.len()); + for (offset, instruction) in instructions.iter().enumerate() { + let output = num_inputs + offset; + let item = match *instruction { + Instruction::Constant(value) => FlatInstruction { + opcode: OpCode::Constant, + output, + left: UNUSED_NODE_ID, + right: UNUSED_NODE_ID, + value, + }, + Instruction::Unary { op, arg } => FlatInstruction { + opcode: OpCode::from_multi_ad(op)?, + output, + left: arg, + right: UNUSED_NODE_ID, + value: 0.0, + }, + Instruction::Binary { op, left, right } => FlatInstruction { + opcode: OpCode::from_multi_ad(op)?, + output, + left, + right, + value: 0.0, + }, + }; + flat.push(item); + } + Ok(flat) + } + + /// Return a zero-copy view of flat instructions with explicit output and argument slots. + #[must_use] + pub fn flat_instructions_slice(&self) -> &[FlatInstruction] { + &self.flat_instructions + } + + /// Return owned flat instructions with explicit output and argument slots. + pub fn flat_instructions(&self) -> Result> { + Ok(self.flat_instructions_slice().to_vec()) + } + + /// Validate whether static graph requirements fit backend capabilities. + pub fn validate_backend_capabilities(&self, capabilities: &BackendCapabilities) -> Result<()> { + if !capabilities.supports_f64 { + return Err(AutodiffError::InvalidGraph { + reason: "backend must support f64 execution", + }); + } + if self.output_nodes.len() > 1 && !capabilities.supports_multi_output { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support multi-output graphs", + }); + } + for instruction in self.flat_instructions_slice() { + if !capabilities.supports_opcode(instruction.opcode) { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support a required opcode", + }); + } + match instruction.opcode.arity() { + 0 if !capabilities.supports_constants => { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support constants", + }); + } + 1 if !capabilities.supports_unary => { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support unary operations", + }); + } + 2 if !capabilities.supports_binary => { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support binary operations", + }); + } + _ => {} + } + } + Ok(()) + } + + /// Return static execution metadata for backend planning. + #[must_use] + pub fn metadata(&self) -> CompiledGraphMetadata { + let mut num_constants = 0; + let mut num_unary = 0; + let mut num_binary = 0; + for instruction in &self.instructions { + match instruction { + Instruction::Constant(_) => num_constants += 1, + Instruction::Unary { .. } => num_unary += 1, + Instruction::Binary { .. } => num_binary += 1, + } + } + CompiledGraphMetadata { + num_inputs: self.num_inputs, + num_outputs: self.output_nodes.len(), + num_instructions: self.instructions.len(), + num_constants, + num_unary, + num_binary, + value_count: self.num_inputs + self.instructions.len(), + is_scalar_output: self.output_nodes.len() == 1, + } + } + + /// Create a reusable workspace sized for this compiled graph. + #[must_use] + pub fn workspace(&self) -> CompiledWorkspace { + let value_count = self.num_inputs + self.instructions.len(); + CompiledWorkspace { + values: Vec::with_capacity(value_count), + cotangents: Vec::with_capacity(value_count), + gradients: Vec::with_capacity(self.num_inputs), + outputs: Vec::with_capacity(self.output_nodes.len()), + } + } + + fn check_input_len(&self, inputs: &[f64]) -> Result<()> { + if inputs.len() == self.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "input length must match compiled graph input count", + }) + } + } + + fn check_batch(&self, batch: BatchInputs<'_>) -> Result<()> { + if batch.input_dim == self.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "batch input_dim must match compiled graph input count", + }) + } + } + + fn gather_unary(values: &[f64], arg: NodeId) -> Result<[f64; 1]> { + let value = values + .get(arg) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: arg, + max_index: values.len().saturating_sub(1), + })?; + Ok([value]) + } + + fn gather_binary(values: &[f64], left: NodeId, right: NodeId) -> Result<[f64; 2]> { + let left_value = values + .get(left) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: left, + max_index: values.len().saturating_sub(1), + })?; + let right_value = values + .get(right) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: right, + max_index: values.len().saturating_sub(1), + })?; + Ok([left_value, right_value]) + } + + fn fill_values(&self, inputs: &[f64], workspace: &mut CompiledWorkspace) -> Result<()> { + self.check_input_len(inputs)?; + workspace.values.clear(); + workspace + .values + .reserve(self.num_inputs + self.instructions.len()); + workspace.values.extend_from_slice(inputs); + + for instruction in &self.instructions { + let value = match *instruction { + Instruction::Constant(value) => value, + Instruction::Unary { op, arg } => { + let args = Self::gather_unary(&workspace.values, arg)?; + op_rules::forward_value(op, &args)? + } + Instruction::Binary { op, left, right } => { + let args = Self::gather_binary(&workspace.values, left, right)?; + op_rules::forward_value(op, &args)? + } + }; + workspace.values.push(value); + } + Ok(()) + } + + /// Compute primary output. + pub fn compute(&self, inputs: &[f64]) -> Result { + let mut workspace = self.workspace(); + self.compute_with_workspace(inputs, &mut workspace) + } + + /// Compute all selected outputs. + pub fn compute_many(&self, inputs: &[f64]) -> Result> { + let mut workspace = self.workspace(); + Ok(self + .compute_many_with_workspace(inputs, &mut workspace)? + .to_vec()) + } + + /// Compute primary output with a reusable workspace. + pub fn compute_with_workspace( + &self, + inputs: &[f64], + workspace: &mut CompiledWorkspace, + ) -> Result { + let outputs = self.compute_many_with_workspace(inputs, workspace)?; + Ok(outputs.first().copied().unwrap_or(0.0)) + } + + /// Compute all selected outputs with a reusable workspace. + pub fn compute_many_with_workspace<'a>( + &self, + inputs: &[f64], + workspace: &'a mut CompiledWorkspace, + ) -> Result<&'a [f64]> { + self.fill_values(inputs, workspace)?; + workspace.outputs.clear(); + for &output in &self.output_nodes { + let value = + workspace + .values + .get(output) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index: workspace.values.len().saturating_sub(1), + })?; + workspace.outputs.push(value); + } + Ok(&workspace.outputs) + } + + fn reverse_for_output(&self, output: NodeId, workspace: &mut CompiledWorkspace) -> Result<()> { + workspace.cotangents.clear(); + workspace.cotangents.resize(workspace.values.len(), 0.0); + if output >= workspace.cotangents.len() { + return Err(AutodiffError::IndexOutOfBounds { + index: output, + max_index: workspace.cotangents.len().saturating_sub(1), + }); + } + workspace.cotangents[output] = 1.0; + + for (offset, instruction) in self.instructions.iter().enumerate().rev() { + let node_id = self.num_inputs + offset; + let current = workspace.cotangents[node_id]; + if current == 0.0 { + continue; + } + match *instruction { + Instruction::Constant(_) => {} + Instruction::Unary { op, arg } => { + let args = Self::gather_unary(&workspace.values, arg)?; + let value = workspace.values[node_id]; + if let op_rules::LocalRule::Unary { dy, .. } = + op_rules::local_rule(op, &args, value)? + { + workspace.cotangents[arg] += current * dy; + } + } + Instruction::Binary { op, left, right } => { + let args = Self::gather_binary(&workspace.values, left, right)?; + let value = workspace.values[node_id]; + if let op_rules::LocalRule::Binary { + dy_left, dy_right, .. + } = op_rules::local_rule(op, &args, value)? + { + workspace.cotangents[left] += current * dy_left; + workspace.cotangents[right] += current * dy_right; + } + } + } + } + Ok(()) + } + + /// Compute primary output and gradient. + pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + let mut workspace = self.workspace(); + let (value, gradient) = self.gradient_with_workspace(inputs, &mut workspace)?; + Ok((value, gradient.to_vec())) + } + + /// Compute primary output and gradient with a reusable workspace. + pub fn gradient_with_workspace<'a>( + &self, + inputs: &[f64], + workspace: &'a mut CompiledWorkspace, + ) -> Result<(f64, &'a [f64])> { + self.fill_values(inputs, workspace)?; + workspace.gradients.clear(); + workspace.gradients.resize(self.num_inputs, 0.0); + let Some(&output) = self.output_nodes.first() else { + return Ok((0.0, &workspace.gradients)); + }; + self.reverse_for_output(output, workspace)?; + workspace + .gradients + .copy_from_slice(&workspace.cotangents[..self.num_inputs]); + Ok((workspace.values[output], &workspace.gradients)) + } + + /// Compute Jacobian for all selected outputs. + pub fn jacobian(&self, inputs: &[f64]) -> Result>> { + let mut workspace = self.workspace(); + self.jacobian_with_workspace(inputs, &mut workspace) + } + + /// Compute Jacobian for all selected outputs with a reusable workspace. + pub fn jacobian_with_workspace( + &self, + inputs: &[f64], + workspace: &mut CompiledWorkspace, + ) -> Result>> { + self.fill_values(inputs, workspace)?; + let mut jacobian = Vec::with_capacity(self.output_nodes.len()); + for &output in &self.output_nodes { + self.reverse_for_output(output, workspace)?; + jacobian.push(workspace.cotangents[..self.num_inputs].to_vec()); + } + Ok(jacobian) + } + + /// Return static compatibility details for a backend. + pub fn backend_support_report(&self, backend: BackendKind) -> Result { + let capabilities = backend.capabilities(); + let supports_required_outputs = + self.output_nodes.len() <= 1 || capabilities.supports_multi_output; + let mut missing_opcodes = Vec::new(); + let mut supports_required_arities = true; + for instruction in self.flat_instructions_slice() { + if !capabilities.supports_opcode(instruction.opcode) + && !missing_opcodes.contains(&instruction.opcode) + { + missing_opcodes.push(instruction.opcode); + } + match instruction.opcode.arity() { + 0 if !capabilities.supports_constants => supports_required_arities = false, + 1 if !capabilities.supports_unary => supports_required_arities = false, + 2 if !capabilities.supports_binary => supports_required_arities = false, + _ => {} + } + } + let runtime_available = backend.runtime_available(); + let supports_required_opcodes = missing_opcodes.is_empty(); + let mut common_reasons = Vec::new(); + if !capabilities.supports_f64 { + common_reasons.push(BackendRejectionReason::MissingF64); + } + if !supports_required_outputs { + common_reasons.push(BackendRejectionReason::UnsupportedOutputs); + } + if !supports_required_opcodes { + common_reasons.push(BackendRejectionReason::UnsupportedOpcodes); + } + if !supports_required_arities { + common_reasons.push(BackendRejectionReason::UnsupportedArities); + } + if !runtime_available { + common_reasons.push(BackendRejectionReason::UnavailableRuntime); + } + + let mut batch_compute_rejection_reasons = common_reasons.clone(); + if !capabilities.supports_batch_compute { + batch_compute_rejection_reasons.push(BackendRejectionReason::NoBatchCompute); + } + let mut batch_gradient_rejection_reasons = common_reasons; + if !capabilities.supports_batch_gradient { + batch_gradient_rejection_reasons.push(BackendRejectionReason::NoBatchGradient); + } + + Ok(BackendSupportReport { + backend, + supports_f64: capabilities.supports_f64, + supports_required_outputs, + supports_required_opcodes, + supports_required_arities, + supports_batch_compute: capabilities.supports_batch_compute, + supports_batch_gradient: capabilities.supports_batch_gradient, + lane_width: backend.lane_width(), + runtime_available, + required_runtime_features: backend.required_runtime_features().to_vec(), + unavailable_runtime_features: backend.unavailable_runtime_features(), + missing_opcodes, + batch_compute_rejection_reasons, + batch_gradient_rejection_reasons, + }) + } + + /// Return static compatibility details for all built-in backends. + pub fn backend_support_reports(&self) -> Result> { + Ok(vec![ + self.backend_support_report(BackendKind::Scalar)?, + self.backend_support_report(BackendKind::MockDeviceCpu)?, + self.backend_support_report(BackendKind::Wgpu)?, + self.backend_support_report(BackendKind::SimdF64x4)?, + self.backend_support_report(BackendKind::SimdF64x2)?, + ]) + } + + /// Return a device-oriented batch buffer plan for a backend. + #[must_use] + pub fn device_batch_plan(&self, backend: BackendKind, batch_size: usize) -> DeviceBatchPlan { + backend.batch_plan(self, batch_size) + } + + /// Allocate mock-device buffers for this compiled graph and batch size. + #[must_use] + pub fn allocate_mock_device_buffers(&self, batch_size: usize) -> DeviceBufferSet { + MockDeviceBackend.allocate_batch_buffers(self, batch_size) + } + + /// Execute batch value computation through mock-device buffers. + pub fn compute_batch_mock_device_into( + &self, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + MockDeviceBackend.compute_batch_with_buffers(self, batch, buffers, output) + } + + /// Execute batch gradient computation through mock-device buffers. + pub fn gradient_batch_mock_device_into( + &self, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + MockDeviceBackend.gradient_batch_with_buffers(self, batch, buffers, output) + } + + #[cfg(feature = "backend-wgpu")] + /// Allocate real WGPU buffers for this compiled graph and batch size. + pub fn allocate_wgpu_buffers( + &self, + backend: &WgpuBackend, + batch_size: usize, + ) -> Result { + backend.allocate_batch_buffers(self, batch_size) + } + + #[cfg(feature = "backend-wgpu")] + /// Execute batch value computation through real WGPU buffers. + pub fn compute_batch_wgpu_into( + &self, + backend: &WgpuBackend, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + backend.compute_batch_with_buffers(self, batch, buffers, output) + } + + #[cfg(feature = "backend-wgpu")] + /// Execute batch gradient computation through real WGPU buffers. + pub fn gradient_batch_wgpu_into( + &self, + backend: &WgpuBackend, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + backend.gradient_batch_with_buffers(self, batch, buffers, output) + } + + #[cfg(feature = "backend-wgpu")] + /// Return whether this graph is statically eligible for the exact-safe native WGPU compute path. + #[must_use] + pub fn supports_native_wgpu_batch_compute(&self, backend: &WgpuBackend) -> bool { + backend.supports_native_batch_compute(self) + } + + #[cfg(feature = "backend-wgpu")] + /// Return whether this graph and concrete batch can use the exact-safe native WGPU compute path. + #[must_use] + pub fn supports_native_wgpu_batch_compute_for_batch( + &self, + backend: &WgpuBackend, + batch: BatchInputs<'_>, + ) -> bool { + backend.supports_native_batch_compute_for_batch(self, batch) + } + + /// Return static compatibility details for the preferred SIMD backend. + pub fn simd_support_report(&self) -> Result { + let f64x4 = self.backend_support_report(BackendKind::SimdF64x4)?; + if f64x4.supports_batch_compute || f64x4.supports_batch_gradient { + return Ok(f64x4); + } + self.backend_support_report(BackendKind::SimdF64x2) + } + + /// Return the preferred backend for batch value computation. + #[must_use] + pub fn recommended_batch_compute_backend(&self) -> BackendKind { + match self.backend_support_report(BackendKind::SimdF64x4) { + Ok(report) if report.can_compute_batch() => BackendKind::SimdF64x4, + _ => match self.backend_support_report(BackendKind::SimdF64x2) { + Ok(report) if report.can_compute_batch() => BackendKind::SimdF64x2, + _ => BackendKind::Scalar, + }, + } + } + + /// Return the preferred backend for batch gradient computation. + #[must_use] + pub fn recommended_batch_gradient_backend(&self) -> BackendKind { + match self.backend_support_report(BackendKind::SimdF64x4) { + Ok(report) if report.can_gradient_batch() => BackendKind::SimdF64x4, + _ => match self.backend_support_report(BackendKind::SimdF64x2) { + Ok(report) if report.can_gradient_batch() => BackendKind::SimdF64x2, + _ => BackendKind::Scalar, + }, + } + } + + /// Compute all selected outputs for a batch of rows. + pub fn compute_batch(&self, batch: BatchInputs<'_>) -> Result { + let mut buffer = BatchValuesBuffer::new(); + self.compute_batch_into(batch, &mut buffer)?; + Ok(BatchValues { + data: buffer.data, + batch_size: buffer.batch_size, + output_dim: buffer.output_dim, + }) + } + + /// Compute all selected outputs for a batch with automatic backend dispatch. + pub fn compute_batch_auto(&self, batch: BatchInputs<'_>) -> Result<(BackendKind, BatchValues)> { + let mut buffer = BatchValuesBuffer::new(); + let backend = self.compute_batch_auto_into(batch, &mut buffer)?; + Ok(( + backend, + BatchValues { + data: buffer.data, + batch_size: buffer.batch_size, + output_dim: buffer.output_dim, + }, + )) + } + + /// Compute all selected outputs into a reusable output buffer. + pub fn compute_batch_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + self.check_batch(batch)?; + let mut workspace = self.workspace(); + let output_dim = self.output_nodes.len(); + buffer.reset(batch.batch_size, output_dim); + for row_index in 0..batch.batch_size { + let outputs = + self.compute_many_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.data.extend_from_slice(outputs); + } + Ok(()) + } + + /// Compute all selected outputs into a reusable buffer with automatic backend dispatch. + pub fn compute_batch_auto_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result { + let backend = self.recommended_batch_compute_backend(); + backend.compute_batch(self, batch, buffer)?; + Ok(backend) + } + + /// Compute primary-output values and gradients for a batch of rows. + pub fn gradient_batch(&self, batch: BatchInputs<'_>) -> Result { + let mut buffer = BatchGradientsBuffer::new(); + self.gradient_batch_into(batch, &mut buffer)?; + Ok(BatchGradients { + values: buffer.values, + gradients: buffer.gradients, + batch_size: buffer.batch_size, + input_dim: buffer.input_dim, + }) + } + + /// Compute primary-output values and gradients for a batch with automatic backend dispatch. + pub fn gradient_batch_auto( + &self, + batch: BatchInputs<'_>, + ) -> Result<(BackendKind, BatchGradients)> { + let mut buffer = BatchGradientsBuffer::new(); + let backend = self.gradient_batch_auto_into(batch, &mut buffer)?; + Ok(( + backend, + BatchGradients { + values: buffer.values, + gradients: buffer.gradients, + batch_size: buffer.batch_size, + input_dim: buffer.input_dim, + }, + )) + } + + /// Compute primary-output values and gradients into a reusable gradient buffer. + pub fn gradient_batch_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + self.check_batch(batch)?; + let mut workspace = self.workspace(); + buffer.reset(batch.batch_size, self.num_inputs); + for row_index in 0..batch.batch_size { + let (value, gradient) = + self.gradient_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.values.push(value); + buffer.gradients.extend_from_slice(gradient); + } + Ok(()) + } + + /// Compute primary-output values and gradients into a reusable buffer with automatic backend dispatch. + pub fn gradient_batch_auto_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result { + let backend = self.recommended_batch_gradient_backend(); + backend.gradient_batch(self, batch, buffer)?; + Ok(backend) + } +} + +fn simd_unsupported_opcode_error() -> AutodiffError { + AutodiffError::InvalidGraph { + reason: "simd backend does not support a required opcode", + } +} + +fn simd_scalar_value(opcode: OpCode, args: &[f64]) -> Result { + let op = opcode + .to_multi_ad() + .ok_or_else(simd_unsupported_opcode_error)?; + op_rules::forward_value(op, args) +} + +fn simd_scalar_first_derivatives(opcode: OpCode, args: &[f64], value: f64) -> Result> { + let op = opcode + .to_multi_ad() + .ok_or_else(simd_unsupported_opcode_error)?; + op_rules::first_derivatives(op, args, value) +} + +fn append_scalar_compute_tail( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + start_row: usize, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + let mut workspace = graph.workspace(); + for row_index in start_row..batch.batch_size { + let outputs = + graph.compute_many_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.data.extend_from_slice(outputs); + } + Ok(()) +} + +fn append_scalar_gradient_tail( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + start_row: usize, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + let mut workspace = graph.workspace(); + for row_index in start_row..batch.batch_size { + let (value, gradient) = + graph.gradient_with_workspace(batch.try_row(row_index)?, &mut workspace)?; + buffer.values.push(value); + buffer.gradients.extend_from_slice(gradient); + } + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn checked_m128_lane(values: &[__m128d], index: NodeId) -> Result<__m128d> { + values + .get(index) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: values.len().saturating_sub(1), + }) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn add_m128_cotangent( + cotangents: &mut [__m128d], + index: NodeId, + contribution: __m128d, +) -> Result<()> { + use std::arch::x86_64::_mm_add_pd; + + let max_index = cotangents.len().saturating_sub(1); + let target = cotangents + .get_mut(index) + .ok_or(AutodiffError::IndexOutOfBounds { index, max_index })?; + *target = _mm_add_pd(*target, contribution); + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn active_m128_contribution(current: __m128d, contribution: __m128d) -> __m128d { + use std::arch::x86_64::{_mm_and_pd, _mm_cmpneq_pd, _mm_setzero_pd}; + + let active = _mm_cmpneq_pd(current, _mm_setzero_pd()); + _mm_and_pd(contribution, active) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_unary(input: __m128d, opcode: OpCode) -> Result<__m128d> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut stored = [0.0_f64; 2]; + _mm_storeu_pd(stored.as_mut_ptr(), input); + let first = simd_scalar_value(opcode, &[stored[0]])?; + let second = simd_scalar_value(opcode, &[stored[1]])?; + Ok(_mm_set_pd(second, first)) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_unary_derivative( + input: __m128d, + output: __m128d, + opcode: OpCode, +) -> Result<__m128d> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut input_values = [0.0_f64; 2]; + let mut output_values = [0.0_f64; 2]; + _mm_storeu_pd(input_values.as_mut_ptr(), input); + _mm_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives(opcode, &[input_values[0]], output_values[0])?[0]; + let second = simd_scalar_first_derivatives(opcode, &[input_values[1]], output_values[1])?[0]; + Ok(_mm_set_pd(second, first)) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_binary( + left: __m128d, + right: __m128d, + opcode: OpCode, +) -> Result<__m128d> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut left_values = [0.0_f64; 2]; + let mut right_values = [0.0_f64; 2]; + _mm_storeu_pd(left_values.as_mut_ptr(), left); + _mm_storeu_pd(right_values.as_mut_ptr(), right); + let first = simd_scalar_value(opcode, &[left_values[0], right_values[0]])?; + let second = simd_scalar_value(opcode, &[left_values[1], right_values[1]])?; + Ok(_mm_set_pd(second, first)) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_scalar_binary_derivatives( + left: __m128d, + right: __m128d, + output: __m128d, + opcode: OpCode, +) -> Result<(__m128d, __m128d)> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + let mut left_values = [0.0_f64; 2]; + let mut right_values = [0.0_f64; 2]; + let mut output_values = [0.0_f64; 2]; + _mm_storeu_pd(left_values.as_mut_ptr(), left); + _mm_storeu_pd(right_values.as_mut_ptr(), right); + _mm_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives( + opcode, + &[left_values[0], right_values[0]], + output_values[0], + )?; + let second = simd_scalar_first_derivatives( + opcode, + &[left_values[1], right_values[1]], + output_values[1], + )?; + Ok(( + _mm_set_pd(second[0], first[0]), + _mm_set_pd(second[1], first[1]), + )) +} + +#[cfg(target_arch = "x86_64")] +unsafe fn simd_f64x2_forward_values( + flat: &[FlatInstruction], + values: &mut Vec<__m128d>, +) -> Result<()> { + use std::arch::x86_64::{ + _mm_add_pd, _mm_and_pd, _mm_andnot_pd, _mm_cmpgt_pd, _mm_div_pd, _mm_mul_pd, _mm_set1_pd, + _mm_setzero_pd, _mm_sqrt_pd, _mm_sub_pd, + }; + + for instruction in flat { + let value = match instruction.opcode { + OpCode::Constant => _mm_set1_pd(instruction.value), + OpCode::Add => _mm_add_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Sub => _mm_sub_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Mul => _mm_mul_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Div => _mm_div_pd( + checked_m128_lane(values, instruction.left)?, + checked_m128_lane(values, instruction.right)?, + ), + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m128_lane(values, instruction.left)?; + let right = checked_m128_lane(values, instruction.right)?; + simd_f64x2_scalar_binary(left, right, instruction.opcode)? + } + OpCode::Neg => _mm_sub_pd( + _mm_setzero_pd(), + checked_m128_lane(values, instruction.left)?, + ), + OpCode::Sqrt => _mm_sqrt_pd(checked_m128_lane(values, instruction.left)?), + OpCode::Relu => { + let input = checked_m128_lane(values, instruction.left)?; + let mask = _mm_cmpgt_pd(input, _mm_setzero_pd()); + _mm_and_pd(input, mask) + } + OpCode::Abs => { + let input = checked_m128_lane(values, instruction.left)?; + _mm_andnot_pd(_mm_set1_pd(-0.0), input) + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input = checked_m128_lane(values, instruction.left)?; + simd_f64x2_scalar_unary(input, instruction.opcode)? + } + }; + values.push(value); + } + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn checked_m256_lane(values: &[__m256d], index: NodeId) -> Result<__m256d> { + values + .get(index) + .copied() + .ok_or(AutodiffError::IndexOutOfBounds { + index, + max_index: values.len().saturating_sub(1), + }) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn add_m256_cotangent( + cotangents: &mut [__m256d], + index: NodeId, + contribution: __m256d, +) -> Result<()> { + use std::arch::x86_64::_mm256_add_pd; + + let max_index = cotangents.len().saturating_sub(1); + let target = cotangents + .get_mut(index) + .ok_or(AutodiffError::IndexOutOfBounds { index, max_index })?; + *target = _mm256_add_pd(*target, contribution); + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn active_m256_contribution(current: __m256d, contribution: __m256d) -> __m256d { + use std::arch::x86_64::{_mm256_and_pd, _mm256_cmp_pd, _mm256_setzero_pd, _CMP_NEQ_UQ}; + + let active = _mm256_cmp_pd(current, _mm256_setzero_pd(), _CMP_NEQ_UQ); + _mm256_and_pd(contribution, active) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_unary(input: __m256d, opcode: OpCode) -> Result<__m256d> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut stored = [0.0_f64; 4]; + _mm256_storeu_pd(stored.as_mut_ptr(), input); + let first = simd_scalar_value(opcode, &[stored[0]])?; + let second = simd_scalar_value(opcode, &[stored[1]])?; + let third = simd_scalar_value(opcode, &[stored[2]])?; + let fourth = simd_scalar_value(opcode, &[stored[3]])?; + Ok(_mm256_set_pd(fourth, third, second, first)) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_unary_derivative( + input: __m256d, + output: __m256d, + opcode: OpCode, +) -> Result<__m256d> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut input_values = [0.0_f64; 4]; + let mut output_values = [0.0_f64; 4]; + _mm256_storeu_pd(input_values.as_mut_ptr(), input); + _mm256_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives(opcode, &[input_values[0]], output_values[0])?[0]; + let second = simd_scalar_first_derivatives(opcode, &[input_values[1]], output_values[1])?[0]; + let third = simd_scalar_first_derivatives(opcode, &[input_values[2]], output_values[2])?[0]; + let fourth = simd_scalar_first_derivatives(opcode, &[input_values[3]], output_values[3])?[0]; + Ok(_mm256_set_pd(fourth, third, second, first)) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_binary( + left: __m256d, + right: __m256d, + opcode: OpCode, +) -> Result<__m256d> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut left_values = [0.0_f64; 4]; + let mut right_values = [0.0_f64; 4]; + _mm256_storeu_pd(left_values.as_mut_ptr(), left); + _mm256_storeu_pd(right_values.as_mut_ptr(), right); + let first = simd_scalar_value(opcode, &[left_values[0], right_values[0]])?; + let second = simd_scalar_value(opcode, &[left_values[1], right_values[1]])?; + let third = simd_scalar_value(opcode, &[left_values[2], right_values[2]])?; + let fourth = simd_scalar_value(opcode, &[left_values[3], right_values[3]])?; + Ok(_mm256_set_pd(fourth, third, second, first)) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_scalar_binary_derivatives( + left: __m256d, + right: __m256d, + output: __m256d, + opcode: OpCode, +) -> Result<(__m256d, __m256d)> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + let mut left_values = [0.0_f64; 4]; + let mut right_values = [0.0_f64; 4]; + let mut output_values = [0.0_f64; 4]; + _mm256_storeu_pd(left_values.as_mut_ptr(), left); + _mm256_storeu_pd(right_values.as_mut_ptr(), right); + _mm256_storeu_pd(output_values.as_mut_ptr(), output); + let first = simd_scalar_first_derivatives( + opcode, + &[left_values[0], right_values[0]], + output_values[0], + )?; + let second = simd_scalar_first_derivatives( + opcode, + &[left_values[1], right_values[1]], + output_values[1], + )?; + let third = simd_scalar_first_derivatives( + opcode, + &[left_values[2], right_values[2]], + output_values[2], + )?; + let fourth = simd_scalar_first_derivatives( + opcode, + &[left_values[3], right_values[3]], + output_values[3], + )?; + Ok(( + _mm256_set_pd(fourth[0], third[0], second[0], first[0]), + _mm256_set_pd(fourth[1], third[1], second[1], first[1]), + )) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn simd_f64x4_forward_values( + flat: &[FlatInstruction], + values: &mut Vec<__m256d>, +) -> Result<()> { + use std::arch::x86_64::{ + _mm256_add_pd, _mm256_and_pd, _mm256_andnot_pd, _mm256_cmp_pd, _mm256_div_pd, + _mm256_mul_pd, _mm256_set1_pd, _mm256_setzero_pd, _mm256_sqrt_pd, _mm256_sub_pd, + _CMP_GT_OQ, + }; + + for instruction in flat { + let value = match instruction.opcode { + OpCode::Constant => _mm256_set1_pd(instruction.value), + OpCode::Add => _mm256_add_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Sub => _mm256_sub_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Mul => _mm256_mul_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Div => _mm256_div_pd( + checked_m256_lane(values, instruction.left)?, + checked_m256_lane(values, instruction.right)?, + ), + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m256_lane(values, instruction.left)?; + let right = checked_m256_lane(values, instruction.right)?; + simd_f64x4_scalar_binary(left, right, instruction.opcode)? + } + OpCode::Neg => _mm256_sub_pd( + _mm256_setzero_pd(), + checked_m256_lane(values, instruction.left)?, + ), + OpCode::Sqrt => _mm256_sqrt_pd(checked_m256_lane(values, instruction.left)?), + OpCode::Relu => { + let input = checked_m256_lane(values, instruction.left)?; + let mask = _mm256_cmp_pd(input, _mm256_setzero_pd(), _CMP_GT_OQ); + _mm256_and_pd(input, mask) + } + OpCode::Abs => { + let input = checked_m256_lane(values, instruction.left)?; + _mm256_andnot_pd(_mm256_set1_pd(-0.0), input) + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input = checked_m256_lane(values, instruction.left)?; + simd_f64x4_scalar_unary(input, instruction.opcode)? + } + }; + values.push(value); + } + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn compute_batch_simd_f64x2( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + use std::arch::x86_64::{_mm_set_pd, _mm_storeu_pd}; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + let output_dim = graph.output_nodes.len(); + let value_count = graph.num_inputs + graph.instructions.len(); + buffer.reset(batch.batch_size, output_dim); + + let mut values: Vec<__m128d> = Vec::with_capacity(value_count); + let mut pair_outputs: Vec<[f64; 2]> = Vec::with_capacity(output_dim); + let mut row_index = 0; + while row_index + 1 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + values.clear(); + for input_index in 0..graph.num_inputs { + unsafe { + values.push(_mm_set_pd(second[input_index], first[input_index])); + } + } + + unsafe { + simd_f64x2_forward_values(flat, &mut values)?; + } + + pair_outputs.clear(); + for &output in &graph.output_nodes { + let lane = checked_m128_lane(&values, output)?; + let mut stored = [0.0_f64; 2]; + unsafe { + _mm_storeu_pd(stored.as_mut_ptr(), lane); + } + pair_outputs.push(stored); + } + for output in &pair_outputs { + buffer.data.push(output[0]); + } + for output in &pair_outputs { + buffer.data.push(output[1]); + } + + row_index += 2; + } + + append_scalar_compute_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn compute_batch_simd_f64x4_impl( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + use std::arch::x86_64::{_mm256_set_pd, _mm256_storeu_pd}; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + let output_dim = graph.output_nodes.len(); + let value_count = graph.num_inputs + graph.instructions.len(); + buffer.reset(batch.batch_size, output_dim); + + let mut values: Vec<__m256d> = Vec::with_capacity(value_count); + let mut quad_outputs: Vec<[f64; 4]> = Vec::with_capacity(output_dim); + let mut row_index = 0; + while row_index + 3 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + let third = batch.try_row(row_index + 2)?; + let fourth = batch.try_row(row_index + 3)?; + values.clear(); + for input_index in 0..graph.num_inputs { + values.push(_mm256_set_pd( + fourth[input_index], + third[input_index], + second[input_index], + first[input_index], + )); + } + + simd_f64x4_forward_values(flat, &mut values)?; + + quad_outputs.clear(); + for &output in &graph.output_nodes { + let lane = checked_m256_lane(&values, output)?; + let mut stored = [0.0_f64; 4]; + _mm256_storeu_pd(stored.as_mut_ptr(), lane); + quad_outputs.push(stored); + } + for lane_index in 0..4 { + for output in &quad_outputs { + buffer.data.push(output[lane_index]); + } + } + + row_index += 4; + } + + append_scalar_compute_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn compute_batch_simd_f64x4( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, +) -> Result<()> { + if !supports_simd_f64x4_runtime() { + return Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }); + } + unsafe { compute_batch_simd_f64x4_impl(graph, batch, buffer) } +} + +#[cfg(not(target_arch = "x86_64"))] +fn compute_batch_simd_f64x4( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchValuesBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }) +} + +#[cfg(not(target_arch = "x86_64"))] +fn compute_batch_simd_f64x2( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchValuesBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd backend requires x86_64 SSE2 support", + }) +} + +#[cfg(target_arch = "x86_64")] +fn gradient_batch_simd_f64x2( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + use std::arch::x86_64::{ + _mm_add_pd, _mm_and_pd, _mm_cmpgt_pd, _mm_cmplt_pd, _mm_div_pd, _mm_mul_pd, _mm_set1_pd, + _mm_set_pd, _mm_setzero_pd, _mm_storeu_pd, _mm_sub_pd, + }; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + buffer.reset(batch.batch_size, graph.num_inputs); + + let value_count = graph.num_inputs + graph.instructions.len(); + let mut values: Vec<__m128d> = Vec::with_capacity(value_count); + let mut cotangents: Vec<__m128d> = Vec::with_capacity(value_count); + let mut gradient_pairs = Vec::with_capacity(graph.num_inputs); + let mut row_index = 0; + while row_index + 1 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + values.clear(); + for input_index in 0..graph.num_inputs { + unsafe { + values.push(_mm_set_pd(second[input_index], first[input_index])); + } + } + + unsafe { + simd_f64x2_forward_values(flat, &mut values)?; + } + + let Some(&output) = graph.output_nodes.first() else { + buffer.values.extend_from_slice(&[0.0, 0.0]); + buffer + .gradients + .resize(buffer.gradients.len() + graph.num_inputs * 2, 0.0); + row_index += 2; + continue; + }; + + cotangents.clear(); + cotangents.resize_with(value_count, || unsafe { _mm_setzero_pd() }); + let max_index = cotangents.len().saturating_sub(1); + *cotangents + .get_mut(output) + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + })? = unsafe { _mm_set1_pd(1.0) }; + + for instruction in flat.iter().rev() { + let current = checked_m128_lane(&cotangents, instruction.output)?; + unsafe { + match instruction.opcode { + OpCode::Constant => {} + OpCode::Add => { + let contribution = active_m128_contribution(current, current); + add_m128_cotangent(&mut cotangents, instruction.left, contribution)?; + add_m128_cotangent(&mut cotangents, instruction.right, contribution)?; + } + OpCode::Sub => { + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, current), + )?; + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution( + current, + _mm_sub_pd(_mm_setzero_pd(), current), + ), + )?; + } + OpCode::Mul => { + let left = checked_m128_lane(&values, instruction.left)?; + let right = checked_m128_lane(&values, instruction.right)?; + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, right)), + )?; + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution(current, _mm_mul_pd(current, left)), + )?; + } + OpCode::Div => { + let left = checked_m128_lane(&values, instruction.left)?; + let right = checked_m128_lane(&values, instruction.right)?; + let right_squared = _mm_mul_pd(right, right); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_div_pd(current, right)), + )?; + let right_contribution = _mm_sub_pd( + _mm_setzero_pd(), + _mm_div_pd(_mm_mul_pd(current, left), right_squared), + ); + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution(current, right_contribution), + )?; + } + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m128_lane(&values, instruction.left)?; + let right = checked_m128_lane(&values, instruction.right)?; + let output_value = checked_m128_lane(&values, instruction.output)?; + let (left_derivative, right_derivative) = + simd_f64x2_scalar_binary_derivatives( + left, + right, + output_value, + instruction.opcode, + )?; + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, left_derivative)), + )?; + add_m128_cotangent( + &mut cotangents, + instruction.right, + active_m128_contribution( + current, + _mm_mul_pd(current, right_derivative), + ), + )?; + } + OpCode::Neg => { + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution( + current, + _mm_sub_pd(_mm_setzero_pd(), current), + ), + )?; + } + OpCode::Sqrt => { + let output_value = checked_m128_lane(&values, instruction.output)?; + let contribution = + _mm_div_pd(_mm_mul_pd(current, _mm_set1_pd(0.5)), output_value); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, contribution), + )?; + } + OpCode::Relu => { + let input_value = checked_m128_lane(&values, instruction.left)?; + let mask = _mm_cmpgt_pd(input_value, _mm_setzero_pd()); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_and_pd(current, mask)), + )?; + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input_value = checked_m128_lane(&values, instruction.left)?; + let output_value = checked_m128_lane(&values, instruction.output)?; + let derivative = simd_f64x2_scalar_unary_derivative( + input_value, + output_value, + instruction.opcode, + )?; + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, derivative)), + )?; + } + OpCode::Abs => { + let input_value = checked_m128_lane(&values, instruction.left)?; + let positive = _mm_cmpgt_pd(input_value, _mm_setzero_pd()); + let negative = _mm_cmplt_pd(input_value, _mm_setzero_pd()); + let sign = _mm_add_pd( + _mm_and_pd(_mm_set1_pd(1.0), positive), + _mm_and_pd(_mm_set1_pd(-1.0), negative), + ); + add_m128_cotangent( + &mut cotangents, + instruction.left, + active_m128_contribution(current, _mm_mul_pd(current, sign)), + )?; + } + } + } + } + + let output_value = checked_m128_lane(&values, output)?; + let mut value_pair = [0.0_f64; 2]; + unsafe { + _mm_storeu_pd(value_pair.as_mut_ptr(), output_value); + } + buffer.values.extend_from_slice(&value_pair); + + gradient_pairs.clear(); + for input_index in 0..graph.num_inputs { + let gradient_lane = checked_m128_lane(&cotangents, input_index)?; + let mut stored = [0.0_f64; 2]; + unsafe { + _mm_storeu_pd(stored.as_mut_ptr(), gradient_lane); + } + gradient_pairs.push(stored); + } + for pair in &gradient_pairs { + buffer.gradients.push(pair[0]); + } + for pair in &gradient_pairs { + buffer.gradients.push(pair[1]); + } + + row_index += 2; + } + + append_scalar_gradient_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx")] +unsafe fn gradient_batch_simd_f64x4_impl( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + use std::arch::x86_64::{ + _mm256_add_pd, _mm256_and_pd, _mm256_cmp_pd, _mm256_div_pd, _mm256_mul_pd, _mm256_set1_pd, + _mm256_set_pd, _mm256_setzero_pd, _mm256_storeu_pd, _mm256_sub_pd, _CMP_GT_OQ, _CMP_LT_OQ, + }; + + graph.check_batch(batch)?; + let flat = graph.flat_instructions_slice(); + buffer.reset(batch.batch_size, graph.num_inputs); + + let value_count = graph.num_inputs + graph.instructions.len(); + let mut values: Vec<__m256d> = Vec::with_capacity(value_count); + let mut cotangents: Vec<__m256d> = Vec::with_capacity(value_count); + let mut gradient_quads = Vec::with_capacity(graph.num_inputs); + let mut row_index = 0; + while row_index + 3 < batch.batch_size { + let first = batch.try_row(row_index)?; + let second = batch.try_row(row_index + 1)?; + let third = batch.try_row(row_index + 2)?; + let fourth = batch.try_row(row_index + 3)?; + values.clear(); + for input_index in 0..graph.num_inputs { + values.push(_mm256_set_pd( + fourth[input_index], + third[input_index], + second[input_index], + first[input_index], + )); + } + + simd_f64x4_forward_values(flat, &mut values)?; + + let Some(&output) = graph.output_nodes.first() else { + buffer.values.extend_from_slice(&[0.0; 4]); + buffer + .gradients + .resize(buffer.gradients.len() + graph.num_inputs * 4, 0.0); + row_index += 4; + continue; + }; + + cotangents.clear(); + cotangents.resize_with(value_count, || _mm256_setzero_pd()); + let max_index = cotangents.len().saturating_sub(1); + *cotangents + .get_mut(output) + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + })? = _mm256_set1_pd(1.0); + + for instruction in flat.iter().rev() { + let current = checked_m256_lane(&cotangents, instruction.output)?; + match instruction.opcode { + OpCode::Constant => {} + OpCode::Add => { + let contribution = active_m256_contribution(current, current); + add_m256_cotangent(&mut cotangents, instruction.left, contribution)?; + add_m256_cotangent(&mut cotangents, instruction.right, contribution)?; + } + OpCode::Sub => { + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, current), + )?; + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution( + current, + _mm256_sub_pd(_mm256_setzero_pd(), current), + ), + )?; + } + OpCode::Mul => { + let left = checked_m256_lane(&values, instruction.left)?; + let right = checked_m256_lane(&values, instruction.right)?; + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, right)), + )?; + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution(current, _mm256_mul_pd(current, left)), + )?; + } + OpCode::Div => { + let left = checked_m256_lane(&values, instruction.left)?; + let right = checked_m256_lane(&values, instruction.right)?; + let right_squared = _mm256_mul_pd(right, right); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_div_pd(current, right)), + )?; + let right_contribution = _mm256_sub_pd( + _mm256_setzero_pd(), + _mm256_div_pd(_mm256_mul_pd(current, left), right_squared), + ); + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution(current, right_contribution), + )?; + } + OpCode::Pow | OpCode::LogAddExp => { + let left = checked_m256_lane(&values, instruction.left)?; + let right = checked_m256_lane(&values, instruction.right)?; + let output_value = checked_m256_lane(&values, instruction.output)?; + let (left_derivative, right_derivative) = simd_f64x4_scalar_binary_derivatives( + left, + right, + output_value, + instruction.opcode, + )?; + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, left_derivative)), + )?; + add_m256_cotangent( + &mut cotangents, + instruction.right, + active_m256_contribution(current, _mm256_mul_pd(current, right_derivative)), + )?; + } + OpCode::Neg => { + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution( + current, + _mm256_sub_pd(_mm256_setzero_pd(), current), + ), + )?; + } + OpCode::Sqrt => { + let output_value = checked_m256_lane(&values, instruction.output)?; + let contribution = + _mm256_div_pd(_mm256_mul_pd(current, _mm256_set1_pd(0.5)), output_value); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, contribution), + )?; + } + OpCode::Relu => { + let input_value = checked_m256_lane(&values, instruction.left)?; + let mask = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_GT_OQ); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_and_pd(current, mask)), + )?; + } + OpCode::Sin + | OpCode::Cos + | OpCode::Tan + | OpCode::Tanh + | OpCode::Log1pExp + | OpCode::Exp + | OpCode::Ln => { + let input_value = checked_m256_lane(&values, instruction.left)?; + let output_value = checked_m256_lane(&values, instruction.output)?; + let derivative = simd_f64x4_scalar_unary_derivative( + input_value, + output_value, + instruction.opcode, + )?; + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, derivative)), + )?; + } + OpCode::Abs => { + let input_value = checked_m256_lane(&values, instruction.left)?; + let positive = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_GT_OQ); + let negative = _mm256_cmp_pd(input_value, _mm256_setzero_pd(), _CMP_LT_OQ); + let sign = _mm256_add_pd( + _mm256_and_pd(_mm256_set1_pd(1.0), positive), + _mm256_and_pd(_mm256_set1_pd(-1.0), negative), + ); + add_m256_cotangent( + &mut cotangents, + instruction.left, + active_m256_contribution(current, _mm256_mul_pd(current, sign)), + )?; + } + } + } + + let output_value = checked_m256_lane(&values, output)?; + let mut value_quad = [0.0_f64; 4]; + _mm256_storeu_pd(value_quad.as_mut_ptr(), output_value); + buffer.values.extend_from_slice(&value_quad); + + gradient_quads.clear(); + for input_index in 0..graph.num_inputs { + let gradient_lane = checked_m256_lane(&cotangents, input_index)?; + let mut stored = [0.0_f64; 4]; + _mm256_storeu_pd(stored.as_mut_ptr(), gradient_lane); + gradient_quads.push(stored); + } + for lane_index in 0..4 { + for quad in &gradient_quads { + buffer.gradients.push(quad[lane_index]); + } + } + + row_index += 4; + } + + append_scalar_gradient_tail(graph, batch, row_index, buffer)?; + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +fn gradient_batch_simd_f64x4( + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + if !supports_simd_f64x4_runtime() { + return Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }); + } + unsafe { gradient_batch_simd_f64x4_impl(graph, batch, buffer) } +} + +#[cfg(not(target_arch = "x86_64"))] +fn gradient_batch_simd_f64x4( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd f64x4 backend requires x86_64 AVX support", + }) +} + +#[cfg(not(target_arch = "x86_64"))] +fn gradient_batch_simd_f64x2( + _graph: &CompiledGraph, + _batch: BatchInputs<'_>, + _buffer: &mut BatchGradientsBuffer, +) -> Result<()> { + Err(AutodiffError::InvalidGraph { + reason: "simd backend requires x86_64 SSE2 support", + }) +} + +impl MockDeviceBackend { + /// Allocate a mock-device buffer set for a compiled graph and batch size. + #[must_use] + pub fn allocate_batch_buffers( + &self, + graph: &CompiledGraph, + batch_size: usize, + ) -> DeviceBufferSet { + DeviceBufferSet::new(self.batch_plan(graph, batch_size)) + } + + /// Execute batch value computation through explicit mock-device transfers. + pub fn compute_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::MockDeviceCpu { + return Err(AutodiffError::InvalidGraph { + reason: "mock execution requires a mock-device buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match mock-device buffer plan", + }); + } + + let transfers = buffers.plan.compute_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "mock compute supports host-to-device input transfers only", + }); + } + buffers.upload(DeviceBufferKind::Inputs, batch.data)?; + } + } + + let device_inputs = buffers.download(DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchValuesBuffer::new(); + ScalarBackend.compute_batch(graph, device_batch, &mut scratch)?; + buffers.upload(DeviceBufferKind::Outputs, &scratch.data)?; + output.reset(batch.batch_size, graph.output_nodes.len()); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + if transfer.buffer != DeviceBufferKind::Outputs { + return Err(AutodiffError::InvalidGraph { + reason: "mock compute supports device-to-host output transfers only", + }); + } + output + .data + .extend_from_slice(&buffers.download(DeviceBufferKind::Outputs)?); + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::MockDeviceCpu, + mode: DeviceExecutionMode::ComputeBatch, + transfers, + used_native_kernel: false, + }) + } + + /// Execute batch gradient computation through explicit mock-device transfers. + pub fn gradient_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::MockDeviceCpu { + return Err(AutodiffError::InvalidGraph { + reason: "mock execution requires a mock-device buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match mock-device buffer plan", + }); + } + + let transfers = buffers.plan.gradient_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "mock gradient supports host-to-device input transfers only", + }); + } + buffers.upload(DeviceBufferKind::Inputs, batch.data)?; + } + } + + let device_inputs = buffers.download(DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchGradientsBuffer::new(); + ScalarBackend.gradient_batch(graph, device_batch, &mut scratch)?; + buffers.upload(DeviceBufferKind::PrimaryValues, &scratch.values)?; + buffers.upload(DeviceBufferKind::Gradients, &scratch.gradients)?; + output.reset(batch.batch_size, graph.num_inputs); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + match transfer.buffer { + DeviceBufferKind::PrimaryValues => output + .values + .extend_from_slice(&buffers.download(DeviceBufferKind::PrimaryValues)?), + DeviceBufferKind::Gradients => output + .gradients + .extend_from_slice(&buffers.download(DeviceBufferKind::Gradients)?), + _ => { + return Err(AutodiffError::InvalidGraph { + reason: + "mock gradient supports primary-value and gradient downloads only", + }); + } + } + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::MockDeviceCpu, + mode: DeviceExecutionMode::GradientBatch, + transfers, + used_native_kernel: false, + }) + } +} + +#[cfg(feature = "backend-wgpu")] +#[inline] +fn wgpu_buffer_size_bytes(len: usize) -> u64 { + let logical = len.saturating_mul(std::mem::size_of::()) as u64; + logical.max(8) +} + +#[cfg(feature = "backend-wgpu")] +fn encode_f64_slice(data: &[f64]) -> Vec { + let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); + for value in data { + bytes.extend_from_slice(&value.to_ne_bytes()); + } + bytes +} + +#[cfg(feature = "backend-wgpu")] +fn decode_f64_bytes(bytes: &[u8], len: usize) -> Result> { + let expected_len = len.saturating_mul(std::mem::size_of::()); + if bytes.len() < expected_len { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu readback length does not match planned buffer length", + }); + } + let mut values = Vec::with_capacity(len); + for chunk in bytes[..expected_len].chunks_exact(std::mem::size_of::()) { + let mut array = [0_u8; std::mem::size_of::()]; + array.copy_from_slice(chunk); + values.push(f64::from_ne_bytes(array)); + } + Ok(values) +} + +#[cfg(feature = "backend-wgpu")] +#[inline] +fn wgpu_buffer_size_bytes_f32(len: usize) -> u64 { + let logical = len.saturating_mul(std::mem::size_of::()) as u64; + logical.max(4) +} + +#[cfg(feature = "backend-wgpu")] +fn encode_f32_slice(data: &[f32]) -> Vec { + let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); + for value in data { + bytes.extend_from_slice(&value.to_ne_bytes()); + } + bytes +} + +#[cfg(feature = "backend-wgpu")] +fn decode_f32_bytes(bytes: &[u8], len: usize) -> Result> { + let expected_len = len.saturating_mul(std::mem::size_of::()); + if bytes.len() < expected_len { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu readback length does not match requested f32 buffer length", + }); + } + let mut values = Vec::with_capacity(len); + for chunk in bytes[..expected_len].chunks_exact(std::mem::size_of::()) { + let mut array = [0_u8; std::mem::size_of::()]; + array.copy_from_slice(chunk); + values.push(f32::from_ne_bytes(array)); + } + Ok(values) +} + +#[cfg(feature = "backend-wgpu")] +fn encode_u32_slice(data: &[u32]) -> Vec { + let mut bytes = Vec::with_capacity(data.len().saturating_mul(std::mem::size_of::())); + for value in data { + bytes.extend_from_slice(&value.to_ne_bytes()); + } + bytes +} + +#[cfg(feature = "backend-wgpu")] +/// Conservative opcode subset currently supported by the exact-safe native WGPU batch-compute path. +pub const WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES: [OpCode; 4] = + [OpCode::Constant, OpCode::Neg, OpCode::Relu, OpCode::Abs]; + +#[cfg(feature = "backend-wgpu")] +fn checked_u32_from_usize(value: usize, reason: &'static str) -> Result { + u32::try_from(value).map_err(|_| AutodiffError::InvalidGraph { reason }) +} + +#[cfg(feature = "backend-wgpu")] +#[inline] +fn is_exact_f32_roundtrip(value: f64) -> bool { + if value.is_nan() { + return false; + } + let narrowed = value as f32; + let widened = f64::from(narrowed); + widened == value && (value != 0.0 || widened.is_sign_negative() == value.is_sign_negative()) +} + +#[cfg(feature = "backend-wgpu")] +fn wgpu_native_exact_safe_supports_opcode(opcode: OpCode) -> bool { + WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES.contains(&opcode) +} + +#[cfg(feature = "backend-wgpu")] +fn wgpu_native_exact_safe_graph(graph: &CompiledGraph) -> bool { + !graph.output_nodes().is_empty() + && graph.flat_instructions_slice().iter().all(|instruction| { + wgpu_native_exact_safe_supports_opcode(instruction.opcode) + && (instruction.opcode != OpCode::Constant + || is_exact_f32_roundtrip(instruction.value)) + }) +} + +#[cfg(feature = "backend-wgpu")] +fn wgpu_native_exact_safe_batch(batch: BatchInputs<'_>) -> bool { + batch.data.iter().copied().all(is_exact_f32_roundtrip) +} + +#[cfg(feature = "backend-wgpu")] +const WGPU_NATIVE_WORDS_PER_INSTRUCTION: usize = 8; + +#[cfg(feature = "backend-wgpu")] +const WGPU_NATIVE_SHADER: &str = r#" +@group(0) @binding(0) var input_data: array; +@group(0) @binding(1) var instruction_words: array; +@group(0) @binding(2) var output_nodes: array; +@group(0) @binding(3) var value_data: array; +@group(0) @binding(4) var output_data: array; +@group(0) @binding(5) var kernel_meta: array; + +fn relu_scalar(x: f32) -> f32 { + if x > 0.0 { + return x; + } + return 0.0; +} + +fn log1p_exp_scalar(x: f32) -> f32 { + if x > 0.0 { + return x + log(1.0 + exp(-x)); + } + return log(1.0 + exp(x)); +} + +fn log_add_exp_scalar(a: f32, b: f32) -> f32 { + var max_value = a; + var min_value = b; + if a < b { + max_value = b; + min_value = a; + } + return max_value + log(1.0 + exp(min_value - max_value)); +} + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let row = gid.x; + let num_inputs = kernel_meta[0u]; + let num_instructions = kernel_meta[1u]; + let num_outputs = kernel_meta[2u]; + let value_count = kernel_meta[3u]; + let batch_size = kernel_meta[4u]; + if row >= batch_size { + return; + } + + let input_base = row * num_inputs; + let value_base = row * value_count; + let output_base = row * num_outputs; + + var input_index = 0u; + loop { + if input_index >= num_inputs { + break; + } + value_data[value_base + input_index] = input_data[input_base + input_index]; + input_index = input_index + 1u; + } + + var instruction_index = 0u; + loop { + if instruction_index >= num_instructions { + break; + } + let word_base = instruction_index * 8u; + let opcode = instruction_words[word_base + 0u]; + let output_index = instruction_words[word_base + 1u]; + let left_index = instruction_words[word_base + 2u]; + let right_index = instruction_words[word_base + 3u]; + let value_bits = instruction_words[word_base + 4u]; + var left_value: f32 = 0.0; + if left_index != 4294967295u { + left_value = value_data[value_base + left_index]; + } + var right_value: f32 = 0.0; + if right_index != 4294967295u { + right_value = value_data[value_base + right_index]; + } + var result_value: f32 = 0.0; + switch opcode { + case 0u: { + result_value = bitcast(value_bits); + } + case 1u: { + result_value = left_value + right_value; + } + case 2u: { + result_value = left_value - right_value; + } + case 3u: { + result_value = left_value * right_value; + } + case 4u: { + result_value = left_value / right_value; + } + case 5u: { + result_value = pow(left_value, right_value); + } + case 6u: { + result_value = sin(left_value); + } + case 7u: { + result_value = cos(left_value); + } + case 8u: { + result_value = tan(left_value); + } + case 9u: { + result_value = tanh(left_value); + } + case 10u: { + result_value = relu_scalar(left_value); + } + case 11u: { + result_value = log1p_exp_scalar(left_value); + } + case 12u: { + result_value = log_add_exp_scalar(left_value, right_value); + } + case 13u: { + result_value = -left_value; + } + case 14u: { + result_value = exp(left_value); + } + case 15u: { + result_value = log(left_value); + } + case 16u: { + result_value = sqrt(left_value); + } + case 17u: { + result_value = abs(left_value); + } + default: { + result_value = 0.0; + } + } + value_data[value_base + output_index] = result_value; + instruction_index = instruction_index + 1u; + } + + var output_index = 0u; + loop { + if output_index >= num_outputs { + break; + } + let node = output_nodes[output_index]; + output_data[output_base + output_index] = value_data[value_base + node]; + output_index = output_index + 1u; + } +} +"#; + +#[cfg(feature = "backend-wgpu")] +impl WgpuBackend { + /// Create the default WGPU backend using device id 0 and automatic transfers. + pub fn new_default() -> Result { + Self::new( + AcceleratorDeviceContext::wgpu(0), + DeviceTransferPolicy::Automatic, + ) + } + + /// Create the WGPU backend skeleton from a context and transfer policy. + pub fn new( + context: AcceleratorDeviceContext, + transfer_policy: DeviceTransferPolicy, + ) -> Result { + Self::from_boundary(GpuBackendBoundary::new(context, transfer_policy)) + } + + /// Create the WGPU backend skeleton from a boundary descriptor. + pub fn from_boundary(boundary: GpuBackendBoundary) -> Result { + if boundary.context.kind != AcceleratorDeviceKind::Wgpu { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu backend requires a WGPU accelerator context", + }); + } + let instance = wgpu::Instance::default(); + let adapter = block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { + power_preference: wgpu::PowerPreference::default(), + force_fallback_adapter: false, + compatible_surface: None, + })) + .map_err(|_| AutodiffError::InvalidGraph { + reason: "wgpu adapter request failed", + })?; + let adapter_info = adapter.get_info(); + let (device, queue) = block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("petite-ad-wgpu"), + required_features: wgpu::Features::empty(), + required_limits: wgpu::Limits::default(), + memory_hints: wgpu::MemoryHints::Performance, + trace: wgpu::Trace::Off, + })) + .map_err(|_| AutodiffError::InvalidGraph { + reason: "wgpu device request failed", + })?; + Ok(Self { + boundary, + device, + queue, + adapter_info, + }) + } + + /// Return the boundary descriptor used for this backend. + #[must_use] + pub fn boundary(&self) -> &GpuBackendBoundary { + &self.boundary + } + + /// Return the accelerator context used for this backend. + #[must_use] + pub fn context(&self) -> &AcceleratorDeviceContext { + &self.boundary.context + } + + /// Return the configured transfer policy. + #[must_use] + pub fn transfer_policy(&self) -> DeviceTransferPolicy { + self.boundary.transfer_policy + } + + /// Return the resolved adapter name. + #[must_use] + pub fn adapter_name(&self) -> &str { + &self.adapter_info.name + } + + /// Return the conservative exact-safe opcode subset used by the native WGPU batch-compute path. + #[must_use] + pub fn native_batch_compute_supported_opcodes() -> &'static [OpCode] { + &WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES + } + + /// Return whether this graph is eligible for the restricted exact-safe native WGPU path. + #[must_use] + pub fn supports_native_batch_compute(&self, graph: &CompiledGraph) -> bool { + wgpu_native_exact_safe_graph(graph) + } + + /// Return whether a concrete batch is eligible for the exact-safe native WGPU path. + #[must_use] + pub fn supports_native_batch_compute_for_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + ) -> bool { + self.supports_native_batch_compute(graph) && wgpu_native_exact_safe_batch(batch) + } + + fn create_initialized_storage_buffer( + &self, + label: &'static str, + bytes: &[u8], + min_size: u64, + ) -> wgpu::Buffer { + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size: min_size.max(bytes.len() as u64).max(4), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + if !bytes.is_empty() { + self.queue.write_buffer(&buffer, 0, bytes); + } + buffer + } + + fn download_raw_buffer(&self, buffer: &wgpu::Buffer, size: u64) -> Result> { + let staging = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("petite-ad-wgpu-readback"), + size, + usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("petite-ad-wgpu-readback-encoder"), + }); + encoder.copy_buffer_to_buffer(buffer, 0, &staging, 0, size); + let submission = self.queue.submit([encoder.finish()]); + let slice = staging.slice(..); + let (tx, rx) = mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |result| { + let _ = tx.send(result); + }); + let _ = self.device.poll(wgpu::PollType::wait_for(submission)); + match rx.recv() { + Ok(Ok(())) => {} + _ => { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu readback mapping failed", + }); + } + } + let view = slice.get_mapped_range(); + let bytes = view.to_vec(); + drop(view); + staging.unmap(); + Ok(bytes) + } + + fn compute_batch_native( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + ) -> Result> { + let metadata = graph.metadata(); + let input_f32: Vec = batch.data.iter().map(|value| *value as f32).collect(); + let output_count = batch.batch_size.saturating_mul(metadata.num_outputs); + let input_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-inputs", + &encode_f32_slice(&input_f32), + wgpu_buffer_size_bytes_f32(input_f32.len()), + ); + let instruction_words = encode_u32_slice(&self.native_instruction_words(graph)?); + let instruction_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-instructions", + &instruction_words, + instruction_words.len() as u64, + ); + let output_nodes: Result> = graph + .output_nodes() + .iter() + .map(|node| { + checked_u32_from_usize(*node, "wgpu native output node index exceeds u32 range") + }) + .collect(); + let output_nodes = output_nodes?; + let output_node_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-output-nodes", + &encode_u32_slice(&output_nodes), + (output_nodes.len().max(1) * std::mem::size_of::()) as u64, + ); + let value_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("petite-ad-wgpu-native-values"), + size: wgpu_buffer_size_bytes_f32(batch.batch_size.saturating_mul(metadata.value_count)), + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("petite-ad-wgpu-native-outputs"), + size: wgpu_buffer_size_bytes_f32(output_count), + usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let meta = [ + checked_u32_from_usize( + metadata.num_inputs, + "wgpu native input dimension exceeds u32 range", + )?, + checked_u32_from_usize( + metadata.num_instructions, + "wgpu native instruction count exceeds u32 range", + )?, + checked_u32_from_usize( + metadata.num_outputs, + "wgpu native output dimension exceeds u32 range", + )?, + checked_u32_from_usize( + metadata.value_count, + "wgpu native value count exceeds u32 range", + )?, + checked_u32_from_usize(batch.batch_size, "wgpu native batch size exceeds u32 range")?, + ]; + let meta_buffer = self.create_initialized_storage_buffer( + "petite-ad-wgpu-native-meta", + &encode_u32_slice(&meta), + (meta.len() * std::mem::size_of::()) as u64, + ); + self.queue.submit([]); + + let shader = self + .device + .create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("petite-ad-wgpu-native-compute-shader"), + source: wgpu::ShaderSource::Wgsl(WGPU_NATIVE_SHADER.into()), + }); + let pipeline = self + .device + .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("petite-ad-wgpu-native-compute-pipeline"), + layout: None, + module: &shader, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + let layout = pipeline.get_bind_group_layout(0); + let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("petite-ad-wgpu-native-bind-group"), + layout: &layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: input_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: instruction_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: output_node_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 3, + resource: value_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 4, + resource: output_buffer.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 5, + resource: meta_buffer.as_entire_binding(), + }, + ], + }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("petite-ad-wgpu-native-compute-encoder"), + }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("petite-ad-wgpu-native-compute-pass"), + timestamp_writes: None, + }); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bind_group, &[]); + let workgroups = checked_u32_from_usize( + batch.batch_size.div_ceil(64), + "wgpu native workgroup count exceeds u32 range", + )?; + pass.dispatch_workgroups(workgroups.max(1), 1, 1); + } + self.queue.submit([encoder.finish()]); + let raw_output = + self.download_raw_buffer(&output_buffer, wgpu_buffer_size_bytes_f32(output_count))?; + let output_f32 = decode_f32_bytes(&raw_output, output_count)?; + Ok(output_f32.into_iter().map(f64::from).collect()) + } + + fn native_instruction_words(&self, graph: &CompiledGraph) -> Result> { + let mut words = Vec::with_capacity( + graph.flat_instructions_slice().len() * WGPU_NATIVE_WORDS_PER_INSTRUCTION, + ); + for instruction in graph.flat_instructions_slice() { + words.push(match instruction.opcode { + OpCode::Constant => 0, + OpCode::Add => 1, + OpCode::Sub => 2, + OpCode::Mul => 3, + OpCode::Div => 4, + OpCode::Pow => 5, + OpCode::Sin => 6, + OpCode::Cos => 7, + OpCode::Tan => 8, + OpCode::Tanh => 9, + OpCode::Relu => 10, + OpCode::Log1pExp => 11, + OpCode::LogAddExp => 12, + OpCode::Neg => 13, + OpCode::Exp => 14, + OpCode::Ln => 15, + OpCode::Sqrt => 16, + OpCode::Abs => 17, + }); + words.push(checked_u32_from_usize( + instruction.output, + "wgpu native instruction output index exceeds u32 range", + )?); + words.push(if instruction.left == UNUSED_NODE_ID { + u32::MAX + } else { + checked_u32_from_usize( + instruction.left, + "wgpu native instruction left index exceeds u32 range", + )? + }); + words.push(if instruction.right == UNUSED_NODE_ID { + u32::MAX + } else { + checked_u32_from_usize( + instruction.right, + "wgpu native instruction right index exceeds u32 range", + )? + }); + words.push((instruction.value as f32).to_bits()); + words.extend_from_slice(&[0, 0, 0]); + } + Ok(words) + } + + /// Allocate a real WGPU buffer set for a compiled graph and batch size. + pub fn allocate_batch_buffers( + &self, + graph: &CompiledGraph, + batch_size: usize, + ) -> Result { + let plan = self.batch_plan(graph, batch_size); + let mut buffers = Vec::with_capacity(plan.buffer_handles.len()); + for handle in &plan.buffer_handles { + let label = format!("petite-ad-wgpu-{:?}", handle.kind); + let buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some(label.as_str()), + size: wgpu_buffer_size_bytes(handle.len), + usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + buffers.push(WgpuBuffer { + handle: *handle, + buffer, + }); + } + Ok(WgpuBufferSet { plan, buffers }) + } + + fn upload_buffer( + &self, + buffers: &WgpuBufferSet, + kind: DeviceBufferKind, + data: &[f64], + ) -> Result<()> { + let buffer = buffers.buffer(kind)?; + if buffer.handle.len != data.len() { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu upload length must match planned buffer length", + }); + } + if data.is_empty() { + return Ok(()); + } + let bytes = encode_f64_slice(data); + self.queue.write_buffer(&buffer.buffer, 0, &bytes); + self.queue.submit([]); + Ok(()) + } + + fn download_buffer(&self, buffers: &WgpuBufferSet, kind: DeviceBufferKind) -> Result> { + let buffer = buffers.buffer(kind)?; + if buffer.handle.len == 0 { + return Ok(Vec::new()); + } + let raw = + self.download_raw_buffer(&buffer.buffer, wgpu_buffer_size_bytes(buffer.handle.len))?; + decode_f64_bytes(&raw, buffer.handle.len) + } + + /// Execute batch value computation through explicit WGPU transfers. + pub fn compute_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::Wgpu { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu execution requires a WGPU buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match WGPU buffer plan", + }); + } + + let transfers = buffers.plan.compute_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu compute supports host-to-device input transfers only", + }); + } + buffers.upload(self, DeviceBufferKind::Inputs, batch.data)?; + } + } + + let used_native_kernel = self.supports_native_batch_compute_for_batch(graph, batch); + let scratch_data = if used_native_kernel { + self.compute_batch_native(graph, batch)? + } else { + let device_inputs = buffers.download(self, DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchValuesBuffer::new(); + ScalarBackend.compute_batch(graph, device_batch, &mut scratch)?; + scratch.data + }; + buffers.upload(self, DeviceBufferKind::Outputs, &scratch_data)?; + output.reset(batch.batch_size, graph.output_nodes.len()); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + if transfer.buffer != DeviceBufferKind::Outputs { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu compute supports device-to-host output transfers only", + }); + } + output + .data + .extend_from_slice(&buffers.download(self, DeviceBufferKind::Outputs)?); + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::Wgpu, + mode: DeviceExecutionMode::ComputeBatch, + transfers, + used_native_kernel, + }) + } + + /// Execute batch gradient computation through explicit WGPU transfers. + pub fn gradient_batch_with_buffers( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffers: &mut WgpuBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + graph.check_batch(batch)?; + if buffers.plan.backend != BackendKind::Wgpu { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu execution requires a WGPU buffer plan", + }); + } + if buffers.plan.batch_size != batch.batch_size || buffers.plan.input_dim != batch.input_dim + { + return Err(AutodiffError::InvalidGraph { + reason: "batch shape must match WGPU buffer plan", + }); + } + + let transfers = buffers.plan.gradient_transfer_plan.clone(); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::HostToDevice { + if transfer.buffer != DeviceBufferKind::Inputs { + return Err(AutodiffError::InvalidGraph { + reason: "wgpu gradient supports host-to-device input transfers only", + }); + } + buffers.upload(self, DeviceBufferKind::Inputs, batch.data)?; + } + } + + let device_inputs = buffers.download(self, DeviceBufferKind::Inputs)?; + let device_batch = BatchInputs::new(&device_inputs, batch.batch_size, batch.input_dim)?; + let mut scratch = BatchGradientsBuffer::new(); + ScalarBackend.gradient_batch(graph, device_batch, &mut scratch)?; + buffers.upload(self, DeviceBufferKind::PrimaryValues, &scratch.values)?; + buffers.upload(self, DeviceBufferKind::Gradients, &scratch.gradients)?; + output.reset(batch.batch_size, graph.num_inputs); + for transfer in &transfers { + if transfer.kind == DeviceTransferKind::DeviceToHost { + match transfer.buffer { + DeviceBufferKind::PrimaryValues => output.values.extend_from_slice( + &buffers.download(self, DeviceBufferKind::PrimaryValues)?, + ), + DeviceBufferKind::Gradients => output + .gradients + .extend_from_slice(&buffers.download(self, DeviceBufferKind::Gradients)?), + _ => { + return Err(AutodiffError::InvalidGraph { + reason: + "wgpu gradient supports primary-value and gradient downloads only", + }); + } + } + } + } + + Ok(DeviceExecutionTrace { + backend: BackendKind::Wgpu, + mode: DeviceExecutionMode::GradientBatch, + transfers, + used_native_kernel: false, + }) + } +} + +impl DeviceBackend for ScalarBackend { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + BackendKind::Scalar + } +} + +impl DeviceBackend for MockDeviceBackend { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + BackendKind::MockDeviceCpu + } +} + +#[cfg(feature = "backend-wgpu")] +impl DeviceBackend for WgpuBackend { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + BackendKind::Wgpu + } +} + +impl DeviceBackend for SimdBackend { + fn backend_kind(&self, graph: &CompiledGraph) -> BackendKind { + graph + .simd_support_report() + .map(|report| report.backend) + .unwrap_or(BackendKind::SimdF64x2) + } +} + +impl DeviceBackend for BackendKind { + fn backend_kind(&self, _graph: &CompiledGraph) -> BackendKind { + *self + } +} + +impl ExecutionBackend for SimdBackend { + fn name(&self) -> &'static str { + "simd" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::simd_f64x2() + } + + fn compute(&self, _graph: &CompiledGraph, _inputs: &[f64]) -> Result { + Err(AutodiffError::InvalidGraph { + reason: "simd backend currently supports batch compute only", + }) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + if let Ok(report) = graph.backend_support_report(BackendKind::SimdF64x4) { + if report.can_compute_batch() { + return compute_batch_simd_f64x4(graph, batch, buffer); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_compute { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch compute on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + compute_batch_simd_f64x2(graph, batch, buffer) + } + + fn gradient(&self, _graph: &CompiledGraph, _inputs: &[f64]) -> Result<(f64, Vec)> { + Err(AutodiffError::InvalidGraph { + reason: "simd backend does not support reverse gradients yet", + }) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + if let Ok(report) = graph.backend_support_report(BackendKind::SimdF64x4) { + if report.can_gradient_batch() { + return gradient_batch_simd_f64x4(graph, batch, buffer); + } + } + let capabilities = self.capabilities(); + if !capabilities.supports_batch_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch gradients on this target", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + gradient_batch_simd_f64x2(graph, batch, buffer) + } +} + +impl ExecutionBackend for MockDeviceBackend { + fn name(&self) -> &'static str { + "mock-device-cpu" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::scalar_f64() + } + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { + ScalarBackend.compute(graph, inputs) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + ScalarBackend.compute_batch(graph, batch, buffer) + } + + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { + ScalarBackend.gradient(graph, inputs) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + ScalarBackend.gradient_batch(graph, batch, buffer) + } +} + +#[cfg(feature = "backend-wgpu")] +impl ExecutionBackend for WgpuBackend { + fn name(&self) -> &'static str { + "wgpu" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::wgpu_f64() + } + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { + let batch = BatchInputs::new(inputs, 1, inputs.len())?; + let mut values = BatchValuesBuffer::new(); + self.compute_batch(graph, batch, &mut values)?; + values + .data + .first() + .copied() + .ok_or(AutodiffError::InvalidGraph { + reason: "wgpu compute did not produce an output value", + }) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + graph.validate_backend_capabilities(&self.capabilities())?; + let mut gpu_buffers = self.allocate_batch_buffers(graph, batch.batch_size)?; + self.compute_batch_with_buffers(graph, batch, &mut gpu_buffers, buffer)?; + Ok(()) + } + + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { + let batch = BatchInputs::new(inputs, 1, inputs.len())?; + let mut gradients = BatchGradientsBuffer::new(); + self.gradient_batch(graph, batch, &mut gradients)?; + let value = gradients + .values + .first() + .copied() + .ok_or(AutodiffError::InvalidGraph { + reason: "wgpu gradient did not produce a primary output value", + })?; + Ok((value, gradients.gradients)) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + graph.validate_backend_capabilities(&self.capabilities())?; + let mut gpu_buffers = self.allocate_batch_buffers(graph, batch.batch_size)?; + self.gradient_batch_with_buffers(graph, batch, &mut gpu_buffers, buffer)?; + Ok(()) + } +} + +impl ExecutionBackend for ScalarBackend { + fn name(&self) -> &'static str { + "scalar" + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::scalar_f64() + } + + fn compute(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result { + graph.validate_backend_capabilities(&self.capabilities())?; + graph.compute(inputs) + } + + fn compute_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + let capabilities = self.capabilities(); + if !capabilities.supports_batch_compute { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch compute", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + graph.compute_batch_into(batch, buffer) + } + + fn gradient(&self, graph: &CompiledGraph, inputs: &[f64]) -> Result<(f64, Vec)> { + let capabilities = self.capabilities(); + if !capabilities.supports_reverse_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support reverse gradients", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + graph.gradient(inputs) + } + + fn gradient_batch( + &self, + graph: &CompiledGraph, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + let capabilities = self.capabilities(); + if !capabilities.supports_batch_gradient { + return Err(AutodiffError::InvalidGraph { + reason: "backend does not support batch gradients", + }); + } + graph.validate_backend_capabilities(&capabilities)?; + graph.gradient_batch_into(batch, buffer) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::approx_eq_eps as approx_eq; + + // --- helpers --- + + fn make_simple_graph() -> CompiledGraph { + let instructions = vec![Instruction::Binary { + op: MultiAD::Add, + left: 0, + right: 1, + }]; + CompiledGraph::new(2, instructions, vec![2]).unwrap() + } + + fn make_multi_output_graph() -> CompiledGraph { + // f(x) = x, outputs: [x, sin(x)] + let instructions = vec![Instruction::Unary { + op: MultiAD::Sin, + arg: 0, + }]; + CompiledGraph::new(1, instructions, vec![0, 1]).unwrap() + } + + // ---- 1. compute / gradient input-length checks ---- + + #[test] + fn test_compiled_compute_wrong_input_length_error() { + let graph = make_simple_graph(); + assert!(graph.compute(&[1.0]).is_err()); + assert!(graph.compute(&[1.0, 2.0, 3.0]).is_err()); + assert!(graph.compute(&[1.0, 2.0]).is_ok()); + } + + #[test] + fn test_compiled_gradient_wrong_input_length_error() { + let graph = make_simple_graph(); + assert!(graph.gradient(&[1.0]).is_err()); + assert!(graph.gradient(&[1.0, 2.0, 3.0]).is_err()); + assert!(graph.gradient(&[1.0, 2.0]).is_ok()); + } + + // ---- 2. batch with odd batch sizes ---- + + #[test] + fn test_compiled_compute_batch_odd_size() { + let graph = make_simple_graph(); + let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let batch = BatchInputs::new(&data, 3, 2).unwrap(); + let result = graph.compute_batch(batch).unwrap(); + assert_eq!(result.batch_size, 3); + assert_eq!(result.output_dim, 1); + assert_eq!(result.data.len(), 3); + assert!(approx_eq(result.data[0], 3.0, 1e-10)); + assert!(approx_eq(result.data[1], 7.0, 1e-10)); + assert!(approx_eq(result.data[2], 11.0, 1e-10)); + } + + #[test] + fn test_compiled_gradient_batch_odd_size() { + let graph = make_simple_graph(); + let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + let batch = BatchInputs::new(&data, 3, 2).unwrap(); + let result = graph.gradient_batch(batch).unwrap(); + assert_eq!(result.batch_size, 3); + assert_eq!(result.input_dim, 2); + assert_eq!(result.values.len(), 3); + assert_eq!(result.gradients.len(), 6); + // f(x,y) = x+y, gradients are always [1, 1] + for i in 0..3 { + assert!(approx_eq(result.gradients[i * 2], 1.0, 1e-10)); + assert!(approx_eq(result.gradients[i * 2 + 1], 1.0, 1e-10)); + } + } + + // ---- 3. large batch ---- + + #[test] + fn test_compiled_compute_batch_large() { + let graph = make_simple_graph(); + let batch_size = 100; + let mut data = Vec::with_capacity(batch_size * 2); + for i in 0..batch_size { + let x = i as f64; + data.push(x); + data.push(x * 2.0); + } + let batch = BatchInputs::new(&data, batch_size, 2).unwrap(); + let result = graph.compute_batch(batch).unwrap(); + assert_eq!(result.batch_size, 100); + assert_eq!(result.data.len(), 100); + for i in 0..batch_size { + assert!(approx_eq(result.data[i], (i as f64) * 3.0, 1e-10)); + } + } + + // ---- 4. multi-output batch ---- + + #[test] + fn test_compiled_compute_batch_many() { + let graph = make_multi_output_graph(); + let data = vec![0.0, 1.0, 2.0]; + let batch = BatchInputs::new(&data, 3, 1).unwrap(); + let result = graph.compute_batch(batch).unwrap(); + assert_eq!(result.batch_size, 3); + assert_eq!(result.output_dim, 2); + assert_eq!(result.data.len(), 6); + // [x, sin(x)] for each row + assert!(approx_eq(result.data[0], 0.0, 1e-10)); + assert!(approx_eq(result.data[1], 0.0_f64.sin(), 1e-10)); + assert!(approx_eq(result.data[2], 1.0, 1e-10)); + assert!(approx_eq(result.data[3], 1.0_f64.sin(), 1e-10)); + assert!(approx_eq(result.data[4], 2.0, 1e-10)); + assert!(approx_eq(result.data[5], 2.0_f64.sin(), 1e-10)); + } + + // ---- 5. CompiledWorkspace reuse ---- + + #[test] + fn test_workspace_reuse_compute() { + let graph = make_simple_graph(); + let mut ws = graph.workspace(); + let v1 = graph.compute_with_workspace(&[1.0, 2.0], &mut ws).unwrap(); + assert!(approx_eq(v1, 3.0, 1e-10)); + let v2 = graph.compute_with_workspace(&[3.0, 4.0], &mut ws).unwrap(); + assert!(approx_eq(v2, 7.0, 1e-10)); + } + + #[test] + fn test_workspace_reuse_gradient() { + let graph = make_simple_graph(); + let mut ws = graph.workspace(); + let (v1, g1) = graph.gradient_with_workspace(&[1.0, 2.0], &mut ws).unwrap(); + assert!(approx_eq(v1, 3.0, 1e-10)); + assert!(approx_eq(g1[0], 1.0, 1e-10)); + assert!(approx_eq(g1[1], 1.0, 1e-10)); + let (v2, g2) = graph.gradient_with_workspace(&[3.0, 4.0], &mut ws).unwrap(); + assert!(approx_eq(v2, 7.0, 1e-10)); + assert!(approx_eq(g2[0], 1.0, 1e-10)); + assert!(approx_eq(g2[1], 1.0, 1e-10)); + } + + // ---- 6. BackendKind ---- + + #[test] + fn test_backend_kind_name_non_empty() { + for kind in &[ + BackendKind::Scalar, + BackendKind::MockDeviceCpu, + BackendKind::Wgpu, + BackendKind::SimdF64x4, + BackendKind::SimdF64x2, + ] { + assert!(!kind.name().is_empty()); + } + } + + #[test] + fn test_backend_kind_capabilities() { + for kind in &[ + BackendKind::Scalar, + BackendKind::MockDeviceCpu, + BackendKind::Wgpu, + BackendKind::SimdF64x4, + BackendKind::SimdF64x2, + ] { + let caps = kind.capabilities(); + // Every built-in backend should support f64. + assert!(caps.supports_f64); + // Every built-in backend should support at least one opcode. + assert!(!caps.supported_opcodes.is_empty()); + } + } + + // ---- 7. BatchInputs ---- + + #[test] + fn test_batch_inputs_try_row_valid() { + let data = vec![1.0, 2.0, 3.0, 4.0]; + let batch = BatchInputs::new(&data, 2, 2).unwrap(); + let row0 = batch.try_row(0).unwrap(); + assert_eq!(row0, &[1.0, 2.0]); + let row1 = batch.try_row(1).unwrap(); + assert_eq!(row1, &[3.0, 4.0]); + } + + #[test] + fn test_batch_inputs_try_row_out_of_bounds() { + let data = vec![1.0, 2.0, 3.0, 4.0]; + let batch = BatchInputs::new(&data, 2, 2).unwrap(); + assert!(batch.try_row(2).is_err()); + assert!(batch.try_row(100).is_err()); + } + + #[test] + fn test_batch_inputs_invalid_shape() { + assert!(BatchInputs::new(&[1.0, 2.0], 2, 2).is_err()); + assert!(BatchInputs::new(&[1.0, 2.0, 3.0], 1, 2).is_err()); + } + + // ---- 8. BatchValues / BatchValuesBuffer ---- + + #[test] + fn test_batch_values_buffer_reuse() { + let graph = make_simple_graph(); + let mut buffer = BatchValuesBuffer::new(); + + let batch1 = BatchInputs::new(&[1.0, 2.0], 1, 2).unwrap(); + graph.compute_batch_into(batch1, &mut buffer).unwrap(); + let values1 = buffer.to_values(); + assert_eq!(values1.batch_size, 1); + assert_eq!(values1.data, &[3.0]); + + let batch2 = BatchInputs::new(&[4.0, 5.0, 6.0, 7.0], 2, 2).unwrap(); + graph.compute_batch_into(batch2, &mut buffer).unwrap(); + let values2 = buffer.to_values(); + assert_eq!(values2.batch_size, 2); + assert_eq!(values2.data, &[9.0, 13.0]); + } + + // ---- 9. BatchGradients / BatchGradientsBuffer ---- + + #[test] + fn test_batch_gradients_buffer_reuse() { + let graph = make_simple_graph(); + let mut buffer = BatchGradientsBuffer::new(); + + let batch1 = BatchInputs::new(&[1.0, 2.0], 1, 2).unwrap(); + graph.gradient_batch_into(batch1, &mut buffer).unwrap(); + let grad1 = buffer.to_gradients(); + assert_eq!(grad1.batch_size, 1); + assert!(approx_eq(grad1.values[0], 3.0, 1e-10)); + assert_eq!(grad1.gradients, &[1.0, 1.0]); + + let batch2 = BatchInputs::new(&[4.0, 5.0, 6.0, 7.0], 2, 2).unwrap(); + graph.gradient_batch_into(batch2, &mut buffer).unwrap(); + let grad2 = buffer.to_gradients(); + assert_eq!(grad2.batch_size, 2); + assert!(approx_eq(grad2.values[0], 9.0, 1e-10)); + assert!(approx_eq(grad2.values[1], 13.0, 1e-10)); + assert_eq!(grad2.gradients, &[1.0, 1.0, 1.0, 1.0]); + } + + // ---- 10. FlatInstruction and OpCode ---- + + #[test] + fn test_flat_instruction_debug() { + let fi = FlatInstruction { + opcode: OpCode::Add, + output: 2, + left: 0, + right: 1, + value: 0.0, + }; + let s = format!("{:?}", fi); + assert!(s.contains("Add")); + assert!(s.contains("2")); + } + + #[test] + fn test_opcode_debug() { + // Every variant should have a non-empty Debug representation. + for op in &[ + OpCode::Constant, + OpCode::Add, + OpCode::Sub, + OpCode::Mul, + OpCode::Div, + OpCode::Pow, + OpCode::Sin, + OpCode::Cos, + OpCode::Tan, + OpCode::Tanh, + OpCode::Relu, + OpCode::Log1pExp, + OpCode::LogAddExp, + OpCode::Neg, + OpCode::Exp, + OpCode::Ln, + OpCode::Sqrt, + OpCode::Abs, + ] { + assert!(!format!("{:?}", op).is_empty()); + } + } + + #[test] + fn test_opcode_arity() { + assert_eq!(OpCode::Constant.arity(), 0); + assert_eq!(OpCode::Sin.arity(), 1); + assert_eq!(OpCode::Neg.arity(), 1); + assert_eq!(OpCode::Add.arity(), 2); + assert_eq!(OpCode::Mul.arity(), 2); + } + + #[test] + fn test_opcode_from_multi_ad() { + assert_eq!(OpCode::from_multi_ad(MultiAD::Add).unwrap(), OpCode::Add); + assert_eq!(OpCode::from_multi_ad(MultiAD::Sin).unwrap(), OpCode::Sin); + assert!(OpCode::from_multi_ad(MultiAD::Inp).is_err()); + } + + // ---- 11. backend support reports ---- + + #[test] + fn test_backend_support_report_scalar() { + let graph = make_simple_graph(); + let report = graph.backend_support_report(BackendKind::Scalar).unwrap(); + assert!(report.supports_f64); + assert!(report.supports_required_opcodes); + assert!(report.can_compute_batch()); + assert!(report.can_gradient_batch()); + assert_eq!(report.backend, BackendKind::Scalar); + } + + #[test] + fn test_simd_support_report() { + let graph = make_simple_graph(); + let report = graph.simd_support_report().unwrap(); + // Report should belong to a SIMD backend. + assert!(matches!( + report.backend, + BackendKind::SimdF64x4 | BackendKind::SimdF64x2 + )); + } + + // ---- 12. Device types ---- + + #[test] + fn test_device_memory_location_variants() { + let host = DeviceMemoryLocation::Host; + let device = DeviceMemoryLocation::Device; + assert_ne!(host, device); + // Equality on same variant. + assert_eq!(host, DeviceMemoryLocation::Host); + assert_eq!(device, DeviceMemoryLocation::Device); + } + + #[test] + fn test_device_batch_plan_creation() { + let graph = make_simple_graph(); + let plan = graph.device_batch_plan(BackendKind::Scalar, 10); + assert_eq!(plan.backend, BackendKind::Scalar); + assert_eq!(plan.batch_size, 10); + assert_eq!(plan.input_dim, 2); + assert_eq!(plan.output_dim, 1); + assert!(!plan.buffers.is_empty()); + assert!(!plan.buffer_handles.is_empty()); + } + + #[test] + fn test_mock_device_allocate_batch_buffers() { + let graph = make_simple_graph(); + let buffers = graph.allocate_mock_device_buffers(5); + assert_eq!(buffers.plan().backend, BackendKind::MockDeviceCpu); + assert_eq!(buffers.plan().batch_size, 5); + assert!(!buffers.buffers().is_empty()); + } + + #[test] + fn test_device_buffer_set_upload_download() { + let graph = make_simple_graph(); + let mut buffers = graph.allocate_mock_device_buffers(3); + let input_data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; + buffers + .upload(DeviceBufferKind::Inputs, &input_data) + .unwrap(); + let downloaded = buffers.download(DeviceBufferKind::Inputs).unwrap(); + assert_eq!(downloaded, input_data); + } + + // ---- 13. graph with all ops ---- + + #[test] + fn test_compiled_graph_all_ops() { + // Build a graph that uses many different ops. + // Graph: 2 inputs (x, y) + // 0: x (input) + // 1: y (input) + // 2: sin(x) [Unary Sin] + // 3: cos(y) [Unary Cos] + // 4: x + y [Binary Add] + // 5: x * y [Binary Mul] + // 6: -x [Unary Neg] + // 7: exp(x) [Unary Exp] + // 8: ln(y) [Unary Ln] + // 9: (x+y) / (x*y) [Binary Div] + // 10: x - y [Binary Sub] + // 11: tanh(x) [Unary Tanh] + // 12: sqrt(abs(x)) [Unary Sqrt] on abs(x)... + // Let's simplify: just use relu, abs on input + let instructions = vec![ + Instruction::Unary { + op: MultiAD::Sin, + arg: 0, + }, + Instruction::Unary { + op: MultiAD::Cos, + arg: 1, + }, + Instruction::Binary { + op: MultiAD::Add, + left: 0, + right: 1, + }, + Instruction::Binary { + op: MultiAD::Mul, + left: 0, + right: 1, + }, + Instruction::Unary { + op: MultiAD::Neg, + arg: 0, + }, + Instruction::Unary { + op: MultiAD::Exp, + arg: 0, + }, + Instruction::Unary { + op: MultiAD::Ln, + arg: 1, + }, + Instruction::Binary { + op: MultiAD::Div, + left: 3, + right: 4, + }, + Instruction::Binary { + op: MultiAD::Sub, + left: 0, + right: 1, + }, + Instruction::Unary { + op: MultiAD::Tanh, + arg: 0, + }, + Instruction::Unary { + op: MultiAD::Relu, + arg: 0, + }, + Instruction::Unary { + op: MultiAD::Abs, + arg: 0, + }, + Instruction::Constant(42.0), + Instruction::Binary { + op: MultiAD::Pow, + left: 0, + right: 1, + }, + Instruction::Unary { + op: MultiAD::Log1pExp, + arg: 0, + }, + Instruction::Binary { + op: MultiAD::LogAddExp, + left: 0, + right: 1, + }, + Instruction::Unary { + op: MultiAD::Sqrt, + arg: 4, + }, + Instruction::Unary { + op: MultiAD::Tan, + arg: 0, + }, + ]; + // Final output: use the Add node (index 4) as primary output. + let graph = CompiledGraph::new(2, instructions, vec![4]).unwrap(); + + // Compute with valid inputs. + let result = graph.compute(&[0.5, 2.0]).unwrap(); + assert!(approx_eq(result, 2.5, 1e-10)); + + // Gradient should also work. + let (val, grad) = graph.gradient(&[0.5, 2.0]).unwrap(); + assert!(approx_eq(val, 2.5, 1e-10)); + assert_eq!(grad.len(), 2); + // d/dx (x+y) = 1, d/dy (x+y) = 1 + assert!(approx_eq(grad[0], 1.0, 1e-10)); + assert!(approx_eq(grad[1], 1.0, 1e-10)); + } +} diff --git a/src/multi/graph.rs b/src/multi/graph.rs new file mode 100644 index 0000000..5b82ee6 --- /dev/null +++ b/src/multi/graph.rs @@ -0,0 +1,6315 @@ +//! Graph and tape APIs for reusable multi-variable autodiff. +//! +//! Unlike the legacy tuple-based representation, this API supports literal +//! constants and returns node handles as you build the graph. +//! +//! # Examples +//! +//! ``` +//! use petite_ad::Graph; +//! +//! let mut graph = Graph::new(2); +//! let x = graph.input(0); +//! let y = graph.input(1); +//! let two = graph.constant(2.0); +//! let xy = graph.mul(x, y); +//! let out = graph.add(xy, two); +//! +//! assert_eq!(out, 4); +//! +//! let value = graph.compute(&[3.0, 4.0]).unwrap(); +//! assert!((value - 14.0).abs() < 1e-10); +//! +//! // Explicitly choose a different output when needed. +//! graph.set_output(xy).unwrap(); +//! let product_only = graph.compute(&[3.0, 4.0]).unwrap(); +//! assert!((product_only - 12.0).abs() < 1e-10); +//! +//! // Or designate multiple outputs and get all values / a Jacobian. +//! graph.set_outputs(&[xy, out]).unwrap(); +//! let values = graph.compute_many(&[3.0, 4.0]).unwrap(); +//! assert_eq!(values.len(), 2); +//! let jacobian = graph.jacobian(&[3.0, 4.0]).unwrap(); +//! assert_eq!(jacobian.len(), 2); +//! ``` + +use std::{ + cell::RefCell, + fmt::Write, + ops::{Add, Div, Mul, Neg, Sub}, + rc::Rc, + sync::Arc, +}; + +use super::compiled::{ + BackendKind, BackendSupportReport, BatchGradients, BatchGradientsBuffer, BatchInputs, + BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata, CompiledWorkspace, + DeviceBatchPlan, DeviceBufferSet, DeviceExecutionTrace, Instruction, +}; +use super::multi_ad::MultiAD; +use super::multi_ad_fr::MultiAD2FR; +use super::multi_ad_rf::MultiAD2RF; +use super::op_rules; +use super::parser; +use super::types::BackwardResultBox; +use crate::{AutodiffError, Result}; + +/// Reusable scratch buffers for repeated [`Tape`] evaluation. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct TapeWorkspace { + values: Vec, + cotangent_values: Vec, + gradients: Vec, +} + +/// A handle to an input or computed node in a graph. +pub type NodeId = usize; + +/// A node in a reusable computation graph. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq)] +pub enum GraphNode { + /// A literal scalar value stored directly in the graph. + Constant(f64), + /// An operation node whose arguments reference prior node ids or inputs. + Operation { op: MultiAD, inputs: Vec }, +} + +/// A reusable multi-variable computation graph. +/// +/// Inputs occupy node ids `0..num_inputs`. Stored nodes are appended after the +/// inputs, so the first computed constant or operation gets id `num_inputs`. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq)] +pub struct Graph { + num_inputs: usize, + nodes: Vec, + output_nodes: Vec, + input_names: Vec>, + output_names: Vec<(NodeId, String)>, + #[cfg_attr(feature = "serde", serde(default))] + parameters: Vec, + #[cfg_attr(feature = "serde", serde(default))] + parameter_names: Vec<(NodeId, String)>, +} + +/// A reusable compiled view of a [`Graph`]. +/// +/// Today this stores the graph structure directly and provides a stable place +/// to grow future precomputation or buffer reuse without changing the public API. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct CompiledArgRange { + start: usize, + len: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum CompiledNode { + Constant(f64), + Operation { + op: MultiAD, + arg_range: CompiledArgRange, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum BackwardLocal { + Unary(f64), + Binary(f64, f64), +} + +#[derive(Debug, Clone, PartialEq)] +enum ExactLocal { + None, + Unary { + parent: NodeId, + dy: f64, + ddy: f64, + }, + Binary { + left: NodeId, + right: NodeId, + dy_left: f64, + dy_right: f64, + ddy_left_left: f64, + ddy_right_right: f64, + ddy_left_right: f64, + }, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Tape { + graph: Graph, + output_indices: Vec, + compiled_nodes: Arc<[CompiledNode]>, + arg_indices: Arc<[usize]>, +} + +/// One component comparison from [`Graph::check_gradient`]. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq)] +pub struct GradientCheckEntry { + /// Input coordinate being checked. + pub index: usize, + /// Reverse-mode gradient value. + pub autodiff: f64, + /// Central finite-difference gradient value. + pub finite_difference: f64, + /// Absolute difference between the two methods. + pub abs_error: f64, +} + +/// Deterministic report returned by [`Graph::check_gradient`]. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq)] +pub struct GradientCheckReport { + /// Whether every component is within tolerance. + pub passed: bool, + /// Requested absolute tolerance. + pub tolerance: f64, + /// Largest absolute error across all entries. + pub max_abs_error: f64, + /// Per-coordinate comparisons. + pub entries: Vec, +} + +/// Domain validation policy for graph evaluation. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DomainPolicy { + /// Preserve raw `f64` behavior. + Unchecked, + /// Reject invalid primal values such as `ln(x <= 0)`. + Checked, + /// Reject invalid primal values and derivative singularities such as `sqrt(0)`. + StrictDerivative, +} + +/// Basic graph statistics for diagnostics and benchmark setup. +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GraphStats { + pub num_inputs: usize, + pub num_constants: usize, + pub num_ops: usize, + pub num_edges: usize, + pub max_depth: usize, + pub op_counts: Vec<(MultiAD, usize)>, +} + +/// Shared expression graph used for operator-overloaded graph construction. +#[derive(Debug, Clone)] +pub struct ExprGraph { + graph: Rc>, +} + +/// A node handle tied to an [`ExprGraph`]. +#[derive(Debug, Clone)] +pub struct ExprNode { + graph: Rc>, + node: NodeId, +} + +#[inline] +fn op_name(op: MultiAD) -> &'static str { + op_rules::op_name(op) +} + +#[inline] +fn expected_arity(op: MultiAD) -> usize { + op_rules::expected_arity(op) +} + +#[inline] +fn compiled_arg_indices(arg_indices_storage: &[usize], arg_range: CompiledArgRange) -> &[usize] { + &arg_indices_storage[arg_range.start..arg_range.start + arg_range.len] +} + +#[inline] +fn with_compiled_arg_values( + arg_indices_storage: &[usize], + arg_range: CompiledArgRange, + values: &[f64], + f: F, +) -> Result +where + F: FnOnce(&[f64]) -> Result, +{ + let arg_indices = compiled_arg_indices(arg_indices_storage, arg_range); + match arg_indices.len() { + 0 => f(&[]), + 1 => { + MultiAD::check_value_index(arg_indices[0], values.len())?; + let args = [values[arg_indices[0]]]; + f(&args) + } + 2 => { + MultiAD::check_value_index(arg_indices[0], values.len())?; + MultiAD::check_value_index(arg_indices[1], values.len())?; + let args = [values[arg_indices[0]], values[arg_indices[1]]]; + f(&args) + } + _ => { + let arg_values = MultiAD::gather_arg_values(arg_indices, values)?; + f(&arg_values) + } + } +} + +#[inline] +fn backward_local(rule: op_rules::LocalRule) -> BackwardLocal { + match rule { + op_rules::LocalRule::Unary { dy, .. } => BackwardLocal::Unary(dy), + op_rules::LocalRule::Binary { + dy_left, dy_right, .. + } => BackwardLocal::Binary(dy_left, dy_right), + } +} + +fn reverse_accumulate_compiled( + num_inputs: usize, + compiled_nodes: &[CompiledNode], + arg_indices_storage: &[usize], + values: &[f64], + cotangent_values: &mut [f64], +) -> Result<()> { + for (offset, node) in compiled_nodes.iter().enumerate().rev() { + let node_id = num_inputs + offset; + let current_cotangent = cotangent_values[node_id]; + if current_cotangent == 0.0 { + continue; + } + + let CompiledNode::Operation { op, arg_range } = *node else { + continue; + }; + let input_indices = compiled_arg_indices(arg_indices_storage, arg_range); + let value = values[node_id]; + + with_compiled_arg_values(arg_indices_storage, arg_range, values, |args| { + match op_rules::local_rule(op, args, value)? { + op_rules::LocalRule::Unary { dy, .. } => { + cotangent_values[input_indices[0]] += current_cotangent * dy; + } + op_rules::LocalRule::Binary { + dy_left, dy_right, .. + } => { + cotangent_values[input_indices[0]] += current_cotangent * dy_left; + cotangent_values[input_indices[1]] += current_cotangent * dy_right; + } + } + Ok(()) + })?; + } + + Ok(()) +} + +impl TapeWorkspace { + /// Create an empty reusable workspace. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Clear all retained buffers while keeping their capacity. + pub fn clear(&mut self) { + self.values.clear(); + self.cotangent_values.clear(); + self.gradients.clear(); + } +} + +impl Graph { + #[inline] + fn output_name_for(&self, node_id: NodeId) -> Option<&str> { + self.output_names + .iter() + .find_map(|(id, name)| (*id == node_id).then_some(name.as_str())) + } + + #[inline] + fn node_label(&self, node_id: NodeId, node: &GraphNode) -> String { + let base = match node { + GraphNode::Constant(value) => format!("Const({value})"), + GraphNode::Operation { op, .. } => format!("{op:?}"), + }; + match self.output_name_for(node_id) { + Some(name) => format!("{name}: {base}"), + None => base, + } + } + + /// Create an empty graph with the given number of inputs. + #[must_use] + pub fn new(num_inputs: usize) -> Self { + Self { + num_inputs, + nodes: Vec::new(), + output_nodes: Vec::new(), + input_names: vec![None; num_inputs], + output_names: Vec::new(), + parameters: Vec::new(), + parameter_names: Vec::new(), + } + } + + /// Convert a legacy tuple graph into a reusable [`Graph`]. + /// + /// Input marker nodes are implicit in this API. Leading `Inp` markers are + /// skipped, while a trailing `Inp(k)` marker is preserved as an explicit + /// output selection for input `k`. + pub fn from_operations(num_inputs: usize, ops: &[(MultiAD, Vec)]) -> Self { + Self::try_from_operations(num_inputs, ops).unwrap_or_else(|_| { + let mut graph = Self::new(num_inputs); + for (op, inputs) in ops { + if *op != MultiAD::Inp { + graph.push_operation(*op, inputs.clone()); + } + } + graph + }) + } + + /// Convert a legacy tuple graph into a reusable [`Graph`] with validation. + pub fn try_from_operations(num_inputs: usize, ops: &[(MultiAD, Vec)]) -> Result { + let mut graph = Self::new(num_inputs); + let mut pending_input_output: Option = None; + + for (op, inputs) in ops { + if *op == MultiAD::Inp { + AutodiffError::check_arity("Inp", 1, inputs.len())?; + if inputs[0] >= num_inputs { + return Err(AutodiffError::IndexOutOfBounds { + index: inputs[0], + max_index: num_inputs.saturating_sub(1), + }); + } + pending_input_output = Some(inputs[0]); + continue; + } + + pending_input_output = None; + graph.try_push_operation(*op, inputs.clone())?; + } + + if let Some(output) = pending_input_output { + graph.set_output(output)?; + } + + Ok(graph) + } + + /// Parse a small expression into a graph. + /// + /// Supported syntax includes named inputs, numeric constants, `+ - * / ^`, + /// parentheses, unary minus, and common unary functions such as `sin`, + /// `exp`, `sqrt`, `tanh`, `relu`, `sigmoid`, `softplus`, `log1p_exp`, and `gelu`. + pub fn parse_expression(expression: &str, input_names: &[&str]) -> Result { + parser::parse_expression(expression, input_names) + } + + /// Return the number of graph inputs. + #[must_use] + pub fn num_inputs(&self) -> usize { + self.num_inputs + } + + /// Return the stored non-input nodes. + #[must_use] + pub fn nodes(&self) -> &[GraphNode] { + &self.nodes + } + + /// Return whether the graph has no stored nodes. + #[must_use] + pub fn is_empty(&self) -> bool { + self.nodes.is_empty() + } + + /// Return the number of stored non-input nodes. + #[must_use] + pub fn len(&self) -> usize { + self.nodes.len() + } + + /// Return the next allocated node id. + #[must_use] + pub fn next_node_id(&self) -> NodeId { + self.num_inputs + self.nodes.len() + } + + /// Return the explicitly selected primary output node, if one is configured. + #[must_use] + pub fn output_node(&self) -> Option { + self.output_nodes.first().copied() + } + + /// Return all explicitly selected output nodes. + #[must_use] + pub fn output_nodes(&self) -> &[NodeId] { + &self.output_nodes + } + + /// Return the effective primary output node used for single-output evaluation. + /// + /// If no explicit outputs have been selected, this falls back to the most + /// recently created node, or the last input for graphs without stored nodes. + #[must_use] + pub fn effective_output_node(&self) -> Option { + self.output_node() + .or_else(|| self.next_node_id().checked_sub(1)) + } + + /// Return the effective output nodes used for multi-output evaluation. + #[must_use] + pub fn effective_output_nodes(&self) -> Vec { + if self.output_nodes.is_empty() { + self.effective_output_node().into_iter().collect() + } else { + self.output_nodes.clone() + } + } + + /// Set a single output node used for evaluation. + pub fn set_output(&mut self, output: NodeId) -> Result<&mut Self> { + self.set_outputs(&[output]) + } + + /// Set multiple output nodes used for vector-output evaluation. + pub fn set_outputs(&mut self, outputs: &[NodeId]) -> Result<&mut Self> { + let max_index = + self.next_node_id() + .checked_sub(1) + .ok_or(AutodiffError::IndexOutOfBounds { + index: outputs.first().copied().unwrap_or(0), + max_index: 0, + })?; + for &output in outputs { + if output > max_index { + return Err(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + }); + } + } + self.output_nodes = outputs.to_vec(); + Ok(self) + } + + /// Add one more output node to the vector-output selection. + pub fn add_output(&mut self, output: NodeId) -> Result<&mut Self> { + let max_index = + self.next_node_id() + .checked_sub(1) + .ok_or(AutodiffError::IndexOutOfBounds { + index: output, + max_index: 0, + })?; + if output > max_index { + return Err(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + }); + } + self.output_nodes.push(output); + Ok(self) + } + + /// Clear any explicit output nodes, restoring implicit last-node behavior. + pub fn clear_output(&mut self) -> &mut Self { + self.output_nodes.clear(); + self + } + + /// Get an input node id. + #[must_use] + pub fn input(&self, input_index: usize) -> NodeId { + input_index + } + + /// Append a literal constant node and return its node id. + pub fn constant(&mut self, value: f64) -> NodeId { + let node_id = self.next_node_id(); + self.nodes.push(GraphNode::Constant(value)); + node_id + } + + /// Append a custom operation node and return its node id. + /// + /// This is a low-level unchecked API. Prefer typed helpers such as + /// [`Graph::sin`] and [`Graph::add`] when possible. Passing `MultiAD::Inp` + /// creates an invalid reusable graph because input nodes are implicit in + /// this API and should be obtained with [`Graph::input`]. + pub fn push_operation(&mut self, op: MultiAD, inputs: Vec) -> NodeId { + let node_id = self.next_node_id(); + self.nodes.push(GraphNode::Operation { op, inputs }); + node_id + } + + /// Append a custom operation after validating arity and input references. + pub fn try_push_operation(&mut self, op: MultiAD, inputs: Vec) -> Result { + if op == MultiAD::Inp { + return Err(AutodiffError::InvalidGraph { + reason: "Graph nodes must not contain input markers", + }); + } + AutodiffError::check_arity(op_name(op), expected_arity(op), inputs.len())?; + let max_index = self.next_node_id().saturating_sub(1); + for &input in &inputs { + if input >= self.next_node_id() { + return Err(AutodiffError::IndexOutOfBounds { + index: input, + max_index, + }); + } + } + Ok(self.push_operation(op, inputs)) + } + + pub fn try_sin(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Sin, vec![arg]) + } + + pub fn try_cos(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Cos, vec![arg]) + } + + pub fn try_tan(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Tan, vec![arg]) + } + + pub fn try_tanh(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Tanh, vec![arg]) + } + + pub fn try_relu(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Relu, vec![arg]) + } + + pub fn try_log1p_exp(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Log1pExp, vec![arg]) + } + + pub fn try_neg(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Neg, vec![arg]) + } + + pub fn try_exp(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Exp, vec![arg]) + } + + pub fn try_ln(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Ln, vec![arg]) + } + + pub fn try_sqrt(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Sqrt, vec![arg]) + } + + pub fn try_abs(&mut self, arg: NodeId) -> Result { + self.try_push_operation(MultiAD::Abs, vec![arg]) + } + + pub fn try_add(&mut self, left: NodeId, right: NodeId) -> Result { + self.try_push_operation(MultiAD::Add, vec![left, right]) + } + + pub fn try_sub(&mut self, left: NodeId, right: NodeId) -> Result { + self.try_push_operation(MultiAD::Sub, vec![left, right]) + } + + pub fn try_mul(&mut self, left: NodeId, right: NodeId) -> Result { + self.try_push_operation(MultiAD::Mul, vec![left, right]) + } + + pub fn try_div(&mut self, left: NodeId, right: NodeId) -> Result { + self.try_push_operation(MultiAD::Div, vec![left, right]) + } + + pub fn try_pow(&mut self, base: NodeId, exp: NodeId) -> Result { + self.try_push_operation(MultiAD::Pow, vec![base, exp]) + } + + pub fn try_log_add_exp(&mut self, left: NodeId, right: NodeId) -> Result { + self.try_push_operation(MultiAD::LogAddExp, vec![left, right]) + } + + /// Append a sine node. + pub fn sin(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Sin, vec![arg]) + } + + /// Append a cosine node. + pub fn cos(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Cos, vec![arg]) + } + + /// Append a tangent node. + pub fn tan(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Tan, vec![arg]) + } + + /// Append a hyperbolic tangent node. + pub fn tanh(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Tanh, vec![arg]) + } + + /// Append a rectified-linear-unit node. + pub fn relu(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Relu, vec![arg]) + } + + /// Append a stable `ln(1 + exp(arg))` node. + pub fn log1p_exp_node(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Log1pExp, vec![arg]) + } + + /// Append a negation node. + pub fn neg(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Neg, vec![arg]) + } + + /// Append an exponential node. + pub fn exp(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Exp, vec![arg]) + } + + /// Append a natural logarithm node. + pub fn ln(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Ln, vec![arg]) + } + + /// Append a square-root node. + pub fn sqrt(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Sqrt, vec![arg]) + } + + /// Append an absolute-value node. + pub fn abs(&mut self, arg: NodeId) -> NodeId { + self.push_operation(MultiAD::Abs, vec![arg]) + } + + /// Append an addition node. + pub fn add(&mut self, left: NodeId, right: NodeId) -> NodeId { + self.push_operation(MultiAD::Add, vec![left, right]) + } + + /// Append a subtraction node. + pub fn sub(&mut self, left: NodeId, right: NodeId) -> NodeId { + self.push_operation(MultiAD::Sub, vec![left, right]) + } + + /// Append a multiplication node. + pub fn mul(&mut self, left: NodeId, right: NodeId) -> NodeId { + self.push_operation(MultiAD::Mul, vec![left, right]) + } + + /// Append a division node. + pub fn div(&mut self, left: NodeId, right: NodeId) -> NodeId { + self.push_operation(MultiAD::Div, vec![left, right]) + } + + /// Append a power node. + pub fn pow(&mut self, base: NodeId, exp: NodeId) -> NodeId { + self.push_operation(MultiAD::Pow, vec![base, exp]) + } + + /// Append a stable binary `ln(exp(left) + exp(right))` node. + pub fn log_add_exp(&mut self, left: NodeId, right: NodeId) -> NodeId { + self.push_operation(MultiAD::LogAddExp, vec![left, right]) + } + + /// Append `node + constant`. + pub fn add_const(&mut self, node: NodeId, value: f64) -> NodeId { + let constant = self.constant(value); + self.add(node, constant) + } + + /// Append `node - constant`. + pub fn sub_const(&mut self, node: NodeId, value: f64) -> NodeId { + let constant = self.constant(value); + self.sub(node, constant) + } + + /// Append `node * constant`. + pub fn mul_const(&mut self, node: NodeId, value: f64) -> NodeId { + let constant = self.constant(value); + self.mul(node, constant) + } + + /// Append `node / constant`. + pub fn div_const(&mut self, node: NodeId, value: f64) -> NodeId { + let constant = self.constant(value); + self.div(node, constant) + } + + /// Append `node.powf(constant)`. + pub fn pow_const(&mut self, node: NodeId, value: f64) -> NodeId { + let constant = self.constant(value); + self.pow(node, constant) + } + + /// Append `node * node`. + pub fn square(&mut self, node: NodeId) -> NodeId { + self.mul(node, node) + } + + /// Append `node * node * node`. + pub fn cube(&mut self, node: NodeId) -> NodeId { + let square = self.square(node); + self.mul(square, node) + } + + /// Append `1 / node`. + pub fn reciprocal(&mut self, node: NodeId) -> NodeId { + let one = self.constant(1.0); + self.div(one, node) + } + + /// Append the logistic sigmoid `1 / (1 + exp(-node))`. + pub fn sigmoid(&mut self, node: NodeId) -> NodeId { + let neg = self.neg(node); + let exp_neg = self.exp(neg); + let denom = self.add_const(exp_neg, 1.0); + self.reciprocal(denom) + } + + /// Append stable `ln(1 + exp(node))`. + pub fn softplus(&mut self, node: NodeId) -> NodeId { + self.log1p_exp(node) + } + + /// Append the tanh-based GELU approximation. + pub fn gelu(&mut self, node: NodeId) -> NodeId { + let x_cubed = self.cube(node); + let inner_cubic = self.mul_const(x_cubed, 0.044_715); + let inner_sum = self.add(node, inner_cubic); + let scaled_inner = self.mul_const(inner_sum, (2.0 / std::f64::consts::PI).sqrt()); + let tanh = self.tanh(scaled_inner); + let one_plus = self.add_const(tanh, 1.0); + let half_x = self.mul_const(node, 0.5); + self.mul(half_x, one_plus) + } + + /// Append softmax nodes for a vector of logits. + pub fn softmax(&mut self, logits: &[NodeId]) -> Vec { + let exp_nodes: Vec = logits.iter().map(|&node| self.exp(node)).collect(); + let Some((&first, rest)) = exp_nodes.split_first() else { + return Vec::new(); + }; + let mut denom = first; + for &node in rest { + denom = self.add(denom, node); + } + exp_nodes + .into_iter() + .map(|node| self.div(node, denom)) + .collect() + } + + /// Sum a slice of nodes. + pub fn sum(&mut self, nodes: &[NodeId]) -> Option { + let (&first, rest) = nodes.split_first()?; + let mut acc = first; + for &node in rest { + acc = self.add(acc, node); + } + Some(acc) + } + + /// Mean of a slice of nodes. + pub fn mean(&mut self, nodes: &[NodeId]) -> Option { + let sum = self.sum(nodes)?; + Some(self.div_const(sum, nodes.len() as f64)) + } + + /// Dot product of two equal-length node slices. + pub fn dot(&mut self, left: &[NodeId], right: &[NodeId]) -> Result { + if left.len() != right.len() || left.is_empty() { + return Err(AutodiffError::InvalidGraph { + reason: "dot inputs must be non-empty and have equal lengths", + }); + } + let products: Vec = left + .iter() + .zip(right.iter()) + .map(|(&lhs, &rhs)| self.mul(lhs, rhs)) + .collect(); + self.sum(&products).ok_or(AutodiffError::InvalidGraph { + reason: "dot product could not be reduced", + }) + } + + /// Sum of squares of a slice of nodes. + pub fn sum_squares(&mut self, nodes: &[NodeId]) -> Option { + let squares: Vec = nodes.iter().map(|&node| self.square(node)).collect(); + self.sum(&squares) + } + + /// Euclidean norm approximation `sqrt(sum_squares(nodes))`. + pub fn norm2(&mut self, nodes: &[NodeId]) -> Option { + let sum_squares = self.sum_squares(nodes)?; + Some(self.sqrt(sum_squares)) + } + + /// Stable sigmoid helper using `0.5 * (tanh(0.5 * x) + 1)`. + pub fn sigmoid_stable(&mut self, node: NodeId) -> NodeId { + let half = self.mul_const(node, 0.5); + let tanh = self.tanh(half); + let shifted = self.add_const(tanh, 1.0); + self.mul_const(shifted, 0.5) + } + + /// Stable `log(1 + exp(x))` helper. + pub fn log1p_exp(&mut self, node: NodeId) -> NodeId { + self.log1p_exp_node(node) + } + + /// Overflow-safe pairwise log-sum-exp approximation. + pub fn logsumexp_approx(&mut self, nodes: &[NodeId]) -> Option { + let (&first, rest) = nodes.split_first()?; + let mut acc = first; + for &node in rest { + acc = self.log_add_exp(acc, node); + } + Some(acc) + } + + /// Overflow-safe softmax approximation using pairwise log-sum-exp. + pub fn stable_softmax_approx(&mut self, logits: &[NodeId]) -> Vec { + let Some(log_denom) = self.logsumexp_approx(logits) else { + return Vec::new(); + }; + logits + .iter() + .map(|&logit| { + let centered = self.sub(logit, log_denom); + self.exp(centered) + }) + .collect() + } + + /// Mean squared error loss. + pub fn mse_loss(&mut self, predictions: &[NodeId], targets: &[NodeId]) -> Result { + if predictions.len() != targets.len() || predictions.is_empty() { + return Err(AutodiffError::InvalidGraph { + reason: "mse_loss inputs must be non-empty and have equal lengths", + }); + } + let squares: Vec = predictions + .iter() + .zip(targets.iter()) + .map(|(&prediction, &target)| { + let diff = self.sub(prediction, target); + self.square(diff) + }) + .collect(); + self.mean(&squares).ok_or(AutodiffError::InvalidGraph { + reason: "mse_loss could not be reduced", + }) + } + + /// Binary cross-entropy loss for probability predictions. + pub fn binary_cross_entropy_loss( + &mut self, + probabilities: &[NodeId], + targets: &[NodeId], + ) -> Result { + if probabilities.len() != targets.len() || probabilities.is_empty() { + return Err(AutodiffError::InvalidGraph { + reason: "binary_cross_entropy_loss inputs must be non-empty and have equal lengths", + }); + } + let mut terms = Vec::with_capacity(probabilities.len()); + for (&probability, &target) in probabilities.iter().zip(targets.iter()) { + let log_p = self.ln(probability); + let one_minus_target = self.sub_const(target, 1.0); + let neg_one_minus_target = self.neg(one_minus_target); + let one_minus_p = self.sub_const(probability, 1.0); + let neg_one_minus_p = self.neg(one_minus_p); + let log_one_minus_p = self.ln(neg_one_minus_p); + let positive_term = self.mul(target, log_p); + let negative_term = self.mul(neg_one_minus_target, log_one_minus_p); + let sum = self.add(positive_term, negative_term); + terms.push(self.neg(sum)); + } + self.mean(&terms).ok_or(AutodiffError::InvalidGraph { + reason: "binary_cross_entropy_loss could not be reduced", + }) + } + + /// Convert this reusable graph into the legacy tuple-based operation form. + /// + /// This succeeds only when the selected output and all nodes needed to reach it + /// can be represented without inline constants. + pub fn to_operations(&self) -> Result)>> { + let effective_output = self.effective_output_node(); + if self.output_nodes.len() > 1 { + return Err(AutodiffError::InvalidGraph { + reason: "legacy tuple graphs support only one output", + }); + } + let last_included_node = + effective_output.and_then(|node| node.checked_sub(self.num_inputs)); + let node_limit = last_included_node.map(|offset| offset + 1).unwrap_or(0); + + let mut ops: Vec<(MultiAD, Vec)> = Vec::with_capacity(node_limit); + for node in self.nodes.iter().take(node_limit) { + match node { + GraphNode::Constant(_) => { + return Err(AutodiffError::InvalidGraph { + reason: "legacy tuple graphs cannot represent constant nodes", + }); + } + GraphNode::Operation { op, inputs } => ops.push((*op, inputs.clone())), + } + } + + if let Some(output) = effective_output { + if output < self.num_inputs { + ops.push((MultiAD::Inp, vec![output])); + } + } + + Ok(ops) + } + + /// Validate graph structure and references. + pub fn validate(&self) -> Result<()> { + let mut next_valid_id = self.num_inputs; + + for node in &self.nodes { + match node { + GraphNode::Constant(_) => { + next_valid_id += 1; + } + GraphNode::Operation { op, inputs } => { + if *op == MultiAD::Inp { + return Err(AutodiffError::InvalidGraph { + reason: "Graph nodes must not contain input markers", + }); + } + AutodiffError::check_arity(op_name(*op), expected_arity(*op), inputs.len())?; + for &input in inputs { + if input >= next_valid_id { + return Err(AutodiffError::IndexOutOfBounds { + index: input, + max_index: next_valid_id.saturating_sub(1), + }); + } + } + next_valid_id += 1; + } + } + } + + let max_index = next_valid_id.saturating_sub(1); + for &output in &self.output_nodes { + if output >= next_valid_id { + return Err(AutodiffError::IndexOutOfBounds { + index: output, + max_index, + }); + } + } + for ¶meter in &self.parameters { + if parameter >= next_valid_id { + return Err(AutodiffError::IndexOutOfBounds { + index: parameter, + max_index, + }); + } + } + + Ok(()) + } + + /// Export the graph as Mermaid flowchart syntax. + #[must_use] + pub fn to_mermaid(&self) -> String { + let mut out = String::from("flowchart LR\n"); + for input_idx in 0..self.num_inputs { + let label = self + .input_name(input_idx) + .map(|name| format!("{name}: Input {input_idx}")) + .unwrap_or_else(|| format!("Input {input_idx}")); + let _ = writeln!(&mut out, " n{input_idx}[\"{label}\"]"); + } + for (offset, node) in self.nodes.iter().enumerate() { + let node_id = self.num_inputs + offset; + let label = if self.output_nodes.contains(&node_id) { + format!("{} [output]", self.node_label(node_id, node)) + } else { + self.node_label(node_id, node) + }; + let _ = writeln!(&mut out, " n{node_id}[\"{label}\"]"); + if let GraphNode::Operation { inputs, .. } = node { + for &input in inputs { + let _ = writeln!(&mut out, " n{input} --> n{node_id}"); + } + } + } + out + } + + /// Export the graph as Graphviz DOT syntax. + #[must_use] + pub fn to_dot(&self) -> String { + let mut out = String::from("digraph Graph {\n rankdir=LR;\n"); + for input_idx in 0..self.num_inputs { + let label = self + .input_name(input_idx) + .map(|name| format!("{name}: Input {input_idx}")) + .unwrap_or_else(|| format!("Input {input_idx}")); + let _ = writeln!(&mut out, " n{input_idx} [label=\"{label}\", shape=box];"); + } + for (offset, node) in self.nodes.iter().enumerate() { + let node_id = self.num_inputs + offset; + let shape = if self.output_nodes.contains(&node_id) { + "doublecircle" + } else { + match node { + GraphNode::Constant(_) => "ellipse", + GraphNode::Operation { .. } => "oval", + } + }; + let label = if self.output_nodes.contains(&node_id) { + format!("{} [output]", self.node_label(node_id, node)) + } else { + self.node_label(node_id, node) + }; + let _ = writeln!( + &mut out, + " n{node_id} [label=\"{label}\", shape={shape}];" + ); + if let GraphNode::Operation { inputs, .. } = node { + for &input in inputs { + let _ = writeln!(&mut out, " n{input} -> n{node_id};"); + } + } + } + out.push_str("}\n"); + out + } + + /// Set a human-readable input name used by graph exporters. + pub fn set_input_name( + &mut self, + input_index: usize, + name: impl Into, + ) -> Result<&mut Self> { + if input_index >= self.num_inputs { + return Err(AutodiffError::IndexOutOfBounds { + index: input_index, + max_index: self.num_inputs.saturating_sub(1), + }); + } + self.input_names[input_index] = Some(name.into()); + Ok(self) + } + + /// Set a human-readable output name used by graph exporters. + pub fn set_output_name( + &mut self, + output: NodeId, + name: impl Into, + ) -> Result<&mut Self> { + self.set_output(output)?; + self.output_names.retain(|(node_id, _)| *node_id != output); + self.output_names.push((output, name.into())); + Ok(self) + } + + /// Return the configured input name, if present. + #[must_use] + pub fn input_name(&self, input_index: usize) -> Option<&str> { + self.input_names + .get(input_index) + .and_then(|name| name.as_deref()) + } + + /// Mark a node as a parameter for optimizer-oriented workflows. + pub fn mark_parameter(&mut self, node: NodeId) -> Result<&mut Self> { + let next_node_id = self.next_node_id(); + let max_index = next_node_id.saturating_sub(1); + if next_node_id == 0 || node >= next_node_id { + return Err(AutodiffError::IndexOutOfBounds { + index: node, + max_index, + }); + } + if !self.parameters.contains(&node) { + self.parameters.push(node); + } + Ok(self) + } + + /// Set a human-readable parameter name. + pub fn set_parameter_name( + &mut self, + node: NodeId, + name: impl Into, + ) -> Result<&mut Self> { + self.mark_parameter(node)?; + self.parameter_names.retain(|(id, _)| *id != node); + self.parameter_names.push((node, name.into())); + Ok(self) + } + + /// Return parameter nodes. + #[must_use] + pub fn parameters(&self) -> &[NodeId] { + &self.parameters + } + + /// Return the configured parameter name, if present. + #[must_use] + pub fn parameter_name(&self, node: NodeId) -> Option<&str> { + self.parameter_names + .iter() + .find(|(id, _)| *id == node) + .map(|(_, name)| name.as_str()) + } + + /// Return all configured parameter names. + #[must_use] + pub fn parameter_names(&self) -> &[(NodeId, String)] { + &self.parameter_names + } + + /// Extract gradients for parameter nodes that are graph inputs. + pub fn parameter_gradient(&self, inputs: &[f64]) -> Result> { + let (_value, gradient) = self.gradient(inputs)?; + let mut result = Vec::with_capacity(self.parameters.len()); + for ¶meter in &self.parameters { + if parameter >= self.num_inputs { + return Err(AutodiffError::InvalidGraph { + reason: "parameter_gradient currently supports input parameters only", + }); + } + result.push((parameter, gradient[parameter])); + } + Ok(result) + } + + /// Return diagnostic graph statistics. + #[must_use] + pub fn stats(&self) -> GraphStats { + let mut num_constants = 0; + let mut num_ops = 0; + let mut num_edges = 0; + let mut depths = vec![0usize; self.next_node_id()]; + let mut op_counts: Vec<(MultiAD, usize)> = Vec::new(); + + for (offset, node) in self.nodes.iter().enumerate() { + let node_id = self.num_inputs + offset; + match node { + GraphNode::Constant(_) => { + num_constants += 1; + depths[node_id] = 0; + } + GraphNode::Operation { op, inputs } => { + num_ops += 1; + num_edges += inputs.len(); + let parent_depth = inputs.iter().map(|&id| depths[id]).max().unwrap_or(0); + depths[node_id] = parent_depth + 1; + if let Some((_, count)) = op_counts.iter_mut().find(|(kind, _)| kind == op) { + *count += 1; + } else { + op_counts.push((*op, 1)); + } + } + } + } + + GraphStats { + num_inputs: self.num_inputs, + num_constants, + num_ops, + num_edges, + max_depth: depths.into_iter().max().unwrap_or(0), + op_counts, + } + } + + /// Return a graph with simple local simplifications applied. + /// + /// Currently performs constant folding and a few algebraic identities with + /// literal `0` and `1` constants. + pub fn simplify(&self) -> Result { + self.validate()?; + let mut simplified = Graph::new(self.num_inputs); + simplified.input_names = self.input_names.clone(); + let mut old_to_new: Vec = (0..self.num_inputs).collect(); + let mut constant_values: Vec> = vec![None; self.num_inputs]; + + for node in &self.nodes { + let new_node = match node { + GraphNode::Constant(value) => { + let id = simplified.constant(*value); + constant_values.push(Some(*value)); + id + } + GraphNode::Operation { op, inputs } => { + let mapped_inputs: Vec = + inputs.iter().map(|&id| old_to_new[id]).collect(); + let input_constants: Vec> = inputs + .iter() + .map(|&id| constant_values.get(id).copied().flatten()) + .collect(); + + let folded = if input_constants.iter().all(Option::is_some) { + let args: Vec = + input_constants.iter().map(|value| value.unwrap()).collect(); + Some(op_rules::forward_value(*op, &args)?) + } else { + None + }; + + if let Some(value) = folded { + let id = simplified.constant(value); + constant_values.push(Some(value)); + id + } else if *op == MultiAD::Add && input_constants.first() == Some(&Some(0.0)) { + constant_values.push(None); + mapped_inputs[1] + } else if *op == MultiAD::Add && input_constants.get(1) == Some(&Some(0.0)) { + constant_values.push(None); + mapped_inputs[0] + } else if *op == MultiAD::Mul && input_constants.first() == Some(&Some(1.0)) { + constant_values.push(None); + mapped_inputs[1] + } else if *op == MultiAD::Mul && input_constants.get(1) == Some(&Some(1.0)) { + constant_values.push(None); + mapped_inputs[0] + } else { + let id = simplified.push_operation(*op, mapped_inputs); + constant_values.push(None); + id + } + } + }; + old_to_new.push(new_node); + } + + let outputs: Vec = self + .effective_output_nodes() + .into_iter() + .map(|id| old_to_new[id]) + .collect(); + simplified.set_outputs(&outputs)?; + for (old_id, name) in &self.output_names { + simplified + .output_names + .push((old_to_new[*old_id], name.clone())); + } + simplified.parameters = self.parameters.iter().map(|&id| old_to_new[id]).collect(); + for (old_id, name) in &self.parameter_names { + simplified + .parameter_names + .push((old_to_new[*old_id], name.clone())); + } + Ok(simplified) + } + + /// Return a new graph containing only nodes reachable from selected outputs. + pub fn prune_to_outputs(&self) -> Result { + self.validate()?; + let outputs = self.effective_output_nodes(); + let mut reachable = vec![false; self.nodes.len()]; + let mut stack = outputs.clone(); + + while let Some(node_id) = stack.pop() { + if node_id < self.num_inputs { + continue; + } + let offset = node_id - self.num_inputs; + if offset >= self.nodes.len() || reachable[offset] { + continue; + } + reachable[offset] = true; + if let GraphNode::Operation { inputs, .. } = &self.nodes[offset] { + stack.extend(inputs.iter().copied()); + } + } + + let mut pruned = Graph::new(self.num_inputs); + pruned.input_names = self.input_names.clone(); + let mut old_to_new: Vec> = vec![None; self.next_node_id()]; + for (input_id, slot) in old_to_new.iter_mut().enumerate().take(self.num_inputs) { + *slot = Some(input_id); + } + + for (offset, node) in self.nodes.iter().enumerate() { + if !reachable[offset] { + continue; + } + let old_id = self.num_inputs + offset; + let new_id = match node { + GraphNode::Constant(value) => pruned.constant(*value), + GraphNode::Operation { op, inputs } => { + let remapped_inputs: Vec = inputs + .iter() + .map(|&id| { + old_to_new[id].ok_or(AutodiffError::InvalidGraph { + reason: "reachable graph remapping failed", + }) + }) + .collect::>>()?; + pruned.push_operation(*op, remapped_inputs) + } + }; + old_to_new[old_id] = Some(new_id); + } + + let remapped_outputs: Vec = outputs + .iter() + .map(|&id| { + old_to_new[id].ok_or(AutodiffError::InvalidGraph { + reason: "output remapping failed", + }) + }) + .collect::>>()?; + pruned.set_outputs(&remapped_outputs)?; + for (old_id, name) in &self.output_names { + if let Some(Some(new_id)) = old_to_new.get(*old_id) { + pruned.output_names.push((*new_id, name.clone())); + } + } + for &old_id in &self.parameters { + if let Some(Some(new_id)) = old_to_new.get(old_id) { + pruned.parameters.push(*new_id); + } + } + for (old_id, name) in &self.parameter_names { + if let Some(Some(new_id)) = old_to_new.get(*old_id) { + pruned.parameter_names.push((*new_id, name.clone())); + } + } + Ok(pruned) + } + + /// Compile the graph into closure-free instruction IR. + pub fn compile_ir(&self) -> Result { + self.validate()?; + let mut instructions = Vec::with_capacity(self.nodes.len()); + for node in &self.nodes { + let instruction = match node { + GraphNode::Constant(value) => Instruction::Constant(*value), + GraphNode::Operation { op, inputs } => match inputs.as_slice() { + [arg] => Instruction::Unary { op: *op, arg: *arg }, + [left, right] => Instruction::Binary { + op: *op, + left: *left, + right: *right, + }, + _ => { + return Err(AutodiffError::InvalidGraph { + reason: "compiled IR supports only unary and binary operations", + }); + } + }, + }; + instructions.push(instruction); + } + CompiledGraph::new(self.num_inputs, instructions, self.effective_output_nodes()) + } + + /// Alias for [`Graph::compile_ir`] using acceleration-oriented naming. + pub fn compile_accelerated(&self) -> Result { + self.compile_ir() + } + + /// Compile the graph into a reusable tape. + #[must_use] + pub fn compile(&self) -> Tape { + let mut compiled_nodes: Vec = Vec::with_capacity(self.nodes.len()); + let mut arg_indices: Vec = Vec::new(); + + for node in &self.nodes { + match node { + GraphNode::Constant(value) => compiled_nodes.push(CompiledNode::Constant(*value)), + GraphNode::Operation { op, inputs } => { + let start = arg_indices.len(); + arg_indices.extend_from_slice(inputs); + compiled_nodes.push(CompiledNode::Operation { + op: *op, + arg_range: CompiledArgRange { + start, + len: inputs.len(), + }, + }); + } + } + } + + Tape { + graph: self.clone(), + output_indices: self.effective_output_nodes(), + compiled_nodes: Arc::from(compiled_nodes), + arg_indices: Arc::from(arg_indices), + } + } + + /// Compile the graph after validating its structure. + pub fn try_compile(&self) -> Result { + self.validate()?; + Ok(self.compile()) + } + + #[inline] + fn check_graph_input_len(&self, inputs: &[f64]) -> Result<()> { + if inputs.len() == self.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "input length must match graph.num_inputs()", + }) + } + } + + fn emit_fr_ops( + &self, + node_id: NodeId, + ops: &mut Vec, + values: &mut Vec, + ) -> Result<()> { + if node_id < self.num_inputs { + ops.push(MultiAD2FR::Inp(node_id)); + return Ok(()); + } + let offset = node_id - self.num_inputs; + let node = self + .nodes + .get(offset) + .ok_or(AutodiffError::IndexOutOfBounds { + index: node_id, + max_index: self.next_node_id().saturating_sub(1), + })?; + match node { + GraphNode::Constant(value) => { + let input_id = self.num_inputs + values.len(); + values.push(*value); + ops.push(MultiAD2FR::Inp(input_id)); + } + GraphNode::Operation { op, inputs } => { + for &input in inputs { + self.emit_fr_ops(input, ops, values)?; + } + ops.push(Self::multi_to_fr(*op)?); + } + } + Ok(()) + } + + fn emit_rf_ops( + &self, + node_id: NodeId, + ops: &mut Vec, + values: &mut Vec, + ) -> Result<()> { + if node_id < self.num_inputs { + ops.push(MultiAD2RF::Inp(node_id)); + return Ok(()); + } + let offset = node_id - self.num_inputs; + let node = self + .nodes + .get(offset) + .ok_or(AutodiffError::IndexOutOfBounds { + index: node_id, + max_index: self.next_node_id().saturating_sub(1), + })?; + match node { + GraphNode::Constant(value) => { + let input_id = self.num_inputs + values.len(); + values.push(*value); + ops.push(MultiAD2RF::Inp(input_id)); + } + GraphNode::Operation { op, inputs } => { + for &input in inputs { + self.emit_rf_ops(input, ops, values)?; + } + ops.push(Self::multi_to_rf(*op)?); + } + } + Ok(()) + } + + fn multi_to_fr(op: MultiAD) -> Result { + Ok(match op { + MultiAD::Sin => MultiAD2FR::Sin, + MultiAD::Cos => MultiAD2FR::Cos, + MultiAD::Tan => MultiAD2FR::Tan, + MultiAD::Neg => MultiAD2FR::Neg, + MultiAD::Exp => MultiAD2FR::Exp, + MultiAD::Ln => MultiAD2FR::Ln, + MultiAD::Sqrt => MultiAD2FR::Sqrt, + MultiAD::Log1pExp => MultiAD2FR::Log1pExp, + MultiAD::Add => MultiAD2FR::Add, + MultiAD::Sub => MultiAD2FR::Sub, + MultiAD::Mul => MultiAD2FR::Mul, + MultiAD::Div => MultiAD2FR::Div, + MultiAD::Pow => MultiAD2FR::Pow, + MultiAD::Tanh | MultiAD::Relu | MultiAD::LogAddExp | MultiAD::Abs | MultiAD::Inp => { + return Err(AutodiffError::InvalidGraph { + reason: + "exact graph Hessian supports only the exact-Hessian smooth operation set", + }); + } + }) + } + + fn multi_to_rf(op: MultiAD) -> Result { + Ok(match op { + MultiAD::Sin => MultiAD2RF::Sin, + MultiAD::Cos => MultiAD2RF::Cos, + MultiAD::Tan => MultiAD2RF::Tan, + MultiAD::Neg => MultiAD2RF::Neg, + MultiAD::Exp => MultiAD2RF::Exp, + MultiAD::Ln => MultiAD2RF::Ln, + MultiAD::Sqrt => MultiAD2RF::Sqrt, + MultiAD::Log1pExp => MultiAD2RF::Log1pExp, + MultiAD::Add => MultiAD2RF::Add, + MultiAD::Sub => MultiAD2RF::Sub, + MultiAD::Mul => MultiAD2RF::Mul, + MultiAD::Div => MultiAD2RF::Div, + MultiAD::Pow => MultiAD2RF::Pow, + MultiAD::Tanh | MultiAD::Relu | MultiAD::LogAddExp | MultiAD::Abs | MultiAD::Inp => { + return Err(AutodiffError::InvalidGraph { + reason: + "exact graph Hessian supports only the exact-Hessian smooth operation set", + }); + } + }) + } + + fn crop_hessian(&self, hessian: Vec>) -> Vec> { + hessian + .into_iter() + .take(self.num_inputs) + .map(|row| row.into_iter().take(self.num_inputs).collect()) + .collect() + } + + fn validate_strict_derivative_domain(&self, inputs: &[f64]) -> Result<()> { + self.check_graph_input_len(inputs)?; + let mut values = inputs.to_vec(); + for node in &self.nodes { + match node { + GraphNode::Constant(value) => values.push(*value), + GraphNode::Operation { + op, + inputs: arg_indices, + } => { + let args = MultiAD::gather_arg_values(arg_indices, &values)?; + match op { + MultiAD::Sqrt if args[0] <= 0.0 => { + return Err(AutodiffError::domain( + "Sqrt", + "input must be positive for derivative evaluation", + )); + } + _ => {} + } + values.push(op.forward_checked(&args)?); + } + } + } + Ok(()) + } + + fn compute_exact_hessian_native(&self, inputs: &[f64]) -> Result>> { + self.check_graph_input_len(inputs)?; + self.validate()?; + let output = self + .effective_output_node() + .ok_or(AutodiffError::EmptyGraph)?; + if self.nodes.iter().any(|node| { + matches!( + node, + GraphNode::Operation { + op: MultiAD::Abs | MultiAD::Relu, + .. + } + ) + }) { + return Err(AutodiffError::InvalidGraph { + reason: "native exact Hessian does not support non-smooth operations", + }); + } + + let value_count = self.next_node_id(); + let mut values = Vec::with_capacity(value_count); + values.extend_from_slice(inputs); + + let mut grads = vec![vec![0.0; self.num_inputs]; value_count]; + for (input_id, grad) in grads.iter_mut().enumerate().take(self.num_inputs) { + grad[input_id] = 1.0; + } + + let mut locals = vec![ExactLocal::None; value_count]; + + for (offset, node) in self.nodes.iter().enumerate() { + let node_id = self.num_inputs + offset; + match node { + GraphNode::Constant(value) => { + values.push(*value); + } + GraphNode::Operation { op, inputs } => { + let arg_values = MultiAD::gather_arg_values(inputs, &values)?; + let value = op_rules::forward_value(*op, &arg_values)?; + let local_rule = op_rules::local_rule(*op, &arg_values, value)?; + values.push(value); + match local_rule { + op_rules::LocalRule::Unary { dy, ddy } => { + let parent = inputs[0]; + let parent_grad = grads[parent].clone(); + for (var, slot) in + grads[node_id].iter_mut().enumerate().take(self.num_inputs) + { + *slot = dy * parent_grad[var]; + } + locals[node_id] = ExactLocal::Unary { parent, dy, ddy }; + } + op_rules::LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + let left = inputs[0]; + let right = inputs[1]; + let left_grad = grads[left].clone(); + let right_grad = grads[right].clone(); + for (var, slot) in + grads[node_id].iter_mut().enumerate().take(self.num_inputs) + { + *slot = dy_left * left_grad[var] + dy_right * right_grad[var]; + } + locals[node_id] = ExactLocal::Binary { + left, + right, + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + }; + } + } + } + } + } + + let mut adjoints = vec![0.0; value_count]; + let mut hessian = vec![vec![0.0; self.num_inputs]; self.num_inputs]; + adjoints[output] = 1.0; + + for node_id in (0..value_count).rev() { + let adjoint = adjoints[node_id]; + if adjoint == 0.0 { + continue; + } + match locals[node_id].clone() { + ExactLocal::None => {} + ExactLocal::Unary { parent, dy, ddy } => { + adjoints[parent] += adjoint * dy; + for row in 0..self.num_inputs { + for col in 0..self.num_inputs { + hessian[row][col] += + adjoint * ddy * grads[parent][row] * grads[parent][col]; + } + } + } + ExactLocal::Binary { + left, + right, + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + adjoints[left] += adjoint * dy_left; + adjoints[right] += adjoint * dy_right; + for row in 0..self.num_inputs { + for col in 0..self.num_inputs { + hessian[row][col] += adjoint + * (ddy_left_left * grads[left][row] * grads[left][col] + + ddy_right_right * grads[right][row] * grads[right][col] + + ddy_left_right + * (grads[left][row] * grads[right][col] + + grads[right][row] * grads[left][col])); + } + } + } + } + } + + Ok(hessian) + } + + /// Compute only the primary output value. + pub fn compute(&self, inputs: &[f64]) -> Result { + self.compile().compute(inputs) + } + + /// Return the preferred backend for batch value computation. + pub fn recommended_batch_compute_backend(&self) -> Result { + Ok(self.compile_ir()?.recommended_batch_compute_backend()) + } + + /// Return the preferred backend for batch gradient computation. + pub fn recommended_batch_gradient_backend(&self) -> Result { + Ok(self.compile_ir()?.recommended_batch_gradient_backend()) + } + + /// Compute all selected outputs for a batch via compiled IR. + pub fn compute_batch(&self, batch: BatchInputs<'_>) -> Result { + self.compile_ir()?.compute_batch(batch) + } + + /// Compute all selected outputs into a reusable output buffer via compiled IR. + pub fn compute_batch_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result<()> { + self.compile_ir()?.compute_batch_into(batch, buffer) + } + + /// Compute all selected outputs for a batch with automatic backend dispatch. + pub fn compute_batch_auto(&self, batch: BatchInputs<'_>) -> Result<(BackendKind, BatchValues)> { + self.compile_ir()?.compute_batch_auto(batch) + } + + /// Compute all selected outputs into a reusable output buffer with automatic backend dispatch. + pub fn compute_batch_auto_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchValuesBuffer, + ) -> Result { + self.compile_ir()?.compute_batch_auto_into(batch, buffer) + } + + /// Compute primary-output values and gradients for a batch via compiled IR. + pub fn gradient_batch(&self, batch: BatchInputs<'_>) -> Result { + self.compile_ir()?.gradient_batch(batch) + } + + /// Compute primary-output values and gradients into a reusable buffer via compiled IR. + pub fn gradient_batch_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result<()> { + self.compile_ir()?.gradient_batch_into(batch, buffer) + } + + /// Compute primary-output values and gradients for a batch with automatic backend dispatch. + pub fn gradient_batch_auto( + &self, + batch: BatchInputs<'_>, + ) -> Result<(BackendKind, BatchGradients)> { + self.compile_ir()?.gradient_batch_auto(batch) + } + + /// Compute primary-output values and gradients into a reusable buffer with automatic backend dispatch. + pub fn gradient_batch_auto_into( + &self, + batch: BatchInputs<'_>, + buffer: &mut BatchGradientsBuffer, + ) -> Result { + self.compile_ir()?.gradient_batch_auto_into(batch, buffer) + } + + /// Return compiled IR metadata for backend planning. + pub fn compiled_metadata(&self) -> Result { + Ok(self.compile_ir()?.metadata()) + } + + /// Return backend compatibility details for a specific backend. + pub fn backend_support_report(&self, backend: BackendKind) -> Result { + self.compile_ir()?.backend_support_report(backend) + } + + /// Return backend compatibility details for all built-in backends. + pub fn backend_support_reports(&self) -> Result> { + self.compile_ir()?.backend_support_reports() + } + + /// Return a device-oriented batch buffer plan for a backend. + pub fn device_batch_plan( + &self, + backend: BackendKind, + batch_size: usize, + ) -> Result { + Ok(self.compile_ir()?.device_batch_plan(backend, batch_size)) + } + + /// Allocate mock-device buffers for this graph and batch size. + pub fn allocate_mock_device_buffers(&self, batch_size: usize) -> Result { + Ok(self.compile_ir()?.allocate_mock_device_buffers(batch_size)) + } + + /// Execute batch value computation through mock-device buffers. + pub fn compute_batch_mock_device_into( + &self, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + self.compile_ir()? + .compute_batch_mock_device_into(batch, buffers, output) + } + + /// Execute batch gradient computation through mock-device buffers. + pub fn gradient_batch_mock_device_into( + &self, + batch: BatchInputs<'_>, + buffers: &mut DeviceBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + self.compile_ir()? + .gradient_batch_mock_device_into(batch, buffers, output) + } + + #[cfg(feature = "backend-wgpu")] + /// Allocate real WGPU buffers for this graph and batch size. + pub fn allocate_wgpu_buffers( + &self, + backend: &crate::WgpuBackend, + batch_size: usize, + ) -> Result { + self.compile_ir()? + .allocate_wgpu_buffers(backend, batch_size) + } + + #[cfg(feature = "backend-wgpu")] + /// Execute batch value computation through real WGPU buffers. + pub fn compute_batch_wgpu_into( + &self, + backend: &crate::WgpuBackend, + batch: BatchInputs<'_>, + buffers: &mut crate::WgpuBufferSet, + output: &mut BatchValuesBuffer, + ) -> Result { + self.compile_ir()? + .compute_batch_wgpu_into(backend, batch, buffers, output) + } + + #[cfg(feature = "backend-wgpu")] + /// Execute batch gradient computation through real WGPU buffers. + pub fn gradient_batch_wgpu_into( + &self, + backend: &crate::WgpuBackend, + batch: BatchInputs<'_>, + buffers: &mut crate::WgpuBufferSet, + output: &mut BatchGradientsBuffer, + ) -> Result { + self.compile_ir()? + .gradient_batch_wgpu_into(backend, batch, buffers, output) + } + + #[cfg(feature = "backend-wgpu")] + /// Return whether this graph is statically eligible for the exact-safe native WGPU compute path. + pub fn supports_native_wgpu_batch_compute(&self, backend: &crate::WgpuBackend) -> Result { + Ok(self + .compile_ir()? + .supports_native_wgpu_batch_compute(backend)) + } + + #[cfg(feature = "backend-wgpu")] + /// Return whether this graph and concrete batch can use the exact-safe native WGPU compute path. + pub fn supports_native_wgpu_batch_compute_for_batch( + &self, + backend: &crate::WgpuBackend, + batch: BatchInputs<'_>, + ) -> Result { + Ok(self + .compile_ir()? + .supports_native_wgpu_batch_compute_for_batch(backend, batch)) + } + + /// Return SIMD backend compatibility details for this graph. + pub fn simd_support_report(&self) -> Result { + self.compile_ir()?.simd_support_report() + } + + /// Create a compiled IR workspace for this graph. + pub fn compiled_workspace(&self) -> Result { + Ok(self.compile_ir()?.workspace()) + } + + /// Compute all selected output values. + pub fn compute_many(&self, inputs: &[f64]) -> Result> { + self.compile().compute_many(inputs) + } + + /// Compute all selected output values with checked-domain validation. + pub fn compute_many_checked(&self, inputs: &[f64]) -> Result> { + self.compile().compute_many_checked(inputs) + } + + /// Compute only the output value with checked-domain validation. + pub fn compute_checked(&self, inputs: &[f64]) -> Result { + self.compile().compute_checked(inputs) + } + + /// Compute the primary output value and gradient closure. + pub fn compute_grad(&self, inputs: &[f64]) -> Result { + self.compile().compute_grad(inputs) + } + + /// Compute the Jacobian for all selected outputs. + pub fn jacobian(&self, inputs: &[f64]) -> Result>> { + self.compile().jacobian(inputs) + } + + /// Compute the Jacobian for all selected outputs with checked-domain validation. + pub fn jacobian_checked(&self, inputs: &[f64]) -> Result>> { + self.compile().jacobian_checked(inputs) + } + + /// Compute all selected output values and their Jacobian. + pub fn value_and_jacobian(&self, inputs: &[f64]) -> Result<(Vec, Vec>)> { + let tape = self.compile(); + Ok((tape.compute_many(inputs)?, tape.jacobian(inputs)?)) + } + + /// Compute all selected output values and their Jacobian with checked domains. + pub fn value_and_jacobian_checked(&self, inputs: &[f64]) -> Result<(Vec, Vec>)> { + let tape = self.compile(); + Ok(( + tape.compute_many_checked(inputs)?, + tape.jacobian_checked(inputs)?, + )) + } + + /// Compute the primary output value and eager gradient. + pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + self.compile().gradient(inputs) + } + + /// Alias for [`Graph::gradient`] with a name common in optimization code. + pub fn value_and_gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + self.gradient(inputs) + } + + /// Compute the primary output value with an explicit domain policy. + pub fn compute_with_domain_policy(&self, inputs: &[f64], policy: DomainPolicy) -> Result { + match policy { + DomainPolicy::Unchecked => self.compute(inputs), + DomainPolicy::Checked => self.compute_checked(inputs), + DomainPolicy::StrictDerivative => { + self.validate_strict_derivative_domain(inputs)?; + self.compute_checked(inputs) + } + } + } + + /// Compute the primary output value and eager gradient with checked-domain validation. + pub fn gradient_checked(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + self.compile().gradient_checked(inputs) + } + + /// Compute the primary output value and gradient with an explicit domain policy. + pub fn gradient_with_domain_policy( + &self, + inputs: &[f64], + policy: DomainPolicy, + ) -> Result<(f64, Vec)> { + match policy { + DomainPolicy::Unchecked => self.gradient(inputs), + DomainPolicy::Checked => self.gradient_checked(inputs), + DomainPolicy::StrictDerivative => { + self.validate_strict_derivative_domain(inputs)?; + self.gradient_checked(inputs) + } + } + } + + /// Compute a finite-difference Hessian using repeated gradient evaluation. + pub fn compute_hessian(&self, inputs: &[f64]) -> Result>> { + self.compile().compute_hessian(inputs) + } + + /// Compute an exact Hessian with a native graph traversal. + /// + /// Constant nodes are handled directly as zero-gradient graph nodes instead + /// of being translated to synthetic inputs. + pub fn exact_hessian_rr(&self, inputs: &[f64]) -> Result>> { + self.compute_exact_hessian_native(inputs) + } + + /// Compute an exact Hessian using the graph-to-RPN FR translation. + pub fn exact_hessian_fr(&self, inputs: &[f64]) -> Result>> { + self.check_graph_input_len(inputs)?; + self.validate()?; + let output = self + .effective_output_node() + .ok_or(AutodiffError::EmptyGraph)?; + let mut ops = Vec::new(); + let mut extra_inputs = Vec::new(); + self.emit_fr_ops(output, &mut ops, &mut extra_inputs)?; + let mut all_inputs = inputs.to_vec(); + all_inputs.extend(extra_inputs); + let hessian = MultiAD2FR::compute_hessian(&ops, &all_inputs)?; + Ok(self.crop_hessian(hessian)) + } + + /// Compute an exact Hessian using the graph-to-RPN RF translation. + pub fn exact_hessian_rf(&self, inputs: &[f64]) -> Result>> { + self.check_graph_input_len(inputs)?; + self.validate()?; + let output = self + .effective_output_node() + .ok_or(AutodiffError::EmptyGraph)?; + let mut ops = Vec::new(); + let mut extra_inputs = Vec::new(); + self.emit_rf_ops(output, &mut ops, &mut extra_inputs)?; + let mut all_inputs = inputs.to_vec(); + all_inputs.extend(extra_inputs); + let hessian = MultiAD2RF::compute_hessian(&ops, &all_inputs)?; + Ok(self.crop_hessian(hessian)) + } + + /// Compute a finite-difference Hessian-vector product. + pub fn hessian_vector_product(&self, inputs: &[f64], vector: &[f64]) -> Result> { + self.check_graph_input_len(inputs)?; + if vector.len() != self.num_inputs { + return Err(AutodiffError::InvalidGraph { + reason: "vector length must match graph.num_inputs()", + }); + } + let epsilon = 1e-5; + let inputs_plus: Vec = inputs + .iter() + .zip(vector) + .map(|(x, v)| x + epsilon * v) + .collect(); + let inputs_minus: Vec = inputs + .iter() + .zip(vector) + .map(|(x, v)| x - epsilon * v) + .collect(); + let (_value_plus, grad_plus) = self.gradient(&inputs_plus)?; + let (_value_minus, grad_minus) = self.gradient(&inputs_minus)?; + Ok(grad_plus + .iter() + .zip(grad_minus) + .map(|(plus, minus)| (plus - minus) / (2.0 * epsilon)) + .collect()) + } + + /// Return non-zero gradient entries using a default `1e-12` threshold. + pub fn gradient_sparse(&self, inputs: &[f64]) -> Result> { + self.gradient_sparse_with_tolerance(inputs, 1e-12) + } + + /// Return gradient entries whose absolute value is greater than `tolerance`. + pub fn gradient_sparse_with_tolerance( + &self, + inputs: &[f64], + tolerance: f64, + ) -> Result> { + let (_value, gradient) = self.gradient(inputs)?; + Ok(gradient + .into_iter() + .enumerate() + .filter(|(_, value)| value.abs() > tolerance) + .collect()) + } + + /// Return non-zero Hessian entries using a default `1e-12` threshold. + pub fn hessian_sparse(&self, inputs: &[f64]) -> Result> { + self.hessian_sparse_with_tolerance(inputs, 1e-12) + } + + /// Return Hessian entries whose absolute value is greater than `tolerance`. + pub fn hessian_sparse_with_tolerance( + &self, + inputs: &[f64], + tolerance: f64, + ) -> Result> { + let hessian = self.compute_hessian(inputs)?; + let mut sparse = Vec::new(); + for (row_index, row) in hessian.iter().enumerate() { + for (col_index, value) in row.iter().enumerate() { + if value.abs() > tolerance { + sparse.push((row_index, col_index, *value)); + } + } + } + Ok(sparse) + } + + /// Compare reverse-mode gradients to central finite differences. + pub fn check_gradient(&self, inputs: &[f64], tolerance: f64) -> Result { + self.check_graph_input_len(inputs)?; + let (_value, autodiff_gradient) = self.gradient(inputs)?; + let epsilon = 1e-6; + let mut entries = Vec::with_capacity(self.num_inputs); + let mut max_abs_error: f64 = 0.0; + + for index in 0..self.num_inputs { + let mut plus = inputs.to_vec(); + let mut minus = inputs.to_vec(); + plus[index] += epsilon; + minus[index] -= epsilon; + let value_plus = self.compute(&plus)?; + let value_minus = self.compute(&minus)?; + let finite_difference = (value_plus - value_minus) / (2.0 * epsilon); + let abs_error = (autodiff_gradient[index] - finite_difference).abs(); + max_abs_error = max_abs_error.max(abs_error); + entries.push(GradientCheckEntry { + index, + autodiff: autodiff_gradient[index], + finite_difference, + abs_error, + }); + } + + Ok(GradientCheckReport { + passed: max_abs_error <= tolerance, + tolerance, + max_abs_error, + entries, + }) + } +} + +impl ExprGraph { + /// Create an expression graph with `num_inputs` input variables. + #[must_use] + pub fn new(num_inputs: usize) -> Self { + Self { + graph: Rc::new(RefCell::new(Graph::new(num_inputs))), + } + } + + /// Return an expression node for an input. + #[must_use] + pub fn input(&self, input_index: usize) -> ExprNode { + let node = self.graph.borrow().input(input_index); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + /// Return an expression node for a literal constant. + pub fn constant(&self, value: f64) -> ExprNode { + let node = self.graph.borrow_mut().constant(value); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + /// Select an expression node as the graph output. + pub fn set_output(&self, expr: &ExprNode) -> Result<()> { + self.graph.borrow_mut().set_output(expr.node)?; + Ok(()) + } + + /// Clone out the underlying reusable graph. + #[must_use] + pub fn graph(&self) -> Graph { + self.graph.borrow().clone() + } +} + +impl ExprNode { + /// Return the underlying node id. + #[must_use] + pub fn node_id(&self) -> NodeId { + self.node + } + + fn same_graph(&self, other: &ExprNode) { + assert!( + Rc::ptr_eq(&self.graph, &other.graph), + "ExprNode graph mismatch" + ); + } + + fn unary(&self, op: MultiAD) -> ExprNode { + let node = self.graph.borrow_mut().push_operation(op, vec![self.node]); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + fn binary(&self, op: MultiAD, other: &ExprNode) -> ExprNode { + self.same_graph(other); + let node = self + .graph + .borrow_mut() + .push_operation(op, vec![self.node, other.node]); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + fn binary_const(&self, op: MultiAD, value: f64) -> ExprNode { + let mut graph = self.graph.borrow_mut(); + let constant = graph.constant(value); + let node = graph.push_operation(op, vec![self.node, constant]); + ExprNode { + graph: Rc::clone(&self.graph), + node, + } + } + + /// Append `sin(self)`. + pub fn sin(&self) -> ExprNode { + self.unary(MultiAD::Sin) + } + + /// Append `cos(self)`. + pub fn cos(&self) -> ExprNode { + self.unary(MultiAD::Cos) + } + + /// Append `exp(self)`. + pub fn exp(&self) -> ExprNode { + self.unary(MultiAD::Exp) + } + + /// Append `ln(self)`. + pub fn ln(&self) -> ExprNode { + self.unary(MultiAD::Ln) + } + + /// Append `sqrt(self)`. + pub fn sqrt(&self) -> ExprNode { + self.unary(MultiAD::Sqrt) + } +} + +impl Add for ExprNode { + type Output = ExprNode; + + fn add(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Add, &rhs) + } +} + +impl Add for ExprNode { + type Output = ExprNode; + + fn add(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Add, rhs) + } +} + +impl Sub for ExprNode { + type Output = ExprNode; + + fn sub(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Sub, &rhs) + } +} + +impl Sub for ExprNode { + type Output = ExprNode; + + fn sub(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Sub, rhs) + } +} + +impl Mul for ExprNode { + type Output = ExprNode; + + fn mul(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Mul, &rhs) + } +} + +impl Mul for ExprNode { + type Output = ExprNode; + + fn mul(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Mul, rhs) + } +} + +impl Div for ExprNode { + type Output = ExprNode; + + fn div(self, rhs: ExprNode) -> Self::Output { + self.binary(MultiAD::Div, &rhs) + } +} + +impl Div for ExprNode { + type Output = ExprNode; + + fn div(self, rhs: f64) -> Self::Output { + self.binary_const(MultiAD::Div, rhs) + } +} + +impl Neg for ExprNode { + type Output = ExprNode; + + fn neg(self) -> Self::Output { + self.unary(MultiAD::Neg) + } +} + +impl TryFrom<&Graph> for Vec<(MultiAD, Vec)> { + type Error = crate::AutodiffError; + + fn try_from(graph: &Graph) -> Result { + graph.to_operations() + } +} + +impl TryFrom for Vec<(MultiAD, Vec)> { + type Error = crate::AutodiffError; + + fn try_from(graph: Graph) -> Result { + graph.to_operations() + } +} + +impl Tape { + #[inline] + fn check_input_len(&self, inputs: &[f64]) -> Result<()> { + if inputs.len() == self.graph.num_inputs { + Ok(()) + } else { + Err(AutodiffError::InvalidGraph { + reason: "input length must match graph.num_inputs()", + }) + } + } + + #[inline] + fn fill_values(&self, inputs: &[f64], workspace: &mut TapeWorkspace) -> Result<()> { + self.fill_values_inner(inputs, workspace, false) + } + + #[inline] + fn fill_values_checked(&self, inputs: &[f64], workspace: &mut TapeWorkspace) -> Result<()> { + self.fill_values_inner(inputs, workspace, true) + } + + fn fill_values_inner( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + checked: bool, + ) -> Result<()> { + self.check_input_len(inputs)?; + workspace.values.clear(); + workspace + .values + .reserve(self.graph.num_inputs + self.graph.nodes.len()); + workspace.values.extend_from_slice(inputs); + + for node in self.compiled_nodes.iter() { + match node { + CompiledNode::Constant(value) => workspace.values.push(*value), + CompiledNode::Operation { op, arg_range } => { + let value = with_compiled_arg_values( + &self.arg_indices, + *arg_range, + &workspace.values, + |args| { + if checked { + op.forward_checked(args) + } else { + op.forward(args) + } + }, + )?; + workspace.values.push(value); + } + } + } + + Ok(()) + } + + /// Return the underlying graph. + #[must_use] + pub fn graph(&self) -> &Graph { + &self.graph + } + + /// Create a reusable workspace sized for this tape. + #[must_use] + pub fn workspace(&self) -> TapeWorkspace { + TapeWorkspace { + values: Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()), + cotangent_values: Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()), + gradients: Vec::with_capacity(self.graph.num_inputs), + } + } + + /// Compute only the primary output value. + pub fn compute(&self, inputs: &[f64]) -> Result { + let mut workspace = self.workspace(); + self.compute_with_workspace(inputs, &mut workspace) + } + + /// Compute all selected output values. + pub fn compute_many(&self, inputs: &[f64]) -> Result> { + let mut workspace = self.workspace(); + self.compute_many_with_workspace(inputs, &mut workspace) + } + + /// Compute all selected output values with checked-domain validation. + pub fn compute_many_checked(&self, inputs: &[f64]) -> Result> { + let mut workspace = self.workspace(); + self.compute_many_with_workspace_checked(inputs, &mut workspace) + } + + /// Compute only the primary output value with checked-domain validation. + pub fn compute_checked(&self, inputs: &[f64]) -> Result { + let mut workspace = self.workspace(); + self.compute_with_workspace_checked(inputs, &mut workspace) + } + + /// Compute only the primary output value using a reusable workspace. + pub fn compute_with_workspace( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result { + self.fill_values(inputs, workspace)?; + Ok(self + .output_indices + .first() + .map(|&index| workspace.values[index]) + .unwrap_or(0.0)) + } + + /// Compute all selected output values using a reusable workspace. + pub fn compute_many_with_workspace( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result> { + self.fill_values(inputs, workspace)?; + Ok(self + .output_indices + .iter() + .map(|&index| workspace.values[index]) + .collect()) + } + + /// Compute all selected output values using a reusable workspace with checked-domain validation. + pub fn compute_many_with_workspace_checked( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result> { + self.fill_values_checked(inputs, workspace)?; + Ok(self + .output_indices + .iter() + .map(|&index| workspace.values[index]) + .collect()) + } + + /// Compute only the primary output value using a reusable workspace with checked-domain validation. + pub fn compute_with_workspace_checked( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result { + self.fill_values_checked(inputs, workspace)?; + Ok(self + .output_indices + .first() + .map(|&index| workspace.values[index]) + .unwrap_or(0.0)) + } + + /// Compute the output value and gradient closure. + pub fn compute_grad(&self, inputs: &[f64]) -> Result { + self.check_input_len(inputs)?; + let mut values: Vec = + Vec::with_capacity(self.graph.num_inputs + self.graph.nodes.len()); + values.extend_from_slice(inputs); + + let mut backward_locals: Vec> = + Vec::with_capacity(self.compiled_nodes.len()); + + for node in self.compiled_nodes.iter() { + match node { + CompiledNode::Constant(value) => { + values.push(*value); + backward_locals.push(None); + } + CompiledNode::Operation { op, arg_range } => { + let (value, backward_local) = + with_compiled_arg_values(&self.arg_indices, *arg_range, &values, |args| { + let value = op.forward(args)?; + let local_rule = op_rules::local_rule(*op, args, value)?; + Ok((value, backward_local(local_rule))) + })?; + values.push(value); + backward_locals.push(Some(backward_local)); + } + } + } + + let final_value = self + .output_indices + .first() + .map(|&index| values[index]) + .unwrap_or(0.0); + let num_inputs = self.graph.num_inputs; + let values_len = values.len(); + let output_index = self.output_indices.first().copied(); + let compiled_nodes = Arc::clone(&self.compiled_nodes); + let arg_indices = Arc::clone(&self.arg_indices); + + let backward_fn = Box::new(move |cotangent: f64| -> Vec { + let Some(final_output_index) = output_index else { + return Vec::new(); + }; + + let mut cotangent_values = vec![0.0; values_len]; + cotangent_values[final_output_index] = cotangent; + + for (offset, backward_local) in backward_locals.iter().enumerate().rev() { + let Some(local) = backward_local else { + continue; + }; + let node_id = num_inputs + offset; + let current_cotangent = cotangent_values[node_id]; + if current_cotangent == 0.0 { + continue; + } + + let CompiledNode::Operation { arg_range, .. } = compiled_nodes[offset] else { + continue; + }; + let input_indices = compiled_arg_indices(&arg_indices, arg_range); + match local { + BackwardLocal::Unary(dy) => { + cotangent_values[input_indices[0]] += current_cotangent * dy; + } + BackwardLocal::Binary(dy_left, dy_right) => { + cotangent_values[input_indices[0]] += current_cotangent * dy_left; + cotangent_values[input_indices[1]] += current_cotangent * dy_right; + } + } + } + + cotangent_values[..num_inputs.min(cotangent_values.len())].to_vec() + }); + + Ok((final_value, backward_fn)) + } + + /// Compute the output value and gradient eagerly using a reusable workspace. + pub fn gradient_with_workspace<'a>( + &self, + inputs: &[f64], + workspace: &'a mut TapeWorkspace, + ) -> Result<(f64, &'a [f64])> { + self.fill_values(inputs, workspace)?; + + workspace.cotangent_values.clear(); + workspace + .cotangent_values + .resize(workspace.values.len(), 0.0); + workspace.gradients.clear(); + workspace.gradients.resize(self.graph.num_inputs, 0.0); + + let Some(output_index) = self.output_indices.first().copied() else { + return Ok((0.0, &workspace.gradients)); + }; + + workspace.cotangent_values[output_index] = 1.0; + reverse_accumulate_compiled( + self.graph.num_inputs, + &self.compiled_nodes, + &self.arg_indices, + &workspace.values, + &mut workspace.cotangent_values, + )?; + + workspace + .gradients + .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); + Ok((workspace.values[output_index], &workspace.gradients)) + } + + /// Compute the output value and gradient eagerly using a reusable workspace with checked-domain validation. + pub fn gradient_with_workspace_checked<'a>( + &self, + inputs: &[f64], + workspace: &'a mut TapeWorkspace, + ) -> Result<(f64, &'a [f64])> { + self.fill_values_checked(inputs, workspace)?; + + workspace.cotangent_values.clear(); + workspace + .cotangent_values + .resize(workspace.values.len(), 0.0); + workspace.gradients.clear(); + workspace.gradients.resize(self.graph.num_inputs, 0.0); + + let Some(output_index) = self.output_indices.first().copied() else { + return Ok((0.0, &workspace.gradients)); + }; + + workspace.cotangent_values[output_index] = 1.0; + reverse_accumulate_compiled( + self.graph.num_inputs, + &self.compiled_nodes, + &self.arg_indices, + &workspace.values, + &mut workspace.cotangent_values, + )?; + + workspace + .gradients + .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); + Ok((workspace.values[output_index], &workspace.gradients)) + } + + /// Compute the output value and gradient eagerly. + pub fn gradient(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + let mut workspace = self.workspace(); + let (value, gradient) = self.gradient_with_workspace(inputs, &mut workspace)?; + Ok((value, gradient.to_vec())) + } + + /// Compute the output value and gradient eagerly with checked-domain validation. + pub fn gradient_checked(&self, inputs: &[f64]) -> Result<(f64, Vec)> { + let mut workspace = self.workspace(); + let (value, gradient) = self.gradient_with_workspace_checked(inputs, &mut workspace)?; + Ok((value, gradient.to_vec())) + } + + /// Compute the Jacobian for all selected outputs. + pub fn jacobian(&self, inputs: &[f64]) -> Result>> { + let mut workspace = self.workspace(); + self.jacobian_with_workspace(inputs, &mut workspace) + } + + /// Compute the Jacobian for all selected outputs with checked-domain validation. + pub fn jacobian_checked(&self, inputs: &[f64]) -> Result>> { + let mut workspace = self.workspace(); + self.jacobian_with_workspace_checked(inputs, &mut workspace) + } + + /// Compute the Jacobian for all selected outputs using a reusable workspace. + pub fn jacobian_with_workspace( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result>> { + self.jacobian_with_workspace_inner(inputs, workspace, false) + } + + /// Compute the Jacobian for all selected outputs using a reusable workspace with checked-domain validation. + pub fn jacobian_with_workspace_checked( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + ) -> Result>> { + self.jacobian_with_workspace_inner(inputs, workspace, true) + } + + fn jacobian_with_workspace_inner( + &self, + inputs: &[f64], + workspace: &mut TapeWorkspace, + checked: bool, + ) -> Result>> { + if checked { + self.fill_values_checked(inputs, workspace)?; + } else { + self.fill_values(inputs, workspace)?; + } + let mut jacobian: Vec> = Vec::with_capacity(self.output_indices.len()); + + for &output_index in &self.output_indices { + workspace.cotangent_values.clear(); + workspace + .cotangent_values + .resize(workspace.values.len(), 0.0); + workspace.gradients.clear(); + workspace.gradients.resize(self.graph.num_inputs, 0.0); + workspace.cotangent_values[output_index] = 1.0; + reverse_accumulate_compiled( + self.graph.num_inputs, + &self.compiled_nodes, + &self.arg_indices, + &workspace.values, + &mut workspace.cotangent_values, + )?; + + workspace + .gradients + .copy_from_slice(&workspace.cotangent_values[..self.graph.num_inputs]); + jacobian.push(workspace.gradients.clone()); + } + + Ok(jacobian) + } + + /// Compute a finite-difference Hessian by differentiating eager gradients. + pub fn compute_hessian(&self, inputs: &[f64]) -> Result>> { + let num_inputs = self.graph.num_inputs; + let epsilon = 1e-5; + let mut hessian = vec![vec![0.0; num_inputs]; num_inputs]; + + if num_inputs == 0 { + let (_value, _grad) = self.gradient(inputs)?; + return Ok(hessian); + } + + let mut plus_workspace = self.workspace(); + let mut minus_workspace = self.workspace(); + + for j in 0..num_inputs { + let mut inputs_plus = inputs.to_vec(); + inputs_plus[j] += epsilon; + + let mut inputs_minus = inputs.to_vec(); + inputs_minus[j] -= epsilon; + + let (_value_plus, grad_plus) = + self.gradient_with_workspace(&inputs_plus, &mut plus_workspace)?; + let (_value_minus, grad_minus) = + self.gradient_with_workspace(&inputs_minus, &mut minus_workspace)?; + + for i in 0..num_inputs { + hessian[i][j] = (grad_plus[i] - grad_minus[i]) / (2.0 * epsilon); + } + } + + Ok(hessian) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::approx_eq_eps as approx_eq; + use crate::{ + BackendCapabilities, BackendKind, ExecutionBackend, OpCode, ScalarBackend, SimdBackend, + }; + + #[test] + fn test_graph_supports_constants() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let two = graph.constant(2.0); + let x_sq = graph.mul(x, x); + graph.add(x_sq, two); + + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 11.0, 1e-10)); + + let (_value, grad_fn) = graph.compute_grad(&[3.0]).unwrap(); + let grad = grad_fn(1.0); + assert_eq!(grad.len(), 1); + assert!(approx_eq(grad[0], 6.0, 1e-10)); + } + + #[test] + fn test_graph_compile_reuses_graph() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + graph.mul(sum, y); + + let tape = graph.compile(); + let value = tape.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 15.0, 1e-10)); + + let (_value, grad_fn) = tape.compute_grad(&[2.0, 3.0]).unwrap(); + let grad = grad_fn(1.0); + assert!(approx_eq(grad[0], 3.0, 1e-10)); + assert!(approx_eq(grad[1], 8.0, 1e-10)); + } + + #[test] + fn test_tape_workspace_eager_gradient() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + graph.mul(sum, y); + + let tape = graph.compile(); + let mut workspace = tape.workspace(); + let (value, grad) = tape + .gradient_with_workspace(&[2.0, 3.0], &mut workspace) + .unwrap(); + assert!(approx_eq(value, 15.0, 1e-10)); + assert!(approx_eq(grad[0], 3.0, 1e-10)); + assert!(approx_eq(grad[1], 8.0, 1e-10)); + + let value_again = tape + .compute_with_workspace(&[4.0, 1.5], &mut workspace) + .unwrap(); + assert!(approx_eq(value_again, 8.25, 1e-10)); + } + + #[test] + fn test_graph_rejects_wrong_input_length() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.mul(x, x); + + assert!(graph.compute(&[]).is_err()); + assert!(graph.compute(&[2.0, 3.0]).is_err()); + assert!(graph.compute_checked(&[2.0, 3.0]).is_err()); + assert!(graph.compute_grad(&[2.0, 3.0]).is_err()); + assert!(graph.jacobian(&[2.0, 3.0]).is_err()); + + let tape = graph.compile(); + assert!(tape.compute(&[2.0, 3.0]).is_err()); + assert!(tape.compute_grad(&[2.0, 3.0]).is_err()); + assert!(tape.gradient(&[2.0, 3.0]).is_err()); + } + + #[test] + fn test_graph_checked_compute_and_gradient() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.ln(x); + + let error = graph.compute_checked(&[0.0]).unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::DomainError { + operation: "Ln", + reason: "input must be positive", + } + ); + + let ok = graph.gradient_checked(&[2.0]).unwrap(); + assert!(approx_eq(ok.0, 2.0_f64.ln(), 1e-10)); + assert!(approx_eq(ok.1[0], 0.5, 1e-10)); + } + + #[test] + fn test_graph_checked_multi_output_and_jacobian() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let ratio = graph.div(x, y); + let log_y = graph.ln(y); + graph.set_outputs(&[ratio, log_y]).unwrap(); + + let values = graph.compute_many_checked(&[4.0, 2.0]).unwrap(); + assert_eq!(values.len(), 2); + assert!(approx_eq(values[0], 2.0, 1e-10)); + assert!(approx_eq(values[1], 2.0_f64.ln(), 1e-10)); + + let jacobian = graph.jacobian_checked(&[4.0, 2.0]).unwrap(); + assert_eq!(jacobian.len(), 2); + assert!(approx_eq(jacobian[0][0], 0.5, 1e-10)); + assert!(approx_eq(jacobian[0][1], -1.0, 1e-10)); + assert!(approx_eq(jacobian[1][0], 0.0, 1e-10)); + assert!(approx_eq(jacobian[1][1], 0.5, 1e-10)); + + let error = graph.compute_many_checked(&[4.0, 0.0]).unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::DomainError { + operation: "Div", + reason: "denominator must be non-zero", + } + ); + } + + #[test] + fn test_tape_checked_workspace() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let ratio = graph.div(x, y); + let log_y = graph.ln(y); + graph.set_outputs(&[ratio, log_y]).unwrap(); + let tape = graph.compile(); + let mut workspace = tape.workspace(); + + let values = tape + .compute_many_with_workspace_checked(&[4.0, 2.0], &mut workspace) + .unwrap(); + assert!(approx_eq(values[0], 2.0, 1e-10)); + assert!(approx_eq(values[1], 2.0_f64.ln(), 1e-10)); + + let jacobian = tape + .jacobian_with_workspace_checked(&[4.0, 2.0], &mut workspace) + .unwrap(); + assert!(approx_eq(jacobian[0][0], 0.5, 1e-10)); + assert!(approx_eq(jacobian[0][1], -1.0, 1e-10)); + assert!(approx_eq(jacobian[1][0], 0.0, 1e-10)); + assert!(approx_eq(jacobian[1][1], 0.5, 1e-10)); + + let error = tape + .compute_with_workspace_checked(&[1.0, 0.0], &mut workspace) + .unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::DomainError { + operation: "Div", + reason: "denominator must be non-zero", + } + ); + } + + #[test] + fn test_graph_from_operations() { + let ops = [ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Add, vec![0, 1]), + (MultiAD::Sin, vec![2]), + ]; + let graph = Graph::from_operations(2, &ops); + let value = graph.compute(&[0.5, 0.25]).unwrap(); + assert!(approx_eq(value, (0.75_f64).sin(), 1e-10)); + } + + #[test] + fn test_graph_from_operations_preserves_trailing_input_output() { + let ops = [(MultiAD::Add, vec![0, 1]), (MultiAD::Inp, vec![0])]; + let graph = Graph::from_operations(2, &ops); + let value = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_graph_from_operations_preserves_input_only_output() { + let ops = [(MultiAD::Inp, vec![0])]; + let graph = Graph::from_operations(2, &ops); + let value = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_graph_to_operations_round_trip() { + let ops = [ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Add, vec![0, 1]), + (MultiAD::Sin, vec![2]), + ]; + let graph = Graph::from_operations(2, &ops); + let converted = graph.to_operations().unwrap(); + assert_eq!( + converted, + vec![(MultiAD::Add, vec![0, 1]), (MultiAD::Sin, vec![2])] + ); + + let value_legacy = MultiAD::compute(&converted, &[0.5, 0.25]).unwrap(); + let value_graph = graph.compute(&[0.5, 0.25]).unwrap(); + assert!(approx_eq(value_legacy, value_graph, 1e-10)); + } + + #[test] + fn test_graph_to_operations_rejects_constants() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let c = graph.constant(2.0); + graph.add(x, c); + + let error = graph.to_operations().unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::InvalidGraph { + reason: "legacy tuple graphs cannot represent constant nodes", + } + ); + } + + #[test] + fn test_graph_validate_accepts_valid_graph() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + graph.sin(sum); + assert!(graph.validate().is_ok()); + } + + #[test] + fn test_graph_validate_rejects_bad_index() { + let mut graph = Graph::new(1); + graph.push_operation(MultiAD::Sin, vec![3]); + let error = graph.validate().unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::IndexOutOfBounds { + index: 3, + max_index: 0, + } + ); + } + + #[test] + fn test_graph_export_formats() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let neg_x = graph.neg(x); + graph.exp(neg_x); + + let mermaid = graph.to_mermaid(); + let dot = graph.to_dot(); + assert!(mermaid.contains("flowchart LR")); + assert!(mermaid.contains("Neg")); + assert!(dot.contains("digraph Graph")); + assert!(dot.contains("Exp")); + } + + #[test] + fn test_graph_explicit_output_uses_selected_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let x_sq = graph.mul(x, x); + let x_cu = graph.mul(x_sq, x); + graph.set_output(x_sq).unwrap(); + + let value = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(value, 4.0, 1e-10)); + + let (_value, grad_fn) = graph.compute_grad(&[2.0]).unwrap(); + let grad = grad_fn(1.0); + assert!(approx_eq(grad[0], 4.0, 1e-10)); + + let ops = graph.to_operations().unwrap(); + let legacy_value = MultiAD::compute(&ops, &[2.0]).unwrap(); + assert!(approx_eq(legacy_value, 4.0, 1e-10)); + + graph.set_output(x_cu).unwrap(); + let value_cu = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(value_cu, 8.0, 1e-10)); + } + + #[test] + fn test_graph_clear_output_restores_last_node_behavior() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let x_sq = graph.mul(x, x); + graph.mul(x_sq, x); + graph.set_output(x_sq).unwrap(); + assert!(approx_eq(graph.compute(&[2.0]).unwrap(), 4.0, 1e-10)); + graph.clear_output(); + assert!(approx_eq(graph.compute(&[2.0]).unwrap(), 8.0, 1e-10)); + } + + #[test] + fn test_graph_validate_rejects_bad_explicit_output() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sin(x); + graph.output_nodes = vec![4]; + let error = graph.validate().unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::IndexOutOfBounds { + index: 4, + max_index: 1, + } + ); + } + + #[test] + fn test_graph_to_operations_with_input_output_marker() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + graph.set_output(x).unwrap(); + + let ops = graph.to_operations().unwrap(); + assert_eq!(ops.last(), Some(&(MultiAD::Inp, vec![0]))); + let value = MultiAD::compute(&ops, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_graph_multi_output_values_and_jacobian() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + let product = graph.mul(x, y); + graph.set_outputs(&[sum, product]).unwrap(); + + let values = graph.compute_many(&[2.0, 3.0]).unwrap(); + assert_eq!(values.len(), 2); + assert!(approx_eq(values[0], 5.0, 1e-10)); + assert!(approx_eq(values[1], 6.0, 1e-10)); + + let jacobian = graph.jacobian(&[2.0, 3.0]).unwrap(); + assert_eq!(jacobian.len(), 2); + assert!(approx_eq(jacobian[0][0], 1.0, 1e-10)); + assert!(approx_eq(jacobian[0][1], 1.0, 1e-10)); + assert!(approx_eq(jacobian[1][0], 3.0, 1e-10)); + assert!(approx_eq(jacobian[1][1], 2.0, 1e-10)); + } + + #[test] + fn test_graph_to_operations_rejects_multi_output() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + let product = graph.mul(x, y); + graph.set_outputs(&[sum, product]).unwrap(); + + let error = graph.to_operations().unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::InvalidGraph { + reason: "legacy tuple graphs support only one output", + } + ); + } + + #[test] + fn test_safe_construction_and_compile() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let x_sq = graph.try_push_operation(MultiAD::Mul, vec![x, x]).unwrap(); + graph.set_output(x_sq).unwrap(); + assert!(graph.try_compile().is_ok()); + assert!(graph.try_push_operation(MultiAD::Inp, vec![0]).is_err()); + assert!(Graph::try_from_operations(1, &[(MultiAD::Sin, vec![2])]).is_err()); + } + + #[test] + fn test_gradient_helpers_and_gradient_check() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.square(x); + + let (value, gradient) = graph.value_and_gradient(&[3.0]).unwrap(); + assert!(approx_eq(value, 9.0, 1e-10)); + assert!(approx_eq(gradient[0], 6.0, 1e-10)); + + let report = graph.check_gradient(&[3.0], 1e-6).unwrap(); + assert!(report.passed); + assert_eq!(report.entries.len(), 1); + } + + #[test] + fn test_constant_and_expression_helpers() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let x_sq = graph.square(x); + let shifted = graph.add_const(x_sq, 2.0); + let sigmoid = graph.sigmoid(x); + graph.set_outputs(&[shifted, sigmoid]).unwrap(); + + let values = graph.compute_many(&[3.0]).unwrap(); + assert!(approx_eq(values[0], 11.0, 1e-10)); + assert!(values[1] > 0.0 && values[1] < 1.0); + } + + #[test] + fn test_graph_builder_constant_node() { + let mut builder = crate::GraphBuilder::new(1); + let x = builder.input_node(0); + let two = builder.constant_node(2.0); + let out = builder.add_node(x, two); + let graph = builder.build_graph_with_output(out).unwrap(); + assert!(approx_eq(graph.compute(&[3.0]).unwrap(), 5.0, 1e-10)); + } + + #[test] + fn test_exact_hessian_hvp_and_sparse_outputs() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let x_sq = graph.square(x); + let y_sq = graph.square(y); + graph.add(x_sq, y_sq); + + let expected = [[2.0, 0.0], [0.0, 2.0]]; + for hessian in [ + graph.exact_hessian_rr(&[1.0, 2.0]).unwrap(), + graph.exact_hessian_fr(&[1.0, 2.0]).unwrap(), + graph.exact_hessian_rf(&[1.0, 2.0]).unwrap(), + ] { + for i in 0..2 { + for j in 0..2 { + assert!(approx_eq(hessian[i][j], expected[i][j], 1e-10)); + } + } + } + + let hvp = graph + .hessian_vector_product(&[1.0, 2.0], &[3.0, 4.0]) + .unwrap(); + assert!(approx_eq(hvp[0], 6.0, 1e-6)); + assert!(approx_eq(hvp[1], 8.0, 1e-6)); + + let sparse_grad = graph.gradient_sparse(&[1.0, 2.0]).unwrap(); + assert_eq!(sparse_grad.len(), 2); + let sparse_hessian = graph.hessian_sparse(&[1.0, 2.0]).unwrap(); + assert_eq!(sparse_hessian.len(), 2); + } + + #[test] + fn test_prune_names_stats_and_exports() { + let mut graph = Graph::new(2); + graph.set_input_name(0, "x").unwrap(); + let x = graph.input(0); + let y = graph.input(1); + let used = graph.square(x); + let _unused = graph.square(y); + graph.set_output_name(used, "objective").unwrap(); + + let pruned = graph.prune_to_outputs().unwrap(); + assert_eq!(pruned.len(), 1); + let stats = graph.stats(); + assert_eq!(stats.num_inputs, 2); + assert_eq!(stats.num_ops, 2); + assert!(graph.to_mermaid().contains("x: Input 0")); + assert!(graph.to_dot().contains("objective")); + } + + #[test] + fn test_expr_graph_operator_overloads() { + let expr_graph = ExprGraph::new(2); + let x = expr_graph.input(0); + let y = expr_graph.input(1); + let out = x.clone().sin() * (x + y) + 2.0; + expr_graph.set_output(&out).unwrap(); + let graph = expr_graph.graph(); + let value = graph.compute(&[0.6, 1.4]).unwrap(); + let expected = 0.6_f64.sin() * 2.0 + 2.0; + assert!(approx_eq(value, expected, 1e-10)); + } + + #[test] + fn test_ml_activations_and_softmax() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let tanh = graph.tanh(x); + let relu = graph.relu(x); + let gelu = graph.gelu(x); + let softmax = graph.softmax(&[x, y]); + graph + .set_outputs(&[tanh, relu, gelu, softmax[0], softmax[1]]) + .unwrap(); + let values = graph.compute_many(&[1.0, 2.0]).unwrap(); + assert!(approx_eq(values[0], 1.0_f64.tanh(), 1e-10)); + assert!(approx_eq(values[1], 1.0, 1e-10)); + assert!(values[2] > 0.0); + assert!(approx_eq(values[3] + values[4], 1.0, 1e-10)); + } + + #[test] + fn test_parser_simplify_vector_and_domain_policy() { + let graph = Graph::parse_expression("sin(x) * (x + y) + 2", &["x", "y"]).unwrap(); + let (values, jacobian) = graph.value_and_jacobian(&[0.6, 1.4]).unwrap(); + assert_eq!(values.len(), 1); + assert_eq!(jacobian[0].len(), 2); + + let mut simplifiable = Graph::new(1); + let x = simplifiable.input(0); + let zero = simplifiable.constant(0.0); + simplifiable.add(x, zero); + let simplified = simplifiable.simplify().unwrap(); + assert!(approx_eq(simplified.compute(&[3.0]).unwrap(), 3.0, 1e-10)); + + let mut strict = Graph::new(1); + let x = strict.input(0); + strict.sqrt(x); + assert!(strict + .gradient_with_domain_policy(&[0.0], DomainPolicy::StrictDerivative) + .is_err()); + } + + #[test] + fn test_compiled_ir_and_batch_match_graph() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + graph.mul(sum, y); + + let compiled = graph.compile_ir().unwrap(); + let value_graph = graph.compute(&[2.0, 3.0]).unwrap(); + let value_ir = compiled.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value_graph, value_ir, 1e-10)); + + let (_value_graph, grad_graph) = graph.gradient(&[2.0, 3.0]).unwrap(); + let (_value_ir, grad_ir) = compiled.gradient(&[2.0, 3.0]).unwrap(); + assert_eq!(grad_graph.len(), grad_ir.len()); + for (left, right) in grad_graph.iter().zip(grad_ir.iter()) { + assert!(approx_eq(*left, *right, 1e-10)); + } + + let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); + let batch_values = compiled.compute_batch(batch).unwrap(); + assert_eq!(batch_values.batch_size, 2); + assert_eq!(batch_values.output_dim, 1); + assert!(approx_eq( + batch_values.data[0], + graph.compute(&[2.0, 3.0]).unwrap(), + 1e-10 + )); + assert!(approx_eq( + batch_values.data[1], + graph.compute(&[4.0, 5.0]).unwrap(), + 1e-10 + )); + + let mut values_buffer = BatchValuesBuffer::new(); + compiled + .compute_batch_into(batch, &mut values_buffer) + .unwrap(); + assert_eq!(values_buffer.data, batch_values.data); + let first_capacity = values_buffer.data.capacity(); + compiled + .compute_batch_into(batch, &mut values_buffer) + .unwrap(); + assert!(values_buffer.data.capacity() >= first_capacity); + + let batch_grad = graph.gradient_batch(batch).unwrap(); + assert_eq!(batch_grad.batch_size, 2); + assert_eq!(batch_grad.input_dim, 2); + assert_eq!(batch_grad.gradients.len(), 4); + + let mut gradients_buffer = BatchGradientsBuffer::new(); + graph + .gradient_batch_into(batch, &mut gradients_buffer) + .unwrap(); + assert_eq!(gradients_buffer.values, batch_grad.values); + assert_eq!(gradients_buffer.gradients, batch_grad.gradients); + + let metadata = graph.compiled_metadata().unwrap(); + assert_eq!(metadata.num_inputs, 2); + assert_eq!(metadata.num_outputs, 1); + assert_eq!(metadata.num_instructions, compiled.instructions().len()); + assert_eq!( + metadata.value_count, + graph.num_inputs() + compiled.instructions().len() + ); + assert!(metadata.is_scalar_output); + + let flat = compiled.flat_instructions().unwrap(); + assert_eq!(flat.len(), compiled.instructions().len()); + assert_eq!(compiled.flat_instructions_slice().len(), flat.len()); + assert_eq!(flat[0].opcode, OpCode::Add); + assert_eq!(flat[0].output, 2); + assert_eq!(flat[0].left, x); + assert_eq!(flat[0].right, y); + assert_eq!(flat[1].opcode, OpCode::Mul); + assert_eq!(flat[1].value, 0.0); + + let scalar_backend = ScalarBackend; + compiled + .validate_backend_capabilities(&scalar_backend.capabilities()) + .unwrap(); + let scalar_value = scalar_backend.compute(&compiled, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(scalar_value, value_ir, 1e-10)); + let mut scalar_values_buffer = BatchValuesBuffer::new(); + scalar_backend + .compute_batch(&compiled, batch, &mut scalar_values_buffer) + .unwrap(); + assert_eq!(scalar_values_buffer.data, batch_values.data); + + let simd_backend = SimdBackend; + if simd_backend.capabilities().supports_batch_compute { + compiled + .validate_backend_capabilities(&simd_backend.capabilities()) + .unwrap(); + let mut simd_values_buffer = BatchValuesBuffer::new(); + simd_backend + .compute_batch(&compiled, batch, &mut simd_values_buffer) + .unwrap(); + assert_eq!(simd_values_buffer.batch_size, batch.batch_size); + assert_eq!(simd_values_buffer.output_dim, batch_values.output_dim); + for (left, right) in simd_values_buffer.data.iter().zip(batch_values.data.iter()) { + assert!(approx_eq(*left, *right, 1e-10)); + } + assert!(simd_backend.gradient(&compiled, &[2.0, 3.0]).is_err()); + } + + let limited_backend = BackendCapabilities { + supports_f64: true, + supports_f32: false, + supports_constants: true, + supports_unary: true, + supports_binary: true, + supports_multi_output: true, + supports_reverse_gradient: true, + supports_batch_compute: true, + supports_batch_gradient: true, + supported_opcodes: vec![OpCode::Add], + }; + assert!(compiled + .validate_backend_capabilities(&limited_backend) + .is_err()); + + let (auto_backend, auto_values) = compiled.compute_batch_auto(batch).unwrap(); + assert_eq!(auto_values.data, batch_values.data); + let mut dispatched_values = BatchValuesBuffer::new(); + auto_backend + .compute_batch(&compiled, batch, &mut dispatched_values) + .unwrap(); + assert_eq!(dispatched_values.data, batch_values.data); + let graph_auto_backend = graph.recommended_batch_compute_backend().unwrap(); + let expected_compute_backend = compiled.recommended_batch_compute_backend(); + assert_eq!(auto_backend, expected_compute_backend); + assert_eq!(graph_auto_backend, expected_compute_backend); + + let (auto_gradient_backend, auto_gradients) = compiled.gradient_batch_auto(batch).unwrap(); + assert_eq!(auto_gradients.values, batch_grad.values); + assert_eq!(auto_gradients.gradients, batch_grad.gradients); + let mut dispatched_gradients = BatchGradientsBuffer::new(); + auto_gradient_backend + .gradient_batch(&compiled, batch, &mut dispatched_gradients) + .unwrap(); + assert_eq!(dispatched_gradients.values, batch_grad.values); + assert_eq!(dispatched_gradients.gradients, batch_grad.gradients); + let graph_auto_gradient_backend = graph.recommended_batch_gradient_backend().unwrap(); + let expected_gradient_backend = compiled.recommended_batch_gradient_backend(); + assert_eq!(auto_gradient_backend, expected_gradient_backend); + assert_eq!(graph_auto_gradient_backend, expected_gradient_backend); + + let plan = graph + .device_batch_plan(expected_compute_backend, batch.batch_size) + .unwrap(); + assert_eq!(plan.backend, expected_compute_backend); + assert_eq!(plan.batch_size, batch.batch_size); + assert_eq!(plan.input_dim, 2); + assert_eq!(plan.output_dim, 1); + assert_eq!(plan.buffers.len(), 5); + assert_eq!(plan.buffer_handles.len(), 5); + assert!(plan.compute_transfer_plan.is_empty()); + assert!(plan.gradient_transfer_plan.is_empty()); + assert_eq!(plan.buffers[3].kind, crate::DeviceBufferKind::PrimaryValues); + assert_eq!(plan.buffers[3].len, batch.batch_size); + } + + #[test] + fn test_auto_backend_reports_track_supported_simd_ops() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let abs_x = graph.abs(x); + let shifted = graph.add_const(abs_x, 1.25); + let exponent = graph.log1p_exp(x); + let powered = graph.pow(shifted, exponent); + let output = graph.log_add_exp(powered, x); + graph.set_output(output).unwrap(); + let compiled = graph.compile_ir().unwrap(); + let batch = BatchInputs::new(&[1.0, 2.0, 3.0], 3, 1).unwrap(); + + let expected_compute_backend = compiled.recommended_batch_compute_backend(); + let expected_gradient_backend = compiled.recommended_batch_gradient_backend(); + assert_eq!( + graph.recommended_batch_compute_backend().unwrap(), + expected_compute_backend + ); + assert_eq!( + graph.recommended_batch_gradient_backend().unwrap(), + expected_gradient_backend + ); + + let report = graph.simd_support_report().unwrap(); + assert!(matches!( + report.backend, + BackendKind::SimdF64x4 | BackendKind::SimdF64x2 + )); + assert_eq!(report.missing_opcodes, Vec::::new()); + assert_eq!(report.lane_width, report.backend.lane_width()); + assert_eq!( + report.can_compute_batch(), + expected_compute_backend != BackendKind::Scalar + ); + assert_eq!( + report.can_gradient_batch(), + expected_gradient_backend != BackendKind::Scalar + ); + + let scalar_report = graph.backend_support_report(BackendKind::Scalar).unwrap(); + assert!(scalar_report.can_compute_batch()); + assert!(scalar_report.can_gradient_batch()); + assert!(scalar_report.runtime_available); + assert_eq!(scalar_report.lane_width, 1); + let mock_report = graph + .backend_support_report(BackendKind::MockDeviceCpu) + .unwrap(); + assert!(mock_report.can_compute_batch()); + assert_eq!(mock_report.backend.name(), "mock-device-cpu"); + let wgpu_report = graph.backend_support_report(BackendKind::Wgpu).unwrap(); + assert_eq!(wgpu_report.backend.name(), "wgpu"); + assert_eq!(wgpu_report.lane_width, 1); + let reports = graph.backend_support_reports().unwrap(); + assert_eq!(reports.len(), 5); + assert_eq!(reports[0].backend, BackendKind::Scalar); + assert_eq!(reports[1].backend, BackendKind::MockDeviceCpu); + assert_eq!(reports[2].backend, BackendKind::Wgpu); + assert_eq!(reports[3].backend, BackendKind::SimdF64x4); + assert_eq!(reports[4].backend, BackendKind::SimdF64x2); + let simd_plan = graph + .device_batch_plan(BackendKind::Scalar, batch.batch_size) + .unwrap(); + assert_eq!(simd_plan.batch_size, batch.batch_size); + assert_eq!(simd_plan.input_dim, 1); + let mock_plan = graph + .device_batch_plan(BackendKind::MockDeviceCpu, batch.batch_size) + .unwrap(); + assert_eq!(mock_plan.backend, BackendKind::MockDeviceCpu); + assert_eq!( + mock_plan.buffer_handles[0].location, + crate::DeviceMemoryLocation::Device + ); + assert_eq!(mock_plan.compute_transfer_plan.len(), 2); + assert_eq!(mock_plan.gradient_transfer_plan.len(), 3); + let wgpu_plan = graph + .device_batch_plan(BackendKind::Wgpu, batch.batch_size) + .unwrap(); + assert_eq!(wgpu_plan.backend, BackendKind::Wgpu); + assert_eq!( + wgpu_plan.buffer_handles[0].location, + crate::DeviceMemoryLocation::Device + ); + assert_eq!(wgpu_plan.compute_transfer_plan.len(), 2); + assert_eq!(wgpu_plan.gradient_transfer_plan.len(), 3); + + let (value_backend, values) = graph.compute_batch_auto(batch).unwrap(); + assert_eq!(value_backend, expected_compute_backend); + assert_eq!(values.data, graph.compute_batch(batch).unwrap().data); + + let mut gradients = BatchGradientsBuffer::new(); + let gradient_backend = graph + .gradient_batch_auto_into(batch, &mut gradients) + .unwrap(); + assert_eq!(gradient_backend, expected_gradient_backend); + assert_eq!( + gradients.gradients, + graph.gradient_batch(batch).unwrap().gradients + ); + } + + #[test] + fn test_mock_device_buffers_execute_transfer_plans() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let product = graph.mul(x, y); + graph.set_output(product).unwrap(); + let compiled = graph.compile_ir().unwrap(); + let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0, 6.0, 7.0], 3, 2).unwrap(); + let mock = crate::MockDeviceBackend; + let mut buffers = mock.allocate_batch_buffers(&compiled, batch.batch_size); + assert_eq!(buffers.plan().backend, BackendKind::MockDeviceCpu); + let input_buffer = buffers.buffer(crate::DeviceBufferKind::Inputs).unwrap(); + assert_eq!(input_buffer.data().len(), batch.data.len()); + assert_eq!(input_buffer.handle().kind, crate::DeviceBufferKind::Inputs); + assert_eq!(buffers.buffers().len(), buffers.plan().buffers.len()); + assert!(buffers + .upload(crate::DeviceBufferKind::Inputs, &[1.0]) + .is_err()); + + let scalar_plan = compiled.device_batch_plan(BackendKind::Scalar, batch.batch_size); + let mut scalar_buffers = crate::DeviceBufferSet::new(scalar_plan); + let mut rejected_values = BatchValuesBuffer::new(); + assert!(mock + .compute_batch_with_buffers(&compiled, batch, &mut scalar_buffers, &mut rejected_values) + .is_err()); + let wrong_batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); + assert!(mock + .compute_batch_with_buffers(&compiled, wrong_batch, &mut buffers, &mut rejected_values) + .is_err()); + + let mut values = BatchValuesBuffer::new(); + let trace = compiled + .compute_batch_mock_device_into(batch, &mut buffers, &mut values) + .unwrap(); + assert_eq!(trace.mode, crate::DeviceExecutionMode::ComputeBatch); + assert_eq!(trace.transfers, buffers.plan().compute_transfer_plan); + assert!(!trace.used_native_kernel); + assert_eq!(values.data, compiled.compute_batch(batch).unwrap().data); + assert_eq!( + buffers.download(crate::DeviceBufferKind::Outputs).unwrap(), + values.data + ); + + let mut graph_buffers = graph + .allocate_mock_device_buffers(batch.batch_size) + .unwrap(); + let mut graph_values = BatchValuesBuffer::new(); + let graph_trace = graph + .compute_batch_mock_device_into(batch, &mut graph_buffers, &mut graph_values) + .unwrap(); + assert_eq!(graph_trace.mode, crate::DeviceExecutionMode::ComputeBatch); + assert!(!graph_trace.used_native_kernel); + assert_eq!(graph_values.data, values.data); + + let mut gradients = BatchGradientsBuffer::new(); + let trace = compiled + .gradient_batch_mock_device_into(batch, &mut buffers, &mut gradients) + .unwrap(); + assert_eq!(trace.mode, crate::DeviceExecutionMode::GradientBatch); + assert_eq!(trace.transfers, buffers.plan().gradient_transfer_plan); + assert!(!trace.used_native_kernel); + let scalar_gradients = compiled.gradient_batch(batch).unwrap(); + assert_eq!(gradients.values, scalar_gradients.values); + assert_eq!(gradients.gradients, scalar_gradients.gradients); + assert_eq!( + buffers + .download(crate::DeviceBufferKind::PrimaryValues) + .unwrap(), + gradients.values + ); + + let boundary = crate::GpuBackendBoundary::new( + crate::AcceleratorDeviceContext::mock_cpu(), + crate::DeviceTransferPolicy::Explicit, + ); + assert!(boundary.unsupported_execution_error::<()>().is_err()); + let cuda = crate::AcceleratorDeviceContext::cuda(2); + assert_eq!(cuda.kind, crate::AcceleratorDeviceKind::Cuda); + assert_eq!(cuda.name, "cuda:2"); + let wgpu = crate::AcceleratorDeviceContext::wgpu(1); + assert_eq!(wgpu.kind, crate::AcceleratorDeviceKind::Wgpu); + assert_eq!(wgpu.name, "wgpu:1"); + } + + #[cfg(feature = "backend-wgpu")] + #[test] + fn test_wgpu_backend_skeleton_buffers_and_execution() { + let boundary = crate::GpuBackendBoundary::new( + crate::AcceleratorDeviceContext::wgpu(0), + crate::DeviceTransferPolicy::Explicit, + ); + let backend = match boundary.initialize_wgpu() { + Ok(backend) => backend, + Err(_) => return, + }; + assert_eq!(backend.context().kind, crate::AcceleratorDeviceKind::Wgpu); + assert_eq!( + backend.transfer_policy(), + crate::DeviceTransferPolicy::Explicit + ); + assert!(!backend.adapter_name().is_empty()); + assert_eq!( + crate::WgpuBackend::native_batch_compute_supported_opcodes(), + crate::WGPU_NATIVE_BATCH_COMPUTE_EXACT_SAFE_OPCODES + ); + + let mut exact_graph = Graph::new(1); + let x = exact_graph.input(0); + let neg_x = exact_graph.neg(x); + let relu_neg_x = exact_graph.relu(neg_x); + let abs_x = exact_graph.abs(x); + exact_graph.set_outputs(&[relu_neg_x, abs_x]).unwrap(); + let compiled = exact_graph.compile_ir().unwrap(); + let batch = BatchInputs::new(&[2.0, -3.0, -0.0], 3, 1).unwrap(); + assert!(compiled.supports_native_wgpu_batch_compute(&backend)); + assert!(compiled.supports_native_wgpu_batch_compute_for_batch(&backend, batch)); + assert!(exact_graph + .supports_native_wgpu_batch_compute(&backend) + .unwrap()); + assert!(exact_graph + .supports_native_wgpu_batch_compute_for_batch(&backend, batch) + .unwrap()); + + let mut buffers = backend + .allocate_batch_buffers(&compiled, batch.batch_size) + .unwrap(); + assert_eq!(buffers.plan().backend, BackendKind::Wgpu); + assert_eq!(buffers.buffers().len(), buffers.plan().buffers.len()); + assert_eq!( + buffers + .buffer(crate::DeviceBufferKind::Inputs) + .unwrap() + .handle() + .location, + crate::DeviceMemoryLocation::Device + ); + assert!(buffers + .upload(&backend, crate::DeviceBufferKind::Inputs, &[1.0]) + .is_err()); + buffers + .upload(&backend, crate::DeviceBufferKind::Inputs, batch.data) + .unwrap(); + assert_eq!( + buffers + .download(&backend, crate::DeviceBufferKind::Inputs) + .unwrap(), + batch.data + ); + + let mut values = BatchValuesBuffer::new(); + let trace = compiled + .compute_batch_wgpu_into(&backend, batch, &mut buffers, &mut values) + .unwrap(); + assert_eq!(trace.backend, BackendKind::Wgpu); + assert_eq!(trace.mode, crate::DeviceExecutionMode::ComputeBatch); + assert_eq!(trace.transfers, buffers.plan().compute_transfer_plan); + assert!(trace.used_native_kernel); + assert_eq!(values.data, compiled.compute_batch(batch).unwrap().data); + assert_eq!( + buffers + .download(&backend, crate::DeviceBufferKind::Outputs) + .unwrap(), + values.data + ); + + let mut gradients = BatchGradientsBuffer::new(); + let gradient_trace = compiled + .gradient_batch_wgpu_into(&backend, batch, &mut buffers, &mut gradients) + .unwrap(); + assert_eq!(gradient_trace.backend, BackendKind::Wgpu); + assert_eq!( + gradient_trace.mode, + crate::DeviceExecutionMode::GradientBatch + ); + assert!(!gradient_trace.used_native_kernel); + let scalar_gradients = compiled.gradient_batch(batch).unwrap(); + assert_eq!(gradients.values, scalar_gradients.values); + assert_eq!(gradients.gradients, scalar_gradients.gradients); + + let mut graph_buffers = exact_graph + .allocate_wgpu_buffers(&backend, batch.batch_size) + .unwrap(); + let mut graph_values = BatchValuesBuffer::new(); + let graph_trace = exact_graph + .compute_batch_wgpu_into(&backend, batch, &mut graph_buffers, &mut graph_values) + .unwrap(); + assert!(graph_trace.used_native_kernel); + assert_eq!(graph_values.data, values.data); + + let inexact_batch = BatchInputs::new(&[0.1, -0.2, 1.5], 3, 1).unwrap(); + assert!(compiled.supports_native_wgpu_batch_compute(&backend)); + assert!(!compiled.supports_native_wgpu_batch_compute_for_batch(&backend, inexact_batch)); + assert!(!exact_graph + .supports_native_wgpu_batch_compute_for_batch(&backend, inexact_batch) + .unwrap()); + let inexact_scalar = compiled.compute_batch(inexact_batch).unwrap(); + let mut inexact_values = BatchValuesBuffer::new(); + let inexact_trace = compiled + .compute_batch_wgpu_into(&backend, inexact_batch, &mut buffers, &mut inexact_values) + .unwrap(); + assert!(!inexact_trace.used_native_kernel); + assert_eq!(inexact_values.data, inexact_scalar.data); + + let mut fallback_graph = Graph::new(2); + let left = fallback_graph.input(0); + let right = fallback_graph.input(1); + let product = fallback_graph.mul(left, right); + fallback_graph.set_output(product).unwrap(); + let fallback_compiled = fallback_graph.compile_ir().unwrap(); + let fallback_batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0, 6.0, 7.0], 3, 2).unwrap(); + assert!(!fallback_compiled.supports_native_wgpu_batch_compute(&backend)); + assert!(!fallback_graph + .supports_native_wgpu_batch_compute(&backend) + .unwrap()); + let mut fallback_buffers = fallback_graph + .allocate_wgpu_buffers(&backend, fallback_batch.batch_size) + .unwrap(); + let mut fallback_values = BatchValuesBuffer::new(); + let fallback_trace = fallback_graph + .compute_batch_wgpu_into( + &backend, + fallback_batch, + &mut fallback_buffers, + &mut fallback_values, + ) + .unwrap(); + assert!(!fallback_trace.used_native_kernel); + assert_eq!( + fallback_values.data, + fallback_compiled + .compute_batch(fallback_batch) + .unwrap() + .data + ); + } + + #[test] + fn test_simd_scalar_unary_fallbacks_and_abs_match_scalar_when_available() { + let compiled = { + let mut graph = Graph::new(1); + let x = graph.input(0); + let sin_x = graph.sin(x); + let cos_sin = graph.cos(sin_x); + let exp_cos = graph.exp(cos_sin); + let shifted = graph.add_const(exp_cos, 1.0); + let ln_shifted = graph.ln(shifted); + let tan_ln = graph.tan(ln_shifted); + let tanh_tan = graph.tanh(tan_ln); + let abs_tanh = graph.abs(tanh_tan); + graph.set_output(abs_tanh).unwrap(); + graph.compile_ir().unwrap() + }; + let batch = BatchInputs::new(&[-0.5, -0.0, 0.2, 0.7, 1.1], 5, 1).unwrap(); + let scalar_values = compiled.compute_batch(batch).unwrap(); + let scalar_gradients = compiled.gradient_batch(batch).unwrap(); + + for backend in [BackendKind::SimdF64x4, BackendKind::SimdF64x2] { + let simd_report = compiled.backend_support_report(backend).unwrap(); + assert!(simd_report.missing_opcodes.is_empty()); + if simd_report.can_compute_batch() { + let mut values = BatchValuesBuffer::new(); + backend + .compute_batch(&compiled, batch, &mut values) + .unwrap(); + assert_eq!(values.data.len(), scalar_values.data.len()); + for (left, right) in values.data.iter().zip(scalar_values.data.iter()) { + assert_eq!(left.to_bits(), right.to_bits()); + } + } + if simd_report.can_gradient_batch() { + let mut gradients = BatchGradientsBuffer::new(); + backend + .gradient_batch(&compiled, batch, &mut gradients) + .unwrap(); + assert_eq!(gradients.values.len(), scalar_gradients.values.len()); + assert_eq!(gradients.gradients.len(), scalar_gradients.gradients.len()); + for (left, right) in gradients.values.iter().zip(scalar_gradients.values.iter()) { + assert_eq!(left.to_bits(), right.to_bits()); + } + for (left, right) in gradients + .gradients + .iter() + .zip(scalar_gradients.gradients.iter()) + { + assert_eq!(left.to_bits(), right.to_bits()); + } + } + } + } + + #[test] + fn test_simd_scalar_binary_fallbacks_match_scalar_when_available() { + let compiled = { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let abs_x = graph.abs(x); + let base = graph.add_const(abs_x, 1.25); + let exponent = graph.log1p_exp(y); + let powered = graph.pow(base, exponent); + let mixed = graph.log_add_exp(powered, y); + let output = graph.tanh(mixed); + graph.set_output(output).unwrap(); + graph.compile_ir().unwrap() + }; + let batch = BatchInputs::new( + &[-0.5, -0.3, -0.0, 0.0, 0.2, 0.4, 0.7, -0.2, 1.1, 0.3], + 5, + 2, + ) + .unwrap(); + let scalar_values = compiled.compute_batch(batch).unwrap(); + let scalar_gradients = compiled.gradient_batch(batch).unwrap(); + + for backend in [BackendKind::SimdF64x4, BackendKind::SimdF64x2] { + let simd_report = compiled.backend_support_report(backend).unwrap(); + assert!(simd_report.missing_opcodes.is_empty()); + if simd_report.can_compute_batch() { + let mut values = BatchValuesBuffer::new(); + backend + .compute_batch(&compiled, batch, &mut values) + .unwrap(); + assert_eq!(values.data.len(), scalar_values.data.len()); + for (left, right) in values.data.iter().zip(scalar_values.data.iter()) { + assert_eq!(left.to_bits(), right.to_bits()); + } + } + if simd_report.can_gradient_batch() { + let mut gradients = BatchGradientsBuffer::new(); + backend + .gradient_batch(&compiled, batch, &mut gradients) + .unwrap(); + assert_eq!(gradients.values.len(), scalar_gradients.values.len()); + assert_eq!(gradients.gradients.len(), scalar_gradients.gradients.len()); + for (left, right) in gradients.values.iter().zip(scalar_gradients.values.iter()) { + assert_eq!(left.to_bits(), right.to_bits()); + } + for (left, right) in gradients + .gradients + .iter() + .zip(scalar_gradients.gradients.iter()) + { + assert_eq!(left.to_bits(), right.to_bits()); + } + } + } + } + + #[test] + fn test_simd_backend_batch_compute_multi_output_and_fallback_support() { + let simd_backend = SimdBackend; + if !simd_backend.capabilities().supports_batch_compute { + return; + } + + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let product = graph.mul(x, y); + let ratio = graph.div(x, y); + let difference = graph.sub(x, y); + let relu = graph.relu(difference); + let sum = graph.add(x, y); + let sqrt = graph.sqrt(sum); + graph.set_outputs(&[product, ratio, relu, sqrt]).unwrap(); + let compiled = graph.compile_ir().unwrap(); + let batch = BatchInputs::new(&[2.0, 1.0, 3.0, 1.5, 4.0, 2.0], 3, 2).unwrap(); + let scalar_values = compiled.compute_batch(batch).unwrap(); + let mut simd_values = BatchValuesBuffer::new(); + simd_backend + .compute_batch(&compiled, batch, &mut simd_values) + .unwrap(); + assert_eq!(simd_values.batch_size, 3); + assert_eq!(simd_values.output_dim, 4); + for (left, right) in simd_values.data.iter().zip(scalar_values.data.iter()) { + assert!(approx_eq(*left, *right, 1e-10)); + } + + let scalar_gradients = compiled.gradient_batch(batch).unwrap(); + let mut simd_gradients = BatchGradientsBuffer::new(); + simd_backend + .gradient_batch(&compiled, batch, &mut simd_gradients) + .unwrap(); + assert_eq!(simd_gradients.batch_size, 3); + assert_eq!(simd_gradients.input_dim, 2); + for (left, right) in simd_gradients + .values + .iter() + .zip(scalar_gradients.values.iter()) + { + assert!(approx_eq(*left, *right, 1e-10)); + } + for (left, right) in simd_gradients + .gradients + .iter() + .zip(scalar_gradients.gradients.iter()) + { + assert!(approx_eq(*left, *right, 1e-10)); + } + + let mut fallback_graph = Graph::new(1); + let z = fallback_graph.input(0); + let abs_z = fallback_graph.abs(z); + let shifted = fallback_graph.add_const(abs_z, 1.25); + let exponent = fallback_graph.log1p_exp(z); + let powered = fallback_graph.pow(shifted, exponent); + let output = fallback_graph.log_add_exp(powered, z); + fallback_graph.set_output(output).unwrap(); + let fallback_compiled = fallback_graph.compile_ir().unwrap(); + assert!(fallback_compiled + .validate_backend_capabilities(&simd_backend.capabilities()) + .is_ok()); + } + + #[test] + fn test_simd_relu_matches_scalar_edge_cases() { + let simd_backend = SimdBackend; + if !simd_backend.capabilities().supports_batch_compute { + return; + } + + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.relu(x); + let compiled = graph.compile_ir().unwrap(); + let inputs = [f64::NAN, -0.0, 0.0, -2.0, 3.0]; + let batch = BatchInputs::new(&inputs, inputs.len(), 1).unwrap(); + + let scalar_values = compiled.compute_batch(batch).unwrap(); + let mut simd_values = BatchValuesBuffer::new(); + simd_backend + .compute_batch(&compiled, batch, &mut simd_values) + .unwrap(); + assert_eq!(simd_values.data.len(), scalar_values.data.len()); + for (left, right) in simd_values.data.iter().zip(scalar_values.data.iter()) { + assert_eq!(left.to_bits(), right.to_bits()); + } + + let scalar_gradients = compiled.gradient_batch(batch).unwrap(); + let mut simd_gradients = BatchGradientsBuffer::new(); + simd_backend + .gradient_batch(&compiled, batch, &mut simd_gradients) + .unwrap(); + for (left, right) in simd_gradients + .values + .iter() + .zip(scalar_gradients.values.iter()) + { + assert_eq!(left.to_bits(), right.to_bits()); + } + for (left, right) in simd_gradients + .gradients + .iter() + .zip(scalar_gradients.gradients.iter()) + { + assert_eq!(left.to_bits(), right.to_bits()); + } + } + + #[test] + fn test_simd_gradient_masks_inactive_nan_lanes() { + let simd_backend = SimdBackend; + if !simd_backend.capabilities().supports_batch_gradient { + return; + } + + let mut relu_sqrt_graph = Graph::new(1); + let x = relu_sqrt_graph.input(0); + let sqrt_x = relu_sqrt_graph.sqrt(x); + relu_sqrt_graph.relu(sqrt_x); + let compiled = relu_sqrt_graph.compile_ir().unwrap(); + let inputs = [-1.0, 4.0, -9.0]; + let batch = BatchInputs::new(&inputs, inputs.len(), 1).unwrap(); + let scalar = compiled.gradient_batch(batch).unwrap(); + let mut simd = BatchGradientsBuffer::new(); + simd_backend + .gradient_batch(&compiled, batch, &mut simd) + .unwrap(); + for (left, right) in simd.gradients.iter().zip(scalar.gradients.iter()) { + assert_eq!(left.to_bits(), right.to_bits()); + } + + let mut unused_div_graph = Graph::new(1); + let x = unused_div_graph.input(0); + let zero = unused_div_graph.constant(0.0); + unused_div_graph.div(x, zero); + unused_div_graph.set_output(x).unwrap(); + let compiled = unused_div_graph.compile_ir().unwrap(); + let inputs = [0.0, 1.0, -1.0]; + let batch = BatchInputs::new(&inputs, inputs.len(), 1).unwrap(); + let scalar = compiled.gradient_batch(batch).unwrap(); + let mut simd = BatchGradientsBuffer::new(); + simd_backend + .gradient_batch(&compiled, batch, &mut simd) + .unwrap(); + assert_eq!(simd.values, scalar.values); + assert_eq!(simd.gradients, scalar.gradients); + } + + #[test] + fn test_reductions_losses_and_parameters() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mark_parameter(x).unwrap(); + graph.set_parameter_name(y, "weight").unwrap(); + assert_eq!(graph.parameters(), &[x, y]); + assert_eq!(graph.parameter_name(y), Some("weight")); + assert_eq!(graph.parameter_names().len(), 1); + + let dot = graph.dot(&[x, y], &[x, y]).unwrap(); + graph.set_output(dot).unwrap(); + let (_value, parameter_grad) = graph.gradient(&[3.0, 4.0]).unwrap(); + let extracted = graph.parameter_gradient(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(extracted[0].1, parameter_grad[0], 1e-10)); + assert!(approx_eq(extracted[1].1, parameter_grad[1], 1e-10)); + + let mut loss_graph = Graph::new(2); + let prediction = loss_graph.input(0); + let target = loss_graph.input(1); + let mse = loss_graph.mse_loss(&[prediction], &[target]).unwrap(); + loss_graph.set_output(mse).unwrap(); + assert!(approx_eq( + loss_graph.compute(&[5.0, 3.0]).unwrap(), + 4.0, + 1e-10 + )); + } + + #[test] + fn test_stable_math_regressions() { + let mut softplus_graph = Graph::new(1); + let x = softplus_graph.input(0); + softplus_graph.log1p_exp(x); + let (value, gradient) = softplus_graph.gradient(&[0.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::LN_2, 1e-12)); + assert!(approx_eq(gradient[0], 0.5, 1e-12)); + assert!(softplus_graph.compute(&[1000.0]).unwrap().is_finite()); + + let mut public_softplus_graph = Graph::new(1); + let x = public_softplus_graph.input(0); + public_softplus_graph.softplus(x); + assert!(public_softplus_graph + .compute(&[1000.0]) + .unwrap() + .is_finite()); + + let mut logsumexp_graph = Graph::new(2); + let a = logsumexp_graph.input(0); + let b = logsumexp_graph.input(1); + let lse = logsumexp_graph.logsumexp_approx(&[a, b]).unwrap(); + logsumexp_graph.set_output(lse).unwrap(); + let lse_value = logsumexp_graph.compute(&[1000.0, 1001.0]).unwrap(); + let expected = 1001.0 + (-1.0_f64).exp().ln_1p(); + assert!(lse_value.is_finite()); + assert!(approx_eq(lse_value, expected, 1e-10)); + + let mut softmax_graph = Graph::new(2); + let a = softmax_graph.input(0); + let b = softmax_graph.input(1); + let softmax = softmax_graph.stable_softmax_approx(&[a, b]); + softmax_graph.set_outputs(&softmax).unwrap(); + let values = softmax_graph.compute_many(&[1000.0, 1001.0]).unwrap(); + assert!(values.iter().all(|value| value.is_finite())); + assert!(approx_eq(values[0] + values[1], 1.0, 1e-12)); + } + + #[test] + fn test_softplus_fr_rf_hessian_lowering() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.softplus(x); + + for input in [0.0, 1000.0] { + let native = graph.exact_hessian_rr(&[input]).unwrap(); + let fr = graph.exact_hessian_fr(&[input]).unwrap(); + let rf = graph.exact_hessian_rf(&[input]).unwrap(); + + assert!(native[0][0].is_finite()); + assert!(fr[0][0].is_finite()); + assert!(rf[0][0].is_finite()); + assert!(approx_eq(fr[0][0], native[0][0], 1e-12)); + assert!(approx_eq(rf[0][0], native[0][0], 1e-12)); + } + + let at_zero = graph.exact_hessian_rr(&[0.0]).unwrap(); + let at_large = graph.exact_hessian_rr(&[1000.0]).unwrap(); + assert!(approx_eq(at_zero[0][0], 0.25, 1e-12)); + assert!(approx_eq(at_large[0][0], 0.0, 1e-12)); + } + + #[test] + fn test_batch_try_row_checks_bounds() { + let batch = BatchInputs::new(&[1.0, 2.0, 3.0, 4.0], 2, 2).unwrap(); + assert_eq!(batch.row(0), &[1.0, 2.0]); + assert_eq!(batch.try_row(0).unwrap(), &[1.0, 2.0]); + assert!(batch.try_row(2).is_err()); + } + + #[cfg(feature = "serde")] + #[test] + fn test_graph_serde_round_trip() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.square(x); + let encoded = serde_json::to_string(&graph).unwrap(); + let decoded: Graph = serde_json::from_str(&encoded).unwrap(); + assert!(approx_eq(decoded.compute(&[3.0]).unwrap(), 9.0, 1e-10)); + } + + #[cfg(feature = "serde")] + #[test] + fn test_graph_serde_accepts_missing_parameter_metadata() { + let encoded = r#"{ + "num_inputs": 1, + "nodes": [{"Operation": {"op": "Sin", "inputs": [0]}}], + "output_nodes": [], + "input_names": [null], + "output_names": [] + }"#; + let decoded: Graph = serde_json::from_str(encoded).unwrap(); + assert!(decoded.parameters().is_empty()); + assert!(approx_eq( + decoded.compute(&[0.5]).unwrap(), + 0.5_f64.sin(), + 1e-10 + )); + } + + // ==================== + // Additional coverage tests for graph.rs uncovered areas + // ==================== + + #[test] + fn test_graph_cos_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.cos(x); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 1.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[0.0]).unwrap(); + assert!(approx_eq(grad[0], 0.0, 1e-10)); + } + + #[test] + fn test_graph_tan_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.tan(x); + let value = graph.compute(&[0.5]).unwrap(); + assert!(approx_eq(value, 0.5_f64.tan(), 1e-10)); + + let (_v, grad) = graph.gradient(&[0.5]).unwrap(); + let expected_grad = 1.0 / 0.5_f64.cos().powi(2); + assert!(approx_eq(grad[0], expected_grad, 1e-8)); + } + + #[test] + fn test_graph_neg_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.neg(x); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, -3.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[3.0]).unwrap(); + assert!(approx_eq(grad[0], -1.0, 1e-10)); + } + + #[test] + fn test_graph_exp_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + let value = graph.compute(&[1.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::E, 1e-10)); + + let (_v, grad) = graph.gradient(&[1.0]).unwrap(); + assert!(approx_eq(grad[0], std::f64::consts::E, 1e-8)); + } + + #[test] + fn test_graph_ln_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.ln(x); + let value = graph.compute(&[std::f64::consts::E]).unwrap(); + assert!(approx_eq(value, 1.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[std::f64::consts::E]).unwrap(); + assert!(approx_eq(grad[0], 1.0 / std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_graph_sqrt_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sqrt(x); + let value = graph.compute(&[4.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[4.0]).unwrap(); + assert!(approx_eq(grad[0], 0.25, 1e-10)); + } + + #[test] + fn test_graph_abs_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.abs(x); + let value = graph.compute(&[-3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[5.0]).unwrap(); + assert!(approx_eq(grad[0], 1.0, 1e-10)); + } + + #[test] + fn test_graph_sub_node() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.sub(x, y); + let value = graph.compute(&[5.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[5.0, 3.0]).unwrap(); + assert!(approx_eq(grad[0], 1.0, 1e-10)); + assert!(approx_eq(grad[1], -1.0, 1e-10)); + } + + #[test] + fn test_graph_div_node() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.div(x, y); + let value = graph.compute(&[6.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[6.0, 3.0]).unwrap(); + assert!(approx_eq(grad[0], 1.0 / 3.0, 1e-10)); + assert!(approx_eq(grad[1], -2.0 / 3.0, 1e-10)); + } + + #[test] + fn test_graph_tanh_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.tanh(x); + let value = graph.compute(&[0.5]).unwrap(); + assert!(approx_eq(value, 0.5_f64.tanh(), 1e-10)); + } + + #[test] + fn test_graph_relu_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.relu(x); + + let pos = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(pos, 3.0, 1e-10)); + + let neg = graph.compute(&[-2.0]).unwrap(); + assert!(approx_eq(neg, 0.0, 1e-10)); + + let (_v, grad) = graph.gradient(&[3.0]).unwrap(); + assert!(approx_eq(grad[0], 1.0, 1e-10)); + } + + #[test] + fn test_graph_log1p_exp_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.log1p_exp(x); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::LN_2, 1e-10)); + } + + #[test] + fn test_graph_log_add_exp_node() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.log_add_exp(x, y); + let value = graph.compute(&[1.0, 2.0]).unwrap(); + let expected = (1.0_f64.exp() + 2.0_f64.exp()).ln(); + assert!(approx_eq(value, expected, 1e-10)); + } + + #[test] + fn test_graph_set_output_errors() { + let mut graph = Graph::new(1); + // Empty graph, no nodes - set_output should error + let result = graph.set_output(0); + assert!(result.is_ok()); // Input 0 is valid + + let result = graph.set_output(1); + assert!(result.is_err()); // No nodes, so index 1 is out of bounds + } + + #[test] + fn test_graph_set_outputs_errors() { + let mut graph = Graph::new(1); + let result = graph.set_outputs(&[1]); + assert!(result.is_err()); + } + + #[test] + fn test_graph_add_output() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + let product = graph.mul(x, y); + + graph.add_output(sum).unwrap(); + graph.add_output(product).unwrap(); + assert_eq!(graph.output_nodes(), &[sum, product]); + + // Error: out of bounds + assert!(graph.add_output(10).is_err()); + } + + #[test] + fn test_graph_add_output_error_empty_graph() { + let mut graph = Graph::new(0); + assert!(graph.add_output(0).is_err()); + } + + #[test] + fn test_graph_clear_output() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let _y = graph.mul(x, x); + graph.set_output(x).unwrap(); + assert_eq!(graph.output_node(), Some(x)); + + graph.clear_output(); + assert!(graph.output_node().is_none()); + + // After clear, effective_output falls back to last node + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 9.0, 1e-10)); + } + + #[test] + fn test_graph_effective_output_node_fallback() { + // No stored nodes, no explicit output + let graph = Graph::new(2); + let effective = graph.effective_output_node(); + // Should fall back to last input + assert_eq!(effective, Some(1)); + } + + #[test] + fn test_graph_effective_output_nodes_empty() { + let graph = Graph::new(2); + let nodes = graph.effective_output_nodes(); + assert_eq!(nodes, vec![1]); + } + + #[test] + fn test_graph_is_empty_and_len() { + let mut graph = Graph::new(2); + assert!(graph.is_empty()); + assert_eq!(graph.len(), 0); + + let x = graph.input(0); + graph.sin(x); + assert!(!graph.is_empty()); + assert_eq!(graph.len(), 1); + } + + #[test] + fn test_graph_next_node_id() { + let mut graph = Graph::new(2); + assert_eq!(graph.next_node_id(), 2); + + let x = graph.input(0); + graph.sin(x); + assert_eq!(graph.next_node_id(), 3); + } + + #[test] + fn test_graph_nodes_access() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sin(x); + assert_eq!(graph.nodes().len(), 1); + } + + #[test] + fn test_graph_num_inputs() { + let graph = Graph::new(3); + assert_eq!(graph.num_inputs(), 3); + } + + #[test] + fn test_graph_compute_checked_valid() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + let value = graph.compute_checked(&[1.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_graph_gradient_checked_valid() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + let (value, grad) = graph.gradient_checked(&[1.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::E, 1e-10)); + assert!(approx_eq(grad[0], std::f64::consts::E, 1e-8)); + } + + #[test] + fn test_graph_compute_many_checked_valid() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.set_outputs(&[x, y]).unwrap(); + let values = graph.compute_many_checked(&[1.0, 2.0]).unwrap(); + assert_eq!(values.len(), 2); + } + + #[test] + fn test_graph_jacobian_checked_valid() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.set_outputs(&[x, y]).unwrap(); + let jacobian = graph.jacobian_checked(&[1.0, 2.0]).unwrap(); + assert_eq!(jacobian.len(), 2); + assert!(approx_eq(jacobian[0][0], 1.0, 1e-10)); + assert!(approx_eq(jacobian[0][1], 0.0, 1e-10)); + } + + #[test] + fn test_graph_value_and_jacobian() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + graph.set_output(sum).unwrap(); + + let (values, jacobian) = graph.value_and_jacobian(&[2.0, 3.0]).unwrap(); + assert_eq!(values.len(), 1); + assert!(approx_eq(values[0], 5.0, 1e-10)); + assert_eq!(jacobian.len(), 1); + } + + #[test] + fn test_graph_value_and_jacobian_checked() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let sum = graph.add(x, y); + graph.set_output(sum).unwrap(); + + let (values, _jacobian) = graph.value_and_jacobian_checked(&[2.0, 3.0]).unwrap(); + assert_eq!(values.len(), 1); + assert!(approx_eq(values[0], 5.0, 1e-10)); + } + + #[test] + fn test_graph_compile_ir_valid() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + + let compiled = graph.compile_ir().unwrap(); + assert_eq!(compiled.num_inputs(), 2); + assert_eq!(compiled.instructions().len(), 1); + } + + #[test] + fn test_graph_compile_accelerated() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + + let compiled = graph.compile_accelerated().unwrap(); + let value = compiled.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 1.0, 1e-10)); + } + + #[test] + fn test_graph_try_compile() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sin(x); + + let tape = graph.try_compile().unwrap(); + let value = tape.compute(&[0.5]).unwrap(); + assert!(approx_eq(value, 0.5_f64.sin(), 1e-10)); + } + + #[test] + fn test_graph_try_compile_invalid() { + let mut graph = Graph::new(1); + graph.push_operation(MultiAD::Sin, vec![5]); // Bad index + assert!(graph.try_compile().is_err()); + } + + #[test] + fn test_graph_from_operations_fallback() { + // Use an invalid operation to trigger the fallback path in from_operations + let ops = [(MultiAD::Sin, vec![0])]; // Valid, should not trigger fallback + let graph = Graph::from_operations(1, &ops); + let value = graph.compute(&[0.5]).unwrap(); + assert!(approx_eq(value, 0.5_f64.sin(), 1e-10)); + } + + #[test] + fn test_graph_try_from_operations_invalid_input_marker() { + let ops = [(MultiAD::Inp, vec![5])]; // Index out of bounds + assert!(Graph::try_from_operations(2, &ops).is_err()); + } + + #[test] + fn test_graph_simplify_mul_identity() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let one = graph.constant(1.0); + graph.mul(x, one); + + let simplified = graph.simplify().unwrap(); + let value = simplified.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_graph_simplify_mul_left_identity() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let one = graph.constant(1.0); + graph.mul(one, x); + + let simplified = graph.simplify().unwrap(); + let value = simplified.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_graph_simplify_add_right_zero() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let zero = graph.constant(0.0); + graph.add(x, zero); + + let simplified = graph.simplify().unwrap(); + let value = simplified.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_graph_simplify_constant_fold() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let two = graph.constant(2.0); + let three = graph.constant(3.0); + let sum = graph.add(two, three); + graph.mul(x, sum); + + let simplified = graph.simplify().unwrap(); + let value = simplified.compute(&[2.0]).unwrap(); + assert!(approx_eq(value, 10.0, 1e-10)); + } + + #[test] + fn test_graph_set_input_name_errors() { + let mut graph = Graph::new(2); + assert!(graph.set_input_name(5, "z").is_err()); + assert!(graph.set_input_name(0, "x").is_ok()); + assert_eq!(graph.input_name(0), Some("x")); + assert_eq!(graph.input_name(1), None); + } + + #[test] + fn test_graph_mark_parameter_errors() { + let mut graph = Graph::new(1); + let x = graph.input(0); + assert!(graph.mark_parameter(5).is_err()); + assert!(graph.mark_parameter(x).is_ok()); + // Duplicate mark should be ok + assert!(graph.mark_parameter(x).is_ok()); + assert_eq!(graph.parameters().len(), 1); + } + + #[test] + fn test_graph_mark_parameter_empty_graph() { + let mut graph = Graph::new(0); + assert!(graph.mark_parameter(0).is_err()); + } + + #[test] + fn test_graph_validate_bad_parameter() { + let mut graph = Graph::new(1); + graph.parameters = vec![5]; + assert!(graph.validate().is_err()); + } + + #[test] + fn test_graph_set_output_name() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let y = graph.sin(x); + graph.set_output_name(y, "sin_x").unwrap(); + assert!(graph.to_mermaid().contains("sin_x")); + assert!(graph.to_dot().contains("sin_x")); + } + + #[test] + fn test_graph_input_name_out_of_bounds() { + let graph = Graph::new(1); + assert_eq!(graph.input_name(5), None); + } + + #[test] + fn test_graph_parameter_gradient_non_input_errors() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let sin_x = graph.sin(x); + graph.mark_parameter(sin_x).unwrap(); + assert!(graph.parameter_gradient(&[1.0]).is_err()); + } + + #[test] + fn test_graph_stats_constants_and_depth() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let c = graph.constant(2.0); + graph.mul(x, c); + + let stats = graph.stats(); + assert_eq!(stats.num_inputs, 1); + assert_eq!(stats.num_constants, 1); + assert_eq!(stats.num_ops, 1); + assert_eq!(stats.num_edges, 2); + assert_eq!(stats.max_depth, 1); + assert_eq!(stats.op_counts.len(), 1); + } + + #[test] + fn test_graph_stats_empty() { + let graph = Graph::new(2); + let stats = graph.stats(); + assert_eq!(stats.num_inputs, 2); + assert_eq!(stats.num_constants, 0); + assert_eq!(stats.num_ops, 0); + assert_eq!(stats.max_depth, 0); + } + + #[test] + fn test_graph_compute_with_domain_policy_unchecked() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.ln(x); + + // Unchecked should not error even for ln(0) + // (Actually ln(0) gives -inf which is a valid f64, but we use a small positive number) + let value = graph + .compute_with_domain_policy(&[1.0], DomainPolicy::Unchecked) + .unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + } + + #[test] + fn test_graph_compute_with_domain_policy_checked() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.ln(x); + + let value = graph + .compute_with_domain_policy(&[1.0], DomainPolicy::Checked) + .unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + + assert!(graph + .compute_with_domain_policy(&[0.0], DomainPolicy::Checked) + .is_err()); + } + + #[test] + fn test_graph_gradient_with_domain_policy() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + + let (val, _grad) = graph + .gradient_with_domain_policy(&[1.0], DomainPolicy::Unchecked) + .unwrap(); + assert!(approx_eq(val, std::f64::consts::E, 1e-10)); + + let (val2, _grad2) = graph + .gradient_with_domain_policy(&[1.0], DomainPolicy::Checked) + .unwrap(); + assert!(approx_eq(val2, std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_graph_check_gradient_pass() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let report = graph.check_gradient(&[3.0, 4.0], 1e-5).unwrap(); + assert!(report.passed); + assert_eq!(report.entries.len(), 2); + assert!(report.max_abs_error < 1e-5); + } + + #[test] + fn test_graph_compute_hessian() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let hessian = graph.compute_hessian(&[1.0, 1.0]).unwrap(); + assert_eq!(hessian.len(), 2); + assert!(approx_eq(hessian[0][0], 0.0, 1e-4)); // d²/dx² + assert!(approx_eq(hessian[0][1], 1.0, 1e-4)); // d²/dxdy + assert!(approx_eq(hessian[1][0], 1.0, 1e-4)); // d²/dydx + assert!(approx_eq(hessian[1][1], 0.0, 1e-4)); // d²/dy² + } + + #[test] + fn test_graph_exact_hessian_with_constants() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let two = graph.constant(2.0); + graph.mul(x, two); + + let hessian = graph.exact_hessian_rr(&[3.0]).unwrap(); + assert!(approx_eq(hessian[0][0], 0.0, 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_errors_for_nonsmooth() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.relu(x); + + assert!(graph.exact_hessian_rr(&[1.0]).is_err()); + } + + #[test] + fn test_graph_exact_hessian_errors_for_abs() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.abs(x); + + assert!(graph.exact_hessian_rr(&[1.0]).is_err()); + } + + #[test] + fn test_graph_exact_hessian_fr_with_constants() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let two = graph.constant(2.0); + graph.mul(x, two); + + let hessian = graph.exact_hessian_fr(&[3.0]).unwrap(); + assert!(approx_eq(hessian[0][0], 0.0, 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_rf_with_constants() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let two = graph.constant(2.0); + graph.mul(x, two); + + let hessian = graph.exact_hessian_rf(&[3.0]).unwrap(); + assert!(approx_eq(hessian[0][0], 0.0, 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_fr_nonsmooth_errors() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.tanh(x); + + assert!(graph.exact_hessian_fr(&[0.0]).is_err()); + } + + #[test] + fn test_graph_exact_hessian_rf_nonsmooth_errors() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.tanh(x); + + assert!(graph.exact_hessian_rf(&[0.0]).is_err()); + } + + #[test] + fn test_graph_binary_cross_entropy_loss() { + let mut graph = Graph::new(2); + let pred = graph.input(0); + let target = graph.input(1); + + // Create probability and target nodes + let sigmoid = graph.sigmoid(pred); + let bce = graph + .binary_cross_entropy_loss(&[sigmoid], &[target]) + .unwrap(); + graph.set_output(bce).unwrap(); + + let value = graph.compute(&[0.0, 1.0]).unwrap(); + // sigmoid(0) = 0.5, target=1 => -ln(0.5) = ln(2) + assert!(approx_eq(value, std::f64::consts::LN_2, 1e-6)); + } + + #[test] + fn test_graph_binary_cross_entropy_empty() { + let mut _graph = Graph::new(1); + assert!(Graph::new(1).binary_cross_entropy_loss(&[], &[]).is_err()); + } + + #[test] + fn test_graph_binary_cross_entropy_mismatched() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + assert!(graph.binary_cross_entropy_loss(&[x], &[y, x]).is_err()); + } + + #[test] + fn test_graph_mse_loss_errors() { + let mut graph = Graph::new(1); + let x = graph.input(0); + assert!(graph.mse_loss(&[], &[]).is_err()); + assert!(graph.mse_loss(&[x], &[]).is_err()); + } + + #[test] + fn test_graph_dot_errors() { + let mut graph = Graph::new(1); + let x = graph.input(0); + assert!(graph.dot(&[], &[]).is_err()); + assert!(graph.dot(&[x], &[]).is_err()); + } + + #[test] + fn test_graph_sum_mean_norm2_empty() { + let mut graph = Graph::new(1); + assert!(graph.sum(&[]).is_none()); + assert!(graph.mean(&[]).is_none()); + assert!(graph.norm2(&[]).is_none()); + assert!(graph.sum_squares(&[]).is_none()); + } + + #[test] + fn test_graph_sum_mean_norm2() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + + let sum = graph.sum(&[x, y]).unwrap(); + let value = graph.set_output(sum).unwrap().compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 7.0, 1e-10)); + + let mean = graph.mean(&[x, y]).unwrap(); + graph.set_output(mean).unwrap(); + let value = graph.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 3.5, 1e-10)); + + let norm2 = graph.norm2(&[x, y]).unwrap(); + graph.set_output(norm2).unwrap(); + let value = graph.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 5.0, 1e-10)); + } + + #[test] + fn test_graph_logsumexp_approx_empty() { + let mut graph = Graph::new(1); + assert!(graph.logsumexp_approx(&[]).is_none()); + } + + #[test] + fn test_graph_stable_softmax_approx_empty() { + let mut graph = Graph::new(1); + let result = graph.stable_softmax_approx(&[]); + assert!(result.is_empty()); + } + + #[test] + fn test_graph_sigmoid_stable() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let sigmoid = graph.sigmoid_stable(x); + graph.set_output(sigmoid).unwrap(); + + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.5, 1e-10)); + + let val_high = graph.compute(&[100.0]).unwrap(); + assert!(val_high > 0.99); + } + + #[test] + fn test_graph_prune_preserves_parameters() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mark_parameter(x).unwrap(); + graph.set_parameter_name(y, "weight").unwrap(); + let used = graph.square(x); + let _unused = graph.square(y); + graph.set_output(used).unwrap(); + + let pruned = graph.prune_to_outputs().unwrap(); + // The pruned graph should still evaluate correctly + let value = pruned.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 9.0, 1e-10)); + } + + #[test] + fn test_graph_compiled_metadata() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + + let metadata = graph.compiled_metadata().unwrap(); + assert_eq!(metadata.num_inputs, 2); + assert_eq!(metadata.num_outputs, 1); + } + + #[test] + fn test_graph_compiled_workspace() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + + let _workspace = graph.compiled_workspace().unwrap(); + } + + #[test] + fn test_graph_hessian_vector_product_errors() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + // Wrong vector length + assert!(graph.hessian_vector_product(&[1.0, 2.0], &[1.0]).is_err()); + } + + #[test] + fn test_graph_compute_batch_and_gradient_batch() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); + + let values = graph.compute_batch(batch).unwrap(); + assert_eq!(values.data.len(), 2); + assert!(approx_eq(values.data[0], 6.0, 1e-10)); + assert!(approx_eq(values.data[1], 20.0, 1e-10)); + + let grads = graph.gradient_batch(batch).unwrap(); + assert_eq!(grads.batch_size, 2); + assert!(approx_eq(grads.values[0], 6.0, 1e-10)); + } + + #[test] + fn test_graph_compute_batch_auto_and_gradient_batch_auto() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); + + let (_backend, values) = graph.compute_batch_auto(batch).unwrap(); + assert_eq!(values.data.len(), 2); + + let (_backend, grads) = graph.gradient_batch_auto(batch).unwrap(); + assert_eq!(grads.values.len(), 2); + } + + #[test] + fn test_graph_compute_batch_into_and_gradient_batch_into() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let batch = BatchInputs::new(&[2.0, 3.0], 1, 2).unwrap(); + + let mut values_buffer = BatchValuesBuffer::new(); + graph.compute_batch_into(batch, &mut values_buffer).unwrap(); + assert!(approx_eq(values_buffer.data[0], 6.0, 1e-10)); + + let mut grad_buffer = BatchGradientsBuffer::new(); + graph.gradient_batch_into(batch, &mut grad_buffer).unwrap(); + assert!(approx_eq(grad_buffer.values[0], 6.0, 1e-10)); + } + + #[test] + fn test_graph_compute_batch_auto_into_and_gradient_batch_auto_into() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let batch = BatchInputs::new(&[2.0, 3.0], 1, 2).unwrap(); + + let mut values_buffer = BatchValuesBuffer::new(); + let backend = graph + .compute_batch_auto_into(batch, &mut values_buffer) + .unwrap(); + assert!(approx_eq(values_buffer.data[0], 6.0, 1e-10)); + + let mut grad_buffer = BatchGradientsBuffer::new(); + let grad_backend = graph + .gradient_batch_auto_into(batch, &mut grad_buffer) + .unwrap(); + assert_eq!(backend, grad_backend); + } + + #[test] + fn test_graph_backend_support_reports() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let reports = graph.backend_support_reports().unwrap(); + assert_eq!(reports.len(), 5); + + let scalar_report = graph.backend_support_report(BackendKind::Scalar).unwrap(); + assert!(scalar_report.can_compute_batch()); + } + + #[test] + fn test_graph_device_batch_plan() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let plan = graph.device_batch_plan(BackendKind::Scalar, 4).unwrap(); + assert_eq!(plan.batch_size, 4); + } + + #[test] + fn test_graph_allocate_mock_device_buffers() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let buffers = graph.allocate_mock_device_buffers(2).unwrap(); + assert_eq!(buffers.buffers().len(), 5); + } + + #[test] + fn test_graph_compute_batch_mock_device_into() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); + let mut buffers = graph.allocate_mock_device_buffers(2).unwrap(); + let mut output = BatchValuesBuffer::new(); + + let trace = graph + .compute_batch_mock_device_into(batch, &mut buffers, &mut output) + .unwrap(); + assert_eq!(trace.mode, crate::DeviceExecutionMode::ComputeBatch); + } + + #[test] + fn test_graph_gradient_batch_mock_device_into() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let batch = BatchInputs::new(&[2.0, 3.0, 4.0, 5.0], 2, 2).unwrap(); + let mut buffers = graph.allocate_mock_device_buffers(2).unwrap(); + let mut output = BatchGradientsBuffer::new(); + + let trace = graph + .gradient_batch_mock_device_into(batch, &mut buffers, &mut output) + .unwrap(); + assert_eq!(trace.mode, crate::DeviceExecutionMode::GradientBatch); + } + + #[test] + fn test_graph_simd_support_report() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let report = graph.simd_support_report().unwrap(); + assert!(report.missing_opcodes.is_empty()); + } + + #[test] + fn test_graph_recommended_batch_backends() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let compute = graph.recommended_batch_compute_backend().unwrap(); + let gradient = graph.recommended_batch_gradient_backend().unwrap(); + assert!( + compute == BackendKind::Scalar + || compute == BackendKind::SimdF64x4 + || compute == BackendKind::SimdF64x2 + ); + assert!( + gradient == BackendKind::Scalar + || gradient == BackendKind::SimdF64x4 + || gradient == BackendKind::SimdF64x2 + ); + } + + #[test] + fn test_graph_to_operations_input_output() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let _y = graph.input(1); + graph.set_output(x).unwrap(); + + let ops = graph.to_operations().unwrap(); + // No operations, just the input marker + assert_eq!(ops, vec![(MultiAD::Inp, vec![0])]); + } + + #[test] + fn test_graph_parse_expression() { + let graph = Graph::parse_expression("x + y", &["x", "y"]).unwrap(); + let value = graph.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 7.0, 1e-10)); + } + + #[test] + fn test_graph_parse_expression_complex() { + let graph = Graph::parse_expression("sin(x) * y + 2", &["x", "y"]).unwrap(); + let value = graph.compute(&[0.5, 3.0]).unwrap(); + let expected = 0.5_f64.sin() * 3.0 + 2.0; + assert!(approx_eq(value, expected, 1e-10)); + } + + #[test] + fn test_graph_to_operations_for_input_only() { + // Empty graph with only inputs - output should be last input + let graph = Graph::new(2); + let ops = graph.to_operations().unwrap(); + assert_eq!(ops, vec![(MultiAD::Inp, vec![1])]); // Falls back to last input + } + + #[test] + fn test_graph_to_operations_no_output() { + // Graph with 0 inputs and 0 nodes + let graph = Graph::new(0); + let ops = graph.to_operations().unwrap(); + assert!(ops.is_empty()); + } + + #[test] + fn test_expr_graph_neg_sub_div() { + let expr = ExprGraph::new(2); + let x = expr.input(0); + let y = expr.input(1); + let neg_x = -x.clone(); + let diff = y.clone() - neg_x; + let quotient = diff / x; + expr.set_output("ient).unwrap(); + let graph = expr.graph(); + let value = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.5, 1e-10)); // (3 - (-2)) / 2 + } + + #[test] + fn test_expr_graph_with_constants() { + let expr = ExprGraph::new(1); + let x = expr.input(0); + let result = x.clone() + 1.0; + let result2 = result * 2.0; + let result3 = result2.clone() - 0.5; + let result4 = result3 / 1.0; + expr.set_output(&result4).unwrap(); + let graph = expr.graph(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 7.5, 1e-10)); // ((3+1)*2 - 0.5)/1 + } + + #[test] + fn test_expr_node_unary_methods() { + let expr = ExprGraph::new(1); + let x = expr.input(0); + let sin_x = x.sin(); + let cos_x = sin_x.cos(); + let exp_x = cos_x.exp(); + let ln_x = exp_x.ln(); + let sqrt_x = ln_x.sqrt(); + expr.set_output(&sqrt_x).unwrap(); + let graph = expr.graph(); + let value = graph.compute(&[0.5]).unwrap(); + let expected = (0.5_f64.sin().cos().exp()).ln().sqrt(); + assert!(approx_eq(value, expected, 1e-8)); + } + + #[test] + fn test_try_from_graph_to_operations() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + + let ops: Vec<(MultiAD, Vec)> = (&graph).try_into().unwrap(); + let value = MultiAD::compute(&ops, &[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 7.0, 1e-10)); + } + + #[test] + fn test_try_from_graph_ref_to_operations() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + + let ops: Vec<(MultiAD, Vec)> = (&graph).try_into().unwrap(); + assert_eq!(ops.len(), 1); + } + + #[test] + fn test_graph_simplify_unary_chain() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let neg_x = graph.neg(x); + let neg_neg = graph.neg(neg_x); + graph.set_output(neg_neg).unwrap(); + + // Simplification won't fold neg(neg(x)) yet, but it should still evaluate correctly + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_graph_simplify_add_left_zero() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let zero = graph.constant(0.0); + graph.add(zero, x); + + let simplified = graph.simplify().unwrap(); + let value = simplified.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_graph_simplify_complex() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let zero = graph.constant(0.0); + let one = graph.constant(1.0); + let x_plus_zero = graph.add(x, zero); + let times_one = graph.mul(x_plus_zero, one); + graph.set_output(times_one).unwrap(); + + let simplified = graph.simplify().unwrap(); + let value = simplified.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_fr_cos() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.cos(x); + + let hessian = graph.exact_hessian_fr(&[0.5]).unwrap(); + // d²/dx² cos(x) = -cos(x) + assert!(approx_eq(hessian[0][0], -0.5_f64.cos(), 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_rf_cos() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.cos(x); + + let hessian = graph.exact_hessian_rf(&[0.5]).unwrap(); + assert!(approx_eq(hessian[0][0], -0.5_f64.cos(), 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_rr_sin() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sin(x); + + let hessian = graph.exact_hessian_rr(&[0.5]).unwrap(); + // d²/dx² sin(x) = -sin(x) + assert!(approx_eq(hessian[0][0], -0.5_f64.sin(), 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_binary_ops() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let product = graph.mul(x, y); + graph.set_output(product).unwrap(); + + let hessian = graph.exact_hessian_rr(&[3.0, 4.0]).unwrap(); + // d²(xy)/dx² = 0, d²(xy)/dxdy = 1, d²(xy)/dy² = 0 + assert!(approx_eq(hessian[0][0], 0.0, 1e-10)); + assert!(approx_eq(hessian[0][1], 1.0, 1e-10)); + assert!(approx_eq(hessian[1][0], 1.0, 1e-10)); + assert!(approx_eq(hessian[1][1], 0.0, 1e-10)); + } + + #[test] + fn test_graph_exact_hessian_div() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.div(x, y); + + let hessian = graph.exact_hessian_rr(&[4.0, 2.0]).unwrap(); + // d²(x/y)/dx² = 0, d²(x/y)/dxdy = -1/y² + assert!(approx_eq(hessian[0][0], 0.0, 1e-10)); + assert!(approx_eq(hessian[0][1], -0.25, 1e-10)); + } + + #[test] + fn test_graph_compute_grad_closure() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let (value, grad_fn) = graph.compute_grad(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 12.0, 1e-10)); + let grad = grad_fn(1.0); + assert!(approx_eq(grad[0], 4.0, 1e-10)); + assert!(approx_eq(grad[1], 3.0, 1e-10)); + } + + #[test] + fn test_graph_compute_grad_multi_output() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.set_outputs(&[x, y]).unwrap(); + + let (value, grad_fn) = graph.compute_grad(&[3.0, 4.0]).unwrap(); + // With multi-output, compute_grad returns the first output + assert!(approx_eq(value, 3.0, 1e-10)); + let grad = grad_fn(1.0); + assert!(approx_eq(grad[0], 1.0, 1e-10)); + assert!(approx_eq(grad[1], 0.0, 1e-10)); + } + + #[test] + fn test_graph_compute_hessian_empty() { + let graph = Graph::new(0); + let hessian = graph.compute_hessian(&[]).unwrap(); + assert_eq!(hessian.len(), 0); + } + + #[test] + fn test_graph_gradient_sparse_with_tolerance() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let sparse = graph + .gradient_sparse_with_tolerance(&[0.0, 4.0], 0.0) + .unwrap(); + // At x=0, dy/dx=4, dy/dy=0 + assert_eq!(sparse.len(), 1); // Only the y gradient (0) should be filtered? No - tolerance is 0 so both + assert_eq!(sparse.len(), 1); // dy/dx=4 > 0, dy/dy=0 == 0 + // Actually at x=0,y=4: dy/dx = y = 4, dy/dy = x = 0 + } + + #[test] + fn test_graph_hessian_sparse() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + + let sparse = graph.hessian_sparse(&[1.0, 2.0]).unwrap(); + // Hessian of x+y is all zeros + assert!(sparse.is_empty()); + } + + #[test] + fn test_graph_hessian_sparse_with_tolerance() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + + let sparse = graph + .hessian_sparse_with_tolerance(&[1.0, 2.0], 0.0) + .unwrap(); + // Hessian of xy: d²/dxdy = 1, d²/dydx = 1, rest are 0 + assert_eq!(sparse.len(), 2); + } + + #[test] + fn test_graph_compile_ir_with_invalid_ops() { + let mut graph = Graph::new(1); + // Push a ternary-like operation which should fail IR compilation + graph.push_operation(MultiAD::Add, vec![0, 0, 0]); // 3 args for Add + assert!(graph.compile_ir().is_err()); + } + + #[test] + fn test_graph_validate_rejects_inp_marker() { + let mut graph = Graph::new(1); + graph.push_operation(MultiAD::Inp, vec![0]); + assert!(graph.validate().is_err()); + } + + #[test] + fn test_graph_validate_rejects_wrong_arity() { + let mut graph = Graph::new(1); + graph.push_operation(MultiAD::Sin, vec![0, 1]); // Sin takes 1 arg + assert!(graph.validate().is_err()); + } + + #[test] + fn test_graph_try_push_operation_errors() { + let mut graph = Graph::new(1); + // Inp marker + assert!(graph.try_push_operation(MultiAD::Inp, vec![0]).is_err()); + // Out of bounds + assert!(graph.try_push_operation(MultiAD::Sin, vec![5]).is_err()); + } + + #[test] + fn test_graph_try_unary_and_binary_errors() { + let mut graph = Graph::new(1); + assert!(graph.try_sin(5).is_err()); + assert!(graph.try_cos(5).is_err()); + assert!(graph.try_tan(5).is_err()); + assert!(graph.try_tanh(5).is_err()); + assert!(graph.try_relu(5).is_err()); + assert!(graph.try_log1p_exp(5).is_err()); + assert!(graph.try_neg(5).is_err()); + assert!(graph.try_exp(5).is_err()); + assert!(graph.try_ln(5).is_err()); + assert!(graph.try_sqrt(5).is_err()); + assert!(graph.try_abs(5).is_err()); + assert!(graph.try_add(5, 0).is_err()); + assert!(graph.try_sub(5, 0).is_err()); + assert!(graph.try_mul(5, 0).is_err()); + assert!(graph.try_div(5, 0).is_err()); + assert!(graph.try_pow(5, 0).is_err()); + assert!(graph.try_log_add_exp(5, 0).is_err()); + } + + #[test] + fn test_graph_constant_and_square_cube_reciprocal() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let sq = graph.square(x); + let cu = graph.cube(x); + let recip = graph.reciprocal(x); + graph.set_outputs(&[sq, cu, recip]).unwrap(); + + let values = graph.compute_many(&[2.0]).unwrap(); + assert!(approx_eq(values[0], 4.0, 1e-10)); + assert!(approx_eq(values[1], 8.0, 1e-10)); + assert!(approx_eq(values[2], 0.5, 1e-10)); + } + + #[test] + fn test_graph_pow_const() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.pow_const(x, 3.0); + let value = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(value, 8.0, 1e-10)); + } + + #[test] + fn test_graph_sub_const_and_div_const() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let subbed = graph.sub_const(x, 1.0); + graph.div_const(subbed, 2.0); + let value = graph.compute(&[5.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_graph_gelu() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.gelu(x); + let value = graph.compute(&[1.0]).unwrap(); + // GELU(1) ≈ 0.8412 + assert!(value > 0.0 && value < 1.0); + } + + #[test] + fn test_graph_node_id_and_output_node() { + let mut graph = Graph::new(2); + let x = graph.input(0); + assert_eq!(x, 0); + + let y = graph.input(1); + assert_eq!(y, 1); + + let sum = graph.add(x, y); + assert_eq!(sum, 2); + + graph.set_output(sum).unwrap(); + assert_eq!(graph.output_node(), Some(2)); + assert_eq!(graph.output_nodes(), &[2]); + } + + #[test] + fn test_graph_to_mermaid_with_output_names() { + let mut graph = Graph::new(1); + graph.set_input_name(0, "x").unwrap(); + let x = graph.input(0); + let sin_x = graph.sin(x); + graph.set_output_name(sin_x, "result").unwrap(); + + let mermaid = graph.to_mermaid(); + assert!(mermaid.contains("x: Input 0")); + assert!(mermaid.contains("result")); + } + + #[test] + fn test_graph_to_dot_with_output_names() { + let mut graph = Graph::new(1); + graph.set_input_name(0, "x").unwrap(); + let x = graph.input(0); + let sin_x = graph.sin(x); + graph.set_output_name(sin_x, "result").unwrap(); + + let dot = graph.to_dot(); + assert!(dot.contains("x: Input 0")); + assert!(dot.contains("result")); + } + + #[test] + fn test_tape_compute_many_with_workspace_checked() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + let tape = graph.compile(); + let mut workspace = tape.workspace(); + + let values = tape + .compute_many_with_workspace_checked(&[1.0], &mut workspace) + .unwrap(); + assert!(approx_eq(values[0], std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_tape_compute_with_workspace_checked() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.exp(x); + let tape = graph.compile(); + let mut workspace = tape.workspace(); + + let value = tape + .compute_with_workspace_checked(&[1.0], &mut workspace) + .unwrap(); + assert!(approx_eq(value, std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_tape_jacobian_with_workspace() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.set_outputs(&[x, y]).unwrap(); + let tape = graph.compile(); + let mut workspace = tape.workspace(); + + let jacobian = tape + .jacobian_with_workspace(&[1.0, 2.0], &mut workspace) + .unwrap(); + assert_eq!(jacobian.len(), 2); + } + + #[test] + fn test_tape_jacobian_with_workspace_checked() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.set_outputs(&[x, y]).unwrap(); + let tape = graph.compile(); + let mut workspace = tape.workspace(); + + let jacobian = tape + .jacobian_with_workspace_checked(&[1.0, 2.0], &mut workspace) + .unwrap(); + assert_eq!(jacobian.len(), 2); + } + + #[test] + fn test_tape_graph_access() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.add(x, y); + let tape = graph.compile(); + assert_eq!(tape.graph().num_inputs(), 2); + } + + #[test] + fn test_tape_compute_hessian() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + let tape = graph.compile(); + + let hessian = tape.compute_hessian(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(hessian[0][1], 1.0, 1e-4)); + assert!(approx_eq(hessian[1][0], 1.0, 1e-4)); + } + + #[test] + fn test_tape_workspace_clear() { + let mut ws = TapeWorkspace::new(); + ws.values.push(1.0); + ws.cotangent_values.push(2.0); + ws.gradients.push(3.0); + ws.clear(); + assert!(ws.values.is_empty()); + assert!(ws.cotangent_values.is_empty()); + assert!(ws.gradients.is_empty()); + } + + #[test] + fn test_graph_value_and_gradient_alias() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.square(x); + + let (v1, g1) = graph.gradient(&[3.0]).unwrap(); + let (v2, g2) = graph.value_and_gradient(&[3.0]).unwrap(); + assert!(approx_eq(v1, v2, 1e-10)); + assert_eq!(g1, g2); + } + + #[test] + fn test_graph_check_gradient_report() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.square(x); + + let report = graph.check_gradient(&[3.0], 1e-5).unwrap(); + assert!(report.passed); + assert_eq!(report.entries.len(), 1); + assert!(report.entries[0].abs_error < 1e-5); + } + + // ==================== + // Coverage: op methods, checked, reductions, hessian, parameter metadata + // ==================== + + #[test] + fn test_graph_pow_forward_and_gradient() { + let mut graph = Graph::new(2); + let base = graph.input(0); + let exp = graph.input(1); + graph.pow(base, exp); + // 2^3 = 8 + let value = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 8.0, 1e-10)); + // dx: 3*2^(3-1)=12, dy: 8*ln(2)≈5.545 + let (_v, grad) = graph.gradient(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(grad[0], 12.0, 1e-8)); + assert!(approx_eq(grad[1], 8.0 * 2.0_f64.ln(), 1e-8)); + } + + #[test] + fn test_graph_pow_const_compute() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.pow_const(x, 4.0); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 81.0, 1e-10)); + } + + #[test] + fn test_graph_softplus_compute_and_gradient() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.softplus(x); + // softplus(0) = ln(2) + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::LN_2, 1e-12)); + // gradient at 0 = sigmoid(0) = 0.5 + let (_v, grad) = graph.gradient(&[0.0]).unwrap(); + assert!(approx_eq(grad[0], 0.5, 1e-12)); + // softplus large x ≈ x + let large_v = graph.compute(&[1000.0]).unwrap(); + assert!((large_v - 1000.0).abs() < 1e-10); + } + + #[test] + fn test_graph_sigmoid_compute_and_gradient() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sigmoid(x); + // sigmoid(0) = 0.5 + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.5, 1e-10)); + // d/dx sigmoid at 0 = sigmoid(0)*(1-sigmoid(0)) = 0.25 + let (_v, grad) = graph.gradient(&[0.0]).unwrap(); + assert!(approx_eq(grad[0], 0.25, 1e-10)); + // sigmoid large x ≈ 1 + let large_v = graph.compute(&[100.0]).unwrap(); + assert!(approx_eq(large_v, 1.0, 1e-10)); + } + + #[test] + fn test_graph_gelu_compute_and_gradient() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.gelu(x); + // GELU(0) = 0 + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-12)); + // GELU(-large) ≈ 0 + let neg_v = graph.compute(&[-100.0]).unwrap(); + assert!((neg_v - 0.0).abs() < 1e-10); + // GELU should be finite and between 0 and x for positive x + let pos_v = graph.compute(&[2.0]).unwrap(); + assert!(pos_v > 0.0 && pos_v < 2.0); + let (_v, grad) = graph.gradient(&[0.0]).unwrap(); + assert!(approx_eq(grad[0], 0.5, 1e-12)); + } + + #[test] + fn test_graph_relu_at_zero() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.relu(x); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + // relu(0) = 0, and standard convention: relu'(0) = 0 + let (_v, grad) = graph.gradient(&[0.0]).unwrap(); + assert!(approx_eq(grad[0], 0.0, 1e-10)); + } + + #[test] + fn test_graph_gradient_checked_domain_error() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.ln(x); + let error = graph.gradient_checked(&[0.0]).unwrap_err(); + assert_eq!( + error, + crate::AutodiffError::DomainError { + operation: "Ln", + reason: "input must be positive", + } + ); + } + + #[test] + fn test_graph_compute_hessian_for_simple_function() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.square(x); + // f(x)=x², f''(x)=2 + let hessian = graph.compute_hessian(&[3.0]).unwrap(); + assert_eq!(hessian.len(), 1); + assert!(approx_eq(hessian[0][0], 2.0, 1e-4)); + } + + #[test] + fn test_graph_compute_hessian_for_empty_graph() { + let graph = Graph::new(0); + let hessian = graph.compute_hessian(&[]).unwrap(); + assert_eq!(hessian.len(), 0); + } + + #[test] + fn test_graph_gradient_sparse_with_tolerance_all_nonzero() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + // At (x=3, y=4): dy/dx=4, dy/dy=3 — both > tolerance 1.0 + let sparse = graph + .gradient_sparse_with_tolerance(&[3.0, 4.0], 1.0) + .unwrap(); + assert_eq!(sparse.len(), 2); + } + + #[test] + fn test_graph_gradient_sparse_with_tolerance_filtered() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + // At (x=0, y=4): dy/dx=4, dy/dy=0 — only dy/dx > tolerance 0.5 + let sparse = graph + .gradient_sparse_with_tolerance(&[0.0, 4.0], 0.5) + .unwrap(); + assert_eq!(sparse.len(), 1); + // dy/dx = y = 4 (input index 0), dy/dy = x = 0 (input index 1, filtered) + assert_eq!(sparse[0].0, 0); + assert!(approx_eq(sparse[0].1, 4.0, 1e-10)); + } + + #[test] + fn test_graph_norm2_single_input() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let norm = graph.norm2(&[x]).unwrap(); + graph.set_output(norm).unwrap(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + let value_neg = graph.compute(&[-4.0]).unwrap(); + assert!(approx_eq(value_neg, 4.0, 1e-10)); + } + + #[test] + fn test_graph_sum_squares_single_input() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let ss = graph.sum_squares(&[x]).unwrap(); + graph.set_output(ss).unwrap(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 9.0, 1e-10)); + } + + #[test] + fn test_graph_dot_with_valid_inputs() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let dot = graph.dot(&[x, y], &[x, y]).unwrap(); + graph.set_output(dot).unwrap(); + // dot([3,4], [3,4]) = 9+16 = 25 + let value = graph.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 25.0, 1e-10)); + // dot([a,b], [c,d]) = ac+bd + let dot2 = graph.dot(&[x, y], &[y, x]).unwrap(); + graph.set_output(dot2).unwrap(); + let value2 = graph.compute(&[2.0, 3.0]).unwrap(); + // 2*3 + 3*2 = 12 + assert!(approx_eq(value2, 12.0, 1e-10)); + } + + #[test] + fn test_graph_dot_with_empty_and_mismatched() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + assert!(graph.dot(&[], &[]).is_err()); + assert!(graph.dot(&[x], &[]).is_err()); + assert!(graph.dot(&[x], &[x, y]).is_err()); + } + + #[test] + fn test_graph_sum_with_multiple_inputs() { + let mut graph = Graph::new(3); + let a = graph.input(0); + let b = graph.input(1); + let c = graph.input(2); + let sum = graph.sum(&[a, b, c]).unwrap(); + graph.set_output(sum).unwrap(); + let value = graph.compute(&[1.0, 2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 6.0, 1e-10)); + } + + #[test] + fn test_graph_mean_with_multiple_inputs() { + let mut graph = Graph::new(3); + let a = graph.input(0); + let b = graph.input(1); + let c = graph.input(2); + let mean = graph.mean(&[a, b, c]).unwrap(); + graph.set_output(mean).unwrap(); + let value = graph.compute(&[1.0, 2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_graph_to_operations_input_only() { + let graph = Graph::new(2); + let ops = graph.to_operations().unwrap(); + assert_eq!(ops, vec![(MultiAD::Inp, vec![1])]); + } + + #[test] + fn test_graph_from_operations_roundtrip() { + let ops = vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), + ]; + let graph = Graph::from_operations(2, &ops); + let value = graph.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, 12.0, 1e-10)); + // Round trip via try_from_operations + let converted = graph.to_operations().unwrap(); + let graph2 = Graph::try_from_operations(2, &converted).unwrap(); + let value2 = graph2.compute(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(value, value2, 1e-10)); + } + + #[test] + fn test_graph_parameter_names_and_metadata() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.set_parameter_name(x, "weight").unwrap(); + graph.set_parameter_name(y, "bias").unwrap(); + + assert_eq!(graph.parameter_name(x), Some("weight")); + assert_eq!(graph.parameter_name(y), Some("bias")); + assert_eq!(graph.parameter_names().len(), 2); + assert!(graph + .parameter_names() + .iter() + .any(|(id, name)| *id == x && name == "weight")); + assert!(graph + .parameter_names() + .iter() + .any(|(id, name)| *id == y && name == "bias")); + // parameter_name for non-parameter returns None + assert_eq!(graph.parameter_name(99), None); + } + + #[test] + fn test_graph_compile_accelerated_full() { + let mut graph = Graph::new(1); + let x = graph.input(0); + graph.sin(x); + let compiled = graph.compile_accelerated().unwrap(); + let value = compiled.compute(&[0.5]).unwrap(); + assert!(approx_eq(value, 0.5_f64.sin(), 1e-10)); + let (_v, grad) = compiled.gradient(&[0.5]).unwrap(); + assert!(approx_eq(grad[0], 0.5_f64.cos(), 1e-8)); + } + + #[test] + fn test_graph_is_empty_and_len_both() { + let mut graph = Graph::new(2); + assert!(graph.is_empty()); + assert_eq!(graph.len(), 0); + + let x = graph.input(0); + graph.sin(x); + assert!(!graph.is_empty()); + assert_eq!(graph.len(), 1); + + graph.cos(x); + assert_eq!(graph.len(), 2); + } + + #[test] + fn test_graph_sum_squares_multiple_inputs() { + let mut graph = Graph::new(3); + let a = graph.input(0); + let b = graph.input(1); + let c = graph.input(2); + let ss = graph.sum_squares(&[a, b, c]).unwrap(); + graph.set_output(ss).unwrap(); + let value = graph.compute(&[1.0, 2.0, 3.0]).unwrap(); + // 1 + 4 + 9 = 14 + assert!(approx_eq(value, 14.0, 1e-10)); + } + + #[test] + fn test_graph_norm2_multiple_inputs() { + let mut graph = Graph::new(3); + let a = graph.input(0); + let b = graph.input(1); + let c = graph.input(2); + let norm = graph.norm2(&[a, b, c]).unwrap(); + graph.set_output(norm).unwrap(); + let value = graph.compute(&[3.0, 4.0, 0.0]).unwrap(); + assert!(approx_eq(value, 5.0, 1e-10)); + } + + #[test] + fn test_graph_hessian_sparse_full() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + let sparse = graph.hessian_sparse(&[2.0, 3.0]).unwrap(); + // Hessian of xy: d²/dxdy=1, d²/dydx=1 — 2 nonzeros + assert_eq!(sparse.len(), 2); + assert!(approx_eq(sparse[0].2, 1.0, 1e-10)); + assert!(approx_eq(sparse[1].2, 1.0, 1e-10)); + } + + #[test] + fn test_graph_hessian_sparse_with_tolerance_full() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + graph.mul(x, y); + // At tolerance 0.0: all nonzero entries returned + let sparse = graph + .hessian_sparse_with_tolerance(&[2.0, 3.0], 0.0) + .unwrap(); + assert_eq!(sparse.len(), 2); + // At tolerance 2.0: no entries survive + let sparse_filtered = graph + .hessian_sparse_with_tolerance(&[2.0, 3.0], 2.0) + .unwrap(); + assert_eq!(sparse_filtered.len(), 0); + } + + #[test] + fn test_graph_try_from_operations_with_inp_error() { + let ops = [(MultiAD::Inp, vec![5])]; + assert!(Graph::try_from_operations(2, &ops).is_err()); + } + + #[test] + fn test_graph_try_from_operations_multi_inp() { + let ops = [(MultiAD::Inp, vec![0, 1])]; + assert!(Graph::try_from_operations(2, &ops).is_err()); + } + + #[test] + fn test_try_sin_and_try_cos() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let s = graph.try_sin(x).unwrap(); + let c = graph.try_cos(x).unwrap(); + graph.set_outputs(&[s, c]).unwrap(); + let v = graph.compute_many(&[1.0]).unwrap(); + assert!(approx_eq(v[0], 1.0_f64.sin(), 1e-10)); + assert!(approx_eq(v[1], 1.0_f64.cos(), 1e-10)); + } + + #[test] + fn test_try_tan_and_try_exp() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let t = graph.try_tan(x).unwrap(); + let e = graph.try_exp(x).unwrap(); + graph.set_outputs(&[t, e]).unwrap(); + let v = graph.compute_many(&[0.5]).unwrap(); + assert!(approx_eq(v[0], 0.5_f64.tan(), 1e-10)); + assert!(approx_eq(v[1], 0.5_f64.exp(), 1e-10)); + } + + #[test] + fn test_try_ln_and_try_sqrt_try_abs() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let l = graph.try_ln(x).unwrap(); + let s = graph.try_sqrt(x).unwrap(); + let a = graph.try_abs(x).unwrap(); + graph.set_outputs(&[l, s, a]).unwrap(); + let v = graph.compute_many(&[4.0]).unwrap(); + assert!(approx_eq(v[0], 4.0_f64.ln(), 1e-10)); + assert!(approx_eq(v[1], 4.0_f64.sqrt(), 1e-10)); + assert!(approx_eq(v[2], 4.0_f64.abs(), 1e-10)); + } + + #[test] + fn test_try_neg_and_try_tanh_try_relu() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let n = graph.try_neg(x).unwrap(); + let t = graph.try_tanh(x).unwrap(); + let r = graph.try_relu(x).unwrap(); + graph.set_outputs(&[n, t, r]).unwrap(); + let v = graph.compute_many(&[2.0]).unwrap(); + assert!(approx_eq(v[0], -2.0, 1e-10)); + assert!(approx_eq(v[1], 2.0_f64.tanh(), 1e-10)); + assert!(approx_eq(v[2], 2.0, 1e-10)); + } + + #[test] + fn test_try_reductions() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let s = graph.sum(&[x, y]).unwrap(); + let m = graph.mean(&[x, y]).unwrap(); + let ss = graph.sum_squares(&[x, y]).unwrap(); + let n = graph.norm2(&[x]).unwrap(); + graph.set_outputs(&[s, m, ss, n]).unwrap(); + let v = graph.compute_many(&[3.0, 4.0]).unwrap(); + assert!(approx_eq(v[0], 7.0, 1e-10)); + assert!(approx_eq(v[1], 3.5, 1e-10)); + assert!(approx_eq(v[2], 25.0, 1e-10)); + assert!(approx_eq(v[3], 3.0, 1e-10)); + } + + #[test] + fn test_try_binary_ops() { + let mut graph = Graph::new(2); + let x = graph.input(0); + let y = graph.input(1); + let a = graph.try_add(x, y).unwrap(); + let s = graph.try_sub(x, y).unwrap(); + let m = graph.try_mul(x, y).unwrap(); + let d = graph.try_div(x, y).unwrap(); + let p = graph.try_pow(x, y).unwrap(); + graph.set_outputs(&[a, s, m, d, p]).unwrap(); + let v = graph.compute_many(&[4.0, 2.0]).unwrap(); + assert!(approx_eq(v[0], 6.0, 1e-10)); + assert!(approx_eq(v[1], 2.0, 1e-10)); + assert!(approx_eq(v[2], 8.0, 1e-10)); + assert!(approx_eq(v[3], 2.0, 1e-10)); + assert!(approx_eq(v[4], 16.0, 1e-10)); + } + + #[test] + fn test_graph_effective_output_nodes_with_set_outputs() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let s = graph.sin(x); + let c = graph.cos(x); + graph.set_outputs(&[s, c]).unwrap(); + let nodes = graph.effective_output_nodes(); + assert_eq!(nodes.len(), 2); + assert!(nodes.contains(&s)); + assert!(nodes.contains(&c)); + } + + #[test] + fn test_graph_clear_output_keeps_last_node() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let s = graph.sin(x); + graph.set_output(s).unwrap(); + assert_eq!(graph.effective_output_node(), Some(s)); + graph.clear_output(); + // After clearing, output goes back to last node + assert_eq!(graph.effective_output_node(), Some(s)); + } + + #[test] + fn test_graph_from_operations_checked() { + let ops = vec![(MultiAD::Inp, vec![0]), (MultiAD::Sin, vec![0])]; + let graph = Graph::try_from_operations(1, &ops).unwrap(); + let v = graph.compute(&[1.0]).unwrap(); + assert!(approx_eq(v, 1.0_f64.sin(), 1e-10)); + } + + #[test] + fn test_graph_from_operations_checked_with_bad_input() { + // Input index out of bounds + let ops = vec![(MultiAD::Inp, vec![99])]; + let result = Graph::try_from_operations(1, &ops); + assert!(result.is_err()); + } + + #[test] + fn test_graph_from_operations_with_constants() { + let mut graph = Graph::new(1); + let _c = graph.constant(5.0); + let ops = graph.to_operations(); + // to_operations should fail when constants are present + assert!(ops.is_err()); + } + + #[test] + fn test_graph_add_const() { + let mut graph = Graph::new(1); + let x = graph.input(0); + let a = graph.add_const(x, 3.0); + graph.set_output(a).unwrap(); + let v = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(v, 5.0, 1e-10)); + let _ = a; + } + + #[test] + fn test_graph_dot_product() { + let mut graph = Graph::new(3); + let a = graph.input(0); + let b = graph.input(1); + let c_ = graph.input(2); + let dot = graph.dot(&[a, b, c_], &[a, b, c_]).unwrap(); + graph.set_output(dot).unwrap(); + let v = graph.compute(&[1.0, 2.0, 3.0]).unwrap(); + assert!(approx_eq(v, 14.0, 1e-10)); + let _ = dot; + } +} diff --git a/src/multi/multi_ad.rs b/src/multi/multi_ad.rs index a681f7b..5614c1d 100644 --- a/src/multi/multi_ad.rs +++ b/src/multi/multi_ad.rs @@ -1,3 +1,4 @@ +use super::op_rules; use super::types::*; use crate::error::{AutodiffError, Result}; @@ -25,6 +26,7 @@ use crate::error::{AutodiffError, Result}; /// println!("f(0.6, 1.4) = {}", value); /// println!("∇f = {:?}", gradients); /// ``` +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum MultiAD { /// Input placeholder - references an input variable @@ -45,6 +47,8 @@ pub enum MultiAD { /// /// # Notes /// - Delegates to `f64::powf()` + /// - Gradients with respect to the exponent use `ln(a)`, so real-valued + /// differentiability requires a positive base for non-constant exponents /// - For `x^n` where n is an integer, consider using repeated multiplication Pow, /// Sine function: sin(x) @@ -65,6 +69,16 @@ pub enum MultiAD { /// - Delegates to `f64::tan()`, which operates in radians /// - Returns very large values near `π/2 + kπ` (asymptotes) Tan, + /// Hyperbolic tangent: tanh(x) + Tanh, + /// Rectified linear unit: max(0, x) with subgradient 0 at x = 0. + Relu, + /// Stable softplus: ln(1 + exp(x)). + Log1pExp, + /// Stable binary log-sum-exp: ln(exp(a) + exp(b)). + LogAddExp, + /// Negation: -x + Neg, /// Exponential function: exp(x) /// /// # Notes @@ -94,187 +108,69 @@ pub enum MultiAD { } impl MultiAD { - /// Get the name of this operation (for error messages and arity checking) - fn op_name(&self) -> &'static str { - match self { - MultiAD::Inp => "Inp", - MultiAD::Add => "Add", - MultiAD::Sub => "Sub", - MultiAD::Mul => "Mul", - MultiAD::Div => "Div", - MultiAD::Pow => "Pow", - MultiAD::Sin => "Sin", - MultiAD::Cos => "Cos", - MultiAD::Tan => "Tan", - MultiAD::Exp => "Exp", - MultiAD::Ln => "Ln", - MultiAD::Sqrt => "Sqrt", - MultiAD::Abs => "Abs", - } + /// Forward pass: compute the output of this operation given inputs + #[inline] + pub(crate) fn forward(&self, args: &[f64]) -> Result { + op_rules::forward_value(*self, args) } - /// Get the expected arity for this operation - fn expected_arity(&self) -> usize { - match self { - MultiAD::Inp - | MultiAD::Sin - | MultiAD::Cos - | MultiAD::Tan - | MultiAD::Exp - | MultiAD::Ln - | MultiAD::Sqrt - | MultiAD::Abs => 1, - MultiAD::Add | MultiAD::Sub | MultiAD::Mul | MultiAD::Div | MultiAD::Pow => 2, - } - } - /// Forward pass: compute the output of this operation given inputs - fn forward(&self, args: &[f64]) -> Result { - Ok(match self { - MultiAD::Inp => { - AutodiffError::check_arity("Inp", 1, args.len())?; - args[0] - } - MultiAD::Sin => { - AutodiffError::check_arity("Sin", 1, args.len())?; - args[0].sin() - } - MultiAD::Cos => { - AutodiffError::check_arity("Cos", 1, args.len())?; - args[0].cos() - } - MultiAD::Tan => { - AutodiffError::check_arity("Tan", 1, args.len())?; - args[0].tan() - } - MultiAD::Exp => { - AutodiffError::check_arity("Exp", 1, args.len())?; - args[0].exp() - } - MultiAD::Ln => { - AutodiffError::check_arity("Ln", 1, args.len())?; - args[0].ln() - } - MultiAD::Sqrt => { - AutodiffError::check_arity("Sqrt", 1, args.len())?; - args[0].sqrt() - } - MultiAD::Abs => { - AutodiffError::check_arity("Abs", 1, args.len())?; - args[0].abs() - } - MultiAD::Add => { - AutodiffError::check_arity("Add", 2, args.len())?; - args[0] + args[1] - } - MultiAD::Sub => { - AutodiffError::check_arity("Sub", 2, args.len())?; - args[0] - args[1] - } - MultiAD::Mul => { - AutodiffError::check_arity("Mul", 2, args.len())?; - args[0] * args[1] - } - MultiAD::Div => { - AutodiffError::check_arity("Div", 2, args.len())?; - args[0] / args[1] - } - MultiAD::Pow => { - AutodiffError::check_arity("Pow", 2, args.len())?; - args[0].powf(args[1]) - } - }) + #[inline] + pub(crate) fn forward_checked(&self, args: &[f64]) -> Result { + op_rules::forward_value_checked(*self, args) } /// Backward pass: compute local gradients ∂output/∂inputs /// Returns a boxed closure that computes gradients given a cotangent value - fn backward_generic(&self, args: &[f64]) -> Result + #[inline] + pub(crate) fn backward_generic(&self, args: &[f64]) -> Result where W: From>, { - AutodiffError::check_arity(self.op_name(), self.expected_arity(), args.len())?; - - let backward_fn: Box Vec> = match self { - MultiAD::Inp => Box::new(|zcotangent: f64| vec![zcotangent]), - MultiAD::Sin => { - let arg_val = args[0]; - Box::new(move |z_cotangent: f64| { - let x_cotangent = z_cotangent * arg_val.cos(); - vec![x_cotangent] - }) - } - MultiAD::Cos => { - let arg_val = args[0]; - Box::new(move |z_cotangent: f64| { - let x_cotangent = z_cotangent * -arg_val.sin(); - vec![x_cotangent] - }) - } - MultiAD::Tan => { - let arg_val = args[0]; - Box::new(move |z_cotangent: f64| { - let x_cotangent = z_cotangent * (1.0 / arg_val.cos().powi(2)); - vec![x_cotangent] - }) - } - MultiAD::Exp => { - let exp_val = args[0].exp(); - Box::new(move |z_cotangent: f64| { - let x_cotangent = z_cotangent * exp_val; - vec![x_cotangent] - }) - } - MultiAD::Ln => { - let arg_val = args[0]; - Box::new(move |z_cotangent: f64| { - let x_cotangent = z_cotangent * (1.0 / arg_val); - vec![x_cotangent] - }) - } - MultiAD::Add => Box::new(|z_cotangent: f64| vec![z_cotangent, z_cotangent]), - MultiAD::Sub => Box::new(|z_cotangent: f64| vec![z_cotangent, -z_cotangent]), - MultiAD::Mul => { - let arg0 = args[0]; - let arg1 = args[1]; - Box::new(move |z_cotangent: f64| vec![z_cotangent * arg1, z_cotangent * arg0]) - } - MultiAD::Div => { - let arg0 = args[0]; - let arg1 = args[1]; - Box::new(move |z_cotangent: f64| { - vec![z_cotangent / arg1, -z_cotangent * arg0 / arg1.powi(2)] - }) - } - MultiAD::Pow => { - let base = args[0]; - let exp = args[1]; - Box::new(move |z_cotangent: f64| { - // d(a^b)/da = b * a^(b-1) - let d_base = z_cotangent * exp * base.powf(exp - 1.0); - // d(a^b)/db = a^b * ln(a) - let d_exp = z_cotangent * base.powf(exp) * base.ln(); - vec![d_base, d_exp] - }) - } - MultiAD::Sqrt => { - let arg_val = args[0]; - Box::new(move |z_cotangent: f64| { - // d(sqrt(x))/dx = 1/(2*sqrt(x)) - let x_cotangent = z_cotangent / (2.0 * arg_val.sqrt()); - vec![x_cotangent] - }) - } - MultiAD::Abs => { - let arg_val = args[0]; - Box::new(move |z_cotangent: f64| { - // d(|x|)/dx = sign(x) where sign(0) = 0 - let sign = if arg_val >= 0.0 { 1.0 } else { -1.0 }; - vec![z_cotangent * sign] - }) - } + let value = op_rules::forward_value(*self, args)?; + let local_grads = match self { + MultiAD::Inp => vec![1.0], + _ => op_rules::first_derivatives(*self, args, value)?, }; + let backward_fn: Box Vec> = Box::new(move |z_cotangent: f64| { + local_grads + .iter() + .map(|grad| z_cotangent * grad) + .collect::>() + }); Ok(W::from(backward_fn)) } + /// Validate that a graph index points to an available value. + #[inline] + pub(crate) fn check_value_index(index: usize, values_len: usize) -> Result<()> { + if index < values_len { + Ok(()) + } else { + Err(AutodiffError::IndexOutOfBounds { + index, + max_index: values_len.saturating_sub(1), + }) + } + } + + /// Validate an input marker operation. + #[inline] + fn check_input_marker(arg_indices: &[usize], input_len: usize) -> Result<()> { + AutodiffError::check_arity("Inp", 1, arg_indices.len())?; + Self::check_value_index(arg_indices[0], input_len) + } + + /// Gather operation arguments after validating all referenced graph indices. + #[inline] + pub(crate) fn gather_arg_values(arg_indices: &[usize], values: &[f64]) -> Result> { + let mut arg_values = Vec::with_capacity(arg_indices.len()); + for &index in arg_indices { + Self::check_value_index(index, values.len())?; + arg_values.push(values[index]); + } + Ok(arg_values) + } + /// Compute forward pass only (no gradient computation). /// /// Evaluates the computational graph to produce the final output value. @@ -298,24 +194,57 @@ impl MultiAD { /// assert!((result - 5.0).abs() < 1e-10); /// ``` #[must_use = "forward computation is expensive; discarding the result is likely a bug"] + #[inline] pub fn compute(exprs: &[(MultiAD, Vec)], inputs: &[f64]) -> Result { - let mut values: Vec = inputs.to_vec(); + // Pre-allocate with estimated capacity + let estimated_size = inputs.len() + exprs.len(); + let mut values: Vec = Vec::with_capacity(estimated_size); + values.extend_from_slice(inputs); + let mut final_output_index = inputs.len().checked_sub(1); for (op, arg_indices) in exprs { if *op == MultiAD::Inp { + Self::check_input_marker(arg_indices, inputs.len())?; + final_output_index = Some(arg_indices[0]); continue; // Input values are already in the values array } - // Gather the argument values from the computation graph - let arg_values: Vec = arg_indices.iter().map(|&i| values[i]).collect(); + // Gather the argument values from the computation graph. + let arg_values = Self::gather_arg_values(arg_indices, &values)?; // Compute this operation let value = op.forward(&arg_values)?; values.push(value); + final_output_index = Some(values.len() - 1); + } + + // Return the final computed value. Empty graph with no inputs is the zero scalar. + Ok(final_output_index.map(|index| values[index]).unwrap_or(0.0)) + } + + /// Compute forward pass with opt-in checked real-domain validation. + #[must_use = "forward computation is expensive; discarding the result is likely a bug"] + #[inline] + pub fn compute_checked(exprs: &[(MultiAD, Vec)], inputs: &[f64]) -> Result { + let estimated_size = inputs.len() + exprs.len(); + let mut values: Vec = Vec::with_capacity(estimated_size); + values.extend_from_slice(inputs); + let mut final_output_index = inputs.len().checked_sub(1); + + for (op, arg_indices) in exprs { + if *op == MultiAD::Inp { + Self::check_input_marker(arg_indices, inputs.len())?; + final_output_index = Some(arg_indices[0]); + continue; + } + + let arg_values = Self::gather_arg_values(arg_indices, &values)?; + let value = op.forward_checked(&arg_values)?; + values.push(value); + final_output_index = Some(values.len() - 1); } - // Return the final computed value - Ok(values.last().copied().unwrap_or(0.0)) + Ok(final_output_index.map(|index| values[index]).unwrap_or(0.0)) } /// Compute forward pass and return gradient function. @@ -357,6 +286,7 @@ impl MultiAD { /// let arc_grad_fn: Arc Vec> = Arc::from(grad_fn); /// ``` #[must_use = "gradient computation is expensive; discarding the result is likely a bug"] + #[inline] pub fn compute_grad_generic( exprs: &[(MultiAD, Vec)], inputs: &[f64], @@ -368,6 +298,7 @@ impl MultiAD { let estimated_size = inputs.len() + exprs.len(); let mut values: Vec = Vec::with_capacity(estimated_size); values.extend_from_slice(inputs); + let mut final_output_index = inputs.len().checked_sub(1); let mut backward_ops: Vec> = Vec::with_capacity(exprs.len()); let mut arg_indices_list: Vec> = Vec::with_capacity(exprs.len()); @@ -375,26 +306,34 @@ impl MultiAD { // Forward pass: compute all values and track backward operations for (op, args) in exprs { if *op == MultiAD::Inp { + Self::check_input_marker(args, inputs.len())?; + final_output_index = Some(args[0]); continue; } - let arg_values: Vec = args.iter().map(|&i| values[i]).collect(); + let arg_values = Self::gather_arg_values(args, &values)?; let value = op.forward(&arg_values)?; values.push(value); + final_output_index = Some(values.len() - 1); // Store the backward operation (which captures necessary values) backward_ops.push(op.backward_generic(&arg_values)?); arg_indices_list.push(args.clone()); } - let final_value = values.last().copied().unwrap_or(0.0); + let final_value = final_output_index.map(|index| values[index]).unwrap_or(0.0); // Clone the data we need for the backward pass let num_inputs = inputs.len(); let values_clone = values; + let final_output_index_clone = final_output_index; let backward_fn = Box::new(move |cotangent: f64| -> Vec { + let Some(final_output_index) = final_output_index_clone else { + return Vec::new(); + }; + let mut cotangent_values = vec![0.0; values_clone.len()]; - cotangent_values[values_clone.len() - 1] = cotangent; + cotangent_values[final_output_index] = cotangent; // Backward pass: propagate cotangents from output to inputs for (i, (backward_op, arg_indices)) in backward_ops @@ -426,4 +365,687 @@ impl MultiAD { ) -> Result { Self::compute_grad_generic::>(exprs, inputs) } + + /// Compute forward pass and return gradient function with checked-domain validation. + #[must_use = "gradient computation is expensive; discarding the result is likely a bug"] + pub fn compute_grad_checked( + exprs: &[(MultiAD, Vec)], + inputs: &[f64], + ) -> Result { + let estimated_size = inputs.len() + exprs.len(); + let mut values: Vec = Vec::with_capacity(estimated_size); + values.extend_from_slice(inputs); + let mut final_output_index = inputs.len().checked_sub(1); + + let mut backward_ops: Vec> = Vec::with_capacity(exprs.len()); + let mut arg_indices_list: Vec> = Vec::with_capacity(exprs.len()); + + for (op, args) in exprs { + if *op == MultiAD::Inp { + Self::check_input_marker(args, inputs.len())?; + final_output_index = Some(args[0]); + continue; + } + let arg_values = Self::gather_arg_values(args, &values)?; + let value = op.forward_checked(&arg_values)?; + values.push(value); + final_output_index = Some(values.len() - 1); + backward_ops.push(op.backward_generic(&arg_values)?); + arg_indices_list.push(args.clone()); + } + + let final_value = final_output_index.map(|index| values[index]).unwrap_or(0.0); + let num_inputs = inputs.len(); + let values_clone = values; + let final_output_index_clone = final_output_index; + + let backward_fn = Box::new(move |cotangent: f64| -> Vec { + let Some(final_output_index) = final_output_index_clone else { + return Vec::new(); + }; + + let mut cotangent_values = vec![0.0; values_clone.len()]; + cotangent_values[final_output_index] = cotangent; + + for (i, (backward_op, arg_indices)) in backward_ops + .iter() + .zip(arg_indices_list.iter()) + .rev() + .enumerate() + { + let output_idx = values_clone.len() - 1 - i; + let current_cotangent_value = cotangent_values[output_idx]; + let argv_cotangents = backward_op(current_cotangent_value); + for (arg_idx, arg_cotangent) in arg_indices.iter().zip(argv_cotangents) { + cotangent_values[*arg_idx] += arg_cotangent; + } + } + + cotangent_values[..num_inputs].to_vec() + }); + + Ok((final_value, backward_fn)) + } + + /// Compute the Hessian matrix using finite differences on the gradient. + /// + /// The Hessian is the matrix of second-order partial derivatives: + /// H\[i\]\[j\] = ∂²f/∂xᵢ∂xⱼ + /// + /// This uses central finite differences on gradients to compute how each + /// gradient component changes with respect to each input. + /// + /// # Arguments + /// + /// * `exprs` - Computational graph as (operation, indices) pairs + /// * `inputs` - Input values to evaluate at + /// + /// # Returns + /// + /// A 2D vector representing the Hessian matrix where result\[i\]\[j\] = ∂²f/∂xᵢ∂xⱼ + /// + /// # Complexity + /// + /// O(n²) where n is the number of inputs. For each of n inputs, this + /// computes two gradient evaluations. + /// + /// # Examples + /// + /// ``` + /// use petite_ad::{MultiAD, multi_ops}; + /// + /// // f(x, y) = x² + y² (Hessian is [[2, 0], [0, 2]]) + /// let exprs = multi_ops![ + /// (inp, 0), (inp, 1), + /// (mul, 0, 0), (mul, 1, 1), + /// (add, 2, 3) + /// ]; + /// let hessian = MultiAD::compute_hessian(&exprs, &[2.0, 3.0]).unwrap(); + /// assert!((hessian[0][0] - 2.0).abs() < 1e-6); + /// assert!((hessian[0][1] - 0.0).abs() < 1e-6); + /// assert!((hessian[1][0] - 0.0).abs() < 1e-6); + /// assert!((hessian[1][1] - 2.0).abs() < 1e-6); + /// ``` + #[must_use = "Hessian computation is expensive; discarding the result is likely a bug"] + pub fn compute_hessian( + exprs: &[(MultiAD, Vec)], + inputs: &[f64], + ) -> Result>> { + let num_inputs = inputs.len(); + let epsilon = 1e-5; + let mut hessian = vec![vec![0.0; num_inputs]; num_inputs]; + + if num_inputs == 0 { + let (_value, _grad_fn) = Self::compute_grad(exprs, inputs)?; + return Ok(hessian); + } + + // For each input variable, compute how the gradient changes using a + // central difference for second-order accuracy on smooth functions. + for j in 0..num_inputs { + let mut inputs_plus = inputs.to_vec(); + inputs_plus[j] += epsilon; + + let mut inputs_minus = inputs.to_vec(); + inputs_minus[j] -= epsilon; + + let (_value_plus, grad_fn_plus) = Self::compute_grad(exprs, &inputs_plus)?; + let grad_plus = grad_fn_plus(1.0); + + let (_value_minus, grad_fn_minus) = Self::compute_grad(exprs, &inputs_minus)?; + let grad_minus = grad_fn_minus(1.0); + + for i in 0..num_inputs { + // ∂²f/∂xᵢ∂xⱼ ≈ (∂f/∂xᵢ(x + εeⱼ) - ∂f/∂xᵢ(x - εeⱼ)) / (2ε) + hessian[i][j] = (grad_plus[i] - grad_minus[i]) / (2.0 * epsilon); + } + } + + Ok(hessian) + } + + /// Compute a single row of the Hessian matrix using finite differences. + /// + /// Computes ∇(∂f/∂xᵢ), which is the i-th row of the Hessian. + /// The j-th element is ∂²f/∂xᵢ∂xⱼ. + /// + /// # Arguments + /// + /// * `exprs` - Computational graph + /// * `inputs` - Input values + /// * `grad_idx` - Index of the gradient component to differentiate + /// + /// # Returns + /// + /// A vector representing the i-th row of the Hessian + /// + /// # Errors + /// + /// Returns `Err(AutodiffError::IndexOutOfBounds)` if `grad_idx` is not a + /// valid input index, or if the graph references unavailable values. + pub fn compute_hessian_row( + exprs: &[(MultiAD, Vec)], + inputs: &[f64], + grad_idx: usize, + ) -> Result> { + let num_inputs = inputs.len(); + Self::check_value_index(grad_idx, num_inputs)?; + + let epsilon = 1e-5; + let mut hessian_row = vec![0.0; num_inputs]; + + // For each input variable, compute how this gradient component changes + // using a central difference for second-order accuracy. + for j in 0..num_inputs { + let mut inputs_plus = inputs.to_vec(); + inputs_plus[j] += epsilon; + + let mut inputs_minus = inputs.to_vec(); + inputs_minus[j] -= epsilon; + + let (_value_plus, grad_fn_plus) = Self::compute_grad(exprs, &inputs_plus)?; + let grad_plus = grad_fn_plus(1.0); + + let (_value_minus, grad_fn_minus) = Self::compute_grad(exprs, &inputs_minus)?; + let grad_minus = grad_fn_minus(1.0); + + // ∂²f/∂x_grad_idx∂xⱼ ≈ (∂f/∂x_grad_idx(x + εeⱼ) - ∂f/∂x_grad_idx(x - εeⱼ)) / (2ε) + hessian_row[j] = (grad_plus[grad_idx] - grad_minus[grad_idx]) / (2.0 * epsilon); + } + + Ok(hessian_row) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::multi_ops; + use crate::test_utils::approx_eq_eps as approx_eq; + use std::sync::Arc; + + // ---- forward / forward_checked ---- + + #[test] + fn test_forward_sin() { + let val = MultiAD::Sin.forward(&[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.sin(), 1e-12)); + } + + #[test] + fn test_forward_cos() { + let val = MultiAD::Cos.forward(&[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.cos(), 1e-12)); + } + + #[test] + fn test_forward_tan() { + let val = MultiAD::Tan.forward(&[0.5]).unwrap(); + assert!(approx_eq(val, 0.5_f64.tan(), 1e-12)); + } + + #[test] + fn test_forward_exp() { + let val = MultiAD::Exp.forward(&[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.exp(), 1e-12)); + } + + #[test] + fn test_forward_ln() { + let val = MultiAD::Ln.forward(&[2.0]).unwrap(); + assert!(approx_eq(val, 2.0_f64.ln(), 1e-12)); + } + + #[test] + fn test_forward_sqrt() { + let val = MultiAD::Sqrt.forward(&[4.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-12)); + } + + #[test] + fn test_forward_neg() { + let val = MultiAD::Neg.forward(&[3.0]).unwrap(); + assert!(approx_eq(val, -3.0, 1e-12)); + } + + #[test] + fn test_forward_abs() { + let val = MultiAD::Abs.forward(&[-5.0]).unwrap(); + assert!(approx_eq(val, 5.0, 1e-12)); + } + + #[test] + fn test_forward_tanh() { + let val = MultiAD::Tanh.forward(&[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.tanh(), 1e-12)); + } + + #[test] + fn test_forward_relu_positive() { + let val = MultiAD::Relu.forward(&[2.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-12)); + } + + #[test] + fn test_forward_relu_negative() { + let val = MultiAD::Relu.forward(&[-2.0]).unwrap(); + assert!(approx_eq(val, 0.0, 1e-12)); + } + + #[test] + fn test_forward_log1p_exp() { + let val = MultiAD::Log1pExp.forward(&[0.0]).unwrap(); + assert!(approx_eq(val, 2.0_f64.ln(), 1e-6)); + } + + #[test] + fn test_forward_log_add_exp() { + let val = MultiAD::LogAddExp.forward(&[0.0, 0.0]).unwrap(); + assert!(approx_eq(val, 2.0_f64.ln(), 1e-6)); + } + + #[test] + fn test_forward_add() { + let val = MultiAD::Add.forward(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(val, 5.0, 1e-12)); + } + + #[test] + fn test_forward_sub() { + let val = MultiAD::Sub.forward(&[5.0, 3.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-12)); + } + + #[test] + fn test_forward_mul() { + let val = MultiAD::Mul.forward(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(val, 6.0, 1e-12)); + } + + #[test] + fn test_forward_div() { + let val = MultiAD::Div.forward(&[6.0, 3.0]).unwrap(); + assert!(approx_eq(val, 2.0, 1e-12)); + } + + #[test] + fn test_forward_pow() { + let val = MultiAD::Pow.forward(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(val, 8.0, 1e-12)); + } + + #[test] + fn test_forward_checked_ln_negative() { + let result = MultiAD::Ln.forward_checked(&[-1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_forward_checked_sqrt_negative() { + let result = MultiAD::Sqrt.forward_checked(&[-1.0]); + assert!(result.is_err()); + } + + // ---- check_value_index / check_input_marker error cases ---- + + #[test] + fn test_check_value_index_ok() { + assert!(MultiAD::check_value_index(0, 5).is_ok()); + assert!(MultiAD::check_value_index(4, 5).is_ok()); + } + + #[test] + fn test_check_value_index_out_of_bounds() { + let result = MultiAD::check_value_index(5, 5); + assert!(result.is_err()); + assert!(matches!( + result.unwrap_err(), + AutodiffError::IndexOutOfBounds { .. } + )); + } + + #[test] + fn test_gather_arg_values_ok() { + let values = vec![1.0, 2.0, 3.0]; + let result = MultiAD::gather_arg_values(&[0, 2], &values).unwrap(); + assert_eq!(result, vec![1.0, 3.0]); + } + + #[test] + fn test_gather_arg_values_out_of_bounds() { + let values = vec![1.0, 2.0]; + let result = MultiAD::gather_arg_values(&[0, 5], &values); + assert!(result.is_err()); + } + + // ---- compute / compute_checked edge cases ---- + + #[test] + fn test_compute_empty_graph_no_inputs() { + // Empty graph with no inputs should return 0.0 + let result = MultiAD::compute(&[], &[]).unwrap(); + assert!(approx_eq(result, 0.0, 1e-12)); + } + + #[test] + fn test_compute_single_input_passthrough() { + // Just an input marker + let exprs = multi_ops![(inp, 0)]; + let result = MultiAD::compute(&exprs, &[42.0]).unwrap(); + assert!(approx_eq(result, 42.0, 1e-12)); + } + + #[test] + fn test_compute_checked_basic() { + let exprs = multi_ops![(inp, 0), (inp, 1), (add, 0, 1)]; + let result = MultiAD::compute_checked(&exprs, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(result, 5.0, 1e-12)); + } + + #[test] + fn test_compute_checked_catches_domain_error() { + // ln(-1) should fail with checked forward + let exprs = multi_ops![(inp, 0), (ln, 0)]; + let result = MultiAD::compute_checked(&exprs, &[-1.0]); + assert!(result.is_err()); + } + + // ---- compute_grad edge cases ---- + + #[test] + fn test_compute_grad_empty_graph() { + let (value, grad_fn) = MultiAD::compute_grad(&[], &[]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-12)); + let grads = grad_fn(1.0); + assert!(grads.is_empty()); + } + + #[test] + fn test_compute_grad_single_input() { + // f(x) = x * x + let exprs = multi_ops![(inp, 0), (mul, 0, 0)]; + let (value, grad_fn) = MultiAD::compute_grad(&exprs, &[3.0]).unwrap(); + assert!(approx_eq(value, 9.0, 1e-10)); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 6.0, 1e-10)); // d/dx(x^2) = 2x = 6 + } + + #[test] + fn test_compute_grad_with_different_cotangent() { + // Test backward with cotangent != 1.0 + let exprs = multi_ops![(inp, 0), (inp, 1), (mul, 0, 1)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[2.0, 3.0]).unwrap(); + let grads = grad_fn(2.0); // cotangent = 2 + // df/dx = y = 3, scaled by 2 = 6 + // df/dy = x = 2, scaled by 2 = 4 + assert!(approx_eq(grads[0], 6.0, 1e-10)); + assert!(approx_eq(grads[1], 4.0, 1e-10)); + } + + #[test] + fn test_compute_grad_generic_arc() { + // Test compute_grad_generic with Arc wrapper + let exprs = multi_ops![(inp, 0), (inp, 1), (add, 0, 1)]; + let (value, grad_fn) = + MultiAD::compute_grad_generic::>(&exprs, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 5.0, 1e-10)); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 1.0, 1e-10)); + assert!(approx_eq(grads[1], 1.0, 1e-10)); + } + + #[test] + fn test_compute_grad_input_marker_error_wrong_arity() { + // Inp with wrong arity + let exprs = vec![(MultiAD::Inp, vec![0, 1])]; + let result = MultiAD::compute_grad(&exprs, &[1.0, 2.0]); + assert!(result.is_err()); + } + + #[test] + fn test_compute_grad_input_marker_error_out_of_bounds() { + // Inp referencing input index beyond inputs length + let exprs = multi_ops![(inp, 5)]; + let result = MultiAD::compute_grad(&exprs, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_compute_grad_index_out_of_bounds_in_op() { + // Operation referencing a non-existent value index + let exprs = vec![(MultiAD::Add, vec![0, 99])]; + let result = MultiAD::compute_grad(&exprs, &[1.0]); + assert!(result.is_err()); + } + + // ---- compute_grad_checked ---- + + #[test] + fn test_compute_grad_checked_basic() { + let exprs = multi_ops![(inp, 0), (inp, 1), (mul, 0, 1)]; + let (value, grad_fn) = MultiAD::compute_grad_checked(&exprs, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 6.0, 1e-10)); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 3.0, 1e-10)); + assert!(approx_eq(grads[1], 2.0, 1e-10)); + } + + #[test] + fn test_compute_grad_checked_empty() { + let (value, grad_fn) = MultiAD::compute_grad_checked(&[], &[]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-12)); + let grads = grad_fn(1.0); + assert!(grads.is_empty()); + } + + #[test] + fn test_compute_grad_checked_domain_error() { + let exprs = multi_ops![(inp, 0), (ln, 0)]; + let result = MultiAD::compute_grad_checked(&exprs, &[-1.0]); + assert!(result.is_err()); + } + + // ---- compute_hessian ---- + + #[test] + fn test_compute_hessian_quadratic() { + // f(x, y) = x^2 + y^2 => H = [[2,0],[0,2]] + let exprs = multi_ops![(inp, 0), (inp, 1), (mul, 0, 0), (mul, 1, 1), (add, 2, 3)]; + let h = MultiAD::compute_hessian(&exprs, &[2.0, 3.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < 1e-4); + assert!((h[0][1]).abs() < 1e-4); + assert!((h[1][0]).abs() < 1e-4); + assert!((h[1][1] - 2.0).abs() < 1e-4); + } + + #[test] + fn test_compute_hessian_empty_inputs() { + let h = MultiAD::compute_hessian(&[], &[]).unwrap(); + assert!(h.is_empty()); + } + + // ---- compute_hessian_row ---- + + #[test] + fn test_compute_hessian_row_quadratic() { + // f(x, y) = x^2 + y^2 => H[0] = [2, 0] + let exprs = multi_ops![(inp, 0), (inp, 1), (mul, 0, 0), (mul, 1, 1), (add, 2, 3)]; + let row = MultiAD::compute_hessian_row(&exprs, &[2.0, 3.0], 0).unwrap(); + assert!((row[0] - 2.0).abs() < 1e-4); + assert!((row[1]).abs() < 1e-4); + } + + #[test] + fn test_compute_hessian_row_out_of_bounds() { + let exprs = multi_ops![(inp, 0), (inp, 1), (add, 0, 1)]; + let result = MultiAD::compute_hessian_row(&exprs, &[1.0, 2.0], 5); + assert!(result.is_err()); + assert!(matches!( + result.unwrap_err(), + AutodiffError::IndexOutOfBounds { .. } + )); + } + + // ---- Various individual op gradient checks ---- + + #[test] + fn test_grad_sub() { + // f(x, y) = x - y => df/dx = 1, df/dy = -1 + let exprs = multi_ops![(inp, 0), (inp, 1), (sub, 0, 1)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[5.0, 3.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 1.0, 1e-10)); + assert!(approx_eq(grads[1], -1.0, 1e-10)); + } + + #[test] + fn test_grad_div() { + // f(x, y) = x / y => df/dx = 1/y, df/dy = -x/y^2 + let exprs = multi_ops![(inp, 0), (inp, 1), (div, 0, 1)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[6.0, 3.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 1.0 / 3.0, 1e-10)); + assert!(approx_eq(grads[1], -6.0 / 9.0, 1e-10)); + } + + #[test] + fn test_grad_sin() { + // f(x) = sin(x) => df/dx = cos(x) + let exprs = multi_ops![(inp, 0), (sin, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[1.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 1.0_f64.cos(), 1e-10)); + } + + #[test] + fn test_grad_cos() { + // f(x) = cos(x) => df/dx = -sin(x) + let exprs = multi_ops![(inp, 0), (cos, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[1.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], -1.0_f64.sin(), 1e-10)); + } + + #[test] + fn test_grad_exp() { + // f(x) = exp(x) => df/dx = exp(x) + let exprs = multi_ops![(inp, 0), (exp, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[1.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 1.0_f64.exp(), 1e-10)); + } + + #[test] + fn test_grad_ln() { + // f(x) = ln(x) => df/dx = 1/x + let exprs = multi_ops![(inp, 0), (ln, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[2.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 0.5, 1e-10)); + } + + #[test] + fn test_grad_sqrt() { + // f(x) = sqrt(x) => df/dx = 1/(2*sqrt(x)) + let exprs = multi_ops![(inp, 0), (sqrt, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[4.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 0.25, 1e-10)); + } + + #[test] + fn test_grad_neg() { + // f(x) = -x => df/dx = -1 + let exprs = multi_ops![(inp, 0), (neg, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[3.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], -1.0, 1e-10)); + } + + #[test] + fn test_grad_abs() { + // f(x) = abs(x) => df/dx = sign(x) + let exprs = multi_ops![(inp, 0), (abs, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[-3.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], -1.0, 1e-10)); + } + + #[test] + fn test_grad_tanh() { + // f(x) = tanh(x) => df/dx = 1 - tanh(x)^2 + let exprs = multi_ops![(inp, 0), (tanh, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[0.5]).unwrap(); + let grads = grad_fn(1.0); + let expected = 1.0 - 0.5_f64.tanh().powi(2); + assert!(approx_eq(grads[0], expected, 1e-10)); + } + + #[test] + fn test_grad_relu_positive() { + let exprs = multi_ops![(inp, 0), (relu, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[2.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 1.0, 1e-10)); + } + + #[test] + fn test_grad_relu_negative() { + let exprs = multi_ops![(inp, 0), (relu, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[-2.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 0.0, 1e-10)); + } + + #[test] + fn test_grad_pow() { + // f(x, y) = x^y => df/dx = y * x^(y-1), df/dy = x^y * ln(x) + let exprs = multi_ops![(inp, 0), (inp, 1), (pow, 0, 1)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[2.0, 3.0]).unwrap(); + let grads = grad_fn(1.0); + assert!(approx_eq(grads[0], 3.0 * 4.0, 1e-10)); // y*x^(y-1) = 3*4 = 12 + assert!(approx_eq(grads[1], 8.0 * 2.0_f64.ln(), 1e-10)); // x^y*ln(x) + } + + #[test] + fn test_grad_log1p_exp() { + // f(x) = log1p_exp(x) = ln(1+exp(x)) => df/dx = sigmoid(x) = exp(x)/(1+exp(x)) + let exprs = multi_ops![(inp, 0), (log1p_exp, 0)]; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[1.0]).unwrap(); + let grads = grad_fn(1.0); + let expected = 1.0_f64.exp() / (1.0 + 1.0_f64.exp()); + assert!(approx_eq(grads[0], expected, 1e-10)); + } + + #[test] + fn test_grad_tan() { + // f(x) = tan(x) => df/dx = 1/cos(x)^2 + let exprs = multi_ops![(inp, 0), (tan, 0)]; + let x = 0.5_f64; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[x]).unwrap(); + let grads = grad_fn(1.0); + let expected = 1.0 / x.cos().powi(2); + assert!(approx_eq(grads[0], expected, 1e-10)); + } + + #[test] + fn test_grad_log_add_exp() { + // f(x, y) = log_add_exp(x, y) = ln(exp(x) + exp(y)) + let exprs = multi_ops![(inp, 0), (inp, 1), (log_add_exp, 0, 1)]; + let x = 1.0_f64; + let y = 2.0_f64; + let (_value, grad_fn) = MultiAD::compute_grad(&exprs, &[x, y]).unwrap(); + let grads = grad_fn(1.0); + let denom = x.exp() + y.exp(); + assert!(approx_eq(grads[0], x.exp() / denom, 1e-10)); + assert!(approx_eq(grads[1], y.exp() / denom, 1e-10)); + } + + #[test] + fn test_compute_input_index_error() { + // Inp with wrong arity (0 args) + let exprs = vec![(MultiAD::Inp, vec![])]; + let result = MultiAD::compute(&exprs, &[1.0]); + assert!(result.is_err()); + } } diff --git a/src/multi/multi_ad_fr.rs b/src/multi/multi_ad_fr.rs new file mode 100644 index 0000000..4bafc85 --- /dev/null +++ b/src/multi/multi_ad_fr.rs @@ -0,0 +1,512 @@ +//! Exact second-order autodiff for multivariate functions using Forward-over-Reverse (FR) mode. +//! +//! This module implements the FR method for computing exact Hessians of +//! multivariate functions f: ℝⁿ → ℝ. +//! +//! # Supported Operations +//! +//! `MultiAD2FR` supports the smooth operation subset shared by the exact Hessian +//! engines: `Inp`, `Sin`, `Cos`, `Tan`, `Neg`, `Exp`, `Ln`, `Sqrt`, `Log1pExp`, +//! `Add`, `Sub`, `Mul`, `Div`, and `Pow`. Non-smooth operations like `Abs` are not supported by +//! the exact Hessian types. For Hessian approximation with the full operation set, +//! use [`MultiAD::compute_hessian`](crate::MultiAD::compute_hessian) +//! (finite-difference based). +//! +//! # Algorithm +//! +//! FR composes forward-mode AD (outer) with reverse-mode AD (inner): +//! +//! 1. **Forward pass** — evaluate the computation graph with **dual numbers** +//! whose tangent component is seeded in direction e_j. Every intermediate +//! value becomes a dual `(val, tan)`. +//! +//! 2. **Reverse pass** — propagate **dual adjoints** backward through the +//! perturbed graph. Each adjoint is a dual `(val, tan)` where: +//! - `val = ∂f/∂node` (standard first-order adjoint) +//! - `tan = ∂²f/(∂node · ∂x_j)` (second-order cross-derivative) +//! +//! The dual adjoint at input node k gives Hessian entry `H[j][k]`. +//! +//! This is repeated for each seed direction j = 0 … n−1, producing the full +//! Hessian column by column. +//! +//! # Accuracy +//! +//! Machine-precision exact (~1e-14 relative error). No finite differences. +//! +//! # Complexity +//! +//! O(n · G) where n = number of inputs, G = graph size. +//! +//! See [docs/multi_ad_hessian.md](../../docs/multi_ad_hessian.md) for complete theory. + +use std::fmt; + +use crate::Result; + +use super::multi_hessian_common::{compute_hessian_dual, OpKind}; + +/// Stack-based operation for multivariate second-order AD. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MultiAD2FR { + /// Input variable by index + Inp(usize), + /// sin(x) + Sin, + /// cos(x) + Cos, + /// tan(x) + Tan, + /// -x + Neg, + /// exp(x) + Exp, + /// ln(x) + Ln, + /// sqrt(x) + Sqrt, + /// stable ln(1 + exp(x)) + Log1pExp, + /// a + b + Add, + /// a - b + Sub, + /// a · b + Mul, + /// a / b + Div, + /// a ^ b + Pow, +} + +impl fmt::Display for MultiAD2FR { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + MultiAD2FR::Inp(idx) => write!(f, "Inp({})", idx), + MultiAD2FR::Sin => write!(f, "Sin"), + MultiAD2FR::Cos => write!(f, "Cos"), + MultiAD2FR::Tan => write!(f, "Tan"), + MultiAD2FR::Neg => write!(f, "Neg"), + MultiAD2FR::Exp => write!(f, "Exp"), + MultiAD2FR::Ln => write!(f, "Ln"), + MultiAD2FR::Sqrt => write!(f, "Sqrt"), + MultiAD2FR::Log1pExp => write!(f, "Log1pExp"), + MultiAD2FR::Add => write!(f, "Add"), + MultiAD2FR::Sub => write!(f, "Sub"), + MultiAD2FR::Mul => write!(f, "Mul"), + MultiAD2FR::Div => write!(f, "Div"), + MultiAD2FR::Pow => write!(f, "Pow"), + } + } +} + +impl From for OpKind { + fn from(op: MultiAD2FR) -> OpKind { + match op { + MultiAD2FR::Inp(k) => OpKind::Inp(k), + MultiAD2FR::Sin => OpKind::Sin, + MultiAD2FR::Cos => OpKind::Cos, + MultiAD2FR::Tan => OpKind::Tan, + MultiAD2FR::Neg => OpKind::Neg, + MultiAD2FR::Exp => OpKind::Exp, + MultiAD2FR::Ln => OpKind::Ln, + MultiAD2FR::Sqrt => OpKind::Sqrt, + MultiAD2FR::Log1pExp => OpKind::Log1pExp, + MultiAD2FR::Add => OpKind::Add, + MultiAD2FR::Sub => OpKind::Sub, + MultiAD2FR::Mul => OpKind::Mul, + MultiAD2FR::Div => OpKind::Div, + MultiAD2FR::Pow => OpKind::Pow, + } + } +} + +impl MultiAD2FR { + /// Compute exact Hessian using Forward-over-Reverse mode. + /// + /// Delegates to the shared dual-number forward + reverse algorithm. + /// + /// # Arguments + /// + /// * `ops` — RPN operation sequence + /// * `x` — input vector (n-dimensional) + /// + /// # Returns + /// + /// Hessian matrix `H` where `H[i][j] = ∂²f/∂xᵢ∂xⱼ`. + /// + /// # Errors + /// + /// Returns `Err(AutodiffError)` if an input index is out of bounds or the + /// RPN expression is malformed. + /// + /// # Examples + /// + /// ``` + /// use petite_ad::MultiAD2FR; + /// + /// // f(x, y) = x·y → Hessian = [[0, 1], [1, 0]] + /// let ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Inp(1), MultiAD2FR::Mul]; + /// let h = MultiAD2FR::compute_hessian(&ops, &[2.0, 3.0]).unwrap(); + /// assert_eq!(h[0][1], 1.0); + /// assert_eq!(h[1][0], 1.0); + /// ``` + pub fn compute_hessian(ops: &[MultiAD2FR], x: &[f64]) -> Result>> { + let ops_repr: Vec = ops.iter().map(|&op| OpKind::from(op)).collect(); + compute_hessian_dual(&ops_repr, x) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + const TOL: f64 = 1e-12; + + #[test] + fn test_product_xy() { + let ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Inp(1), MultiAD2FR::Mul]; + let h = MultiAD2FR::compute_hessian(&ops, &[2.0, 3.0]).unwrap(); + assert!((h[0][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + assert!((h[0][1] - 1.0).abs() < TOL); + assert!((h[1][0] - 1.0).abs() < TOL); + } + + #[test] + fn test_simple_quadratic() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Inp(0), + MultiAD2FR::Mul, + MultiAD2FR::Inp(1), + MultiAD2FR::Inp(1), + MultiAD2FR::Mul, + MultiAD2FR::Add, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < TOL); + assert!((h[1][1] - 2.0).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + } + + #[test] + fn test_empty_ops() { + let h = MultiAD2FR::compute_hessian(&[], &[1.0, 2.0]).unwrap(); + assert_eq!(h.len(), 2); + assert!(h.iter().all(|row| row.iter().all(|&v| v == 0.0))); + } + + #[test] + fn test_sin() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Inp(1), + MultiAD2FR::Add, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.sin())).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_cos() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Cos, + MultiAD2FR::Inp(1), + MultiAD2FR::Add, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.cos())).abs() < TOL); + } + + #[test] + fn test_exp() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Exp, + MultiAD2FR::Inp(1), + MultiAD2FR::Add, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 1.0_f64.exp()).abs() < TOL); + } + + #[test] + fn test_three_variables_quadratic() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Inp(0), + MultiAD2FR::Mul, + MultiAD2FR::Inp(1), + MultiAD2FR::Inp(1), + MultiAD2FR::Mul, + MultiAD2FR::Add, + MultiAD2FR::Inp(2), + MultiAD2FR::Inp(2), + MultiAD2FR::Mul, + MultiAD2FR::Add, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0, 3.0]).unwrap(); + for (i, row) in h.iter().enumerate().take(3) { + assert!((row[i] - 2.0).abs() < TOL, "H[{}][{}] = {}", i, i, row[i]); + } + assert!((h[0][1]).abs() < TOL); + assert!((h[0][2]).abs() < TOL); + assert!((h[1][2]).abs() < TOL); + } + + #[test] + fn test_sin_plus_exp() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Inp(1), + MultiAD2FR::Exp, + MultiAD2FR::Add, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.sin())).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1] - 2.0_f64.exp()).abs() < TOL); + } + + #[test] + fn test_sin_times_cos() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Inp(1), + MultiAD2FR::Cos, + MultiAD2FR::Mul, + ]; + let x = 1.0; + let y = 2.0; + let h = MultiAD2FR::compute_hessian(&ops, &[x, y]).unwrap(); + let ex00 = -x.sin() * y.cos(); + let ex01 = -x.cos() * y.sin(); + assert!((h[0][0] - ex00).abs() < TOL, "H[0][0] = {}", h[0][0]); + assert!((h[0][1] - ex01).abs() < TOL, "H[0][1] = {}", h[0][1]); + assert!((h[1][0] - ex01).abs() < TOL, "H[1][0] = {}", h[1][0]); + assert!((h[1][1] - ex00).abs() < TOL, "H[1][1] = {}", h[1][1]); + } + + #[test] + fn test_exp_times_exp() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Exp, + MultiAD2FR::Inp(1), + MultiAD2FR::Exp, + MultiAD2FR::Mul, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + let expected = (1.0_f64.exp()) * (2.0_f64.exp()); + assert!((h[0][0] - expected).abs() < TOL); + assert!((h[0][1] - expected).abs() < TOL); + assert!((h[1][1] - expected).abs() < TOL); + } + + #[test] + fn test_exp_times_sin_single_var() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Exp, + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Mul, + ]; + let x = 1.0; + let h = MultiAD2FR::compute_hessian(&ops, &[x]).unwrap(); + let expected = 2.0 * x.exp() * x.cos(); + assert!((h[0][0] - expected).abs() < TOL, "H[0][0] = {}", h[0][0]); + } + + #[test] + fn test_sum_squared() { + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Inp(1), + MultiAD2FR::Add, + MultiAD2FR::Inp(0), + MultiAD2FR::Inp(1), + MultiAD2FR::Add, + MultiAD2FR::Mul, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < TOL); + assert!((h[0][1] - 2.0).abs() < TOL); + assert!((h[1][1] - 2.0).abs() < TOL); + } + + #[test] + fn test_tan_and_sub() { + let x = 0.3_f64; + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Tan, + MultiAD2FR::Inp(1), + MultiAD2FR::Sub, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[x, 2.0]).unwrap(); + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + assert!((h[0][0] - expected).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_div_ln_sqrt() { + let x = 4.0_f64; + let y = 2.0_f64; + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sqrt, + MultiAD2FR::Inp(1), + MultiAD2FR::Ln, + MultiAD2FR::Div, + ]; + let h = MultiAD2FR::compute_hessian(&ops, &[x, y]).unwrap(); + let sqrt_x = x.sqrt(); + let expected_xx = -1.0 / (4.0 * x * sqrt_x * y.ln()); + let expected_xy = -1.0 / (2.0 * sqrt_x * y * y.ln().powi(2)); + let expected_yy = + 2.0 * sqrt_x / (y * y * y.ln().powi(3)) + sqrt_x / (y * y * y.ln().powi(2)); + assert!((h[0][0] - expected_xx).abs() < 1e-10); + assert!((h[0][1] - expected_xy).abs() < 1e-10); + assert!((h[1][0] - expected_xy).abs() < 1e-10); + assert!((h[1][1] - expected_yy).abs() < 1e-10); + } + + #[test] + fn test_pow_and_neg() { + let x = 2.0_f64; + let y = 3.0_f64; + let pow_ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Inp(1), MultiAD2FR::Pow]; + let h = MultiAD2FR::compute_hessian(&pow_ops, &[x, y]).unwrap(); + let expected_xx = y * (y - 1.0) * x.powf(y - 2.0); + let expected_xy = x.powf(y - 1.0) * (1.0 + y * x.ln()); + let expected_yy = x.powf(y) * x.ln().powi(2); + assert!((h[0][0] - expected_xx).abs() < 1e-10); + assert!((h[0][1] - expected_xy).abs() < 1e-10); + assert!((h[1][0] - expected_xy).abs() < 1e-10); + assert!((h[1][1] - expected_yy).abs() < 1e-10); + + let neg_ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Sin, MultiAD2FR::Neg]; + let h_neg = MultiAD2FR::compute_hessian(&neg_ops, &[0.5]).unwrap(); + assert!((h_neg[0][0] - 0.5_f64.sin()).abs() < 1e-10); + } + + // ---- Additional coverage tests ---- + + #[test] + fn test_display_format() { + assert_eq!(format!("{}", MultiAD2FR::Inp(2)), "Inp(2)"); + assert_eq!(format!("{}", MultiAD2FR::Sin), "Sin"); + assert_eq!(format!("{}", MultiAD2FR::Cos), "Cos"); + assert_eq!(format!("{}", MultiAD2FR::Tan), "Tan"); + assert_eq!(format!("{}", MultiAD2FR::Neg), "Neg"); + assert_eq!(format!("{}", MultiAD2FR::Exp), "Exp"); + assert_eq!(format!("{}", MultiAD2FR::Ln), "Ln"); + assert_eq!(format!("{}", MultiAD2FR::Sqrt), "Sqrt"); + assert_eq!(format!("{}", MultiAD2FR::Log1pExp), "Log1pExp"); + assert_eq!(format!("{}", MultiAD2FR::Add), "Add"); + assert_eq!(format!("{}", MultiAD2FR::Sub), "Sub"); + assert_eq!(format!("{}", MultiAD2FR::Mul), "Mul"); + assert_eq!(format!("{}", MultiAD2FR::Div), "Div"); + assert_eq!(format!("{}", MultiAD2FR::Pow), "Pow"); + } + + #[test] + fn test_display_all_variants() { + // Exercise Display for all remaining variants through compute_hessian + // (which covers From for OpKind for those variants) + let ops = vec![ + MultiAD2FR::Inp(0), + MultiAD2FR::Sin, + MultiAD2FR::Cos, + MultiAD2FR::Tan, + MultiAD2FR::Neg, + MultiAD2FR::Exp, + MultiAD2FR::Ln, + MultiAD2FR::Sqrt, + MultiAD2FR::Log1pExp, + MultiAD2FR::Add, + MultiAD2FR::Sub, + MultiAD2FR::Mul, + MultiAD2FR::Div, + MultiAD2FR::Pow, + ]; + // Just check display doesn't panic + for op in &ops { + let _s = format!("{}", op); + } + } + + #[test] + fn test_hessian_sub_op() { + // f(x, y) = x - y → H = [[0, 0], [0, 0]] + let ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Inp(1), MultiAD2FR::Sub]; + let h = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0]).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_log1p_exp_hessian() { + // f(x) = log1p_exp(x) = ln(1+exp(x)) + // f''(x) = exp(x)/(1+exp(x))^2 = sigmoid(x)*(1-sigmoid(x)) + let ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Log1pExp]; + let x = 1.0_f64; + let h = MultiAD2FR::compute_hessian(&ops, &[x]).unwrap(); + let sigmoid = x.exp() / (1.0 + x.exp()); + let expected = sigmoid * (1.0 - sigmoid); + assert!((h[0][0] - expected).abs() < 1e-10, "H[0][0] = {}", h[0][0]); + } + + #[test] + fn test_input_index_out_of_bounds() { + let ops = vec![MultiAD2FR::Inp(5)]; + let result = MultiAD2FR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_unary_missing_operand() { + let ops = vec![MultiAD2FR::Sin]; + let result = MultiAD2FR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_binary_missing_right_operand() { + let ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Add]; + let result = MultiAD2FR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_binary_missing_left_operand() { + let ops = vec![MultiAD2FR::Add]; + let result = MultiAD2FR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_extra_items_on_stack() { + let ops = vec![MultiAD2FR::Inp(0), MultiAD2FR::Inp(1)]; + let result = MultiAD2FR::compute_hessian(&ops, &[1.0, 2.0]); + assert!(result.is_err()); + } +} diff --git a/src/multi/multi_ad_rf.rs b/src/multi/multi_ad_rf.rs new file mode 100644 index 0000000..66aa171 --- /dev/null +++ b/src/multi/multi_ad_rf.rs @@ -0,0 +1,522 @@ +//! Exact second-order autodiff for multivariate functions using Reverse-over-Forward (RF) mode. +//! +//! This module implements the RF method for computing exact Hessians of +//! multivariate functions f: ℝⁿ → ℝ. +//! +//! # Supported Operations +//! +//! `MultiAD2RF` supports the smooth operation subset shared by the exact Hessian +//! engines: `Inp`, `Sin`, `Cos`, `Tan`, `Neg`, `Exp`, `Ln`, `Sqrt`, `Log1pExp`, +//! `Add`, `Sub`, `Mul`, `Div`, and `Pow`. Non-smooth operations like `Abs` are not supported by +//! the exact Hessian types. For Hessian approximation with the full operation set, +//! use [`MultiAD::compute_hessian`](crate::MultiAD::compute_hessian) +//! (finite-difference based). +//! +//! # Algorithm +//! +//! RF composes reverse-mode AD (outer) with forward-mode AD (inner): +//! +//! 1. **Forward pass** — evaluate the computation graph with **dual numbers** +//! whose tangent component is seeded in direction e_j. This computes the +//! directional derivative D_{e_j} f for each node. +//! +//! 2. **Reverse pass** — propagate **dual adjoints** backward to differentiate +//! the directional derivative computation. Each adjoint is a dual +//! `(val, tan)` where: +//! - `val = ∂f/∂node` (standard first-order adjoint) +//! - `tan = ∂²f/(∂node · ∂x_j)` (second-order cross-derivative) +//! +//! The dual adjoint at input node k gives Hessian entry `H[j][k]`. +//! +//! # FR vs RF for scalar functions +//! +//! For scalar functions f: ℝⁿ → ℝ, FR and RF produce **identical computations**: +//! both require a dual-number forward pass followed by a dual-adjoint reverse pass +//! for each seed direction. The distinction is conceptual: +//! +//! | Aspect | FR | RF | +//! |--------|----|----| +//! | Inner AD | Reverse (gradient) | Forward (directional derivative) | +//! | Outer AD | Forward (differentiate gradient) | Reverse (differentiate directional deriv) | +//! | Implementation | Dual forward + dual reverse | Dual forward + dual reverse | +//! +//! The distinction becomes meaningful for vector-valued functions f: ℝⁿ → ℝᵐ (m > 1). +//! +//! # Accuracy +//! +//! Machine-precision exact (~1e-14 relative error). No finite differences. +//! +//! # Complexity +//! +//! O(n · G) where n = number of inputs, G = graph size. +//! +//! See [docs/multi_ad_hessian.md](../../docs/multi_ad_hessian.md) for complete theory. + +use std::fmt; + +use crate::Result; + +use super::multi_hessian_common::{compute_hessian_dual, OpKind}; + +/// Stack-based operation for multivariate second-order AD. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MultiAD2RF { + /// Input variable by index + Inp(usize), + /// sin(x) + Sin, + /// cos(x) + Cos, + /// tan(x) + Tan, + /// -x + Neg, + /// exp(x) + Exp, + /// ln(x) + Ln, + /// sqrt(x) + Sqrt, + /// stable ln(1 + exp(x)) + Log1pExp, + /// a + b + Add, + /// a - b + Sub, + /// a · b + Mul, + /// a / b + Div, + /// a ^ b + Pow, +} + +impl fmt::Display for MultiAD2RF { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + MultiAD2RF::Inp(idx) => write!(f, "Inp({})", idx), + MultiAD2RF::Sin => write!(f, "Sin"), + MultiAD2RF::Cos => write!(f, "Cos"), + MultiAD2RF::Tan => write!(f, "Tan"), + MultiAD2RF::Neg => write!(f, "Neg"), + MultiAD2RF::Exp => write!(f, "Exp"), + MultiAD2RF::Ln => write!(f, "Ln"), + MultiAD2RF::Sqrt => write!(f, "Sqrt"), + MultiAD2RF::Log1pExp => write!(f, "Log1pExp"), + MultiAD2RF::Add => write!(f, "Add"), + MultiAD2RF::Sub => write!(f, "Sub"), + MultiAD2RF::Mul => write!(f, "Mul"), + MultiAD2RF::Div => write!(f, "Div"), + MultiAD2RF::Pow => write!(f, "Pow"), + } + } +} + +impl From for OpKind { + fn from(op: MultiAD2RF) -> OpKind { + match op { + MultiAD2RF::Inp(k) => OpKind::Inp(k), + MultiAD2RF::Sin => OpKind::Sin, + MultiAD2RF::Cos => OpKind::Cos, + MultiAD2RF::Tan => OpKind::Tan, + MultiAD2RF::Neg => OpKind::Neg, + MultiAD2RF::Exp => OpKind::Exp, + MultiAD2RF::Ln => OpKind::Ln, + MultiAD2RF::Sqrt => OpKind::Sqrt, + MultiAD2RF::Log1pExp => OpKind::Log1pExp, + MultiAD2RF::Add => OpKind::Add, + MultiAD2RF::Sub => OpKind::Sub, + MultiAD2RF::Mul => OpKind::Mul, + MultiAD2RF::Div => OpKind::Div, + MultiAD2RF::Pow => OpKind::Pow, + } + } +} + +impl MultiAD2RF { + /// Compute exact Hessian using Reverse-over-Forward mode. + /// + /// Delegates to the shared dual-number forward + reverse algorithm, + /// which implements the RF composition for scalar functions. + /// + /// # Arguments + /// + /// * `ops` — RPN operation sequence + /// * `x` — input vector (n-dimensional) + /// + /// # Returns + /// + /// Hessian matrix `H` where `H[i][j] = ∂²f/∂xᵢ∂xⱼ`. + /// + /// # Errors + /// + /// Returns `Err(AutodiffError)` if an input index is out of bounds or the + /// RPN expression is malformed. + /// + /// # Examples + /// + /// ``` + /// use petite_ad::MultiAD2RF; + /// + /// // f(x, y) = x·y → Hessian = [[0, 1], [1, 0]] + /// let ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Inp(1), MultiAD2RF::Mul]; + /// let h = MultiAD2RF::compute_hessian(&ops, &[2.0, 3.0]).unwrap(); + /// assert_eq!(h[0][1], 1.0); + /// assert_eq!(h[1][0], 1.0); + /// ``` + pub fn compute_hessian(ops: &[MultiAD2RF], x: &[f64]) -> Result>> { + let ops_repr: Vec = ops.iter().map(|&op| OpKind::from(op)).collect(); + compute_hessian_dual(&ops_repr, x) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + const TOL: f64 = 1e-12; + + #[test] + fn test_product_xy() { + let ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Inp(1), MultiAD2RF::Mul]; + let h = MultiAD2RF::compute_hessian(&ops, &[2.0, 3.0]).unwrap(); + assert!((h[0][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + assert!((h[0][1] - 1.0).abs() < TOL); + assert!((h[1][0] - 1.0).abs() < TOL); + } + + #[test] + fn test_simple_quadratic() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Inp(0), + MultiAD2RF::Mul, + MultiAD2RF::Inp(1), + MultiAD2RF::Inp(1), + MultiAD2RF::Mul, + MultiAD2RF::Add, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < TOL); + assert!((h[1][1] - 2.0).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + } + + #[test] + fn test_empty_ops() { + let h = MultiAD2RF::compute_hessian(&[], &[1.0, 2.0]).unwrap(); + assert_eq!(h.len(), 2); + assert!(h.iter().all(|row| row.iter().all(|&v| v == 0.0))); + } + + #[test] + fn test_sin() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Inp(1), + MultiAD2RF::Add, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.sin())).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_cos() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Cos, + MultiAD2RF::Inp(1), + MultiAD2RF::Add, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.cos())).abs() < TOL); + } + + #[test] + fn test_exp() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Exp, + MultiAD2RF::Inp(1), + MultiAD2RF::Add, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 1.0_f64.exp()).abs() < TOL); + } + + #[test] + fn test_three_variables_quadratic() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Inp(0), + MultiAD2RF::Mul, + MultiAD2RF::Inp(1), + MultiAD2RF::Inp(1), + MultiAD2RF::Mul, + MultiAD2RF::Add, + MultiAD2RF::Inp(2), + MultiAD2RF::Inp(2), + MultiAD2RF::Mul, + MultiAD2RF::Add, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0, 3.0]).unwrap(); + for (i, row) in h.iter().enumerate().take(3) { + assert!((row[i] - 2.0).abs() < TOL, "H[{}][{}] = {}", i, i, row[i]); + } + assert!((h[0][1]).abs() < TOL); + assert!((h[0][2]).abs() < TOL); + assert!((h[1][2]).abs() < TOL); + } + + #[test] + fn test_sin_plus_exp() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Inp(1), + MultiAD2RF::Exp, + MultiAD2RF::Add, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.sin())).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1] - 2.0_f64.exp()).abs() < TOL); + } + + #[test] + fn test_sin_times_cos() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Inp(1), + MultiAD2RF::Cos, + MultiAD2RF::Mul, + ]; + let x = 1.0; + let y = 2.0; + let h = MultiAD2RF::compute_hessian(&ops, &[x, y]).unwrap(); + let ex00 = -x.sin() * y.cos(); + let ex01 = -x.cos() * y.sin(); + assert!((h[0][0] - ex00).abs() < TOL, "H[0][0] = {}", h[0][0]); + assert!((h[0][1] - ex01).abs() < TOL, "H[0][1] = {}", h[0][1]); + assert!((h[1][0] - ex01).abs() < TOL, "H[1][0] = {}", h[1][0]); + assert!((h[1][1] - ex00).abs() < TOL, "H[1][1] = {}", h[1][1]); + } + + #[test] + fn test_exp_times_exp() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Exp, + MultiAD2RF::Inp(1), + MultiAD2RF::Exp, + MultiAD2RF::Mul, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + let expected = (1.0_f64.exp()) * (2.0_f64.exp()); + assert!((h[0][0] - expected).abs() < TOL); + assert!((h[0][1] - expected).abs() < TOL); + assert!((h[1][1] - expected).abs() < TOL); + } + + #[test] + fn test_exp_times_sin_single_var() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Exp, + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Mul, + ]; + let x = 1.0; + let h = MultiAD2RF::compute_hessian(&ops, &[x]).unwrap(); + let expected = 2.0 * x.exp() * x.cos(); + assert!((h[0][0] - expected).abs() < TOL, "H[0][0] = {}", h[0][0]); + } + + #[test] + fn test_sum_squared() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Inp(1), + MultiAD2RF::Add, + MultiAD2RF::Inp(0), + MultiAD2RF::Inp(1), + MultiAD2RF::Add, + MultiAD2RF::Mul, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < TOL); + assert!((h[0][1] - 2.0).abs() < TOL); + assert!((h[1][1] - 2.0).abs() < TOL); + } + + #[test] + fn test_tan_and_sub() { + let x = 0.3_f64; + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Tan, + MultiAD2RF::Inp(1), + MultiAD2RF::Sub, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[x, 2.0]).unwrap(); + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + assert!((h[0][0] - expected).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_div_ln_sqrt() { + let x = 4.0_f64; + let y = 2.0_f64; + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sqrt, + MultiAD2RF::Inp(1), + MultiAD2RF::Ln, + MultiAD2RF::Div, + ]; + let h = MultiAD2RF::compute_hessian(&ops, &[x, y]).unwrap(); + let sqrt_x = x.sqrt(); + let expected_xx = -1.0 / (4.0 * x * sqrt_x * y.ln()); + let expected_xy = -1.0 / (2.0 * sqrt_x * y * y.ln().powi(2)); + let expected_yy = + 2.0 * sqrt_x / (y * y * y.ln().powi(3)) + sqrt_x / (y * y * y.ln().powi(2)); + assert!((h[0][0] - expected_xx).abs() < 1e-10); + assert!((h[0][1] - expected_xy).abs() < 1e-10); + assert!((h[1][0] - expected_xy).abs() < 1e-10); + assert!((h[1][1] - expected_yy).abs() < 1e-10); + } + + #[test] + fn test_pow_and_neg() { + let x = 2.0_f64; + let y = 3.0_f64; + let pow_ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Inp(1), MultiAD2RF::Pow]; + let h = MultiAD2RF::compute_hessian(&pow_ops, &[x, y]).unwrap(); + let expected_xx = y * (y - 1.0) * x.powf(y - 2.0); + let expected_xy = x.powf(y - 1.0) * (1.0 + y * x.ln()); + let expected_yy = x.powf(y) * x.ln().powi(2); + assert!((h[0][0] - expected_xx).abs() < 1e-10); + assert!((h[0][1] - expected_xy).abs() < 1e-10); + assert!((h[1][0] - expected_xy).abs() < 1e-10); + assert!((h[1][1] - expected_yy).abs() < 1e-10); + + let neg_ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Sin, MultiAD2RF::Neg]; + let h_neg = MultiAD2RF::compute_hessian(&neg_ops, &[0.5]).unwrap(); + assert!((h_neg[0][0] - 0.5_f64.sin()).abs() < 1e-10); + } + + // ---- Additional coverage tests ---- + + #[test] + fn test_display_format() { + assert_eq!(format!("{}", MultiAD2RF::Inp(2)), "Inp(2)"); + assert_eq!(format!("{}", MultiAD2RF::Sin), "Sin"); + assert_eq!(format!("{}", MultiAD2RF::Cos), "Cos"); + assert_eq!(format!("{}", MultiAD2RF::Tan), "Tan"); + assert_eq!(format!("{}", MultiAD2RF::Neg), "Neg"); + assert_eq!(format!("{}", MultiAD2RF::Exp), "Exp"); + assert_eq!(format!("{}", MultiAD2RF::Ln), "Ln"); + assert_eq!(format!("{}", MultiAD2RF::Sqrt), "Sqrt"); + assert_eq!(format!("{}", MultiAD2RF::Log1pExp), "Log1pExp"); + assert_eq!(format!("{}", MultiAD2RF::Add), "Add"); + assert_eq!(format!("{}", MultiAD2RF::Sub), "Sub"); + assert_eq!(format!("{}", MultiAD2RF::Mul), "Mul"); + assert_eq!(format!("{}", MultiAD2RF::Div), "Div"); + assert_eq!(format!("{}", MultiAD2RF::Pow), "Pow"); + } + + #[test] + fn test_display_all_variants() { + let ops = vec![ + MultiAD2RF::Inp(0), + MultiAD2RF::Sin, + MultiAD2RF::Cos, + MultiAD2RF::Tan, + MultiAD2RF::Neg, + MultiAD2RF::Exp, + MultiAD2RF::Ln, + MultiAD2RF::Sqrt, + MultiAD2RF::Log1pExp, + MultiAD2RF::Add, + MultiAD2RF::Sub, + MultiAD2RF::Mul, + MultiAD2RF::Div, + MultiAD2RF::Pow, + ]; + for op in &ops { + let _s = format!("{}", op); + } + } + + #[test] + fn test_hessian_sub_op() { + // f(x, y) = x - y → H = [[0, 0], [0, 0]] + let ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Inp(1), MultiAD2RF::Sub]; + let h = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0]).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_log1p_exp_hessian() { + // f(x) = log1p_exp(x) = ln(1+exp(x)) + // f''(x) = exp(x)/(1+exp(x))^2 + let ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Log1pExp]; + let x = 1.0_f64; + let h = MultiAD2RF::compute_hessian(&ops, &[x]).unwrap(); + let sigmoid = x.exp() / (1.0 + x.exp()); + let expected = sigmoid * (1.0 - sigmoid); + assert!((h[0][0] - expected).abs() < 1e-10, "H[0][0] = {}", h[0][0]); + } + + #[test] + fn test_input_index_out_of_bounds() { + let ops = vec![MultiAD2RF::Inp(5)]; + let result = MultiAD2RF::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_unary_missing_operand() { + let ops = vec![MultiAD2RF::Sin]; + let result = MultiAD2RF::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_binary_missing_right_operand() { + let ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Add]; + let result = MultiAD2RF::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_binary_missing_both_operands() { + let ops = vec![MultiAD2RF::Add]; + let result = MultiAD2RF::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_extra_items_on_stack() { + let ops = vec![MultiAD2RF::Inp(0), MultiAD2RF::Inp(1)]; + let result = MultiAD2RF::compute_hessian(&ops, &[1.0, 2.0]); + assert!(result.is_err()); + } +} diff --git a/src/multi/multi_ad_rr.rs b/src/multi/multi_ad_rr.rs new file mode 100644 index 0000000..49f5479 --- /dev/null +++ b/src/multi/multi_ad_rr.rs @@ -0,0 +1,879 @@ +//! Exact second-order autodiff for multivariate functions using Reverse-over-Reverse (RR) mode. +//! +//! This module implements the RR method for computing exact Hessians of +//! multivariate functions f: ℝⁿ → ℝ. +//! +//! # Supported Operations +//! +//! `MultiAD2RR` supports the smooth operation subset shared by the exact Hessian +//! engines: `Inp`, `Sin`, `Cos`, `Tan`, `Neg`, `Exp`, `Ln`, `Sqrt`, `Add`, `Sub`, +//! `Mul`, `Div`, and `Pow`. Non-smooth operations like `Abs` are not supported by +//! the exact Hessian types. For Hessian approximation with the full operation set, +//! use [`MultiAD::compute_hessian`](crate::MultiAD::compute_hessian) +//! (finite-difference based). +//! +//! # Algorithm +//! +//! The RR method applies reverse-mode AD with second-order chain rule: +//! +//! 1. **Forward pass** — evaluate the computation graph, storing at each node: +//! - The node's scalar value +//! - The node's **gradient vector** `g = ∂node/∂x` (one entry per input variable) +//! - Local first and second derivatives of the operation +//! +//! 2. **Reverse pass** — starting from the output, accumulate: +//! - **Adjoint** (first-order): standard reverse-mode +//! - **Hessian** (second-order): for each operation z = op(u, v): +//! ```text +//! H += a * ddy_uu * g_u (x) g_u^T +//! H += a * ddy_vv * g_v (x) g_v^T +//! H += a * ddy_uv * (g_u (x) g_v^T + g_v (x) g_u^T) +//! ``` +//! +//! See [docs/multi_ad_hessian.md](../../docs/multi_ad_hessian.md) for complete theory. + +use std::fmt; + +use crate::error::{AutodiffError, Result}; + +use super::multi_ad::MultiAD; +use super::op_rules::{self, LocalRule}; + +/// Stack-based operation for multivariate second-order AD. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum MultiAD2RR { + /// Input variable by index + Inp(usize), + /// sin(x) + Sin, + /// cos(x) + Cos, + /// tan(x) + Tan, + /// -x + Neg, + /// exp(x) + Exp, + /// ln(x) + Ln, + /// sqrt(x) + Sqrt, + /// a + b + Add, + /// a - b + Sub, + /// a · b + Mul, + /// a / b + Div, + /// a ^ b + Pow, +} + +impl fmt::Display for MultiAD2RR { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + MultiAD2RR::Inp(idx) => write!(f, "Inp({})", idx), + MultiAD2RR::Sin => write!(f, "Sin"), + MultiAD2RR::Cos => write!(f, "Cos"), + MultiAD2RR::Tan => write!(f, "Tan"), + MultiAD2RR::Neg => write!(f, "Neg"), + MultiAD2RR::Exp => write!(f, "Exp"), + MultiAD2RR::Ln => write!(f, "Ln"), + MultiAD2RR::Sqrt => write!(f, "Sqrt"), + MultiAD2RR::Add => write!(f, "Add"), + MultiAD2RR::Sub => write!(f, "Sub"), + MultiAD2RR::Mul => write!(f, "Mul"), + MultiAD2RR::Div => write!(f, "Div"), + MultiAD2RR::Pow => write!(f, "Pow"), + } + } +} + +// --------------------------------------------------------------------------- +// Internal data structures +// --------------------------------------------------------------------------- + +/// A node in the computation graph. +struct StackNode { + value: f64, + /// ∂(this node)/∂x_j for each input variable x_j + grad: Vec, +} + +/// Local derivative information stored during the forward pass. +enum LocalDerivs { + Unary { + parent: usize, + dy: f64, // dz/du + ddy: f64, // d²z/du² + }, + Binary { + left: usize, + right: usize, + dy_left: f64, // dz/du + dy_right: f64, // dz/dv + ddy_left_left: f64, // d²z/du² + ddy_right_right: f64, // d²z/dv² + ddy_left_right: f64, // d²z/(du dv) + }, +} + +// --------------------------------------------------------------------------- +// Implementation +// --------------------------------------------------------------------------- + +#[inline(always)] +fn as_multiad(op: MultiAD2RR) -> Option { + match op { + MultiAD2RR::Inp(_) => None, + MultiAD2RR::Sin => Some(MultiAD::Sin), + MultiAD2RR::Cos => Some(MultiAD::Cos), + MultiAD2RR::Tan => Some(MultiAD::Tan), + MultiAD2RR::Neg => Some(MultiAD::Neg), + MultiAD2RR::Exp => Some(MultiAD::Exp), + MultiAD2RR::Ln => Some(MultiAD::Ln), + MultiAD2RR::Sqrt => Some(MultiAD::Sqrt), + MultiAD2RR::Add => Some(MultiAD::Add), + MultiAD2RR::Sub => Some(MultiAD::Sub), + MultiAD2RR::Mul => Some(MultiAD::Mul), + MultiAD2RR::Div => Some(MultiAD::Div), + MultiAD2RR::Pow => Some(MultiAD::Pow), + } +} + +impl MultiAD2RR { + /// Compute exact Hessian using Reverse-over-Reverse mode. + /// + /// # Arguments + /// + /// * `ops` — RPN operation sequence + /// * `x` — input vector (n-dimensional) + /// + /// # Returns + /// + /// Hessian matrix `H` where `H[i][j] = ∂²f/∂xᵢ∂xⱼ`. + /// + /// # Errors + /// + /// Returns `Err(AutodiffError)` if an input index is out of bounds or the + /// RPN expression is malformed. + /// + /// # Accuracy + /// + /// Machine-precision exact (~1e-14 relative error). + /// + /// # Complexity + /// + /// - Time: O(G · n²) where G = graph size, n = number of inputs + /// - Space: O(G · n) for per-node gradient vectors + /// + /// # Examples + /// + /// ``` + /// use petite_ad::MultiAD2RR; + /// + /// // f(x, y) = x² + y² → Hessian = [[2, 0], [0, 2]] + /// let ops = vec![ + /// MultiAD2RR::Inp(0), MultiAD2RR::Inp(0), MultiAD2RR::Mul, + /// MultiAD2RR::Inp(1), MultiAD2RR::Inp(1), MultiAD2RR::Mul, + /// MultiAD2RR::Add, + /// ]; + /// let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + /// assert!((h[0][0] - 2.0).abs() < 1e-12); + /// assert!((h[1][1] - 2.0).abs() < 1e-12); + /// ``` + pub fn compute_hessian(ops: &[MultiAD2RR], x: &[f64]) -> Result>> { + let n_vars = x.len(); + if ops.is_empty() { + return Ok(vec![vec![0.0; n_vars]; n_vars]); + } + + // ================================================================ + // FORWARD PASS: value + per-node gradient vector + // ================================================================ + let mut nodes: Vec = Vec::new(); + let mut local_derivs: Vec> = Vec::new(); + let mut eval_stack: Vec = Vec::new(); + + for &op in ops { + match op { + MultiAD2RR::Inp(k) => { + if k >= n_vars { + return Err(AutodiffError::IndexOutOfBounds { + index: k, + max_index: n_vars.saturating_sub(1), + }); + } + let idx = nodes.len(); + let mut grad = vec![0.0; n_vars]; + grad[k] = 1.0; + nodes.push(StackNode { value: x[k], grad }); + local_derivs.push(None); + eval_stack.push(idx); + } + MultiAD2RR::Sin + | MultiAD2RR::Cos + | MultiAD2RR::Tan + | MultiAD2RR::Neg + | MultiAD2RR::Exp + | MultiAD2RR::Ln + | MultiAD2RR::Sqrt => { + let p = eval_stack.pop().ok_or(AutodiffError::InvalidGraph { + reason: "unary operation missing operand", + })?; + let op = as_multiad(op).expect("non-input op"); + let args = [nodes[p].value]; + let value = op_rules::forward_value(op, &args)?; + let rule = op_rules::local_rule(op, &args, value)?; + let idx = nodes.len(); + let (dy, ddy) = match rule { + LocalRule::Unary { dy, ddy } => (dy, ddy), + LocalRule::Binary { .. } => unreachable!("unary op must have unary rule"), + }; + let grad: Vec = nodes[p].grad.iter().map(|&g| dy * g).collect(); + nodes.push(StackNode { value, grad }); + local_derivs.push(Some(LocalDerivs::Unary { parent: p, dy, ddy })); + eval_stack.push(idx); + } + MultiAD2RR::Add + | MultiAD2RR::Sub + | MultiAD2RR::Mul + | MultiAD2RR::Div + | MultiAD2RR::Pow => { + let r = eval_stack.pop().ok_or(AutodiffError::InvalidGraph { + reason: "binary operation missing right operand", + })?; + let l = eval_stack.pop().ok_or(AutodiffError::InvalidGraph { + reason: "binary operation missing left operand", + })?; + let op = as_multiad(op).expect("non-input op"); + let args = [nodes[l].value, nodes[r].value]; + let value = op_rules::forward_value(op, &args)?; + let rule = op_rules::local_rule(op, &args, value)?; + let idx = nodes.len(); + let (dy_left, dy_right, ddy_left_left, ddy_right_right, ddy_left_right) = + match rule { + LocalRule::Unary { .. } => { + unreachable!("binary op must have binary rule") + } + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => ( + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + ), + }; + let grad: Vec = nodes[l] + .grad + .iter() + .zip(&nodes[r].grad) + .map(|(a, b)| dy_left * a + dy_right * b) + .collect(); + nodes.push(StackNode { value, grad }); + local_derivs.push(Some(LocalDerivs::Binary { + left: l, + right: r, + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + })); + eval_stack.push(idx); + } + } + } + + if eval_stack.len() != 1 { + return Err(AutodiffError::InvalidGraph { + reason: "RPN expression must leave exactly one output on the stack", + }); + } + + // ================================================================ + // REVERSE PASS: adjoint + Hessian accumulation + // ================================================================ + let mut adjoint: Vec = vec![0.0; nodes.len()]; + let mut hessian: Vec> = vec![vec![0.0; n_vars]; n_vars]; + + adjoint[nodes.len() - 1] = 1.0; + + for i in (0..nodes.len()).rev() { + let a = adjoint[i]; + + if let Some(ref d) = local_derivs[i] { + match d { + LocalDerivs::Unary { parent, dy, ddy } => { + adjoint[*parent] += a * dy; + + // H += a · ddy · g_parent ⊗ g_parent + let g = &nodes[*parent].grad; + for vi in 0..n_vars { + let gi = g[vi]; + if gi == 0.0 { + continue; + } + // Precompute factor outside the vj loop + let factor = a * ddy * gi; + for vj in vi..n_vars { + let contrib = factor * g[vj]; + hessian[vi][vj] += contrib; + if vi != vj { + hessian[vj][vi] += contrib; + } + } + } + } + LocalDerivs::Binary { + left, + right, + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + adjoint[*left] += a * dy_left; + adjoint[*right] += a * dy_right; + + let gl = &nodes[*left].grad; + let gr = &nodes[*right].grad; + + // H += a · ddy_ll · g_left ⊗ g_left + if *ddy_left_left != 0.0 { + for vi in 0..n_vars { + let gi = gl[vi]; + if gi == 0.0 { + continue; + } + for vj in 0..n_vars { + hessian[vi][vj] += a * ddy_left_left * gi * gl[vj]; + } + } + } + + // H += a · ddy_rr · g_right ⊗ g_right + if *ddy_right_right != 0.0 { + for vi in 0..n_vars { + let gi = gr[vi]; + if gi == 0.0 { + continue; + } + for vj in 0..n_vars { + hessian[vi][vj] += a * ddy_right_right * gi * gr[vj]; + } + } + } + + // H += a · ddy_lr · (g_left ⊗ g_right + g_right ⊗ g_left) + if *ddy_left_right != 0.0 { + for vi in 0..n_vars { + for vj in 0..n_vars { + hessian[vi][vj] += + a * ddy_left_right * (gl[vi] * gr[vj] + gr[vi] * gl[vj]); + } + } + } + } + } + } + } + + Ok(hessian) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + /// Absolute tolerance for exact (non-finite-difference) methods. + const TOL: f64 = 1e-12; + + #[test] + fn test_product_xy() { + // f = x·y → H = [[0,1],[1,0]] + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Inp(1), MultiAD2RR::Mul]; + let h = MultiAD2RR::compute_hessian(&ops, &[2.0, 3.0]).unwrap(); + assert_eq!(h[0][0], 0.0); + assert_eq!(h[1][1], 0.0); + assert_eq!(h[0][1], 1.0); + assert_eq!(h[1][0], 1.0); + } + + #[test] + fn test_simple_quadratic() { + // f = x² + y² → H = [[2,0],[0,2]] + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(0), + MultiAD2RR::Mul, + MultiAD2RR::Inp(1), + MultiAD2RR::Inp(1), + MultiAD2RR::Mul, + MultiAD2RR::Add, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < TOL); + assert!((h[1][1] - 2.0).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + } + + #[test] + fn test_empty_ops() { + let h = MultiAD2RR::compute_hessian(&[], &[1.0, 2.0]).unwrap(); + assert_eq!(h.len(), 2); + assert!(h.iter().all(|row| row.iter().all(|&v| v == 0.0))); + } + + #[test] + fn test_sin() { + // f = sin(x) + y → H[0][0] = -sin(x), rest 0 + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Sin, + MultiAD2RR::Inp(1), + MultiAD2RR::Add, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.sin())).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_cos() { + // f = cos(x) + y → H[0][0] = -cos(x) + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Cos, + MultiAD2RR::Inp(1), + MultiAD2RR::Add, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.cos())).abs() < TOL); + } + + #[test] + fn test_exp() { + // f = exp(x) + y → H[0][0] = exp(x) + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Exp, + MultiAD2RR::Inp(1), + MultiAD2RR::Add, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - 1.0_f64.exp()).abs() < TOL); + } + + #[test] + fn test_three_variables_quadratic() { + // f = x²+y²+z² → H = diag(2,2,2) + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(0), + MultiAD2RR::Mul, + MultiAD2RR::Inp(1), + MultiAD2RR::Inp(1), + MultiAD2RR::Mul, + MultiAD2RR::Add, + MultiAD2RR::Inp(2), + MultiAD2RR::Inp(2), + MultiAD2RR::Mul, + MultiAD2RR::Add, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0, 3.0]).unwrap(); + for (i, row) in h.iter().enumerate().take(3) { + assert!((row[i] - 2.0).abs() < TOL, "H[{}][{}] = {}", i, i, row[i]); + } + assert!((h[0][1]).abs() < TOL); + assert!((h[0][2]).abs() < TOL); + assert!((h[1][2]).abs() < TOL); + } + + #[test] + fn test_sin_plus_exp() { + // f = sin(x) + exp(y) → H = [[-sin(x), 0], [0, exp(y)]] + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Sin, + MultiAD2RR::Inp(1), + MultiAD2RR::Exp, + MultiAD2RR::Add, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0] - (-1.0_f64.sin())).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1] - 2.0_f64.exp()).abs() < TOL); + } + + // ---- New correctness tests (non-trivial compositions) ---- + + #[test] + fn test_sin_times_cos() { + // f = sin(x)·cos(y) + // H[0][0] = -sin(x)·cos(y) + // H[1][1] = -sin(x)·cos(y) + // H[0][1] = H[1][0] = -cos(x)·sin(y) + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Sin, + MultiAD2RR::Inp(1), + MultiAD2RR::Cos, + MultiAD2RR::Mul, + ]; + let x = 1.0; + let y = 2.0; + let h = MultiAD2RR::compute_hessian(&ops, &[x, y]).unwrap(); + + let ex00 = -x.sin() * y.cos(); + let ex11 = ex00; + let ex01 = -x.cos() * y.sin(); + + assert!( + (h[0][0] - ex00).abs() < TOL, + "H[0][0] = {} (expected {})", + h[0][0], + ex00 + ); + assert!( + (h[1][1] - ex11).abs() < TOL, + "H[1][1] = {} (expected {})", + h[1][1], + ex11 + ); + assert!( + (h[0][1] - ex01).abs() < TOL, + "H[0][1] = {} (expected {})", + h[0][1], + ex01 + ); + assert!( + (h[1][0] - ex01).abs() < TOL, + "H[1][0] = {} (expected {})", + h[1][0], + ex01 + ); + } + + #[test] + fn test_exp_times_exp() { + // f = exp(x)·exp(y) = exp(x+y) + // H[i][j] = exp(x+y) for all i,j + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Exp, + MultiAD2RR::Inp(1), + MultiAD2RR::Exp, + MultiAD2RR::Mul, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + let expected = (1.0_f64.exp()) * (2.0_f64.exp()); + + assert!( + (h[0][0] - expected).abs() < TOL, + "H[0][0] = {} (expected {})", + h[0][0], + expected + ); + assert!( + (h[0][1] - expected).abs() < TOL, + "H[0][1] = {} (expected {})", + h[0][1], + expected + ); + assert!( + (h[1][1] - expected).abs() < TOL, + "H[1][1] = {} (expected {})", + h[1][1], + expected + ); + } + + #[test] + fn test_exp_times_sin_single_var() { + // f(x) = exp(x)·sin(x) (single variable) + // f''(x) = 2·exp(x)·cos(x) + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Exp, + MultiAD2RR::Inp(0), + MultiAD2RR::Sin, + MultiAD2RR::Mul, + ]; + let x = 1.0; + let h = MultiAD2RR::compute_hessian(&ops, &[x]).unwrap(); + let expected = 2.0 * x.exp() * x.cos(); + + assert!( + (h[0][0] - expected).abs() < TOL, + "H[0][0] = {} (expected {})", + h[0][0], + expected + ); + } + + #[test] + fn test_sum_squared() { + // f = (x+y)·(x+y) = x²+2xy+y² → H = [[2,2],[2,2]] + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(1), + MultiAD2RR::Add, + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(1), + MultiAD2RR::Add, + MultiAD2RR::Mul, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]).unwrap(); + + assert!( + (h[0][0] - 2.0).abs() < TOL, + "H[0][0] = {} (expected 2.0)", + h[0][0] + ); + assert!( + (h[0][1] - 2.0).abs() < TOL, + "H[0][1] = {} (expected 2.0)", + h[0][1] + ); + assert!( + (h[1][1] - 2.0).abs() < TOL, + "H[1][1] = {} (expected 2.0)", + h[1][1] + ); + } + + #[test] + fn test_tan_and_sub() { + let x = 0.3_f64; + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Tan, + MultiAD2RR::Inp(1), + MultiAD2RR::Sub, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[x, 2.0]).unwrap(); + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + assert!((h[0][0] - expected).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_div_ln_sqrt() { + let x = 4.0_f64; + let y = 2.0_f64; + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Sqrt, + MultiAD2RR::Inp(1), + MultiAD2RR::Ln, + MultiAD2RR::Div, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[x, y]).unwrap(); + let sqrt_x = x.sqrt(); + let expected_xx = -1.0 / (4.0 * x * sqrt_x * y.ln()); + let expected_xy = -1.0 / (2.0 * sqrt_x * y * y.ln().powi(2)); + let expected_yy = + 2.0 * sqrt_x / (y * y * y.ln().powi(3)) + sqrt_x / (y * y * y.ln().powi(2)); + assert!((h[0][0] - expected_xx).abs() < 1e-10); + assert!((h[0][1] - expected_xy).abs() < 1e-10); + assert!((h[1][0] - expected_xy).abs() < 1e-10); + assert!((h[1][1] - expected_yy).abs() < 1e-10); + } + + #[test] + fn test_pow_and_neg() { + let x = 2.0_f64; + let y = 3.0_f64; + let pow_ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Inp(1), MultiAD2RR::Pow]; + let h = MultiAD2RR::compute_hessian(&pow_ops, &[x, y]).unwrap(); + let expected_xx = y * (y - 1.0) * x.powf(y - 2.0); + let expected_xy = x.powf(y - 1.0) * (1.0 + y * x.ln()); + let expected_yy = x.powf(y) * x.ln().powi(2); + assert!((h[0][0] - expected_xx).abs() < 1e-10); + assert!((h[0][1] - expected_xy).abs() < 1e-10); + assert!((h[1][0] - expected_xy).abs() < 1e-10); + assert!((h[1][1] - expected_yy).abs() < 1e-10); + + let neg_ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Sin, MultiAD2RR::Neg]; + let h_neg = MultiAD2RR::compute_hessian(&neg_ops, &[0.5]).unwrap(); + assert!((h_neg[0][0] - 0.5_f64.sin()).abs() < 1e-10); + } + + // ---- Error path tests ---- + + #[test] + fn test_empty_ops_returns_zero_hessian() { + let h = MultiAD2RR::compute_hessian(&[], &[]).unwrap(); + assert!(h.is_empty()); + } + + #[test] + fn test_single_var_hessian() { + // f(x) = x^2 => f''(x) = 2 + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Inp(0), MultiAD2RR::Mul]; + let h = MultiAD2RR::compute_hessian(&ops, &[3.0]).unwrap(); + assert!((h[0][0] - 2.0).abs() < TOL, "H[0][0] = {}", h[0][0]); + } + + #[test] + fn test_input_index_out_of_bounds() { + let ops = vec![MultiAD2RR::Inp(5)]; + let result = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]); + assert!(result.is_err()); + } + + #[test] + fn test_unary_missing_operand() { + let ops = vec![MultiAD2RR::Sin]; + let result = MultiAD2RR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_binary_missing_operand() { + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Add]; + let result = MultiAD2RR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } + + #[test] + fn test_extra_items_on_stack() { + // Two inputs pushed, no reduction => stack has 2 items + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Inp(1)]; + let result = MultiAD2RR::compute_hessian(&ops, &[1.0, 2.0]); + assert!(result.is_err()); + } + + #[test] + fn test_display_format() { + assert_eq!(format!("{}", MultiAD2RR::Inp(3)), "Inp(3)"); + assert_eq!(format!("{}", MultiAD2RR::Sin), "Sin"); + assert_eq!(format!("{}", MultiAD2RR::Cos), "Cos"); + assert_eq!(format!("{}", MultiAD2RR::Tan), "Tan"); + assert_eq!(format!("{}", MultiAD2RR::Neg), "Neg"); + assert_eq!(format!("{}", MultiAD2RR::Exp), "Exp"); + assert_eq!(format!("{}", MultiAD2RR::Ln), "Ln"); + assert_eq!(format!("{}", MultiAD2RR::Sqrt), "Sqrt"); + assert_eq!(format!("{}", MultiAD2RR::Add), "Add"); + assert_eq!(format!("{}", MultiAD2RR::Sub), "Sub"); + assert_eq!(format!("{}", MultiAD2RR::Mul), "Mul"); + assert_eq!(format!("{}", MultiAD2RR::Div), "Div"); + assert_eq!(format!("{}", MultiAD2RR::Pow), "Pow"); + } + + // ---- Additional operation tests for uncovered branches ---- + + #[test] + fn test_hessian_single_ln() { + // f(x) = ln(x), f''(x) = -1/x² + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Ln]; + let x = 2.0_f64; + let h = MultiAD2RR::compute_hessian(&ops, &[x]).unwrap(); + let expected = -1.0 / (x * x); + assert!((h[0][0] - expected).abs() < TOL); + } + + #[test] + fn test_hessian_single_sqrt() { + // f(x) = sqrt(x), f''(x) = -1/(4 * x * sqrt(x)) + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Sqrt]; + let x = 4.0_f64; + let h = MultiAD2RR::compute_hessian(&ops, &[x]).unwrap(); + let expected = -1.0 / (4.0 * x * x.sqrt()); + assert!((h[0][0] - expected).abs() < TOL); + } + + #[test] + fn test_hessian_single_neg() { + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Neg]; + let h = MultiAD2RR::compute_hessian(&ops, &[3.0]).unwrap(); + assert!((h[0][0]).abs() < TOL); + } + + #[test] + fn test_hessian_single_tan() { + let x = 0.4_f64; + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Tan]; + let h = MultiAD2RR::compute_hessian(&ops, &[x]).unwrap(); + let sec_sq = 1.0 / x.cos().powi(2); + let expected = 2.0 * sec_sq * x.tan(); + assert!((h[0][0] - expected).abs() < TOL); + } + + #[test] + fn test_hessian_single_cos() { + let x = 1.0_f64; + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Cos]; + let h = MultiAD2RR::compute_hessian(&ops, &[x]).unwrap(); + assert!((h[0][0] - (-x.cos())).abs() < TOL); + } + + #[test] + fn test_hessian_composition_neg_sin() { + // f(x) = -sin(x), f''(x) = sin(x) + let ops = vec![MultiAD2RR::Inp(0), MultiAD2RR::Sin, MultiAD2RR::Neg]; + let x = 1.0_f64; + let h = MultiAD2RR::compute_hessian(&ops, &[x]).unwrap(); + assert!((h[0][0] - x.sin()).abs() < TOL); + } + + #[test] + fn test_hessian_sub_and_ln() { + // f(x, y) = x - ln(y) → H = [[0, 0], [0, 1/y²]] + let y = 3.0_f64; + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Inp(1), + MultiAD2RR::Ln, + MultiAD2RR::Sub, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[2.0, y]).unwrap(); + assert!((h[0][0]).abs() < TOL); + assert!((h[0][1]).abs() < TOL); + assert!((h[1][0]).abs() < TOL); + assert!((h[1][1] - 1.0 / (y * y)).abs() < TOL); + } + + #[test] + fn test_hessian_mul_and_cos() { + // f(x, y) = cos(x) * y → H = [[-cos(x)*y, -sin(x)], [-sin(x), 0]] + let x = 1.0_f64; + let y = 2.0_f64; + let ops = vec![ + MultiAD2RR::Inp(0), + MultiAD2RR::Cos, + MultiAD2RR::Inp(1), + MultiAD2RR::Mul, + ]; + let h = MultiAD2RR::compute_hessian(&ops, &[x, y]).unwrap(); + assert!((h[0][0] - (-x.cos() * y)).abs() < TOL); + assert!((h[0][1] - (-x.sin())).abs() < TOL); + assert!((h[1][0] - (-x.sin())).abs() < TOL); + assert!((h[1][1]).abs() < TOL); + } + + #[test] + fn test_binary_missing_left_operand() { + let ops = vec![MultiAD2RR::Mul]; + let result = MultiAD2RR::compute_hessian(&ops, &[1.0]); + assert!(result.is_err()); + } +} diff --git a/src/multi/multi_fn.rs b/src/multi/multi_fn.rs index 9a0f60d..f00fef7 100644 --- a/src/multi/multi_fn.rs +++ b/src/multi/multi_fn.rs @@ -45,7 +45,7 @@ pub trait MultiFn { } println!("\nForward pass only:"); - println!("f({:?}) = {}", &self.inputs(), result); + println!("f({:?}) = {}", self.inputs(), result); // Forward + backward (automatic differentiation) let (value, backprop_fn) = self.compute_with_gradients().unwrap(); @@ -54,7 +54,7 @@ pub trait MultiFn { } let grads = backprop_fn(1.0); println!("\nForward + backward (automatic differentiation):"); - println!("f({:?}) = {}", &self.inputs(), value); + println!("f({:?}) = {}", self.inputs(), value); for (i, grad) in grads.iter().enumerate() { println!("∂f/∂x{} = {}", i + 1, grad); } @@ -62,7 +62,7 @@ pub trait MultiFn { // Verify against analytical solution let expected_grad = self.expected_gradients(); println!("\nAnalytical gradients:"); - println!("∂f/∂x₁, ∂f/∂x₂, ... = {:?}", &expected_grad); + println!("∂f/∂x₁, ∂f/∂x₂, ... = {:?}", expected_grad); println!("\nGradient differences:"); for (expected, auto) in expected_grad.iter().zip(grads.iter()) { @@ -76,3 +76,67 @@ pub trait MultiFn { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::approx_eq_eps as approx_eq; + use std::sync::OnceLock; + + /// f(x, y) = x * y, ∇f = [y, x] + struct MulFn; + + impl MultiFn for MulFn { + fn inputs(&self) -> Vec { + vec![2.0, 3.0] + } + fn graph(&self) -> &'static GraphType { + static G: OnceLock)>> = OnceLock::new(); + G.get_or_init(|| { + vec![ + (MultiAD::Inp, vec![0]), + (MultiAD::Inp, vec![1]), + (MultiAD::Mul, vec![0, 1]), + ] + }) + } + fn expected_value(&self) -> f64 { + 6.0 + } + fn expected_gradients(&self) -> Vec { + vec![3.0, 2.0] + } + } + + #[test] + fn test_multi_fn_compute_with_gradients() { + let f = MulFn; + let (value, backprop) = f.compute_with_gradients().unwrap(); + assert!( + approx_eq(value, f.expected_value(), 1e-10), + "value mismatch" + ); + let grads = backprop(1.0); + let expected = f.expected_gradients(); + for (a, b) in grads.iter().zip(expected.iter()) { + assert!((a - b).abs() < 1e-10, "gradient mismatch"); + } + } + + #[test] + fn test_multi_fn_demonstrate_with_assert() { + MulFn.demonstrate(true); + } + + #[test] + fn test_multi_fn_demonstrate_without_assert() { + MulFn.demonstrate(false); + } + + #[test] + fn test_multi_fn_compute_forward_only() { + let f = MulFn; + let result = f.compute().unwrap(); + assert!((result - f.expected_value()).abs() < 1e-10); + } +} diff --git a/src/multi/multi_hessian_common.rs b/src/multi/multi_hessian_common.rs new file mode 100644 index 0000000..88639ab --- /dev/null +++ b/src/multi/multi_hessian_common.rs @@ -0,0 +1,783 @@ +//! Shared types and the dual-number Hessian algorithm for multivariate AD. +//! +//! This module provides: +//! - [`OpKind`]: an operation descriptor used by FR/RF to avoid enum duplication +//! - [`Dual`]: dual numbers for forward-mode AD +//! - [`compute_hessian_dual`]: the shared FR/RF Hessian algorithm +//! +//! For scalar functions f: ℝⁿ → ℝ, Forward-over-Reverse (FR) and Reverse-over-Forward (RF) +//! are equivalent in implementation. Both work by: +//! 1. Forward pass with dual numbers (value + tangent in seed direction e_j) +//! 2. Reverse pass with **dual adjoints** (val = ∂f/∂node, tan = ∂²f/∂node∂x_j) +//! +//! The dual adjoint at each input node's tangent gives a column of the Hessian. + +use crate::error::{AutodiffError, Result}; + +use super::multi_ad::MultiAD; +use super::op_rules::{self, LocalRule}; + +/// An operation descriptor for the stack-based RPN computation graph. +/// Used by FR and RF Hessian methods. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum OpKind { + /// Input variable by index + Inp(usize), + /// Unary sine + Sin, + /// Unary cosine + Cos, + /// Unary tangent + Tan, + /// Unary negation + Neg, + /// Unary exponential + Exp, + /// Unary natural logarithm + Ln, + /// Unary square root + Sqrt, + /// Unary stable softplus: ln(1 + exp(x)) + Log1pExp, + /// Binary addition + Add, + /// Binary subtraction + Sub, + /// Binary multiplication + Mul, + /// Binary division + Div, + /// Binary power + Pow, +} + +// --------------------------------------------------------------------------- +// Dual number +// --------------------------------------------------------------------------- + +/// A dual number `(val, tan)` for forward-mode automatic differentiation. +/// +/// `val` is the function value; `tan` is the directional tangent +/// (∂/∂x_j in the current seed direction). +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct Dual { + pub val: f64, + pub tan: f64, +} + +impl Dual { + /// Create a seed variable: `(x, 1)` — tangent = 1 in the seed direction. + #[inline(always)] + pub fn variable(val: f64) -> Self { + Dual { val, tan: 1.0 } + } + + /// Create a constant: `(x, 0)` — tangent = 0 (independent of seed direction). + #[inline(always)] + pub fn constant(val: f64) -> Self { + Dual { val, tan: 0.0 } + } +} + +impl std::ops::Add for Dual { + type Output = Dual; + #[inline(always)] + fn add(self, rhs: Dual) -> Dual { + Dual { + val: self.val + rhs.val, + tan: self.tan + rhs.tan, + } + } +} + +impl std::ops::Sub for Dual { + type Output = Dual; + #[inline(always)] + fn sub(self, rhs: Dual) -> Dual { + Dual { + val: self.val - rhs.val, + tan: self.tan - rhs.tan, + } + } +} + +impl std::ops::Mul for Dual { + type Output = Dual; + #[inline(always)] + fn mul(self, rhs: Dual) -> Dual { + Dual { + val: self.val * rhs.val, + tan: self.val * rhs.tan + self.tan * rhs.val, + } + } +} + +impl std::ops::Div for Dual { + type Output = Dual; + #[inline(always)] + fn div(self, rhs: Dual) -> Dual { + let denom = rhs.val * rhs.val; + Dual { + val: self.val / rhs.val, + tan: (self.tan * rhs.val - self.val * rhs.tan) / denom, + } + } +} + +// --------------------------------------------------------------------------- +// Dual adjoint accumulation +// --------------------------------------------------------------------------- + +/// Accumulate a dual adjoint: `target += source * local_deriv`. +/// +/// Using the product rule on dual numbers: +/// `(a.val + a.tan·ε) × (b.val + b.tan·ε) = a.val·b.val + (a.val·b.tan + a.tan·b.val)·ε` +#[inline(always)] +pub(crate) fn dual_adj_accum(target: &mut Dual, source: Dual, local_deriv: Dual) { + target.val += source.val * local_deriv.val; + target.tan += source.val * local_deriv.tan + source.tan * local_deriv.val; +} + +// --------------------------------------------------------------------------- +// Node metadata (internal to the forward pass) +// --------------------------------------------------------------------------- + +/// Tracks how each node in the computation graph was created and the +/// associated local derivatives needed for the reverse pass. +#[derive(Clone, Copy)] +enum NodeKind { + Input { + var_idx: usize, + }, + Unary { + parent: usize, + op: MultiAD, + }, + Binary { + left: usize, + right: usize, + op: MultiAD, + }, +} + +// --------------------------------------------------------------------------- +// Core algorithm: dual forward + dual reverse +// --------------------------------------------------------------------------- + +#[inline(always)] +fn opkind_to_multiad(op: OpKind) -> Option { + match op { + OpKind::Inp(_) => None, + OpKind::Sin => Some(MultiAD::Sin), + OpKind::Cos => Some(MultiAD::Cos), + OpKind::Tan => Some(MultiAD::Tan), + OpKind::Neg => Some(MultiAD::Neg), + OpKind::Exp => Some(MultiAD::Exp), + OpKind::Ln => Some(MultiAD::Ln), + OpKind::Sqrt => Some(MultiAD::Sqrt), + OpKind::Log1pExp => Some(MultiAD::Log1pExp), + OpKind::Add => Some(MultiAD::Add), + OpKind::Sub => Some(MultiAD::Sub), + OpKind::Mul => Some(MultiAD::Mul), + OpKind::Div => Some(MultiAD::Div), + OpKind::Pow => Some(MultiAD::Pow), + } +} + +#[inline(always)] +fn apply_dual(op: MultiAD, args: &[Dual]) -> Result { + let arg_values: Vec = args.iter().map(|arg| arg.val).collect(); + let arg_tangents: Vec = args.iter().map(|arg| arg.tan).collect(); + let value = op_rules::forward_value(op, &arg_values)?; + let tangent = op_rules::directional_tangent(op, &arg_values, &arg_tangents, value)?; + Ok(Dual { + val: value, + tan: tangent, + }) +} + +/// Validate input indices and RPN stack shape before seed-direction evaluation. +fn validate_rpn_structure(ops: &[OpKind], n_vars: usize) -> Result<()> { + let mut stack_size: usize = 0; + + for &op in ops { + match op { + OpKind::Inp(k) => { + if k >= n_vars { + return Err(AutodiffError::IndexOutOfBounds { + index: k, + max_index: n_vars.saturating_sub(1), + }); + } + stack_size += 1; + } + OpKind::Sin + | OpKind::Cos + | OpKind::Tan + | OpKind::Neg + | OpKind::Exp + | OpKind::Ln + | OpKind::Sqrt + | OpKind::Log1pExp => { + if stack_size == 0 { + return Err(AutodiffError::InvalidGraph { + reason: "unary operation missing operand", + }); + } + } + OpKind::Add | OpKind::Sub | OpKind::Mul | OpKind::Div | OpKind::Pow => { + if stack_size < 2 { + return Err(AutodiffError::InvalidGraph { + reason: "binary operation missing operand", + }); + } + stack_size -= 1; + } + } + } + + if stack_size != 1 { + return Err(AutodiffError::InvalidGraph { + reason: "RPN expression must leave exactly one output on the stack", + }); + } + + Ok(()) +} + +/// Compute the exact Hessian using dual-number forward and reverse passes. +/// +/// For each seed direction `e_j` (j = 0 … n−1): +/// +/// 1. **Forward pass** — evaluate the graph with dual numbers whose tangent +/// component is seeded in direction `e_j`. Every intermediate value becomes +/// a [`Dual`]. +/// +/// 2. **Reverse pass** — propagate **dual adjoints** backward. Each adjoint +/// is a `Dual { val, tan }` where: +/// - `val = ∂f/∂node` (standard first-order adjoint) +/// - `tan = ∂²f/(∂node · ∂x_j)` (second-order cross-derivative) +/// +/// At input nodes the tangent of the adjoint equals `H[j][k]`. +/// +/// # Accuracy +/// +/// Machine-precision exact (up to floating-point rounding, ~1e-14 relative error). +/// +/// # Complexity +/// +/// O(n · G) where n = number of inputs and G = graph size (number of ops). +/// Each of the n seed directions requires one forward pass and one reverse pass. +pub(crate) fn compute_hessian_dual(ops: &[OpKind], x: &[f64]) -> Result>> { + let n_vars = x.len(); + if ops.is_empty() { + return Ok(vec![vec![0.0; n_vars]; n_vars]); + } + + validate_rpn_structure(ops, n_vars)?; + + let mut hessian = vec![vec![0.0; n_vars]; n_vars]; + + for (j, hessian_row) in hessian.iter_mut().enumerate().take(n_vars) { + // ================================================================ + // FORWARD PASS: dual values in direction e_j + // ================================================================ + let mut dual_values: Vec = Vec::new(); + let mut node_kinds: Vec = Vec::new(); + let mut eval_stack: Vec = Vec::new(); + + for &op in ops { + match op { + OpKind::Inp(k) => { + if k >= n_vars { + return Err(AutodiffError::IndexOutOfBounds { + index: k, + max_index: n_vars.saturating_sub(1), + }); + } + let idx = dual_values.len(); + dual_values.push(if k == j { + Dual::variable(x[k]) + } else { + Dual::constant(x[k]) + }); + node_kinds.push(NodeKind::Input { var_idx: k }); + eval_stack.push(idx); + } + OpKind::Sin + | OpKind::Cos + | OpKind::Tan + | OpKind::Neg + | OpKind::Exp + | OpKind::Ln + | OpKind::Sqrt + | OpKind::Log1pExp => { + let p = eval_stack.pop().ok_or(AutodiffError::InvalidGraph { + reason: "unary operation missing operand", + })?; + let op = opkind_to_multiad(op).expect("non-input op"); + let idx = dual_values.len(); + dual_values.push(apply_dual(op, &[dual_values[p]])?); + node_kinds.push(NodeKind::Unary { parent: p, op }); + eval_stack.push(idx); + } + OpKind::Add | OpKind::Sub | OpKind::Mul | OpKind::Div | OpKind::Pow => { + let r = eval_stack.pop().ok_or(AutodiffError::InvalidGraph { + reason: "binary operation missing right operand", + })?; + let l = eval_stack.pop().ok_or(AutodiffError::InvalidGraph { + reason: "binary operation missing left operand", + })?; + let op = opkind_to_multiad(op).expect("non-input op"); + let idx = dual_values.len(); + dual_values.push(apply_dual(op, &[dual_values[l], dual_values[r]])?); + node_kinds.push(NodeKind::Binary { + left: l, + right: r, + op, + }); + eval_stack.push(idx); + } + } + } + + if eval_stack.len() != 1 { + return Err(AutodiffError::InvalidGraph { + reason: "RPN expression must leave exactly one output on the stack", + }); + } + + // ================================================================ + // REVERSE PASS: dual adjoints + // ================================================================ + let mut adj: Vec = vec![Dual::constant(0.0); dual_values.len()]; + // Seed: ∂f/∂f = 1, ∂²f/(∂f·∂x_j) = 0 + adj[dual_values.len() - 1] = Dual::constant(1.0); + + for i in (0..dual_values.len()).rev() { + let a = adj[i]; + let nk = node_kinds[i]; // Copy (NodeKind is Copy) + match nk { + NodeKind::Input { var_idx: k } => { + // Extract Hessian entry: ∂²f/∂x_j∂x_k + hessian_row[k] += a.tan; + } + NodeKind::Unary { parent, op } => { + let u = dual_values[parent]; + let current = dual_values[i]; + let rule = op_rules::local_rule(op, &[u.val], current.val)?; + let dual_deriv = match rule { + LocalRule::Unary { dy, ddy } => Dual { + val: dy, + tan: ddy * u.tan, + }, + LocalRule::Binary { .. } => unreachable!("unary node must have unary rule"), + }; + dual_adj_accum(&mut adj[parent], a, dual_deriv); + } + NodeKind::Binary { left, right, op } => { + let u = dual_values[left]; + let v = dual_values[right]; + let current = dual_values[i]; + let rule = op_rules::local_rule(op, &[u.val, v.val], current.val)?; + match rule { + LocalRule::Unary { .. } => { + unreachable!("binary node must have binary rule") + } + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + let dz_du = Dual { + val: dy_left, + tan: ddy_left_left * u.tan + ddy_left_right * v.tan, + }; + let dz_dv = Dual { + val: dy_right, + tan: ddy_left_right * u.tan + ddy_right_right * v.tan, + }; + dual_adj_accum(&mut adj[left], a, dz_du); + dual_adj_accum(&mut adj[right], a, dz_dv); + } + } + } + } + } + } + + Ok(hessian) +} + +#[cfg(test)] +mod tests { + use super::*; + + // ---- Dual number arithmetic tests ---- + + #[test] + fn test_dual_variable() { + let d = Dual::variable(3.0); + assert_eq!(d.val, 3.0); + assert_eq!(d.tan, 1.0); + } + + #[test] + fn test_dual_constant() { + let d = Dual::constant(3.0); + assert_eq!(d.val, 3.0); + assert_eq!(d.tan, 0.0); + } + + #[test] + fn test_dual_add() { + let a = Dual { val: 1.0, tan: 2.0 }; + let b = Dual { val: 3.0, tan: 4.0 }; + let c = a + b; + assert_eq!(c.val, 4.0); + assert_eq!(c.tan, 6.0); + } + + #[test] + fn test_dual_sub() { + let a = Dual { val: 5.0, tan: 2.0 }; + let b = Dual { val: 3.0, tan: 1.0 }; + let c = a - b; + assert_eq!(c.val, 2.0); + assert_eq!(c.tan, 1.0); + } + + #[test] + fn test_dual_mul() { + let a = Dual { val: 2.0, tan: 3.0 }; + let b = Dual { val: 4.0, tan: 5.0 }; + let c = a * b; + assert_eq!(c.val, 8.0); + assert_eq!(c.tan, 2.0 * 5.0 + 3.0 * 4.0); // a.val*b.tan + a.tan*b.val + } + + #[test] + fn test_dual_div() { + let a = Dual { val: 6.0, tan: 2.0 }; + let b = Dual { val: 3.0, tan: 1.0 }; + let c = a / b; + assert_eq!(c.val, 2.0); + // tan = (a.tan*b.val - a.val*b.tan) / b.val^2 + let expected_tan = (2.0 * 3.0 - 6.0 * 1.0) / 9.0; + assert!((c.tan - expected_tan).abs() < 1e-12); + } + + #[test] + fn test_dual_adj_accum() { + let mut target = Dual { val: 1.0, tan: 2.0 }; + let source = Dual { val: 3.0, tan: 4.0 }; + let deriv = Dual { val: 5.0, tan: 6.0 }; + dual_adj_accum(&mut target, source, deriv); + // val += source.val * deriv.val = 1 + 3*5 = 16 + assert_eq!(target.val, 16.0); + // tan += source.val * deriv.tan + source.tan * deriv.val = 2 + 3*6 + 4*5 = 2+18+20 = 40 + assert_eq!(target.tan, 40.0); + } + + // ---- OpKind to MultiAD conversion ---- + + #[test] + fn test_opkind_to_multiad() { + assert!(opkind_to_multiad(OpKind::Inp(0)).is_none()); + assert_eq!(opkind_to_multiad(OpKind::Sin), Some(MultiAD::Sin)); + assert_eq!(opkind_to_multiad(OpKind::Cos), Some(MultiAD::Cos)); + assert_eq!(opkind_to_multiad(OpKind::Tan), Some(MultiAD::Tan)); + assert_eq!(opkind_to_multiad(OpKind::Neg), Some(MultiAD::Neg)); + assert_eq!(opkind_to_multiad(OpKind::Exp), Some(MultiAD::Exp)); + assert_eq!(opkind_to_multiad(OpKind::Ln), Some(MultiAD::Ln)); + assert_eq!(opkind_to_multiad(OpKind::Sqrt), Some(MultiAD::Sqrt)); + assert_eq!(opkind_to_multiad(OpKind::Log1pExp), Some(MultiAD::Log1pExp)); + assert_eq!(opkind_to_multiad(OpKind::Add), Some(MultiAD::Add)); + assert_eq!(opkind_to_multiad(OpKind::Sub), Some(MultiAD::Sub)); + assert_eq!(opkind_to_multiad(OpKind::Mul), Some(MultiAD::Mul)); + assert_eq!(opkind_to_multiad(OpKind::Div), Some(MultiAD::Div)); + assert_eq!(opkind_to_multiad(OpKind::Pow), Some(MultiAD::Pow)); + } + + // ---- apply_dual tests ---- + + #[test] + fn test_apply_dual_sin() { + let result = apply_dual(MultiAD::Sin, &[Dual::variable(1.0)]).unwrap(); + assert!((result.val - 1.0_f64.sin()).abs() < 1e-12); + assert!((result.tan - 1.0_f64.cos()).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_add() { + let a = Dual::variable(1.0); + let b = Dual::constant(2.0); + let result = apply_dual(MultiAD::Add, &[a, b]).unwrap(); + assert_eq!(result.val, 3.0); + assert_eq!(result.tan, 1.0); // Only a contributes to tangent + } + + // ---- validate_rpn_structure tests ---- + + #[test] + fn test_validate_rpn_valid() { + let ops = vec![OpKind::Inp(0), OpKind::Inp(1), OpKind::Mul]; + assert!(validate_rpn_structure(&ops, 2).is_ok()); + } + + #[test] + fn test_validate_rpn_input_out_of_bounds() { + let ops = vec![OpKind::Inp(5)]; + let result = validate_rpn_structure(&ops, 2); + assert!(result.is_err()); + } + + #[test] + fn test_validate_rpn_unary_missing_operand() { + let ops = vec![OpKind::Sin]; + let result = validate_rpn_structure(&ops, 1); + assert!(result.is_err()); + } + + #[test] + fn test_validate_rpn_binary_missing_operand() { + let ops = vec![OpKind::Inp(0), OpKind::Add]; + let result = validate_rpn_structure(&ops, 1); + assert!(result.is_err()); + } + + #[test] + fn test_validate_rpn_extra_stack_items() { + let ops = vec![OpKind::Inp(0), OpKind::Inp(1)]; + let result = validate_rpn_structure(&ops, 2); + assert!(result.is_err()); + } + + #[test] + fn test_validate_rpn_empty_stack_at_end() { + let ops: Vec = vec![]; + let result = validate_rpn_structure(&ops, 0); + assert!(result.is_err()); + } + + // ---- compute_hessian_dual edge cases ---- + + #[test] + fn test_compute_hessian_dual_empty_ops() { + let h = compute_hessian_dual(&[], &[1.0, 2.0]).unwrap(); + assert_eq!(h.len(), 2); + assert!(h.iter().all(|row| row.iter().all(|&v| v == 0.0))); + } + + #[test] + fn test_compute_hessian_dual_empty_ops_no_vars() { + let h = compute_hessian_dual(&[], &[]).unwrap(); + assert!(h.is_empty()); + } + + #[test] + fn test_compute_hessian_dual_log1p_exp() { + // f(x) = log1p_exp(x) + let ops = vec![OpKind::Inp(0), OpKind::Log1pExp]; + let x = 1.0_f64; + let h = compute_hessian_dual(&ops, &[x]).unwrap(); + let sigmoid = x.exp() / (1.0 + x.exp()); + let expected = sigmoid * (1.0 - sigmoid); + assert!((h[0][0] - expected).abs() < 1e-10, "H[0][0] = {}", h[0][0]); + } + + #[test] + fn test_compute_hessian_dual_input_out_of_bounds() { + let ops = vec![OpKind::Inp(10), OpKind::Sin]; + let result = compute_hessian_dual(&ops, &[1.0]); + assert!(result.is_err()); + } + + // ---- apply_dual with more operations ---- + + #[test] + fn test_apply_dual_cos() { + let x = 1.0_f64; + let result = apply_dual(MultiAD::Cos, &[Dual::variable(x)]).unwrap(); + assert!((result.val - x.cos()).abs() < 1e-12); + assert!((result.tan - (-x.sin())).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_exp() { + let x = 1.0_f64; + let result = apply_dual(MultiAD::Exp, &[Dual::variable(x)]).unwrap(); + assert!((result.val - x.exp()).abs() < 1e-12); + assert!((result.tan - x.exp()).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_ln() { + let x = 2.0_f64; + let result = apply_dual(MultiAD::Ln, &[Dual::variable(x)]).unwrap(); + assert!((result.val - x.ln()).abs() < 1e-12); + assert!((result.tan - 1.0 / x).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_neg() { + let result = apply_dual(MultiAD::Neg, &[Dual::variable(5.0)]).unwrap(); + assert_eq!(result.val, -5.0); + assert_eq!(result.tan, -1.0); + } + + #[test] + fn test_apply_dual_sqrt() { + let x = 4.0_f64; + let result = apply_dual(MultiAD::Sqrt, &[Dual::variable(x)]).unwrap(); + assert!((result.val - 2.0).abs() < 1e-12); + assert!((result.tan - 0.25).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_tan() { + let x = 0.5_f64; + let result = apply_dual(MultiAD::Tan, &[Dual::variable(x)]).unwrap(); + assert!((result.val - x.tan()).abs() < 1e-12); + let expected_tan = 1.0 / x.cos().powi(2); + assert!((result.tan - expected_tan).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_log1pexp() { + let x = 1.0_f64; + let result = apply_dual(MultiAD::Log1pExp, &[Dual::variable(x)]).unwrap(); + assert!((result.val - (1.0 + x.exp()).ln()).abs() < 1e-12); + // derivative = sigmoid(x) = exp(x)/(1+exp(x)) + let sigmoid = x.exp() / (1.0 + x.exp()); + assert!((result.tan - sigmoid).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_sub() { + let a = Dual::variable(5.0); + let b = Dual::constant(3.0); + let result = apply_dual(MultiAD::Sub, &[a, b]).unwrap(); + assert_eq!(result.val, 2.0); + assert_eq!(result.tan, 1.0); + } + + #[test] + fn test_apply_dual_mul_two_vars() { + let a = Dual::variable(3.0); + let b = Dual::variable(4.0); + let result = apply_dual(MultiAD::Mul, &[a, b]).unwrap(); + assert_eq!(result.val, 12.0); + // d/dx[x*y] with both variable = y + x = 4 + 3 = 7 + assert!((result.tan - 7.0).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_div() { + let a = Dual::variable(6.0); + let b = Dual::constant(3.0); + let result = apply_dual(MultiAD::Div, &[a, b]).unwrap(); + assert_eq!(result.val, 2.0); + assert!((result.tan - 1.0 / 3.0).abs() < 1e-12); + } + + #[test] + fn test_apply_dual_pow() { + let a = Dual::variable(2.0); + let b = Dual::constant(3.0); + let result = apply_dual(MultiAD::Pow, &[a, b]).unwrap(); + assert!((result.val - 8.0).abs() < 1e-12); + // d/dx[x^3] = 3*x^2 = 12 + assert!((result.tan - 12.0).abs() < 1e-12); + } + + // ---- compute_hessian_dual with more operation combos ---- + + #[test] + fn test_compute_hessian_dual_sub() { + let ops = vec![OpKind::Inp(0), OpKind::Inp(1), OpKind::Sub]; + let h = compute_hessian_dual(&ops, &[1.0, 2.0]).unwrap(); + assert!((h[0][0]).abs() < 1e-12); + assert!((h[0][1]).abs() < 1e-12); + assert!((h[1][0]).abs() < 1e-12); + assert!((h[1][1]).abs() < 1e-12); + } + + #[test] + fn test_compute_hessian_dual_div() { + // f(x, y) = x / y → H = [[0, -1/y²], [-1/y², 2*x/y³]] + let x = 6.0_f64; + let y = 3.0_f64; + let ops = vec![OpKind::Inp(0), OpKind::Inp(1), OpKind::Div]; + let h = compute_hessian_dual(&ops, &[x, y]).unwrap(); + assert!((h[0][0]).abs() < TOL); + assert!((h[0][1] - (-1.0 / (y * y))).abs() < TOL); + assert!((h[1][0] - (-1.0 / (y * y))).abs() < TOL); + assert!((h[1][1] - (2.0 * x / (y * y * y))).abs() < TOL); + } + + #[test] + fn test_compute_hessian_dual_ln() { + // f(x) = ln(x), f''(x) = -1/x² + let x = 2.0_f64; + let ops = vec![OpKind::Inp(0), OpKind::Ln]; + let h = compute_hessian_dual(&ops, &[x]).unwrap(); + assert!((h[0][0] - (-1.0 / (x * x))).abs() < TOL); + } + + #[test] + fn test_compute_hessian_dual_pow() { + // f(x, y) = x^y with y constant (=3) + // We can't make y constant in OpKind directly, so test with Inp for both + let x = 2.0_f64; + let y = 3.0_f64; + let ops = vec![OpKind::Inp(0), OpKind::Inp(1), OpKind::Pow]; + let h = compute_hessian_dual(&ops, &[x, y]).unwrap(); + let expected_xx = y * (y - 1.0) * x.powf(y - 2.0); + assert!((h[0][0] - expected_xx).abs() < TOL); + } + + #[test] + fn test_compute_hessian_dual_sqrt() { + let x = 4.0_f64; + let ops = vec![OpKind::Inp(0), OpKind::Sqrt]; + let h = compute_hessian_dual(&ops, &[x]).unwrap(); + let expected = -1.0 / (4.0 * x * x.sqrt()); + assert!((h[0][0] - expected).abs() < TOL); + } + + #[test] + fn test_compute_hessian_dual_cos_sin() { + // f(x) = cos(sin(x)) + let x = 1.0_f64; + let ops = vec![OpKind::Inp(0), OpKind::Sin, OpKind::Cos]; + let h = compute_hessian_dual(&ops, &[x]).unwrap(); + assert!(h[0][0].is_finite()); + } + + // ---- validate_rpn_structure edge cases ---- + + #[test] + fn test_validate_rpn_binary_missing_both() { + let ops = vec![OpKind::Mul]; + let result = validate_rpn_structure(&ops, 1); + assert!(result.is_err()); + } + + #[test] + fn test_validate_rpn_empty_vars() { + let ops = vec![OpKind::Inp(0), OpKind::Sin]; + let result = validate_rpn_structure(&ops, 1); + assert!(result.is_ok()); + } + + #[test] + fn test_compute_hessian_dual_empty_single_var() { + let h = compute_hessian_dual(&[], &[5.0]).unwrap(); + assert_eq!(h, vec![vec![0.0]]); + } + + const TOL: f64 = 1e-12; +} diff --git a/src/multi/op_rules.rs b/src/multi/op_rules.rs new file mode 100644 index 0000000..27ece17 --- /dev/null +++ b/src/multi/op_rules.rs @@ -0,0 +1,1110 @@ +//! Shared local operation rules for multivariate autodiff. +//! +//! This module centralizes scalar values plus first- and second-order local +//! derivatives for [`super::multi_ad::MultiAD`] operations so first-order, +//! forward-mode, and exact Hessian implementations can reuse the same formulas. + +use super::multi_ad::MultiAD; +use crate::{AutodiffError, Result}; + +#[inline(always)] +fn stable_sigmoid(x: f64) -> f64 { + if x >= 0.0 { + 1.0 / (1.0 + (-x).exp()) + } else { + let exp_x = x.exp(); + exp_x / (1.0 + exp_x) + } +} + +#[inline(always)] +fn stable_log1p_exp(x: f64) -> f64 { + if x > 0.0 { + x + (-x).exp().ln_1p() + } else { + x.exp().ln_1p() + } +} + +#[inline(always)] +fn stable_logaddexp(left: f64, right: f64) -> f64 { + if left == f64::NEG_INFINITY && right == f64::NEG_INFINITY { + f64::NEG_INFINITY + } else if left == f64::INFINITY || right == f64::INFINITY { + f64::INFINITY + } else { + let max_value = left.max(right); + max_value + ((left - max_value).exp() + (right - max_value).exp()).ln() + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) enum LocalRule { + Unary { + dy: f64, + ddy: f64, + }, + Binary { + dy_left: f64, + dy_right: f64, + ddy_left_left: f64, + ddy_right_right: f64, + ddy_left_right: f64, + }, +} + +#[inline(always)] +pub(crate) fn op_name(op: MultiAD) -> &'static str { + match op { + MultiAD::Inp => "Inp", + MultiAD::Add => "Add", + MultiAD::Sub => "Sub", + MultiAD::Mul => "Mul", + MultiAD::Div => "Div", + MultiAD::Pow => "Pow", + MultiAD::Sin => "Sin", + MultiAD::Cos => "Cos", + MultiAD::Tan => "Tan", + MultiAD::Tanh => "Tanh", + MultiAD::Relu => "Relu", + MultiAD::Log1pExp => "Log1pExp", + MultiAD::LogAddExp => "LogAddExp", + MultiAD::Neg => "Neg", + MultiAD::Exp => "Exp", + MultiAD::Ln => "Ln", + MultiAD::Sqrt => "Sqrt", + MultiAD::Abs => "Abs", + } +} + +#[inline(always)] +pub(crate) fn expected_arity(op: MultiAD) -> usize { + match op { + MultiAD::Inp + | MultiAD::Sin + | MultiAD::Cos + | MultiAD::Tan + | MultiAD::Tanh + | MultiAD::Relu + | MultiAD::Log1pExp + | MultiAD::Neg + | MultiAD::Exp + | MultiAD::Ln + | MultiAD::Sqrt + | MultiAD::Abs => 1, + MultiAD::Add + | MultiAD::Sub + | MultiAD::Mul + | MultiAD::Div + | MultiAD::Pow + | MultiAD::LogAddExp => 2, + } +} + +#[inline(always)] +pub(crate) fn check_arity(op: MultiAD, actual: usize) -> Result<()> { + AutodiffError::check_arity(op_name(op), expected_arity(op), actual) +} + +#[inline(always)] +pub(crate) fn check_domain(op: MultiAD, args: &[f64]) -> Result<()> { + check_arity(op, args.len())?; + match op { + MultiAD::Div if args[1] == 0.0 => { + Err(AutodiffError::domain("Div", "denominator must be non-zero")) + } + MultiAD::Ln if args[0] <= 0.0 => Err(AutodiffError::domain("Ln", "input must be positive")), + MultiAD::Sqrt if args[0] < 0.0 => { + Err(AutodiffError::domain("Sqrt", "input must be non-negative")) + } + MultiAD::Pow if args[0] <= 0.0 => Err(AutodiffError::domain( + "Pow", + "base must be positive in checked mode", + )), + _ => Ok(()), + } +} + +#[inline(always)] +pub(crate) fn forward_value(op: MultiAD, args: &[f64]) -> Result { + check_arity(op, args.len())?; + Ok(match op { + MultiAD::Inp => args[0], + MultiAD::Sin => args[0].sin(), + MultiAD::Cos => args[0].cos(), + MultiAD::Tan => args[0].tan(), + MultiAD::Tanh => args[0].tanh(), + MultiAD::Relu => args[0].max(0.0), + MultiAD::Log1pExp => stable_log1p_exp(args[0]), + MultiAD::Neg => -args[0], + MultiAD::Exp => args[0].exp(), + MultiAD::Ln => args[0].ln(), + MultiAD::Sqrt => args[0].sqrt(), + MultiAD::Abs => args[0].abs(), + MultiAD::Add => args[0] + args[1], + MultiAD::Sub => args[0] - args[1], + MultiAD::Mul => args[0] * args[1], + MultiAD::Div => args[0] / args[1], + MultiAD::Pow => args[0].powf(args[1]), + MultiAD::LogAddExp => stable_logaddexp(args[0], args[1]), + }) +} + +#[inline(always)] +pub(crate) fn forward_value_checked(op: MultiAD, args: &[f64]) -> Result { + check_domain(op, args)?; + forward_value(op, args) +} + +#[inline(always)] +pub(crate) fn local_rule(op: MultiAD, args: &[f64], value: f64) -> Result { + check_arity(op, args.len())?; + Ok(match op { + MultiAD::Sin => LocalRule::Unary { + dy: args[0].cos(), + ddy: -args[0].sin(), + }, + MultiAD::Cos => LocalRule::Unary { + dy: -args[0].sin(), + ddy: -args[0].cos(), + }, + MultiAD::Tan => { + let sec_sq = 1.0 / args[0].cos().powi(2); + LocalRule::Unary { + dy: sec_sq, + ddy: 2.0 * sec_sq * args[0].tan(), + } + } + MultiAD::Tanh => { + let tanh = value; + let dy = 1.0 - tanh * tanh; + LocalRule::Unary { + dy, + ddy: -2.0 * tanh * dy, + } + } + MultiAD::Relu => { + let dy = if args[0] > 0.0 { 1.0 } else { 0.0 }; + LocalRule::Unary { dy, ddy: 0.0 } + } + MultiAD::Log1pExp => { + let dy = stable_sigmoid(args[0]); + LocalRule::Unary { + dy, + ddy: dy * (1.0 - dy), + } + } + MultiAD::Neg => LocalRule::Unary { dy: -1.0, ddy: 0.0 }, + MultiAD::Exp => LocalRule::Unary { + dy: value, + ddy: value, + }, + MultiAD::Ln => LocalRule::Unary { + dy: 1.0 / args[0], + ddy: -1.0 / args[0].powi(2), + }, + MultiAD::Sqrt => LocalRule::Unary { + dy: 1.0 / (2.0 * value), + ddy: -1.0 / (4.0 * args[0] * value), + }, + MultiAD::Abs => { + let sign = if args[0] > 0.0 { + 1.0 + } else if args[0] < 0.0 { + -1.0 + } else { + 0.0 + }; + LocalRule::Unary { dy: sign, ddy: 0.0 } + } + MultiAD::Add => LocalRule::Binary { + dy_left: 1.0, + dy_right: 1.0, + ddy_left_left: 0.0, + ddy_right_right: 0.0, + ddy_left_right: 0.0, + }, + MultiAD::Sub => LocalRule::Binary { + dy_left: 1.0, + dy_right: -1.0, + ddy_left_left: 0.0, + ddy_right_right: 0.0, + ddy_left_right: 0.0, + }, + MultiAD::Mul => LocalRule::Binary { + dy_left: args[1], + dy_right: args[0], + ddy_left_left: 0.0, + ddy_right_right: 0.0, + ddy_left_right: 1.0, + }, + MultiAD::Div => LocalRule::Binary { + dy_left: 1.0 / args[1], + dy_right: -args[0] / args[1].powi(2), + ddy_left_left: 0.0, + ddy_right_right: 2.0 * args[0] / args[1].powi(3), + ddy_left_right: -1.0 / args[1].powi(2), + }, + MultiAD::Pow => LocalRule::Binary { + dy_left: args[1] * args[0].powf(args[1] - 1.0), + dy_right: if args[0] == 0.0 { + 0.0 + } else { + value * args[0].ln() + }, + ddy_left_left: args[1] * (args[1] - 1.0) * args[0].powf(args[1] - 2.0), + ddy_right_right: if args[0] == 0.0 { + 0.0 + } else { + value * args[0].ln().powi(2) + }, + ddy_left_right: if args[0] == 0.0 { + args[0].powf(args[1] - 1.0) + } else { + args[0].powf(args[1] - 1.0) * (1.0 + args[1] * args[0].ln()) + }, + }, + MultiAD::LogAddExp => { + let (left_weight, right_weight) = if value == f64::NEG_INFINITY { + (0.5, 0.5) + } else if value == f64::INFINITY { + match (args[0] == f64::INFINITY, args[1] == f64::INFINITY) { + (true, true) => (0.5, 0.5), + (true, false) => (1.0, 0.0), + (false, true) => (0.0, 1.0), + (false, false) => (0.5, 0.5), + } + } else { + ((args[0] - value).exp(), (args[1] - value).exp()) + }; + LocalRule::Binary { + dy_left: left_weight, + dy_right: right_weight, + ddy_left_left: left_weight * (1.0 - left_weight), + ddy_right_right: right_weight * (1.0 - right_weight), + ddy_left_right: -left_weight * right_weight, + } + } + MultiAD::Inp => { + return Err(AutodiffError::InvalidGraph { + reason: "input markers do not have local derivative rules", + }); + } + }) +} + +#[inline(always)] +pub(crate) fn first_derivatives(op: MultiAD, args: &[f64], value: f64) -> Result> { + Ok(match local_rule(op, args, value)? { + LocalRule::Unary { dy, .. } => vec![dy], + LocalRule::Binary { + dy_left, dy_right, .. + } => vec![dy_left, dy_right], + }) +} + +#[inline(always)] +pub(crate) fn directional_tangent( + op: MultiAD, + args: &[f64], + tangents: &[f64], + value: f64, +) -> Result { + let first = first_derivatives(op, args, value)?; + Ok(first.iter().zip(tangents.iter()).map(|(d, t)| d * t).sum()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::approx_eq_eps as approx_eq; + + const EPS: f64 = 1e-10; + + // ─── forward_value ──────────────────────────────────────────────── + + #[test] + fn test_forward_inp() { + assert_eq!(forward_value(MultiAD::Inp, &[2.5]).unwrap(), 2.5); + } + + #[test] + fn test_forward_sin() { + let val = forward_value(MultiAD::Sin, &[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.sin(), EPS)); + } + + #[test] + fn test_forward_cos() { + let val = forward_value(MultiAD::Cos, &[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.cos(), EPS)); + } + + #[test] + fn test_forward_tan() { + let val = forward_value(MultiAD::Tan, &[0.5]).unwrap(); + assert!(approx_eq(val, 0.5_f64.tan(), EPS)); + } + + #[test] + fn test_forward_tanh() { + let val = forward_value(MultiAD::Tanh, &[1.0]).unwrap(); + assert!(approx_eq(val, 1.0_f64.tanh(), EPS)); + } + + #[test] + fn test_forward_relu_positive() { + assert_eq!(forward_value(MultiAD::Relu, &[3.0]).unwrap(), 3.0); + } + + #[test] + fn test_forward_relu_negative() { + assert_eq!(forward_value(MultiAD::Relu, &[-2.0]).unwrap(), 0.0); + } + + #[test] + fn test_forward_relu_zero() { + assert_eq!(forward_value(MultiAD::Relu, &[0.0]).unwrap(), 0.0); + } + + #[test] + fn test_forward_log1p_exp() { + // x > 0 branch + let val = forward_value(MultiAD::Log1pExp, &[5.0]).unwrap(); + assert!(approx_eq(val, stable_log1p_exp(5.0), EPS)); + // x <= 0 branch + let val2 = forward_value(MultiAD::Log1pExp, &[-5.0]).unwrap(); + assert!(approx_eq(val2, stable_log1p_exp(-5.0), EPS)); + } + + #[test] + fn test_forward_neg() { + assert_eq!(forward_value(MultiAD::Neg, &[7.0]).unwrap(), -7.0); + } + + #[test] + fn test_forward_exp() { + let val = forward_value(MultiAD::Exp, &[1.0]).unwrap(); + assert!(approx_eq(val, std::f64::consts::E, EPS)); + } + + #[test] + fn test_forward_ln() { + let val = forward_value(MultiAD::Ln, &[std::f64::consts::E]).unwrap(); + assert!(approx_eq(val, 1.0, EPS)); + } + + #[test] + fn test_forward_sqrt() { + let val = forward_value(MultiAD::Sqrt, &[4.0]).unwrap(); + assert!(approx_eq(val, 2.0, EPS)); + } + + #[test] + fn test_forward_abs() { + assert_eq!(forward_value(MultiAD::Abs, &[-3.0]).unwrap(), 3.0); + assert_eq!(forward_value(MultiAD::Abs, &[3.0]).unwrap(), 3.0); + } + + #[test] + fn test_forward_add() { + assert_eq!(forward_value(MultiAD::Add, &[2.0, 3.0]).unwrap(), 5.0); + } + + #[test] + fn test_forward_sub() { + assert_eq!(forward_value(MultiAD::Sub, &[2.0, 3.0]).unwrap(), -1.0); + } + + #[test] + fn test_forward_mul() { + assert_eq!(forward_value(MultiAD::Mul, &[2.0, 3.0]).unwrap(), 6.0); + } + + #[test] + fn test_forward_div() { + let val = forward_value(MultiAD::Div, &[6.0, 3.0]).unwrap(); + assert!(approx_eq(val, 2.0, EPS)); + } + + #[test] + fn test_forward_pow() { + let val = forward_value(MultiAD::Pow, &[2.0, 3.0]).unwrap(); + assert!(approx_eq(val, 8.0, EPS)); + } + + #[test] + fn test_forward_logaddexp_normal() { + let val = forward_value(MultiAD::LogAddExp, &[1.0, 2.0]).unwrap(); + assert!(approx_eq(val, (1.0_f64.exp() + 2.0_f64.exp()).ln(), 1e-9)); + } + + // ─── arity checks ──────────────────────────────────────────────── + + #[test] + fn test_forward_wrong_arity() { + assert!(forward_value(MultiAD::Sin, &[1.0, 2.0]).is_err()); + assert!(forward_value(MultiAD::Add, &[1.0]).is_err()); + } + + // ─── forward_value_checked / check_domain ──────────────────────── + + #[test] + fn test_checked_div_by_zero() { + assert!(forward_value_checked(MultiAD::Div, &[1.0, 0.0]).is_err()); + } + + #[test] + fn test_checked_ln_non_positive() { + assert!(forward_value_checked(MultiAD::Ln, &[0.0]).is_err()); + assert!(forward_value_checked(MultiAD::Ln, &[-1.0]).is_err()); + } + + #[test] + fn test_checked_sqrt_negative() { + assert!(forward_value_checked(MultiAD::Sqrt, &[-1.0]).is_err()); + } + + #[test] + fn test_checked_pow_non_positive_base() { + assert!(forward_value_checked(MultiAD::Pow, &[0.0, 2.0]).is_err()); + assert!(forward_value_checked(MultiAD::Pow, &[-1.0, 2.0]).is_err()); + } + + #[test] + fn test_checked_ok() { + assert!(forward_value_checked(MultiAD::Sin, &[1.0]).is_ok()); + assert!(forward_value_checked(MultiAD::Div, &[1.0, 2.0]).is_ok()); + assert!(forward_value_checked(MultiAD::Ln, &[1.0]).is_ok()); + assert!(forward_value_checked(MultiAD::Sqrt, &[0.0]).is_ok()); + assert!(forward_value_checked(MultiAD::Pow, &[2.0, 3.0]).is_ok()); + } + + // ─── local_rule — unary ops ───────────────────────────────────── + + #[test] + fn test_local_rule_sin() { + let rule = local_rule(MultiAD::Sin, &[1.0], 1.0_f64.sin()).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, 1.0_f64.cos(), EPS)); + assert!(approx_eq(ddy, -1.0_f64.sin(), EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_cos() { + let rule = local_rule(MultiAD::Cos, &[1.0], 1.0_f64.cos()).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, -1.0_f64.sin(), EPS)); + assert!(approx_eq(ddy, -1.0_f64.cos(), EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_tan() { + let x = 0.5; + let rule = local_rule(MultiAD::Tan, &[x], x.tan()).unwrap(); + let sec_sq = 1.0 / x.cos().powi(2); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, sec_sq, EPS)); + assert!(approx_eq(ddy, 2.0 * sec_sq * x.tan(), EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_tanh() { + let x: f64 = 0.7; + let tanh_x = x.tanh(); + let rule = local_rule(MultiAD::Tanh, &[x], tanh_x).unwrap(); + let dy = 1.0 - tanh_x * tanh_x; + match rule { + LocalRule::Unary { dy: dy_val, ddy } => { + assert!(approx_eq(dy_val, dy, EPS)); + assert!(approx_eq(ddy, -2.0 * tanh_x * dy, EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_relu() { + // x > 0 + let rule = local_rule(MultiAD::Relu, &[2.0], 2.0).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert_eq!(dy, 1.0); + assert_eq!(ddy, 0.0); + } + _ => panic!("expected Unary"), + } + // x < 0 + let rule2 = local_rule(MultiAD::Relu, &[-1.0], 0.0).unwrap(); + match rule2 { + LocalRule::Unary { dy, .. } => assert_eq!(dy, 0.0), + _ => panic!("expected Unary"), + } + // x == 0 + let rule3 = local_rule(MultiAD::Relu, &[0.0], 0.0).unwrap(); + match rule3 { + LocalRule::Unary { dy, .. } => assert_eq!(dy, 0.0), + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_log1p_exp() { + let x = 2.0; + let value = stable_log1p_exp(x); + let rule = local_rule(MultiAD::Log1pExp, &[x], value).unwrap(); + let sigma = stable_sigmoid(x); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, sigma, EPS)); + assert!(approx_eq(ddy, sigma * (1.0 - sigma), EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_neg() { + let rule = local_rule(MultiAD::Neg, &[5.0], -5.0).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert_eq!(dy, -1.0); + assert_eq!(ddy, 0.0); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_exp() { + let x: f64 = 1.5; + let value = x.exp(); + let rule = local_rule(MultiAD::Exp, &[x], value).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, value, EPS)); + assert!(approx_eq(ddy, value, EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_ln() { + let x: f64 = 3.0; + let value = x.ln(); + let rule = local_rule(MultiAD::Ln, &[x], value).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, 1.0 / x, EPS)); + assert!(approx_eq(ddy, -1.0 / (x * x), EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_sqrt() { + let x: f64 = 9.0; + let value = x.sqrt(); // 3.0 + let rule = local_rule(MultiAD::Sqrt, &[x], value).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert!(approx_eq(dy, 1.0 / (2.0 * value), EPS)); + assert!(approx_eq(ddy, -1.0 / (4.0 * x * value), EPS)); + } + _ => panic!("expected Unary"), + } + } + + #[test] + fn test_local_rule_abs() { + // positive + let rule = local_rule(MultiAD::Abs, &[3.0], 3.0).unwrap(); + match rule { + LocalRule::Unary { dy, ddy } => { + assert_eq!(dy, 1.0); + assert_eq!(ddy, 0.0); + } + _ => panic!(), + } + // negative + let rule2 = local_rule(MultiAD::Abs, &[-3.0], 3.0).unwrap(); + match rule2 { + LocalRule::Unary { dy, .. } => assert_eq!(dy, -1.0), + _ => panic!(), + } + // zero + let rule3 = local_rule(MultiAD::Abs, &[0.0], 0.0).unwrap(); + match rule3 { + LocalRule::Unary { dy, .. } => assert_eq!(dy, 0.0), + _ => panic!(), + } + } + + #[test] + fn test_local_rule_inp_errors() { + assert!(local_rule(MultiAD::Inp, &[1.0], 1.0).is_err()); + } + + // ─── local_rule — binary ops ───────────────────────────────────── + + #[test] + fn test_local_rule_add() { + let rule = local_rule(MultiAD::Add, &[2.0, 3.0], 5.0).unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert_eq!(dy_left, 1.0); + assert_eq!(dy_right, 1.0); + assert_eq!(ddy_left_left, 0.0); + assert_eq!(ddy_right_right, 0.0); + assert_eq!(ddy_left_right, 0.0); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_sub() { + let rule = local_rule(MultiAD::Sub, &[2.0, 3.0], -1.0).unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert_eq!(dy_left, 1.0); + assert_eq!(dy_right, -1.0); + assert_eq!(ddy_left_left, 0.0); + assert_eq!(ddy_right_right, 0.0); + assert_eq!(ddy_left_right, 0.0); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_mul() { + let rule = local_rule(MultiAD::Mul, &[2.0, 3.0], 6.0).unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert_eq!(dy_left, 3.0); // args[1] + assert_eq!(dy_right, 2.0); // args[0] + assert_eq!(ddy_left_left, 0.0); + assert_eq!(ddy_right_right, 0.0); + assert_eq!(ddy_left_right, 1.0); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_div() { + let a = 6.0; + let b = 3.0; + let rule = local_rule(MultiAD::Div, &[a, b], a / b).unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert!(approx_eq(dy_left, 1.0 / b, EPS)); + assert!(approx_eq(dy_right, -a / b.powi(2), EPS)); + assert_eq!(ddy_left_left, 0.0); + assert!(approx_eq(ddy_right_right, 2.0 * a / b.powi(3), EPS)); + assert!(approx_eq(ddy_left_right, -1.0 / b.powi(2), EPS)); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_pow_normal() { + let base: f64 = 2.0; + let exp: f64 = 3.0; + let value = base.powf(exp); // 8.0 + let rule = local_rule(MultiAD::Pow, &[base, exp], value).unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert!(approx_eq(dy_left, exp * base.powf(exp - 1.0), EPS)); + assert!(approx_eq(dy_right, value * base.ln(), EPS)); + assert!(approx_eq( + ddy_left_left, + exp * (exp - 1.0) * base.powf(exp - 2.0), + EPS + )); + assert!(approx_eq(ddy_right_right, value * base.ln().powi(2), EPS)); + assert!(approx_eq( + ddy_left_right, + base.powf(exp - 1.0) * (1.0 + exp * base.ln()), + EPS + )); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_pow_zero_base() { + // base == 0.0 triggers special branches + let rule = local_rule(MultiAD::Pow, &[0.0, 2.0], 0.0_f64.powf(2.0)).unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + // dy_left: exp * base^(exp-1) = 2 * 0^1 = 0 + assert_eq!(dy_left, 0.0); + assert_eq!(dy_right, 0.0); // base == 0.0 branch + // ddy_left_left: exp*(exp-1)*base^(exp-2) = 2*1*0^0 = 2*1*1 = 2 + assert!(approx_eq(ddy_left_left, 2.0, EPS)); + assert_eq!(ddy_right_right, 0.0); + // ddy_left_right: args[0].powf(args[1]-1) = 0^(1) = 0 + assert!(approx_eq(ddy_left_right, 0.0_f64.powf(2.0 - 1.0), EPS)); + } + _ => panic!("expected Binary"), + } + } + + // ─── local_rule — LogAddExp edge cases ─────────────────────────── + + #[test] + fn test_local_rule_logaddexp_normal() { + let a = 1.0; + let b = 2.0; + let value = stable_logaddexp(a, b); + let rule = local_rule(MultiAD::LogAddExp, &[a, b], value).unwrap(); + let lw = (a - value).exp(); + let rw = (b - value).exp(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert!(approx_eq(dy_left, lw, 1e-9)); + assert!(approx_eq(dy_right, rw, 1e-9)); + assert!(approx_eq(ddy_left_left, lw * (1.0 - lw), 1e-9)); + assert!(approx_eq(ddy_right_right, rw * (1.0 - rw), 1e-9)); + assert!(approx_eq(ddy_left_right, -lw * rw, 1e-9)); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_logaddexp_neg_inf() { + let rule = local_rule( + MultiAD::LogAddExp, + &[f64::NEG_INFINITY, f64::NEG_INFINITY], + f64::NEG_INFINITY, + ) + .unwrap(); + match rule { + LocalRule::Binary { + dy_left, + dy_right, + ddy_left_left, + ddy_right_right, + ddy_left_right, + } => { + assert_eq!(dy_left, 0.5); + assert_eq!(dy_right, 0.5); + assert_eq!(ddy_left_left, 0.25); + assert_eq!(ddy_right_right, 0.25); + assert_eq!(ddy_left_right, -0.25); + } + _ => panic!("expected Binary"), + } + } + + #[test] + fn test_local_rule_logaddexp_both_inf() { + let rule = local_rule( + MultiAD::LogAddExp, + &[f64::INFINITY, f64::INFINITY], + f64::INFINITY, + ) + .unwrap(); + match rule { + LocalRule::Binary { + dy_left, dy_right, .. + } => { + assert_eq!(dy_left, 0.5); + assert_eq!(dy_right, 0.5); + } + _ => panic!(), + } + } + + #[test] + fn test_local_rule_logaddexp_left_inf() { + let rule = local_rule(MultiAD::LogAddExp, &[f64::INFINITY, 1.0], f64::INFINITY).unwrap(); + match rule { + LocalRule::Binary { + dy_left, dy_right, .. + } => { + assert_eq!(dy_left, 1.0); + assert_eq!(dy_right, 0.0); + } + _ => panic!(), + } + } + + #[test] + fn test_local_rule_logaddexp_right_inf() { + let rule = local_rule(MultiAD::LogAddExp, &[1.0, f64::INFINITY], f64::INFINITY).unwrap(); + match rule { + LocalRule::Binary { + dy_left, dy_right, .. + } => { + assert_eq!(dy_left, 0.0); + assert_eq!(dy_right, 1.0); + } + _ => panic!(), + } + } + + // ─── directional_tangent ───────────────────────────────────────── + + #[test] + fn test_directional_tangent_sin() { + let x: f64 = 1.0; + let value = x.sin(); + let t = directional_tangent(MultiAD::Sin, &[x], &[0.5], value).unwrap(); + assert!(approx_eq(t, x.cos() * 0.5, EPS)); + } + + #[test] + fn test_directional_tangent_add() { + let t = directional_tangent(MultiAD::Add, &[2.0, 3.0], &[1.0, -1.0], 5.0).unwrap(); + // dy_left * t_left + dy_right * t_right = 1*1 + 1*(-1) = 0 + assert!(approx_eq(t, 0.0, EPS)); + } + + #[test] + fn test_directional_tangent_mul() { + let t = directional_tangent(MultiAD::Mul, &[2.0, 3.0], &[1.0, 2.0], 6.0).unwrap(); + // dy_left*1 + dy_right*2 = 3*1 + 2*2 = 7 + assert!(approx_eq(t, 7.0, EPS)); + } + + #[test] + fn test_directional_tangent_exp() { + let x: f64 = 1.0; + let value = x.exp(); + let t = directional_tangent(MultiAD::Exp, &[x], &[2.0], value).unwrap(); + assert!(approx_eq(t, value * 2.0, EPS)); + } + + #[test] + fn test_directional_tangent_div() { + let a = 6.0; + let b = 3.0; + let t = directional_tangent(MultiAD::Div, &[a, b], &[1.0, 1.0], a / b).unwrap(); + // (1/3)*1 + (-6/9)*1 = 1/3 - 2/3 = -1/3 + assert!(approx_eq(t, 1.0 / 3.0 - 2.0 / 3.0, EPS)); + } + + // ─── first_derivatives ────────────────────────────────────────── + + #[test] + fn test_first_derivatives_unary() { + let derivs = first_derivatives(MultiAD::Neg, &[5.0], -5.0).unwrap(); + assert_eq!(derivs, vec![-1.0]); + } + + #[test] + fn test_first_derivatives_binary() { + let derivs = first_derivatives(MultiAD::Mul, &[2.0, 3.0], 6.0).unwrap(); + assert_eq!(derivs, vec![3.0, 2.0]); + } + + // ─── op_name / expected_arity / check_arity ───────────────────── + + #[test] + fn test_op_name_all() { + assert_eq!(op_name(MultiAD::Inp), "Inp"); + assert_eq!(op_name(MultiAD::Add), "Add"); + assert_eq!(op_name(MultiAD::Sub), "Sub"); + assert_eq!(op_name(MultiAD::Mul), "Mul"); + assert_eq!(op_name(MultiAD::Div), "Div"); + assert_eq!(op_name(MultiAD::Pow), "Pow"); + assert_eq!(op_name(MultiAD::Sin), "Sin"); + assert_eq!(op_name(MultiAD::Cos), "Cos"); + assert_eq!(op_name(MultiAD::Tan), "Tan"); + assert_eq!(op_name(MultiAD::Tanh), "Tanh"); + assert_eq!(op_name(MultiAD::Relu), "Relu"); + assert_eq!(op_name(MultiAD::Log1pExp), "Log1pExp"); + assert_eq!(op_name(MultiAD::LogAddExp), "LogAddExp"); + assert_eq!(op_name(MultiAD::Neg), "Neg"); + assert_eq!(op_name(MultiAD::Exp), "Exp"); + assert_eq!(op_name(MultiAD::Ln), "Ln"); + assert_eq!(op_name(MultiAD::Sqrt), "Sqrt"); + assert_eq!(op_name(MultiAD::Abs), "Abs"); + } + + #[test] + fn test_expected_arity() { + assert_eq!(expected_arity(MultiAD::Sin), 1); + assert_eq!(expected_arity(MultiAD::Cos), 1); + assert_eq!(expected_arity(MultiAD::Tan), 1); + assert_eq!(expected_arity(MultiAD::Tanh), 1); + assert_eq!(expected_arity(MultiAD::Relu), 1); + assert_eq!(expected_arity(MultiAD::Log1pExp), 1); + assert_eq!(expected_arity(MultiAD::Neg), 1); + assert_eq!(expected_arity(MultiAD::Exp), 1); + assert_eq!(expected_arity(MultiAD::Ln), 1); + assert_eq!(expected_arity(MultiAD::Sqrt), 1); + assert_eq!(expected_arity(MultiAD::Abs), 1); + assert_eq!(expected_arity(MultiAD::Inp), 1); + assert_eq!(expected_arity(MultiAD::Add), 2); + assert_eq!(expected_arity(MultiAD::Sub), 2); + assert_eq!(expected_arity(MultiAD::Mul), 2); + assert_eq!(expected_arity(MultiAD::Div), 2); + assert_eq!(expected_arity(MultiAD::Pow), 2); + assert_eq!(expected_arity(MultiAD::LogAddExp), 2); + } + + #[test] + fn test_check_arity_ok() { + assert!(check_arity(MultiAD::Sin, 1).is_ok()); + assert!(check_arity(MultiAD::Add, 2).is_ok()); + } + + #[test] + fn test_check_arity_fail() { + assert!(check_arity(MultiAD::Sin, 2).is_err()); + assert!(check_arity(MultiAD::Add, 1).is_err()); + } + + // ─── stable helpers ───────────────────────────────────────────── + + #[test] + fn test_stable_sigmoid() { + // x >= 0 + assert!(approx_eq(stable_sigmoid(0.0), 0.5, EPS)); + assert!(approx_eq( + stable_sigmoid(10.0), + 1.0 / (1.0 + (-10.0_f64).exp()), + EPS + )); + // x < 0 + let exp_x = (-5.0_f64).exp(); + assert!(approx_eq(stable_sigmoid(-5.0), exp_x / (1.0 + exp_x), EPS)); + } + + #[test] + fn test_stable_log1p_exp_positive() { + // x > 0 branch + let x: f64 = 5.0; + let expected = x + (-x).exp().ln_1p(); + assert!(approx_eq(stable_log1p_exp(x), expected, EPS)); + } + + #[test] + fn test_stable_log1p_exp_nonpositive() { + // x <= 0 branch + let x: f64 = -3.0; + let expected = x.exp().ln_1p(); + assert!(approx_eq(stable_log1p_exp(x), expected, EPS)); + } + + #[test] + fn test_stable_logaddexp_both_neg_inf() { + assert_eq!( + stable_logaddexp(f64::NEG_INFINITY, f64::NEG_INFINITY), + f64::NEG_INFINITY + ); + } + + #[test] + fn test_stable_logaddexp_one_inf() { + assert_eq!(stable_logaddexp(f64::INFINITY, 1.0), f64::INFINITY); + assert_eq!(stable_logaddexp(1.0, f64::INFINITY), f64::INFINITY); + } + + #[test] + fn test_stable_logaddexp_normal() { + let val = stable_logaddexp(1.0, 2.0); + let max_v = 2.0_f64; + let expected = max_v + ((1.0 - max_v).exp() + (2.0 - max_v).exp()).ln(); + assert!(approx_eq(val, expected, EPS)); + } + + // ─── forward_value edge cases ──────────────────────────────────── + + #[test] + fn test_forward_div_by_zero_unchecked() { + // forward_value does NOT check domain; f64 division yields inf + let val = forward_value(MultiAD::Div, &[1.0, 0.0]).unwrap(); + assert!(val.is_infinite()); + } + + #[test] + fn test_forward_ln_zero_unchecked() { + // forward_value does NOT check domain + let val = forward_value(MultiAD::Ln, &[0.0]).unwrap(); + assert!(val.is_infinite() || val.is_nan()); + } + + #[test] + fn test_forward_sqrt_zero() { + let val = forward_value(MultiAD::Sqrt, &[0.0]).unwrap(); + assert_eq!(val, 0.0); + } + + #[test] + fn test_forward_log1p_exp_large_positive() { + // For large x, log1p(exp(x)) ≈ x + let val = forward_value(MultiAD::Log1pExp, &[100.0]).unwrap(); + assert!(approx_eq(val, 100.0, 1e-6)); + } + + #[test] + fn test_forward_log1p_exp_large_negative() { + // For very negative x, log1p(exp(x)) ≈ exp(x) ≈ 0 + let val = forward_value(MultiAD::Log1pExp, &[-100.0]).unwrap(); + assert!(val < 1e-40); + } +} diff --git a/src/multi/parser.rs b/src/multi/parser.rs new file mode 100644 index 0000000..6e02cd2 --- /dev/null +++ b/src/multi/parser.rs @@ -0,0 +1,635 @@ +//! Small expression parser for `Graph` construction. +//! +//! Supported grammar is intentionally small and predictable: +//! numbers, named inputs, `+ - * / ^`, parentheses, unary minus, and functions +//! `sin`, `cos`, `tan`, `tanh`, `relu`, `exp`, `ln`, `sqrt`, `abs`, `sigmoid`, +//! `softplus`, `log1p_exp`, and `gelu`. + +use crate::{AutodiffError, Graph, NodeId, Result}; + +#[derive(Debug, Clone, PartialEq)] +enum Token { + Number(f64), + Ident(String), + Plus, + Minus, + Star, + Slash, + Caret, + LParen, + RParen, +} + +pub(crate) fn parse_expression(expression: &str, input_names: &[&str]) -> Result { + let tokens = tokenize(expression)?; + let mut parser = Parser { + tokens, + pos: 0, + graph: Graph::new(input_names.len()), + input_names: input_names.iter().map(|name| (*name).to_string()).collect(), + }; + for (index, name) in input_names.iter().enumerate() { + parser.graph.set_input_name(index, *name)?; + } + let output = parser.parse_expr()?; + if parser.pos != parser.tokens.len() { + return Err(AutodiffError::InvalidGraph { + reason: "unexpected trailing expression tokens", + }); + } + parser.graph.set_output(output)?; + Ok(parser.graph) +} + +fn tokenize(input: &str) -> Result> { + let mut chars = input.chars().peekable(); + let mut tokens = Vec::new(); + while let Some(&ch) = chars.peek() { + match ch { + ' ' | '\t' | '\n' | '\r' => { + chars.next(); + } + '0'..='9' | '.' => { + let mut text = String::new(); + let mut previous_was_exponent = false; + while let Some(&next) = chars.peek() { + if next.is_ascii_digit() || next == '.' || next == 'e' || next == 'E' { + previous_was_exponent = next == 'e' || next == 'E'; + text.push(next); + chars.next(); + } else if (next == '-' || next == '+') && previous_was_exponent { + previous_was_exponent = false; + text.push(next); + chars.next(); + } else { + break; + } + } + let value = text + .parse::() + .map_err(|_| AutodiffError::InvalidGraph { + reason: "invalid numeric literal", + })?; + tokens.push(Token::Number(value)); + } + 'a'..='z' | 'A'..='Z' | '_' => { + let mut ident = String::new(); + while let Some(&next) = chars.peek() { + if next.is_ascii_alphanumeric() || next == '_' { + ident.push(next); + chars.next(); + } else { + break; + } + } + tokens.push(Token::Ident(ident)); + } + '+' => { + chars.next(); + tokens.push(Token::Plus); + } + '-' => { + chars.next(); + tokens.push(Token::Minus); + } + '*' => { + chars.next(); + tokens.push(Token::Star); + } + '/' => { + chars.next(); + tokens.push(Token::Slash); + } + '^' => { + chars.next(); + tokens.push(Token::Caret); + } + '(' => { + chars.next(); + tokens.push(Token::LParen); + } + ')' => { + chars.next(); + tokens.push(Token::RParen); + } + _ => { + return Err(AutodiffError::InvalidGraph { + reason: "unsupported expression character", + }); + } + } + } + Ok(tokens) +} + +struct Parser { + tokens: Vec, + pos: usize, + graph: Graph, + input_names: Vec, +} + +impl Parser { + fn peek(&self) -> Option<&Token> { + self.tokens.get(self.pos) + } + + fn bump(&mut self) -> Option { + let token = self.tokens.get(self.pos).cloned(); + self.pos += usize::from(token.is_some()); + token + } + + fn parse_expr(&mut self) -> Result { + self.parse_add_sub() + } + + fn parse_add_sub(&mut self) -> Result { + let mut left = self.parse_mul_div()?; + loop { + match self.peek() { + Some(Token::Plus) => { + self.bump(); + let right = self.parse_mul_div()?; + left = self.graph.add(left, right); + } + Some(Token::Minus) => { + self.bump(); + let right = self.parse_mul_div()?; + left = self.graph.sub(left, right); + } + _ => return Ok(left), + } + } + } + + fn parse_mul_div(&mut self) -> Result { + let mut left = self.parse_pow()?; + loop { + match self.peek() { + Some(Token::Star) => { + self.bump(); + let right = self.parse_pow()?; + left = self.graph.mul(left, right); + } + Some(Token::Slash) => { + self.bump(); + let right = self.parse_pow()?; + left = self.graph.div(left, right); + } + _ => return Ok(left), + } + } + } + + fn parse_pow(&mut self) -> Result { + let left = self.parse_unary()?; + if matches!(self.peek(), Some(Token::Caret)) { + self.bump(); + let right = self.parse_pow()?; + Ok(self.graph.pow(left, right)) + } else { + Ok(left) + } + } + + fn parse_unary(&mut self) -> Result { + if matches!(self.peek(), Some(Token::Minus)) { + self.bump(); + let node = self.parse_unary()?; + Ok(self.graph.neg(node)) + } else { + self.parse_primary() + } + } + + fn parse_primary(&mut self) -> Result { + match self.bump() { + Some(Token::Number(value)) => Ok(self.graph.constant(value)), + Some(Token::Ident(name)) => { + if matches!(self.peek(), Some(Token::LParen)) { + self.bump(); + let arg = self.parse_expr()?; + if !matches!(self.bump(), Some(Token::RParen)) { + return Err(AutodiffError::InvalidGraph { + reason: "function call missing closing parenthesis", + }); + } + self.apply_function(&name, arg) + } else { + let index = self + .input_names + .iter() + .position(|candidate| candidate == &name) + .ok_or(AutodiffError::InvalidGraph { + reason: "unknown expression identifier", + })?; + Ok(self.graph.input(index)) + } + } + Some(Token::LParen) => { + let node = self.parse_expr()?; + if !matches!(self.bump(), Some(Token::RParen)) { + return Err(AutodiffError::InvalidGraph { + reason: "expression missing closing parenthesis", + }); + } + Ok(node) + } + _ => Err(AutodiffError::InvalidGraph { + reason: "expected expression primary", + }), + } + } + + fn apply_function(&mut self, name: &str, arg: NodeId) -> Result { + Ok(match name { + "sin" => self.graph.sin(arg), + "cos" => self.graph.cos(arg), + "tan" => self.graph.tan(arg), + "tanh" => self.graph.tanh(arg), + "relu" => self.graph.relu(arg), + "exp" => self.graph.exp(arg), + "ln" | "log" => self.graph.ln(arg), + "sqrt" => self.graph.sqrt(arg), + "abs" => self.graph.abs(arg), + "sigmoid" => self.graph.sigmoid(arg), + "softplus" | "log1p_exp" => self.graph.log1p_exp(arg), + "gelu" => self.graph.gelu(arg), + _ => { + return Err(AutodiffError::InvalidGraph { + reason: "unknown expression function", + }); + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::approx_eq_eps as approx_eq; + + // ---- Valid parse + evaluate round trips ---- + + #[test] + fn test_parse_single_input_addition() { + // f(x) = x + x + let graph = parse_expression("x + x", &["x"]).unwrap(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 6.0, 1e-10)); + } + + #[test] + fn test_parse_two_variables_mul() { + // f(x, y) = x * y + let graph = parse_expression("x * y", &["x", "y"]).unwrap(); + let value = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 6.0, 1e-10)); + } + + #[test] + fn test_parse_subtraction_and_division() { + // f(a, b) = (a - b) / b + let graph = parse_expression("(a - b) / b", &["a", "b"]).unwrap(); + let value = graph.compute(&[5.0, 2.0]).unwrap(); + assert!(approx_eq(value, 1.5, 1e-10)); + } + + #[test] + fn test_parse_power() { + // f(x, y) = x ^ y + let graph = parse_expression("x ^ y", &["x", "y"]).unwrap(); + let value = graph.compute(&[2.0, 3.0]).unwrap(); + assert!(approx_eq(value, 8.0, 1e-10)); + } + + #[test] + fn test_parse_unary_minus() { + // f(x) = -x + 5 + let graph = parse_expression("-x + 5", &["x"]).unwrap(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_parse_nested_unary_minus() { + // f(x) = --x (double negation) + let graph = parse_expression("--x", &["x"]).unwrap(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_parse_numeric_constant() { + // Use a non-PI approximate constant + let graph = parse_expression("2.71", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 2.71, 1e-10)); + } + + #[test] + fn test_parse_scientific_notation() { + // f() = 1e-3 + let graph = parse_expression("1e-3", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 1e-3, 1e-10)); + } + + #[test] + fn test_parse_scientific_notation_positive_exponent() { + // f() = 2E+5 + let graph = parse_expression("2E+5", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 2e5, 1e-8)); + } + + #[test] + fn test_parse_function_sin() { + // f(x) = sin(x) + let graph = parse_expression("sin(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + } + + #[test] + fn test_parse_function_cos() { + // f(x) = cos(x) + let graph = parse_expression("cos(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 1.0, 1e-10)); + } + + #[test] + fn test_parse_function_tan() { + // f(x) = tan(x) + let graph = parse_expression("tan(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + } + + #[test] + fn test_parse_function_tanh() { + // f(x) = tanh(x) + let graph = parse_expression("tanh(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + } + + #[test] + fn test_parse_function_relu() { + // f(x) = relu(x) + let graph = parse_expression("relu(x)", &["x"]).unwrap(); + let value_pos = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(value_pos, 2.0, 1e-10)); + } + + #[test] + fn test_parse_function_relu_negative() { + let graph = parse_expression("relu(x)", &["x"]).unwrap(); + let value_neg = graph.compute(&[-2.0]).unwrap(); + assert!(approx_eq(value_neg, 0.0, 1e-10)); + } + + #[test] + fn test_parse_function_exp() { + // f(x) = exp(x) + let graph = parse_expression("exp(x)", &["x"]).unwrap(); + let value = graph.compute(&[1.0]).unwrap(); + assert!(approx_eq(value, std::f64::consts::E, 1e-10)); + } + + #[test] + fn test_parse_function_ln() { + // f(x) = ln(x) + let graph = parse_expression("ln(x)", &["x"]).unwrap(); + let value = graph.compute(&[std::f64::consts::E]).unwrap(); + assert!(approx_eq(value, 1.0, 1e-10)); + } + + #[test] + fn test_parse_function_log_alias() { + // f(x) = log(x) (alias for ln) + let graph = parse_expression("log(x)", &["x"]).unwrap(); + let value = graph.compute(&[std::f64::consts::E]).unwrap(); + assert!(approx_eq(value, 1.0, 1e-10)); + } + + #[test] + fn test_parse_function_sqrt() { + // f(x) = sqrt(x) + let graph = parse_expression("sqrt(x)", &["x"]).unwrap(); + let value = graph.compute(&[4.0]).unwrap(); + assert!(approx_eq(value, 2.0, 1e-10)); + } + + #[test] + fn test_parse_function_abs() { + // f(x) = abs(x) + let graph = parse_expression("abs(x)", &["x"]).unwrap(); + let value = graph.compute(&[-3.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_parse_function_sigmoid() { + // f(x) = sigmoid(x) + let graph = parse_expression("sigmoid(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.5, 1e-10)); + } + + #[test] + fn test_parse_function_softplus() { + // f(x) = softplus(x) (alias for log1p_exp) + let graph = parse_expression("softplus(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 2.0_f64.ln(), 1e-6)); + } + + #[test] + fn test_parse_function_log1p_exp() { + // f(x) = log1p_exp(x) + let graph = parse_expression("log1p_exp(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 2.0_f64.ln(), 1e-6)); + } + + #[test] + fn test_parse_function_gelu() { + // f(x) = gelu(x) + let graph = parse_expression("gelu(x)", &["x"]).unwrap(); + let value = graph.compute(&[0.0]).unwrap(); + assert!(approx_eq(value, 0.0, 1e-10)); + } + + #[test] + fn test_parse_complex_expression() { + // f(x, y) = sin(x) * (x + y) + let graph = parse_expression("sin(x) * (x + y)", &["x", "y"]).unwrap(); + let value = graph.compute(&[0.6, 1.4]).unwrap(); + let expected = 0.6_f64.sin() * (0.6 + 1.4); + assert!(approx_eq(value, expected, 1e-10)); + } + + #[test] + fn test_parse_chained_functions() { + // f(x) = sin(cos(exp(x))) + let graph = parse_expression("sin(cos(exp(x)))", &["x"]).unwrap(); + let value = graph.compute(&[0.5]).unwrap(); + let expected = 0.5_f64.exp().cos().sin(); + assert!(approx_eq(value, expected, 1e-10)); + } + + #[test] + fn test_parse_whitespace_handling() { + // f(x, y) = x+y with various whitespace + let graph = parse_expression(" x + y ", &["x", "y"]).unwrap(); + let value = graph.compute(&[1.0, 2.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_parse_operator_precedence() { + // 2 + 3 * 4 = 14 (not 20) + let graph = parse_expression("2 + 3 * 4", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 14.0, 1e-10)); + } + + #[test] + fn test_parse_deeply_nested_parens() { + // f(x) = ((x + 1)) + let graph = parse_expression("((x + 1))", &["x"]).unwrap(); + let value = graph.compute(&[2.0]).unwrap(); + assert!(approx_eq(value, 3.0, 1e-10)); + } + + #[test] + fn test_parse_mul_div_precedence() { + // 6 / 2 * 3 = 9 (left-to-right) + let graph = parse_expression("6 / 2 * 3", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 9.0, 1e-10)); + } + + #[test] + fn test_parse_add_sub_precedence() { + // 10 - 3 - 2 = 5 (left-to-right) + let graph = parse_expression("10 - 3 - 2", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 5.0, 1e-10)); + } + + #[test] + fn test_parse_float_literal() { + // f() = 0.5 + let graph = parse_expression("0.5", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 0.5, 1e-10)); + } + + #[test] + fn test_parse_dot_literal() { + // f() = .5 (starting with dot) + let graph = parse_expression(".5", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 0.5, 1e-10)); + } + + #[test] + fn test_parse_negative_times_variable() { + // f(x) = -2 * x + let graph = parse_expression("-2 * x", &["x"]).unwrap(); + let value = graph.compute(&[3.0]).unwrap(); + assert!(approx_eq(value, -6.0, 1e-10)); + } + + #[test] + fn test_parse_power_right_associative() { + // 2 ^ 3 ^ 2 = 2^(3^2) = 512 + let graph = parse_expression("2 ^ 3 ^ 2", &[]).unwrap(); + let value = graph.compute(&[]).unwrap(); + assert!(approx_eq(value, 512.0, 1e-10)); + } + + // ---- Error cases ---- + + #[test] + fn test_parse_error_unsupported_character() { + let result = parse_expression("x # y", &["x", "y"]); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!(matches!(err, AutodiffError::InvalidGraph { .. })); + } + + #[test] + fn test_parse_error_unknown_identifier() { + let result = parse_expression("x + z", &["x", "y"]); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + matches!(err, AutodiffError::InvalidGraph { reason } if reason.contains("unknown expression identifier")) + ); + } + + #[test] + fn test_parse_error_unknown_function() { + let result = parse_expression("foo(x)", &["x"]); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + matches!(err, AutodiffError::InvalidGraph { reason } if reason.contains("unknown expression function")) + ); + } + + #[test] + fn test_parse_error_function_missing_closing_paren() { + let result = parse_expression("sin(x", &["x"]); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + matches!(err, AutodiffError::InvalidGraph { reason } if reason.contains("function call missing closing parenthesis")) + ); + } + + #[test] + fn test_parse_error_expression_missing_closing_paren() { + let result = parse_expression("(x + y", &["x", "y"]); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + matches!(err, AutodiffError::InvalidGraph { reason } if reason.contains("expression missing closing parenthesis")) + ); + } + + #[test] + fn test_parse_error_trailing_tokens() { + let result = parse_expression("x + y z", &["x", "y"]); + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + matches!(err, AutodiffError::InvalidGraph { reason } if reason.contains("unexpected trailing expression tokens")) + ); + } + + #[test] + fn test_parse_error_expected_primary() { + // Just a plus sign with nothing before it + let result = parse_expression("+", &["x"]); + assert!(result.is_err()); + } + + #[test] + fn test_parse_error_invalid_numeric_literal() { + // "1e" is not a valid float + let result = parse_expression("1e", &[]); + assert!(result.is_err()); + } + + #[test] + fn test_parse_empty_expression() { + let result = parse_expression("", &[]); + assert!(result.is_err()); + } +} diff --git a/src/multi/tests.rs b/src/multi/tests.rs index bfd439b..c53813d 100644 --- a/src/multi/tests.rs +++ b/src/multi/tests.rs @@ -389,3 +389,145 @@ fn test_sqrt_and_mul_chain() { assert!(approx_eq(grads[0], 5.0 / (2.0 * 16.0_f64.sqrt()), 1e-10)); assert!(approx_eq(grads[1], 4.0, 1e-10)); } + +#[test] +fn test_hessian_quadratic() { + // Test f(x, y) = x² + y² + // Hessian is [[2, 0], [0, 2]] + let exprs = &multi_ops![(inp, 0), (inp, 1), (mul, 0, 0), (mul, 1, 1), (add, 2, 3)]; + let inputs = &[2.0, 3.0]; + + let hessian = MultiAD::compute_hessian(exprs, inputs).unwrap(); + + assert_eq!(hessian.len(), 2); + assert_eq!(hessian[0].len(), 2); + assert_eq!(hessian[1].len(), 2); + + // ∂²f/∂x² = 2, ∂²f/∂x∂y = 0, ∂²f/∂y∂x = 0, ∂²f/∂y² = 2 + // Using 1e-6 tolerance for finite difference approximation + assert!(approx_eq(hessian[0][0], 2.0, 1e-6)); + assert!(approx_eq(hessian[0][1], 0.0, 1e-6)); + assert!(approx_eq(hessian[1][0], 0.0, 1e-6)); + assert!(approx_eq(hessian[1][1], 2.0, 1e-6)); +} + +#[test] +fn test_hessian_exp_sin() { + // Test f(x) = exp(sin(x)) + // First derivative: f'(x) = exp(sin(x)) * cos(x) + // Second derivative: f''(x) = exp(sin(x)) * cos²(x) - exp(sin(x)) * sin(x) + let exprs = &multi_ops![(inp, 0), (sin, 0), (exp, 1)]; + let x = 0.5; + let inputs = &[x]; + let hessian = MultiAD::compute_hessian(exprs, inputs).unwrap(); + + let expected_hessian = x.sin().exp() * x.cos().powi(2) - x.sin().exp() * x.sin(); + // Using 1e-4 tolerance for finite difference approximation + assert!(approx_eq(hessian[0][0], expected_hessian, 1e-4)); +} + +#[test] +fn test_hessian_complex_function() { + // Test f(x, y) = sin(x) * (x + y) + let exprs = &multi_ops![(inp, 0), (inp, 1), (add, 0, 1), (sin, 0), (mul, 2, 3)]; + let x = 0.6; + let y = 1.4; + let inputs = &[x, y]; + let hessian = MultiAD::compute_hessian(exprs, inputs).unwrap(); + + // Analytical second derivatives: + // f(x, y) = sin(x) * (x + y) + // ∂f/∂x = cos(x) * (x + y) + sin(x) + // ∂f/∂y = sin(x) + // ∂²f/∂x² = -sin(x) * (x + y) + cos(x) + cos(x) = -sin(x) * (x + y) + 2*cos(x) + // ∂²f/∂x∂y = cos(x) + // ∂²f/∂y∂x = cos(x) + // ∂²f/∂y² = 0 + + let expected_dxx = -x.sin() * (x + y) + 2.0 * x.cos(); + let expected_dxy = x.cos(); + let expected_dyx = x.cos(); + let expected_dyy = 0.0; + + // Using 1e-4 tolerance for finite difference approximation + assert!(approx_eq(hessian[0][0], expected_dxx, 1e-4)); + assert!(approx_eq(hessian[0][1], expected_dxy, 1e-4)); + assert!(approx_eq(hessian[1][0], expected_dyx, 1e-4)); + assert!(approx_eq(hessian[1][1], expected_dyy, 1e-4)); +} + +#[test] +fn test_hessian_single_variable() { + // Test f(x) = x³ + // Hessian is [[6x]] + let exprs = &multi_ops![(inp, 0), (mul, 0, 0), (mul, 1, 0)]; + let x = 3.0; + let inputs = &[x]; + + let hessian = MultiAD::compute_hessian(exprs, inputs).unwrap(); + + // f''(x) = 6x = 18 + // Using 1e-4 tolerance for finite difference approximation (higher for second derivatives) + assert_eq!(hessian.len(), 1); + assert_eq!(hessian[0].len(), 1); + assert!(approx_eq(hessian[0][0], 18.0, 1e-4)); +} + +#[test] +fn test_hessian_quadratic_form() { + // Test f(x, y, z) = x² + xy + y² + yz + z² + // Hessian is [[2, 1, 0], [1, 2, 1], [0, 1, 2]] + let exprs = &multi_ops![ + (inp, 0), + (inp, 1), + (inp, 2), + (mul, 0, 0), + (mul, 0, 1), + (mul, 1, 1), + (mul, 1, 2), + (mul, 2, 2), + (add, 3, 4), + (add, 5, 6), + (add, 7, 8), + (add, 9, 10) + ]; + let inputs = &[1.0, 2.0, 3.0]; + let hessian = MultiAD::compute_hessian(exprs, inputs).unwrap(); + + assert_eq!(hessian.len(), 3); + assert_eq!(hessian[0].len(), 3); + assert_eq!(hessian[1].len(), 3); + assert_eq!(hessian[2].len(), 3); + + // Row 0: [2, 1, 0] + assert!(approx_eq(hessian[0][0], 2.0, 1e-6)); + assert!(approx_eq(hessian[0][1], 1.0, 1e-6)); + assert!(approx_eq(hessian[0][2], 0.0, 1e-6)); + + // Row 1: [1, 2, 1] + assert!(approx_eq(hessian[1][0], 1.0, 1e-6)); + assert!(approx_eq(hessian[1][1], 2.0, 1e-6)); + assert!(approx_eq(hessian[1][2], 1.0, 1e-6)); + + // Row 2: [0, 1, 2] + assert!(approx_eq(hessian[2][0], 0.0, 1e-6)); + assert!(approx_eq(hessian[2][1], 1.0, 1e-6)); + assert!(approx_eq(hessian[2][2], 2.0, 1e-6)); +} + +#[test] +fn test_hessian_row() { + // Test compute_hessian_row function + let exprs = &multi_ops![(inp, 0), (inp, 1), (mul, 0, 0), (mul, 1, 1), (add, 2, 3)]; + let inputs = &[2.0, 3.0]; + + // Test first row (derivatives of ∂f/∂x) + let row0 = MultiAD::compute_hessian_row(exprs, inputs, 0).unwrap(); + assert!(approx_eq(row0[0], 2.0, 1e-6)); + assert!(approx_eq(row0[1], 0.0, 1e-6)); + + // Test second row (derivatives of ∂f/∂y) + let row1 = MultiAD::compute_hessian_row(exprs, inputs, 1).unwrap(); + assert!(approx_eq(row1[0], 0.0, 1e-6)); + assert!(approx_eq(row1[1], 2.0, 1e-6)); +} diff --git a/src/optim.rs b/src/optim.rs new file mode 100644 index 0000000..30d732c --- /dev/null +++ b/src/optim.rs @@ -0,0 +1,179 @@ +//! Small optimizer utilities independent from graph construction. + +use crate::{AutodiffError, Result}; + +fn check_lengths(params: &[f64], grads: &[f64]) -> Result<()> { + if params.len() == grads.len() { + Ok(()) + } else { + Err(AutodiffError::InvalidArguments { + reason: "parameter and gradient lengths must match", + }) + } +} + +/// Plain gradient descent optimizer. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct GradientDescent { + /// Learning rate multiplier. + pub learning_rate: f64, +} + +impl GradientDescent { + /// Apply one in-place gradient descent step. + pub fn step(&self, params: &mut [f64], grads: &[f64]) -> Result<()> { + check_lengths(params, grads)?; + for (param, grad) in params.iter_mut().zip(grads.iter()) { + *param -= self.learning_rate * grad; + } + Ok(()) + } +} + +/// Adam optimizer with explicit state. +#[derive(Debug, Clone, PartialEq)] +pub struct Adam { + pub learning_rate: f64, + pub beta1: f64, + pub beta2: f64, + pub epsilon: f64, + m: Vec, + v: Vec, + step: usize, +} + +impl Adam { + /// Create Adam state for `parameter_count` scalar parameters. + /// + /// # Panics + /// + /// Panics if beta1, beta2 are outside [0, 1) or epsilon <= 0. + #[must_use] + pub fn new(parameter_count: usize, learning_rate: f64) -> Self { + Self::with_params(parameter_count, learning_rate, 0.9, 0.999, 1e-8) + } + + /// Create Adam state with custom hyperparameters. + /// + /// # Panics + /// + /// Panics if beta1, beta2 are outside [0, 1) or epsilon <= 0. + #[must_use] + pub fn with_params( + parameter_count: usize, + learning_rate: f64, + beta1: f64, + beta2: f64, + epsilon: f64, + ) -> Self { + assert!( + (0.0..1.0).contains(&beta1), + "beta1 must be in [0, 1), got {}", + beta1 + ); + assert!( + (0.0..1.0).contains(&beta2), + "beta2 must be in [0, 1), got {}", + beta2 + ); + assert!(epsilon > 0.0, "epsilon must be > 0, got {}", epsilon); + Self { + learning_rate, + beta1, + beta2, + epsilon, + m: vec![0.0; parameter_count], + v: vec![0.0; parameter_count], + step: 0, + } + } + + /// Apply one in-place Adam update. + pub fn step(&mut self, params: &mut [f64], grads: &[f64]) -> Result<()> { + check_lengths(params, grads)?; + check_lengths(&self.m, grads)?; + self.step += 1; + let bias1 = 1.0 - self.beta1.powf(self.step as f64); + let bias2 = 1.0 - self.beta2.powf(self.step as f64); + + for index in 0..params.len() { + self.m[index] = self.beta1 * self.m[index] + (1.0 - self.beta1) * grads[index]; + self.v[index] = + self.beta2 * self.v[index] + (1.0 - self.beta2) * grads[index] * grads[index]; + let m_hat = self.m[index] / bias1; + let v_hat = self.v[index] / bias2; + params[index] -= self.learning_rate * m_hat / (v_hat.sqrt() + self.epsilon); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_gradient_descent_step() { + let optimizer = GradientDescent { learning_rate: 0.1 }; + let mut params = [1.0, 2.0]; + optimizer.step(&mut params, &[0.5, -0.5]).unwrap(); + assert!((params[0] - 0.95).abs() < 1e-12); + assert!((params[1] - 2.05).abs() < 1e-12); + } + + #[test] + fn test_adam_step_changes_params() { + let mut optimizer = Adam::new(1, 0.1); + let mut params = [0.0]; + optimizer.step(&mut params, &[-1.0]).unwrap(); + assert!(params[0] > 0.0); + } + + // --- Additional tests for uncovered lines --- + + #[test] + fn test_gradient_descent_length_mismatch() { + let optimizer = GradientDescent { learning_rate: 0.1 }; + let mut params = [1.0]; + let result = optimizer.step(&mut params, &[1.0, 2.0]); + assert!(result.is_err()); + } + + #[test] + fn test_adam_new_defaults() { + let adam = Adam::new(3, 0.01); + assert_eq!(adam.learning_rate, 0.01); + assert!((adam.beta1 - 0.9).abs() < 1e-10); + assert!((adam.beta2 - 0.999).abs() < 1e-10); + assert!((adam.epsilon - 1e-8).abs() < 1e-20); + assert_eq!(adam.m.len(), 3); + assert_eq!(adam.v.len(), 3); + assert_eq!(adam.step, 0); + } + + #[test] + fn test_adam_multiple_steps() { + let mut adam = Adam::new(1, 0.1); + let mut params = [1.0]; + // Run several steps to exercise powf with increasing step count + for _ in 0..5 { + adam.step(&mut params, &[1.0]).unwrap(); + } + // After 5 steps of gradient=1.0 starting from 1.0, params should decrease + assert!(params[0] < 1.0); + } + + #[test] + fn test_adam_length_mismatch_params() { + let mut adam = Adam::new(2, 0.1); + let mut params = [1.0]; + assert!(adam.step(&mut params, &[1.0]).is_err()); + } + + #[test] + fn test_adam_length_mismatch_grads() { + let mut adam = Adam::new(1, 0.1); + let mut params = [1.0]; + assert!(adam.step(&mut params, &[1.0, 2.0]).is_err()); + } +} diff --git a/src/tests_comprehensive.rs b/src/tests_comprehensive.rs new file mode 100644 index 0000000..163c797 --- /dev/null +++ b/src/tests_comprehensive.rs @@ -0,0 +1,897 @@ +//! Comprehensive tests for all features and edge cases. +//! This module ensures very high code coverage. + +use crate::error::AutodiffError; +use crate::mono::types::Dual; +use crate::test_utils::approx_eq_eps as approx_eq; +use crate::{mono_ops, mono_ops_fr, mono_ops_rf, mono_ops_rr, multi_ops}; +use crate::{ + GraphBuilder, MonoAD, MonoAD2FR, MonoAD2RF, MonoAD2RR, MultiAD, MultiAD2FR, MultiAD2RF, + MultiAD2RR, +}; + +// ============================================================================ +// Error Handling Tests +// ============================================================================ + +#[test] +fn test_error_display() { + let err = AutodiffError::arity("Sin", 1, 2); + let msg = format!("{}", err); + assert!(msg.contains("Sin")); + assert!(msg.contains("expected 1")); + assert!(msg.contains("got 2")); + + let err = AutodiffError::EmptyGraph; + let msg = format!("{}", err); + assert!(msg.contains("empty")); + + let err = AutodiffError::IndexOutOfBounds { + index: 5, + max_index: 3, + }; + let msg = format!("{}", err); + assert!(msg.contains("5")); + assert!(msg.contains("3")); + + let err = AutodiffError::InvalidGraph { + reason: "missing operand", + }; + let msg = format!("{}", err); + assert!(msg.contains("invalid")); + assert!(msg.contains("missing operand")); + + let err = AutodiffError::domain("Ln", "input must be positive"); + let msg = format!("{}", err); + assert!(msg.contains("Domain error")); + assert!(msg.contains("Ln")); +} + +#[test] +fn test_error_send_sync() { + fn assert_send_sync() {} + assert_send_sync::(); +} + +#[test] +fn test_check_arity_ok() { + assert!(AutodiffError::check_arity("Test", 2, 2).is_ok()); +} + +#[test] +fn test_check_arity_err() { + let result = AutodiffError::check_arity("Test", 2, 3); + assert!(result.is_err()); + match result { + Err(AutodiffError::ArityError { + operation, + expected, + actual, + }) => { + assert_eq!(operation, "Test"); + assert_eq!(expected, 2); + assert_eq!(actual, 3); + } + _ => panic!("Expected ArityError"), + } +} + +// ============================================================================ +// MonoAD Unary Operation Tests +// ============================================================================ + +#[test] +fn test_mono_neg_forward() { + let ops = &[MonoAD::Neg]; + let result = MonoAD::compute(ops, 5.0); + assert!(approx_eq(result, -5.0, 1e-10)); + + let result = MonoAD::compute(ops, -3.0); + assert!(approx_eq(result, 3.0, 1e-10)); +} + +#[test] +fn test_mono_neg_gradient() { + let ops = &[MonoAD::Neg]; + let (_value, backprop) = MonoAD::compute_grad(ops, 5.0); + assert!(approx_eq(backprop(1.0), -1.0, 1e-10)); + assert!(approx_eq(backprop(2.0), -2.0, 1e-10)); +} + +#[test] +fn test_mono_neg_chain() { + let ops = &[MonoAD::Sin, MonoAD::Neg]; + let x: f64 = 1.0; + let (value, backprop) = MonoAD::compute_grad(ops, x); + + assert!(approx_eq(value, -x.sin(), 1e-10)); + assert!(approx_eq(backprop(1.0), -x.cos(), 1e-10)); +} + +#[test] +fn test_mono_neg_hessian() { + let ops = &[MonoAD::Neg]; + let hessian = MonoAD::compute_hessian(ops, 5.0); + assert!(approx_eq(hessian, 0.0, 1e-10)); +} + +#[test] +fn test_mono_tan_forward_and_grad() { + let x: f64 = 0.5; + let ops = &[MonoAD::Tan]; + let (value, backprop) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(value, x.tan(), 1e-10)); + assert!(approx_eq(backprop(1.0), 1.0 / x.cos().powi(2), 1e-10)); +} + +#[test] +fn test_mono_ln_forward_and_grad() { + let x: f64 = 2.0; + let ops = &[MonoAD::Ln]; + let (value, backprop) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(value, x.ln(), 1e-10)); + assert!(approx_eq(backprop(1.0), 1.0 / x, 1e-10)); +} + +#[test] +fn test_mono_sqrt_forward_and_grad() { + let x: f64 = 9.0; + let ops = &[MonoAD::Sqrt]; + let (value, backprop) = MonoAD::compute_grad(ops, x); + assert!(approx_eq(value, 3.0, 1e-10)); + assert!(approx_eq(backprop(1.0), 1.0 / 6.0, 1e-10)); +} + +#[test] +fn test_mono_abs_forward_and_grad() { + let ops = &[MonoAD::Abs]; + let (value_pos, backprop_pos) = MonoAD::compute_grad(ops, 2.0); + assert!(approx_eq(value_pos, 2.0, 1e-10)); + assert!(approx_eq(backprop_pos(1.0), 1.0, 1e-10)); + + let (value_zero, backprop_zero) = MonoAD::compute_grad(ops, 0.0); + assert!(approx_eq(value_zero, 0.0, 1e-10)); + assert!(approx_eq(backprop_zero(1.0), 0.0, 1e-10)); +} + +#[test] +fn test_mono_checked_domain_errors() { + assert!(MonoAD::compute_checked(&[MonoAD::Ln], 2.0).is_ok()); + assert!(matches!( + MonoAD::compute_checked(&[MonoAD::Ln], 0.0), + Err(AutodiffError::DomainError { + operation: "Ln", + .. + }) + )); + assert!(MonoAD::compute_checked(&[MonoAD::Sqrt], 0.0).is_ok()); + assert!(matches!( + MonoAD::compute_checked(&[MonoAD::Sqrt], -1.0), + Err(AutodiffError::DomainError { + operation: "Sqrt", + .. + }) + )); +} + +#[test] +fn test_mono_checked_grad_and_hessian() { + let (value, grad_fn) = MonoAD::compute_grad_checked(&[MonoAD::Ln], 2.0).unwrap(); + assert!(approx_eq(value, 2.0_f64.ln(), 1e-10)); + assert!(approx_eq(grad_fn(1.0), 0.5, 1e-10)); + + assert!(MonoAD::compute_grad_checked(&[MonoAD::Sqrt], -1.0).is_err()); + assert!(MonoAD::compute_hessian_checked(&[MonoAD::Ln], 2.0).is_ok()); + assert!(MonoAD::compute_hessian_checked(&[MonoAD::Sqrt], 1.0).is_ok()); + assert!(MonoAD::compute_hessian_checked(&[MonoAD::Sqrt], 0.0).is_err()); +} + +#[test] +fn test_mono_exact_checked_domain_errors() { + assert!(MonoAD2RR::compute_checked(&[MonoAD2RR::Ln], 2.0).is_ok()); + assert!(MonoAD2RR::compute_grad_checked(&[MonoAD2RR::Sqrt], 4.0).is_ok()); + assert!(MonoAD2RR::compute_hessian_checked(&[MonoAD2RR::Sqrt], 4.0).is_ok()); + assert!(MonoAD2FR::compute_checked(&[MonoAD2FR::Ln], 2.0).is_ok()); + assert!(MonoAD2RF::compute_checked(&[MonoAD2RF::Sqrt], 0.0).is_ok()); + + assert!(matches!( + MonoAD2RR::compute_checked(&[MonoAD2RR::Ln], 0.0), + Err(AutodiffError::DomainError { + operation: "Ln", + .. + }) + )); + assert!(matches!( + MonoAD2FR::compute_grad_checked(&[MonoAD2FR::Sqrt], -1.0), + Err(AutodiffError::DomainError { + operation: "Sqrt", + .. + }) + )); + assert!(matches!( + MonoAD2RF::compute_hessian_checked(&[MonoAD2RF::Ln], -1.0), + Err(AutodiffError::DomainError { + operation: "Ln", + .. + }) + )); +} + +#[test] +fn test_mono_ops_macro_extended() { + let ops = mono_ops![tan, neg, ln, sqrt, abs]; + assert_eq!(ops.len(), 5); +} + +// ============================================================================ +// MultiAD Missing Operations Tests (Div, Sub, Tan, Ln) +// ============================================================================ + +#[test] +fn test_multi_div_forward() { + let exprs = &multi_ops![(inp, 0), (inp, 1), (div, 0, 1)]; + let inputs = &[10.0, 2.0]; + let result = MultiAD::compute(exprs, inputs).unwrap(); + assert!(approx_eq(result, 5.0, 1e-10)); +} + +#[test] +fn test_multi_div_gradient() { + let exprs = &multi_ops![(inp, 0), (inp, 1), (div, 0, 1)]; + let x = 10.0; + let y = 2.0; + let inputs = &[x, y]; + + let (_value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + assert!(approx_eq(grads[0], 1.0 / y, 1e-10)); + assert!(approx_eq(grads[1], -x / (y * y), 1e-10)); +} + +#[test] +fn test_multi_sub_forward() { + let exprs = &multi_ops![(inp, 0), (inp, 1), (sub, 0, 1)]; + let inputs = &[5.0, 3.0]; + let result = MultiAD::compute(exprs, inputs).unwrap(); + assert!(approx_eq(result, 2.0, 1e-10)); +} + +#[test] +fn test_multi_sub_gradient() { + let exprs = &multi_ops![(inp, 0), (inp, 1), (sub, 0, 1)]; + let inputs = &[5.0, 3.0]; + + let (_value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + assert!(approx_eq(grads[0], 1.0, 1e-10)); + assert!(approx_eq(grads[1], -1.0, 1e-10)); +} + +#[test] +fn test_multi_tan_forward() { + let exprs = &multi_ops![(tan, 0)]; + let x: f64 = 0.5; + let inputs = &[x]; + let result = MultiAD::compute(exprs, inputs).unwrap(); + assert!(approx_eq(result, x.tan(), 1e-10)); +} + +#[test] +fn test_multi_tan_gradient() { + let exprs = &multi_ops![(tan, 0)]; + let x: f64 = 0.5; + let inputs = &[x]; + + let (_value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + let expected_grad = 1.0 / x.cos().powi(2); + assert!(approx_eq(grads[0], expected_grad, 1e-10)); +} + +#[test] +fn test_multi_checked_domain_errors() { + let ln_exprs = &multi_ops![(ln, 0)]; + let ln_error = MultiAD::compute_checked(ln_exprs, &[0.0]).unwrap_err(); + assert_eq!( + ln_error, + AutodiffError::DomainError { + operation: "Ln", + reason: "input must be positive", + } + ); + + let div_exprs = &multi_ops![(inp, 0), (inp, 1), (div, 0, 1)]; + let div_error = MultiAD::compute_checked(div_exprs, &[1.0, 0.0]).unwrap_err(); + assert_eq!( + div_error, + AutodiffError::DomainError { + operation: "Div", + reason: "denominator must be non-zero", + } + ); +} + +#[test] +fn test_multi_checked_gradient_errors() { + let exprs = &multi_ops![(inp, 0), (inp, 1), (pow, 0, 1)]; + let result = MultiAD::compute_grad_checked(exprs, &[-1.0, 2.0]); + assert!(matches!( + result, + Err(AutodiffError::DomainError { + operation: "Pow", + reason: "base must be positive in checked mode", + }) + )); +} + +#[test] +fn test_multi_ln_forward() { + let exprs = &multi_ops![(ln, 0)]; + let x: f64 = 2.0; + let inputs = &[x]; + let result = MultiAD::compute(exprs, inputs).unwrap(); + assert!(approx_eq(result, x.ln(), 1e-10)); +} + +#[test] +fn test_multi_ln_gradient() { + let exprs = &multi_ops![(ln, 0)]; + let x: f64 = 2.0; + let inputs = &[x]; + + let (_value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + assert!(approx_eq(grads[0], 1.0 / x, 1e-10)); +} + +// ============================================================================ +// MultiAD Edge Cases +// ============================================================================ + +#[test] +fn test_multi_empty_graph() { + let exprs: &[(MultiAD, Vec)] = &[]; + let inputs: &[f64] = &[]; + let result = MultiAD::compute(exprs, inputs).unwrap(); + assert_eq!(result, 0.0); +} + +#[test] +fn test_multi_single_input() { + // Inp operation alone returns the input value at that index. + let exprs = &multi_ops![(inp, 0)]; + let inputs = &[5.0]; + let result = MultiAD::compute(exprs, inputs).unwrap(); + assert_eq!(result, 5.0); +} + +#[test] +fn test_multi_input_marker_selects_requested_input() { + let exprs = &multi_ops![(inp, 0)]; + let inputs = &[5.0, 7.0]; + + let (value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + assert_eq!(value, 5.0); + assert_eq!(grads, vec![1.0, 0.0]); +} + +#[test] +fn test_multi_complex_chain() { + let exprs = &multi_ops![ + (inp, 0), + (inp, 1), + (sub, 0, 1), + (sin, 2), + (div, 0, 1), + (ln, 4), + (mul, 3, 5), + ]; + + let x: f64 = 3.0; + let y: f64 = 2.0; + let inputs = &[x, y]; + + let (value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + let expected_value: f64 = (x - y).sin() * (x / y).ln(); + assert!(approx_eq(value, expected_value, 1e-10)); + + assert!(grads[0].is_finite()); + assert!(grads[1].is_finite()); +} + +// ============================================================================ +// GraphBuilder Comprehensive Tests +// ============================================================================ + +#[test] +fn test_builder_all_operations() { + let graph = GraphBuilder::new(3) + .input(0) + .sin(0) + .cos(0) + .tan(0) + .exp(0) + .ln(0) + .sqrt(0) + .abs(0) + .add(0, 1) + .sub(0, 1) + .mul(0, 1) + .div(0, 1) + .pow(0, 1) + .build(); + + assert!(!graph.is_empty()); + assert_eq!(graph.len(), 13); +} + +#[test] +fn test_builder_chained_complex() { + let graph = GraphBuilder::new(2) + .sin(0) + .cos(1) + .tan(0) + .exp(1) + .ln(0) + .sqrt(1) + .abs(0) + .sub(0, 1) + .div(2, 3) + .build(); + + let inputs = &[0.5, 0.8]; + let result = MultiAD::compute(&graph, inputs); + assert!(result.is_ok()); +} + +// ============================================================================ +// Second-Order Methods Tests +// ============================================================================ + +#[test] +fn test_mono_ad2rr_compute() { + let ops = &[MonoAD2RR::Sin, MonoAD2RR::Exp]; + let x: f64 = 0.5; + let value = MonoAD2RR::compute(ops, x); + assert!(approx_eq(value, x.sin().exp(), 1e-10)); +} + +#[test] +fn test_mono_ad2rr_tan_hessian() { + let ops = &[MonoAD2RR::Tan]; + let x: f64 = 0.4; + let sec_sq = 1.0 / x.cos().powi(2); + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 2.0 * sec_sq * x.tan(), 1e-10)); +} + +#[test] +fn test_mono_ad2rr_ln_hessian() { + let ops = &[MonoAD2RR::Ln]; + let x: f64 = 1.7; + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, -1.0 / x.powi(2), 1e-10)); +} + +#[test] +fn test_mono_ad2rr_sqrt_hessian() { + let ops = &[MonoAD2RR::Sqrt]; + let x: f64 = 2.5; + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, -1.0 / (4.0 * x * x.sqrt()), 1e-10)); +} + +#[test] +fn test_mono_ad2rr_abs_hessian() { + let ops = &[MonoAD2RR::Abs]; + for x in [-2.5, 0.0, 2.5] { + let hessian = MonoAD2RR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, 1e-10)); + } +} + +#[test] +fn test_mono_ad2rr_compute_grad() { + let ops = &[MonoAD2RR::Sin]; + let x: f64 = 0.5; + let (value, backprop) = MonoAD2RR::compute_grad(ops, x); + assert!(approx_eq(value, x.sin(), 1e-10)); + assert!(approx_eq(backprop(1.0), x.cos(), 1e-10)); +} + +#[test] +fn test_mono_ad2fr_compute() { + let ops = &[MonoAD2FR::Sin, MonoAD2FR::Exp]; + let x: f64 = 0.5; + let value = MonoAD2FR::compute(ops, x); + assert!(approx_eq(value, x.sin().exp(), 1e-10)); +} + +#[test] +fn test_mono_ad2fr_tan_hessian() { + let ops = &[MonoAD2FR::Tan]; + let x: f64 = 0.4; + let sec_sq = 1.0 / x.cos().powi(2); + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 2.0 * sec_sq * x.tan(), 1e-10)); +} + +#[test] +fn test_mono_ad2fr_ln_hessian() { + let ops = &[MonoAD2FR::Ln]; + let x: f64 = 1.7; + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, -1.0 / x.powi(2), 1e-10)); +} + +#[test] +fn test_mono_ad2fr_sqrt_hessian() { + let ops = &[MonoAD2FR::Sqrt]; + let x: f64 = 2.5; + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, -1.0 / (4.0 * x * x.sqrt()), 1e-10)); +} + +#[test] +fn test_mono_ad2fr_abs_hessian() { + let ops = &[MonoAD2FR::Abs]; + for x in [-2.5, 0.0, 2.5] { + let hessian = MonoAD2FR::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, 1e-10)); + } +} + +#[test] +fn test_mono_ad2fr_compute_grad() { + let ops = &[MonoAD2FR::Sin]; + let x: f64 = 0.5; + let (value, backprop) = MonoAD2FR::compute_grad(ops, x); + assert!(approx_eq(value, x.sin(), 1e-10)); + assert!(approx_eq(backprop(1.0), x.cos(), 1e-10)); +} + +#[test] +fn test_mono_ad2rf_compute() { + let ops = &[MonoAD2RF::Sin, MonoAD2RF::Exp]; + let x: f64 = 0.5; + let value = MonoAD2RF::compute(ops, x); + assert!(approx_eq(value, x.sin().exp(), 1e-10)); +} + +#[test] +fn test_mono_ad2rf_tan_hessian() { + let ops = &[MonoAD2RF::Tan]; + let x: f64 = 0.4; + let sec_sq = 1.0 / x.cos().powi(2); + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, 2.0 * sec_sq * x.tan(), 1e-10)); +} + +#[test] +fn test_mono_ad2rf_ln_hessian() { + let ops = &[MonoAD2RF::Ln]; + let x: f64 = 1.7; + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, -1.0 / x.powi(2), 1e-10)); +} + +#[test] +fn test_mono_ad2rf_sqrt_hessian() { + let ops = &[MonoAD2RF::Sqrt]; + let x: f64 = 2.5; + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, -1.0 / (4.0 * x * x.sqrt()), 1e-10)); +} + +#[test] +fn test_mono_ad2rf_abs_hessian() { + let ops = &[MonoAD2RF::Abs]; + for x in [-2.5, 0.0, 2.5] { + let hessian = MonoAD2RF::compute_hessian(ops, x); + assert!(approx_eq(hessian, 0.0, 1e-10)); + } +} + +#[test] +fn test_mono_ad2rf_compute_grad() { + let ops = &[MonoAD2RF::Sin]; + let x: f64 = 0.5; + let (value, backprop) = MonoAD2RF::compute_grad(ops, x); + assert!(approx_eq(value, x.sin(), 1e-10)); + assert!(approx_eq(backprop(1.0), x.cos(), 1e-10)); +} + +// ============================================================================ +// Dual Type Tests +// ============================================================================ + +#[test] +fn test_dual_variable() { + let d = Dual::variable(3.0); + assert!(approx_eq(d.val, 3.0, 1e-10)); + assert!(approx_eq(d.tan, 1.0, 1e-10)); +} + +#[test] +fn test_dual_constant() { + let d = Dual::constant(3.0); + assert!(approx_eq(d.val, 3.0, 1e-10)); + assert!(approx_eq(d.tan, 0.0, 1e-10)); +} + +// ============================================================================ +// MonoFn Trait Tests - covered in mono::tests module +// since MF1-4 are private implementation details + +// ============================================================================ +// MultiFn Trait Tests +// ============================================================================ + +// Note: F1, F2, F3 are private modules, tested through multi::tests + +// MultiFn trait tests are covered in multi::tests module +// since F1, F2, F3 are private implementation details + +// ============================================================================ +// Arity Error Tests +// ============================================================================ + +#[test] +fn test_multi_arity_errors() { + let exprs = &[(MultiAD::Sin, vec![0, 1])]; + let inputs = &[1.0, 2.0]; + let result = MultiAD::compute(exprs, inputs); + assert!(result.is_err()); + + let exprs = &[(MultiAD::Add, vec![0])]; + let result = MultiAD::compute(exprs, inputs); + assert!(result.is_err()); +} + +// ============================================================================ +// Macros Tests +// ============================================================================ + +#[test] +fn test_mono_ops_rr_macro() { + let ops = mono_ops_rr![sin, cos, tan, exp, neg, ln, sqrt, abs]; + assert_eq!(ops.len(), 8); +} + +#[test] +fn test_mono_ops_fr_macro() { + let ops = mono_ops_fr![sin, cos, tan, exp, neg, ln, sqrt, abs]; + assert_eq!(ops.len(), 8); +} + +#[test] +fn test_mono_ops_rf_macro() { + let ops = mono_ops_rf![sin, cos, tan, exp, neg, ln, sqrt, abs]; + assert_eq!(ops.len(), 8); +} + +// ============================================================================ +// Edge Cases Tests +// ============================================================================ + +#[test] +fn test_mono_empty_hessian() { + let ops: &[MonoAD] = &[]; + let hessian = MonoAD::compute_hessian(ops, 5.0); + assert_eq!(hessian, 0.0); +} + +#[test] +fn test_mono2rr_empty_hessian() { + let ops: &[MonoAD2RR] = &[]; + let hessian = MonoAD2RR::compute_hessian(ops, 5.0); + assert_eq!(hessian, 0.0); +} + +#[test] +fn test_mono2fr_empty_hessian() { + let ops: &[MonoAD2FR] = &[]; + let hessian = MonoAD2FR::compute_hessian(ops, 5.0); + assert_eq!(hessian, 0.0); +} + +#[test] +fn test_mono2rf_empty_hessian() { + let ops: &[MonoAD2RF] = &[]; + let hessian = MonoAD2RF::compute_hessian(ops, 5.0); + assert_eq!(hessian, 0.0); +} + +#[test] +fn test_multi_hessian_empty() { + let exprs: &[(MultiAD, Vec)] = &[]; + let inputs = &[1.0, 2.0]; + let hessian = MultiAD::compute_hessian(exprs, inputs).unwrap(); + assert_eq!(hessian.len(), 2); + assert_eq!(hessian[0][0], 0.0); +} + +#[test] +fn test_multi_hessian_row() { + let exprs = &multi_ops![(inp, 0), (inp, 1), (mul, 0, 0), (mul, 1, 1), (add, 2, 3)]; + let inputs = &[2.0, 3.0]; + + let row0 = MultiAD::compute_hessian_row(exprs, inputs, 0).unwrap(); + assert!(approx_eq(row0[0], 2.0, 1e-6)); + assert!(approx_eq(row0[1], 0.0, 1e-6)); + + let row1 = MultiAD::compute_hessian_row(exprs, inputs, 1).unwrap(); + assert!(approx_eq(row1[0], 0.0, 1e-6)); + assert!(approx_eq(row1[1], 2.0, 1e-6)); +} + +#[test] +fn test_deep_composition() { + let ops = mono_ops![sin, sin, sin, sin, sin]; + let x: f64 = 0.5; + + let (value, backprop) = MonoAD::compute_grad(&ops, x); + let expected = x.sin().sin().sin().sin().sin(); + assert!(approx_eq(value, expected, 1e-10)); + + let grad = backprop(1.0); + assert!(grad.is_finite()); + + let hessian: f64 = MonoAD::compute_hessian(&ops, x); + assert!(hessian.is_finite()); +} + +#[test] +fn test_abs_at_zero() { + // Test abs at x=0 using the explicit subgradient convention sign(0) = 0. + let exprs = &multi_ops![(abs, 0)]; + let inputs = &[0.0]; + + let (_value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + assert_eq!(grads[0], 0.0); +} + +#[test] +fn test_multi_invalid_index_returns_error() { + let exprs = &[(MultiAD::Sin, vec![10])]; + let inputs = &[0.5]; + + let compute_result = MultiAD::compute(exprs, inputs); + assert!(matches!( + compute_result, + Err(AutodiffError::IndexOutOfBounds { index: 10, .. }) + )); + + let grad_result = MultiAD::compute_grad(exprs, inputs); + assert!(matches!( + grad_result, + Err(AutodiffError::IndexOutOfBounds { index: 10, .. }) + )); +} + +#[test] +fn test_multi_empty_compute_grad_zero_inputs_is_safe() { + let exprs: &[(MultiAD, Vec)] = &[]; + let inputs: &[f64] = &[]; + + let (value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + assert_eq!(value, 0.0); + assert!(grads.is_empty()); +} + +#[test] +fn test_multi_hessian_row_invalid_index_returns_error() { + let exprs = &multi_ops![(inp, 0), (mul, 0, 0)]; + let result = MultiAD::compute_hessian_row(exprs, &[2.0], 1); + + assert!(matches!( + result, + Err(AutodiffError::IndexOutOfBounds { index: 1, .. }) + )); +} + +#[test] +fn test_multi_hessian_zero_inputs_validates_graph() { + let exprs = &[(MultiAD::Sin, vec![0])]; + let result = MultiAD::compute_hessian(exprs, &[]); + + assert!(matches!( + result, + Err(AutodiffError::IndexOutOfBounds { index: 0, .. }) + )); +} + +#[test] +fn test_multi_exact_hessian_invalid_input_returns_error() { + let rr_result = MultiAD2RR::compute_hessian(&[MultiAD2RR::Inp(1)], &[1.0]); + let fr_result = MultiAD2FR::compute_hessian(&[MultiAD2FR::Inp(1)], &[1.0]); + let rf_result = MultiAD2RF::compute_hessian(&[MultiAD2RF::Inp(1)], &[1.0]); + + assert!(matches!( + rr_result, + Err(AutodiffError::IndexOutOfBounds { index: 1, .. }) + )); + assert!(matches!( + fr_result, + Err(AutodiffError::IndexOutOfBounds { index: 1, .. }) + )); + assert!(matches!( + rf_result, + Err(AutodiffError::IndexOutOfBounds { index: 1, .. }) + )); +} + +#[test] +fn test_multi_exact_hessian_zero_inputs_validates_graph() { + let rr_result = MultiAD2RR::compute_hessian(&[MultiAD2RR::Inp(0)], &[]); + let fr_result = MultiAD2FR::compute_hessian(&[MultiAD2FR::Inp(0)], &[]); + let rf_result = MultiAD2RF::compute_hessian(&[MultiAD2RF::Inp(0)], &[]); + + assert!(matches!( + rr_result, + Err(AutodiffError::IndexOutOfBounds { index: 0, .. }) + )); + assert!(matches!( + fr_result, + Err(AutodiffError::IndexOutOfBounds { index: 0, .. }) + )); + assert!(matches!( + rf_result, + Err(AutodiffError::IndexOutOfBounds { index: 0, .. }) + )); +} + +#[test] +fn test_multi_exact_hessian_malformed_rpn_returns_error() { + let rr_result = MultiAD2RR::compute_hessian(&[MultiAD2RR::Sin], &[1.0]); + let fr_result = MultiAD2FR::compute_hessian(&[MultiAD2FR::Sin], &[1.0]); + let rf_result = MultiAD2RF::compute_hessian(&[MultiAD2RF::Sin], &[1.0]); + + assert!(matches!(rr_result, Err(AutodiffError::InvalidGraph { .. }))); + assert!(matches!(fr_result, Err(AutodiffError::InvalidGraph { .. }))); + assert!(matches!(rf_result, Err(AutodiffError::InvalidGraph { .. }))); +} + +#[test] +fn test_builder_input_marker_does_not_shift_indices() { + let graph = GraphBuilder::new(1).input(0).sin(0).build(); + let result = MultiAD::compute(&graph, &[0.5]).unwrap(); + + assert!(approx_eq(result, 0.5_f64.sin(), 1e-10)); +} + +#[test] +fn test_mixed_operations() { + let exprs = &multi_ops![ + (inp, 0), + (inp, 1), + (add, 0, 1), + (sub, 0, 1), + (mul, 2, 3), + (div, 0, 1), + (pow, 0, 1), + ]; + + let x: f64 = 2.0; + let y: f64 = 1.5; + let inputs = &[x, y]; + + let (value, backprop) = MultiAD::compute_grad(exprs, inputs).unwrap(); + let grads = backprop(1.0); + + let _value: f64 = value; + assert!(value.is_finite()); + assert!(grads[0].is_finite()); + assert!(grads[1].is_finite()); +}