Enzyme.jl: Unable to differentiate through Interpolator construction

[kylebeggs]

Summary

Enzyme.jl cannot differentiate through Interpolator(points, values) construction in loss functions. This blocks a common use case of differentiating RBF interpolation w.r.t. input values.

Minimal Reproducible Example

using RadialBasisFunctions
using StaticArrays
import Enzyme

N = 30
points = [SVector{2}(rand(), rand()) for _ in 1:N]
values = sin.(getindex.(points, 1))
eval_points = [SVector{2}(rand(), rand()) for _ in 1:10]

function loss_interp(v)
    interp = Interpolator(points, v)
    result = interp(eval_points)
    return sum(result .^ 2)
end

dv = zeros(N)
Enzyme.autodiff(Enzyme.Reverse, loss_interp, Enzyme.Active, Enzyme.Duplicated(values, dv))

Error Messages

Without custom rules (HEAD of enzyme branch)

ERROR: IllegalTypeAnalysisException: Enzyme compilation failed due to illegal type analysis.
 This usually indicates the use of a Union type, which is not fully supported...
 Failure within method: MethodInstance for LinearAlgebra.var"#_factorize#135"(...)

The issue is that Interpolator internally calls factorize which returns a Union type (BunchKaufman or Diagonal).

Setting Enzyme.API.strictAliasing!(false) causes Julia to crash/segfault.

With custom EnzymeRules (attempted fix)

Adding custom augmented_primal and reverse rules for Interpolator construction leads to:

ERROR: AugmentedRuleReturnError: Incorrect return type for primal and shadow configuration...
  expected : Interpolator{..., MB} where MB<:MonomialBasis
  found    : Interpolator{..., MonomialBasis{2, 2, ...}}

Enzyme requires exact type matching for custom rules, but the rule returns a concrete type while Julia's type inference produces an abstract type.

Root Causes

1. LinearAlgebra.factorize Union return type: The _factorize function can return either BunchKaufman or Diagonal, which Enzyme cannot handle
2. Custom rule type matching: When defining custom rules to bypass the factorize issue, Enzyme's strict type matching requirements conflict with Julia's type inference

Environment

• Julia: 1.10.10 and 1.11.8
• Enzyme.jl: 0.13.129
• EnzymeCore.jl: 0.8.18
• RadialBasisFunctions.jl: 0.3.0 (enzyme branch)

Workaround

Currently, use ChainRules/Mooncake for Interpolator construction differentiation:

import DifferentiationInterface as DI
import Mooncake

backend = DI.AutoMooncake(; config=nothing)
grad = DI.gradient(loss_interp, backend, values)

Related Issues

• Enzyme.jl #2699 (Julia 1.12+ compatibility)

Potential Fixes

1. Add an Enzyme-compatible implementation of Interpolator that avoids factorize (use explicit LU or Cholesky instead)
2. File an Enzyme.jl issue about relaxing type matching requirements for custom rules
3. Rely on ChainRulesCore rrules which Enzyme can use on Julia < 1.12

🤖 Generated with Claude Code

[kylebeggs]

Minimal Working Example (Simplified)

The issue is upstream in Enzyme.jl - it cannot handle factorize() because it returns a Union type.

using LinearAlgebra
import Enzyme

# This FAILS - factorize returns Union type
function loss_with_factorize(b::Vector{Float64})
    A = [4.0 1.0; 1.0 3.0]
    F = factorize(Symmetric(A))
    x = F \ b
    return sum(x.^2)
end

# This WORKS - cholesky returns concrete type
function loss_with_cholesky(b::Vector{Float64})
    A = [4.0 1.0; 1.0 3.0]
    F = cholesky(Symmetric(A))
    x = F \ b
    return sum(x.^2)
end

b = [1.0, 2.0]

# Test factorize - FAILS
db = zeros(2)
Enzyme.autodiff(Enzyme.Reverse, loss_with_factorize, Enzyme.Active, Enzyme.Duplicated(b, db))
# ERROR: IllegalTypeAnalysisException

# Test cholesky - WORKS
db = zeros(2)
Enzyme.autodiff(Enzyme.Reverse, loss_with_cholesky, Enzyme.Active, Enzyme.Duplicated(b, db))
# Success! Gradient: [-0.066, 0.446]

Output:

Enzyme version: 0.13.129
Julia version: 1.11.8

=== Test 1: factorize (Union return type) ===
✗ Failed: IllegalTypeAnalysisException

=== Test 2: cholesky (concrete return type) ===
✓ Success! Gradient: [-0.06611570247933886, 0.44628099173553726]

Potential Fix

Since RBF collocation matrices are symmetric positive definite, we could use cholesky() instead of factorize() in the Interpolator constructor. This would make it compatible with Enzyme.

🤖 Generated with Claude Code

[kylebeggs]

Complete Issue Analysis

After extensive testing, here's a comprehensive analysis of Enzyme.jl limitations affecting RBF interpolation:

Test Results

============================================================
Enzyme.jl Linear Algebra MWE
Enzyme version: 0.13.129
Julia version: 1.11.8
============================================================

### ISSUE 1: Symmetric matrix \ ###
  Symmetric \ b: ✗ FAILS: IllegalTypeAnalysisException
  Matrix \ b: ✓ WORKS

### ISSUE 2: Block matrix construction ###
  Literal construction: ✓ WORKS
  Block [Φ P; P' c]: ✗ FAILS: IllegalTypeAnalysisException

### ISSUE 3: Specific factorization methods ###
  factorize(Symmetric): ✗ FAILS: IllegalTypeAnalysisException
  cholesky(Symmetric): ✓ WORKS
  bunchkaufman(Symmetric): ✗ FAILS: EnzymeNoDerivativeError
  lu(Matrix): ✗ FAILS: EnzymeNoDerivativeError

Root Causes

1. Symmetric \ b fails because it internally calls factorize() which returns a Union type (BunchKaufman or Cholesky or Diagonal, etc.)

2. Block matrix construction fails - [A B; C D] syntax causes Enzyme compilation errors even when the resulting matrix is identical to literal construction

3. RBF collocation matrices are indefinite (not positive definite) due to the polynomial augmentation block structure with zeros on diagonal - this means cholesky() cannot be used directly

Why Interpolator Fails

The Interpolator constructor uses:

A = Symmetric(zeros(...))  # Symmetric wrapper
_build_collocation_matrix!(A, ...)  # Fill entries
b = vcat(y, zeros(...))  # vcat works!
w = A \ b  # FAILS - Symmetric \ b uses factorize()

Potential Fixes

1. Drop Symmetric wrapper and use plain Matrix - Matrix \ b works
2. Use explicit LU decomposition - but lu(A) \ b fails too (no Enzyme rule)
3. Use Mooncake/ChainRules - already working via DI.AutoMooncake

MWE file: dev/enzyme_mwe.jl

🤖 Generated with Claude Code

[kylebeggs]

Update from PR #83

PR #83 (merged) replaced the ChainRulesCore extension with native Enzyme and Mooncake rules. Here's the current status of Interpolator construction differentiation:

Mooncake: Fully working

A native rrule!! for Interpolator(x, y, basis) was added in the Mooncake extension. It runs the constructor opaquely in the forward pass and uses the implicit function theorem in the backward pass to propagate cotangents back to y. This makes the workaround path from the original issue first-class:

import DifferentiationInterface as DI
import Mooncake
backend = DI.AutoMooncake(; config=nothing)
grad = DI.gradient(loss_interp, backend, values)

Enzyme: Still blocked

The Enzyme extension added rules for Interpolator evaluation (single point and batch), but not for Interpolator construction. The core issue remains: Symmetric \ b internally calls factorize() which returns a Union type that Enzyme cannot handle (IllegalTypeAnalysisException).

The Enzyme construction path would require either:

1. An explicit EnzymeRule for the Interpolator constructor (matching the Mooncake approach)
2. Upstream Enzyme.jl fix for factorize Union return types
3. Replacing factorize with a concrete factorization (but RBF collocation matrices are indefinite, so cholesky doesn't work)

Recommendation

Keep this issue open — the Enzyme constructor path is still broken. The Mooncake path is fully functional as a workaround.