Skip to content

feat: RBFLayer for Lux.jl, fix Gaussian/IMQ ε validation#109

Merged
kylebeggs merged 20 commits intomainfrom
lux
Mar 15, 2026
Merged

feat: RBFLayer for Lux.jl, fix Gaussian/IMQ ε validation#109
kylebeggs merged 20 commits intomainfrom
lux

Conversation

@kylebeggs
Copy link
Member

@kylebeggs kylebeggs commented Mar 12, 2026

Summary

  • feat: RBFLayer — a Lux.jl-compatible RBF network layer via LuxCore package extension
  • fix: Gaussian and IMQ constructors now correctly reject ε=0 (was all(ε .< 0), now ε <= 0)
  • refactor: Gaussian/IMQ basis functions split into rbf(r²) and rbf(x, xᵢ) overloads to eliminate redundant distance computation
  • docs: Julia best practices added to CLAUDE.md
  • chore: Version bump 0.3.00.4.0

kylebeggs and others added 9 commits March 12, 2026 10:48
- Add scalar-distance callable `(rbf::T)(r2)` to Gaussian, IMQ, PHS1/3/5/7;
  delegate point-pair callables to it so the formula lives in one place
- Add `src/lux.jl` with `RBFLayer` (Lux.jl-compatible RBF network layer):
  learnable centers, per-center shape parameters via softplus, all six
  RBF types supported; uses the new scalar callables to avoid duplication
- Use `basis_type::Type{B}` struct field instead of `basis::Symbol`,
  enforced by the type system; `_rbf_activate` and `_has_shape_parameter`
  dispatch on types rather than `Val{:symbol}`
- Add docs guide `docs/src/guides/lux.md` covering theory, training example,
  and comparison with MLP; wire into docs nav
- Add `test/lux.jl` with 122 tests covering construction, initialization,
  forward pass, activation math, and pretty printing

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The validation used `all(ε .< 0)` which only rejected negative values,
allowing ε=0 through. Changed to `ε <= 0` to correctly reject zero.
Move Lux integration from a hard dependency to a package extension that
loads only when LuxCore is available. Adds RBFLayer struct in main
package with LuxCore trait implementations in the extension.
The v0.4.0 compat entry caused Downgrade CI to fail with KeyError
since the version isn't in the General registry yet. The docs project
uses the local dev version via path, so this bound is unnecessary.
The docs project has heavy deps (CairoMakie, Lux, Enzyme, etc.) whose
minimum compat versions conflict when downgraded together. Docs builds
documentation, not package tests, so downgrade testing is irrelevant.
@codecov
Copy link

codecov bot commented Mar 12, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines Coverage Δ
src/basis/gaussian.jl 100.00% <100.00%> (ø)
src/basis/inverse_multiquadric.jl 100.00% <100.00%> (ø)
src/basis/polyharmonic_spline.jl 100.00% <100.00%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

- Add [compat] section to test/Project.toml to prevent resolver from
  picking ancient broken versions of test dependencies
- Switch downgrade mode from alldeps to deps to focus on direct dep bounds
- Expand skip list to include test-only and extension deps
- Replace julia-runtest with manual test runner (Pkg.develop + include)
  to avoid "can not merge projects" error on Julia < 1.12
- Update benchmark julia compat from 1.9 to 1.10
The downgrade-compat action was crashing with KeyError on the
RadialBasisFunctions UUID because it tried to resolve the dev'd package
from the registry. Remove it from test/Project.toml [compat] and add it
to the workflow skip list.
@github-actions
Copy link
Contributor

github-actions bot commented Mar 12, 2026

Benchmark Results

main 9aef9cf... main / 9aef9cf...
Directional 2.45 ± 0.15 ms 2.37 ± 0.13 ms 1.03 ± 0.085
Directional (per point) 2.43 ± 0.14 ms 2.39 ± 0.14 ms 1.02 ± 0.085
Gradient 8.82 ± 0.37 ms 8.8 ± 0.37 ms 1 ± 0.06
MonomialBasis/dim=1/deg=0 0.0464 ± 0.014 μs 0.0453 ± 0.013 μs 1.02 ± 0.43
MonomialBasis/dim=1/deg=1 0.0781 ± 0.022 μs 0.0757 ± 0.022 μs 1.03 ± 0.41
MonomialBasis/dim=1/deg=2 0.0891 ± 0.022 μs 0.0869 ± 0.022 μs 1.03 ± 0.36
MonomialBasis/dim=2/deg=0 0.0336 ± 0.00063 μs 0.0345 ± 0.0017 μs 0.973 ± 0.051
MonomialBasis/dim=2/deg=1 0.0333 ± 0.014 μs 0.0343 ± 0.014 μs 0.969 ± 0.57
MonomialBasis/dim=2/deg=2 0.0406 ± 0.014 μs 0.041 ± 0.014 μs 0.99 ± 0.48
MonomialBasis/dim=3/deg=0 0.0376 ± 0.014 μs 0.0359 ± 0.014 μs 1.05 ± 0.56
MonomialBasis/dim=3/deg=1 0.0466 ± 0.014 μs 0.0407 ± 0.014 μs 1.15 ± 0.52
MonomialBasis/dim=3/deg=2 0.0438 ± 0.014 μs 0.0477 ± 0.013 μs 0.918 ± 0.39
Partial 2.58 ± 0.15 ms 2.53 ± 0.13 ms 1.02 ± 0.079
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∂ 9.89 ± 0.12 ns 9.6 ± 0.15 ns 1.03 ± 0.02
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∂² 10 ± 0.25 ns 10.2 ± 0.07 ns 0.985 ± 0.026
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∇ 17.1 ± 0.06 ns 17.1 ± 0.06 ns 1 ± 0.005
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∇² 18.6 ± 0.05 ns 18.1 ± 0.06 ns 1.03 ± 0.0044
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∂ 9.89 ± 0.14 ns 9.68 ± 0.06 ns 1.02 ± 0.016
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∂² 10.2 ± 0.16 ns 10.1 ± 0.13 ns 1.01 ± 0.02
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∇ 17 ± 0.06 ns 17 ± 0.061 ns 1 ± 0.005
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∇² 18.6 ± 0.05 ns 18.1 ± 0.06 ns 1.03 ± 0.0044
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∂ 9.85 ± 0.12 ns 9.68 ± 0.07 ns 1.02 ± 0.014
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∂² 10.1 ± 0.2 ns 10.1 ± 0.21 ns 1 ± 0.029
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∇ 17.1 ± 0.06 ns 17 ± 0.061 ns 1.01 ± 0.0051
RBF/Gaussian, exp(-(ε*r)²)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∇² 18.6 ± 0.29 ns 18.1 ± 0.061 ns 1.03 ± 0.016
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∂ 6.32 ± 0.011 ns 6.32 ± 0.02 ns 1 ± 0.0036
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∂² 14.2 ± 0.2 ns 14.2 ± 0.02 ns 0.999 ± 0.014
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∇ 8.55 ± 0.11 ns 8.65 ± 0.31 ns 0.988 ± 0.038
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 0/0/∇² 15.7 ± 0.13 ns 15.8 ± 0.089 ns 0.991 ± 0.01
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∂ 6.32 ± 0.01 ns 6.32 ± 0.01 ns 1 ± 0.0022
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∂² 14.1 ± 0.2 ns 14.2 ± 0.02 ns 0.991 ± 0.014
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∇ 8.63 ± 0.11 ns 8.61 ± 0.3 ns 1 ± 0.037
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 1/1/∇² 15.7 ± 0.12 ns 15.8 ± 0.089 ns 0.992 ± 0.0095
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∂ 6.32 ± 0.07 ns 6.32 ± 0.01 ns 1 ± 0.011
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∂² 14.1 ± 0.2 ns 14.2 ± 0.02 ns 0.991 ± 0.014
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∇ 8.55 ± 0.11 ns 8.63 ± 0.3 ns 0.992 ± 0.037
RBF/Inverse Multiquadrics, 1/sqrt((r*ε)²+1)
├─Shape factor: ε = 1
└─Polynomial augmentation: degree 2/2/∇² 15.7 ± 0.13 ns 15.8 ± 0.089 ns 0.992 ± 0.01
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 0/0/∂ 3.4 ± 0.03 ns 3.73 ± 0.01 ns 0.911 ± 0.0084
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 0/0/∂² 4.7 ± 0.01 ns 4.7 ± 0.01 ns 1 ± 0.003
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 0/0/∇ 5.65 ± 0.02 ns 5.61 ± 0.031 ns 1.01 ± 0.0066
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 0/0/∇² 3.42 ± 0.001 ns 3.11 ± 0 ns 1.1 ± 0.00032
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 1/1/∂ 3.41 ± 0.031 ns 3.73 ± 0.01 ns 0.914 ± 0.0087
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 1/1/∂² 4.7 ± 0.01 ns 4.7 ± 0.009 ns 1 ± 0.0029
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 1/1/∇ 5.65 ± 0.02 ns 5.61 ± 0.031 ns 1.01 ± 0.0066
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 1/1/∇² 3.42 ± 0.01 ns 3.11 ± 0.001 ns 1.1 ± 0.0032
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 2/2/∂ 3.4 ± 0.039 ns 3.73 ± 0.01 ns 0.911 ± 0.011
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 2/2/∂² 4.7 ± 0.01 ns 4.7 ± 0.009 ns 1 ± 0.0029
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 2/2/∇ 5.65 ± 0.02 ns 5.61 ± 0.049 ns 1.01 ± 0.0095
RBF/Polyharmonic spline (r³)
└─Polynomial augmentation: degree 2/2/∇² 3.42 ± 0.01 ns 3.11 ± 0.001 ns 1.1 ± 0.0032
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 0/0/∂ 4.27 ± 0.01 ns 4.27 ± 0.009 ns 1 ± 0.0032
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 0/0/∂² 5.58 ± 0.011 ns 5.56 ± 0.011 ns 1 ± 0.0028
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 0/0/∇ 6.85 ± 0.01 ns 7.11 ± 0.3 ns 0.963 ± 0.041
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 0/0/∇² 4.27 ± 0.01 ns 4.27 ± 0.01 ns 1 ± 0.0033
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 1/1/∂ 4.27 ± 0.01 ns 4.27 ± 0.01 ns 1 ± 0.0033
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 1/1/∂² 5.58 ± 0.24 ns 5.56 ± 0.02 ns 1 ± 0.043
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 1/1/∇ 6.85 ± 0.02 ns 7.11 ± 0.27 ns 0.963 ± 0.037
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 1/1/∇² 4.27 ± 0.01 ns 4.27 ± 0.01 ns 1 ± 0.0033
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 2/2/∂ 4.27 ± 0.01 ns 4.27 ± 0.01 ns 1 ± 0.0033
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 2/2/∂² 5.58 ± 0.24 ns 5.56 ± 0.02 ns 1 ± 0.043
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 2/2/∇ 6.85 ± 0.02 ns 7.11 ± 0.14 ns 0.963 ± 0.019
RBF/Polyharmonic spline (r¹)
└─Polynomial augmentation: degree 2/2/∇² 4.27 ± 0.01 ns 4.27 ± 0.01 ns 1 ± 0.0033
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 0/0/∂ 4.65 ± 0.011 ns 5.26 ± 0.01 ns 0.884 ± 0.0027
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 0/0/∂² 5.26 ± 0.01 ns 4.96 ± 0.01 ns 1.06 ± 0.0029
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 0/0/∇ 6.19 ± 0.02 ns 6.1 ± 0.041 ns 1.01 ± 0.0076
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 0/0/∇² 3.11 ± 0.01 ns 3.11 ± 0.01 ns 1 ± 0.0046
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 1/1/∂ 4.65 ± 0.01 ns 5.26 ± 0.01 ns 0.884 ± 0.0025
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 1/1/∂² 5.26 ± 0.01 ns 4.96 ± 0.011 ns 1.06 ± 0.0031
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 1/1/∇ 6.19 ± 0.011 ns 6.1 ± 0.041 ns 1.01 ± 0.0071
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 1/1/∇² 3.11 ± 0.01 ns 3.11 ± 0.01 ns 1 ± 0.0046
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 2/2/∂ 4.65 ± 0.01 ns 5.26 ± 0.01 ns 0.884 ± 0.0025
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 2/2/∂² 5.26 ± 0.01 ns 4.96 ± 0.011 ns 1.06 ± 0.0031
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 2/2/∇ 6.19 ± 0.02 ns 6.1 ± 0.041 ns 1.01 ± 0.0076
RBF/Polyharmonic spline (r⁵)
└─Polynomial augmentation: degree 2/2/∇² 3.11 ± 0.01 ns 3.11 ± 0.01 ns 1 ± 0.0046
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 0/0/∂ 10.4 ± 0.099 ns 10.1 ± 0.13 ns 1.02 ± 0.016
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 0/0/∂² 4.75 ± 0.04 ns 4.96 ± 0.01 ns 0.958 ± 0.0083
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 0/0/∇ 12.6 ± 0.27 ns 12.5 ± 0.09 ns 1.01 ± 0.023
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 0/0/∇² 8.14 ± 0.051 ns 8.05 ± 0.079 ns 1.01 ± 0.012
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 1/1/∂ 10.4 ± 0.08 ns 10.1 ± 0.14 ns 1.02 ± 0.016
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 1/1/∂² 4.73 ± 0.04 ns 4.96 ± 0.01 ns 0.954 ± 0.0083
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 1/1/∇ 12.6 ± 0.091 ns 12.5 ± 0.081 ns 1.01 ± 0.0098
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 1/1/∇² 8.14 ± 0.039 ns 8.05 ± 0.072 ns 1.01 ± 0.01
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 2/2/∂ 10.3 ± 0.12 ns 10.1 ± 0.17 ns 1.02 ± 0.021
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 2/2/∂² 4.75 ± 0.04 ns 4.96 ± 0.01 ns 0.958 ± 0.0083
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 2/2/∇ 12.6 ± 0.061 ns 12.5 ± 0.081 ns 1.01 ± 0.0082
RBF/Polyharmonic spline (r⁷)
└─Polynomial augmentation: degree 2/2/∇² 8.14 ± 0.05 ns 8.05 ± 0.04 ns 1.01 ± 0.008
time_to_load 0.788 ± 0.0057 s 0.785 ± 0.0028 s 1 ± 0.0081

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

Mooncake 0.4.x uses Core.Intrinsics.arraylen which was removed in
modern Julia, causing precompilation failures during downgrade resolution.
Remove test/ from downgrade CI projects and trim skip list to only
stdlibs and weakdeps. Test-only deps (DifferentiationInterface,
ForwardDiff, etc.) aren't part of the public API contract and were
causing repeated CI failures on Julia 1.10 minimum-version resolution.
Lux now requires all layers in Chain to be <: AbstractLuxLayer.
Move LuxCore from weakdep to direct dep (lightweight, just abstract
types) so RBFLayer can inherit from AbstractLuxLayer at definition.
Inline the former extension code into src/lux.jl since the extension
trigger is no longer valid with LuxCore as a direct dep.
…ux example

Zygote can't differentiate through in-place mutations in _pairwise_sq_euclidean,
breaking the docs build. Switch to DifferentiationInterface with Mooncake backend
which handles mutations correctly.
Mooncake.Tangent doesn't implement `length`, which Optimisers.update
needs, breaking the @example blocks in docs CI. Enzyme works on 1.12
and handles this correctly with Const function annotation.
Widen threshold from 1e-10 to 1e-9 to match the max_error bound.
Actual value on 1.11 is 1.027e-10, just barely over the old limit.
LuxCore was a direct dependency used solely for AbstractLuxLayer. Moving
the Lux code to a package extension avoids loading LuxCore for users who
don't need it, following the existing pattern (Enzyme, Mooncake).
Add Lux integration to README feature list and docs links. Update
CLAUDE.md architecture section to list LuxCoreExt. Fix docs build by
adding LuxCore to docs/Project.toml and using const assignment for
RBFLayer in @example blocks. Relax docs compat bounds.
@kylebeggs kylebeggs merged commit e7af955 into main Mar 15, 2026
26 checks passed
@kylebeggs kylebeggs deleted the lux branch March 15, 2026 15:47
@kylebeggs
Copy link
Member Author

closes #45

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant