Skip to content

Generate device-compatible FullSpecialize initialization maps - #5045

Draft
ChrisRackauckas-Claude wants to merge 7 commits into
SciML:masterfrom
ChrisRackauckas-Claude:agent/full-specialize-initialization-maps
Draft

Generate device-compatible FullSpecialize initialization maps#5045
ChrisRackauckas-Claude wants to merge 7 commits into
SciML:masterfrom
ChrisRackauckas-Claude:agent/full-specialize-initialization-maps

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Aug 30, 2026

Copy link
Copy Markdown
Member

Ignore this PR until it has been reviewed by @ChrisRackauckas.

What changed and why

Explicit FullSpecialize problem construction now generates each initialization state/parameter transfer map as one RuntimeGeneratedFunction. The generated maps drop their expression bodies, use static output buffers, and avoid capturing the dynamic InitializationMap/closure graph that blocked device-compatible problem capture in #5043.

The AutoSpecialize/AutoDespecialize path is unchanged. Specialization is propagated through NonlinearLeastSquaresProblem and SteadyStateProblem construction, and the generated parameter map handles DiffCache, immutable nonnumeric buffers, and array-valued callable parameters such as the namespaced Lux function used by ModelingToolkitNeuralNets.

Verification

Failing before

On unmodified upstream/master at f9131984d9, the issue reproducer reported:

SPEC=SciMLBase.FullSpecialize
UMAP_TYPE=ModelingToolkitBase.InitializationMap{true, ...}
UMAP_ISBITS=false

The new regression predicates therefore failed before the implementation: the state map was neither a RuntimeGeneratedFunction nor isbits and retained the dynamic wrapper graph.

The existing DiffCache regression initially discriminated two parameter-map bugs. Before support was added it failed with:

ArgumentError: FullSpecialize initialization maps do not yet support DiffCache parameters.

After removing that rejection but before correcting immutable nonnumeric-buffer reconstruction, it failed with:

IndexCache Test | 66 pass, 2 errors
BoundsError: attempt to access DataType at index [1]

The latter came from querying fieldtype(C, i) while reconstructing N; this PR queries fieldtype(N, i).

The namespaced array-valued callable reproducer on the prior feature commit b3256600f9 failed with:

ArgumentError: FullSpecialize initialization maps do not yet support array-valued parameter storage slots.
Test Summary: FullSpecialize array-valued storage slot | 1 Error

That is the same storage shape exercised by ModelingToolkitNeuralNets' lux_apply parameter.

Passing after

The focused issue regression executes both generated maps and checks their types, dropped expression bodies, static outputs, and equality with AutoSpecialize output:

FullSpecialize regression assertions passed

The array-valued callable regression passes and confirms that the generated parameter map preserves the exact function object:

grouped[5] = [[array_block₊array_fn]]
symbolic_type(...) = ArraySymbolic()
Test Summary: array callable | 3 pass | 3 total | 14.9s

The existing DiffCache gradient test passes:

IndexCache Test | 68 pass | 68 total | 59.7s

The rebased Initialization group was run through the repository's ModelingToolkitBase test runner using the resolved package-test environment:

$ GROUP=Initialization julia +1.12 --startup-file=no --project=<resolved-test-env> lib/ModelingToolkitBase/test/runtests.jl
Test Summary:     | Pass  Total      Time
Guess Propagation |   11     11  10m20.0s
Test Summary:             | Pass  Broken  Total      Time
InitializationSystem Test |  680      12    692  15m37.2s
Test Summary:       | Pass  Total     Time
Initial Values Test |   65     65  2m38.0s

The complete ModelingToolkitNeuralNets test suite against this branch passed:

Basic | 10 pass, 2 broken, 12 total
MTK model macro compatibility | 1 pass
Symbolic Neural Network Macro | 37 pass
Neural Network Parameter Metadata | 55 pass
Reported issues | 7 pass
Testing ModelingToolkitNeuralNets tests passed

Local hygiene checks completed with exit code 0 and no findings:

Runic.jl v1.10 --check <five modified files>
git diff --check upstream/master...HEAD
git diff --unified=0 upstream/master...HEAD | typos -

Known base-branch failures

GROUP=InterfaceII passed the corrected IndexCache test, then reached the independently failing SDEs + jumps with symbolic tstops test:

JumpSystem Test | 4202 pass | 2 fail | 4204 total

A clean-base reproduction and bisect identified the ModelingToolkit dependency-widening commit and reduced the doubled callback to a standalone non-MTK reproducer. It is tracked at SciML/OrdinaryDiffEq.jl#4439. InterfaceII is not claimed as passing.

Root QA on the rebased branch passed 19/20 checks. Its only failure was the same clean-base public-reexport drift (Unknown, scalarize, shape, unwrap); method ambiguity, dependency, compat, piracy, ExplicitImports, and public-documentation checks passed. ModelingToolkitBase QA separately retains the known 263 JET findings tracked at #4958. QA is not claimed as passing.

CI-detected mutable-buffer regression

The first CI run found that generated SVector parameter buffers broke the existing mutating SciMLStructures interface:

MTKParameters Test | 118 pass | 6 errors | 124 total
setindex!(::SVector, ...) is not defined

Commit 93854e0 preserves static sizing while selecting MVector for mutable isbits prototypes, SizedVector for mutable non-isbits prototypes, and SVector for immutable prototypes. The focused tunable/constant replace! reproducer exits successfully, and the full group now reports:

SymbolicIndexingInterface | 1984 pass | 1984 total | 8m06.7s
Testing ModelingToolkit tests passed

CI-detected non-isbits mutable-buffer regression

The subsequent FMI LTS job exercised mutable parameter buffers containing vectors and FMI functors. MVector cannot implement setindex! for non-isbits element types, so the two co-simulation cases failed before commit 613bdd1:

FMI | 56 pass | 2 errors | 58 total
setindex!() with non-isbitstype eltype is not supported by StaticArrays. Consider using SizedArray.

The helper now selects immutable SVector storage, mutable MVector storage for isbits elements, and mutable SizedVector storage backed by Vector for non-isbits elements. The exact FMI group now passes locally:

$ GROUP=FMI julia +1.12 --startup-file=no --project=. -e "using Pkg; Pkg.test()"
FMI | 60 pass | 60 total | 15m53.9s
Testing ModelingToolkit tests passed

The final Initialization group also passes:

$ GROUP=Initialization julia +1.12 --startup-file=no --project=. -e "using Pkg; Pkg.test()"
Initialization | 870 pass | 12 broken | 882 total | 36m22.1s
Testing ModelingToolkit tests passed

\n## Not verified

  • No GPU test was run locally; CPU regressions verify the structural requirements (RuntimeGeneratedFunction, isbits map object, dropped body, and static buffers).
  • Downstreams other than ModelingToolkitNeuralNets and allowed-to-fail GPU jobs were not run locally.
  • Documentation was not built because this changes no public API, docstring, or documentation file.

Fixes #5043.

Related downstream workaround: SciML/DiffEqGPU.jl#516.

🤖 Generated with Codex (version unknown) (model: unknown). Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Plan/status: generate single dropped-body RuntimeGeneratedFunctions only for explicit FullSpecialize; preserve AutoSpecialize and AutoDespecialize behavior; verify map structure and output against existing maps; run Initialization and local hygiene checks; track the independently reproduced upstream/master QA failures in the required parallel investigation before final handoff.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

CI Runic failure is pre-existing formatter-version drift in three files untouched by this PR. The mechanical cleanup is isolated in draft PR #5046. This feature PR should be rebased after that cleanup lands; no formatting-only changes will be mixed into this behavior diff.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Follow-up: the isolated formatting PR #5046 now passes CI Runic, Runic Suggestions, and spelling. This confirms the formatter failure here is resolved by that base cleanup rather than a change to this feature diff.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Follow-up commit b3256600f9d2f0dae9ebdbb25f0841ff375a03c8 addresses the CI-discovered DiffCache path. The original run failed with the explicit unsupported-parameter ArgumentError; removing that guard exposed an existing immutable nonnumeric-buffer typo (fieldtype(C, i) while reconstructing N) and produced 66 passes / 2 errors in IndexCache. Correcting it to fieldtype(N, i) produced IndexCache Test | 68 pass | 68 total locally. The broader InterfaceII run later hit two unrelated stochastic jump assertions after 4,202 passes; clean-base reproduction/bisect is running separately.

ChrisRackauckas and others added 3 commits August 31, 2026 00:02
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
@ChrisRackauckas-Claude
ChrisRackauckas-Claude force-pushed the agent/full-specialize-initialization-maps branch from b325660 to d3c8a49 Compare August 31, 2026 04:02
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Added array-valued callable-parameter support in d3c8a49.

Failure on the preceding implementation (b3256600f9):

ArgumentError: FullSpecialize initialization maps do not yet support array-valued parameter storage slots.
Test Summary: FullSpecialize array-valued storage slot | 1 Error

The same namespaced ArraySymbolic reproduction now passes:

Test Summary: array callable | 3 pass | 3 total | 14.9s

The complete ModelingToolkitNeuralNets suite also passes against this branch (10 + 1 + 37 + 55 + 7 passing assertions, with its 2 pre-existing broken tests). This directly covers the lux_apply parameter shape that failed in https://github.com/SciML/ModelingToolkit.jl/actions/runs/33300901287/job/99228760344.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

CI exposed a mutable-buffer regression in the first generated-parameter implementation: https://github.com/SciML/ModelingToolkit.jl/actions/runs/33355753893/job/99377920719 failed six MTKParameters cases because SciMLStructures.replace! attempted setindex! on generated SVector tunable/constant buffers.

Fixed in 93854e0 by selecting MVector when the source buffer is mutable while retaining SVector for immutable prototypes. The generated map itself remains an isbits, body-dropped RuntimeGeneratedFunction.

Local verification after the fix:

Focused Tunable/Constants replace! reproducer: exit 0
SymbolicIndexingInterface | 1984 pass | 1984 total | 8m06.7s
Testing ModelingToolkit tests passed
Runic.jl v1.10, typos, git diff --check: exit 0

Use SizedVector-backed storage when a mutable FullSpecialize parameter buffer contains non-isbits values, while retaining MVector for isbits element types.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Commit 613bdd1 fixes the non-isbits mutable-buffer regression found by the FMI LTS job: https://github.com/SciML/ModelingToolkit.jl/actions/runs/33358934228/job/99386604733.

Failing before (CI):

FMI | 56 pass | 2 errors | 58 total
setindex!() with non-isbitstype eltype is not supported by StaticArrays. Consider using SizedArray.

Passing after (local):

GROUP=FMI ... Pkg.test()
FMI | 60 pass | 60 total | 15m53.9s
Testing ModelingToolkit tests passed

GROUP=Initialization ... Pkg.test()
Initialization | 870 pass | 12 broken | 882 total | 36m22.1s
Testing ModelingToolkit tests passed

Runic 1.10, typos over the diff, and git diff --check also exited successfully.

Assert the declared array shape instead of the SymbolicIndexingInterface classification enum, which differs across supported dependency versions while the generated map behavior remains identical.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Commit 010bc30 addresses the downgraded Initialization failure: https://github.com/SciML/ModelingToolkit.jl/actions/runs/33366256108/job/99407718377.

The failed job used SymbolicIndexingInterface 0.3.46, where the callable parameter is classified as ScalarSymbolic despite retaining its declared [1:2] shape. The generated map already returned the exact callable; the failing enum assertion encoded a newer dependency detail.

Local Julia 1.10 reproduction pinned the relevant CI floors (ArrayInterface 7.28.0, SciMLBase 3.48.0, StaticArraysCore 1.4.3, SymbolicUtils 4.37.0, SymbolicIndexingInterface 0.3.46, Symbolics 7.37.0):

classifier=ScalarSymbolic
shape/generated-map/callable-identity assertions: 3/3 pass

The same semantic assertions also pass 3/3 on the current dependency stack. Runic 1.10, typos, and git diff --check pass for the follow-up diff.

Read the callable parameter shape from SymbolicUtils directly so the regression works across supported versions that do not re-export shape through ModelingToolkitBase.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: OpenAI Codex <noreply@openai.com>
Agent-Harness: Codex (version unknown)
Agent-Model: unknown
Agent-Session: local session ID 01a04feb-43da-7ce1-99e4-a68438c69833
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

Follow-up commit 258f72c fixes the remaining downgrade error from https://github.com/SciML/ModelingToolkit.jl/actions/runs/33387712747/job/99474377835. SymbolicUtils 4.37 provides shape at its owning module but does not re-export it through ModelingToolkitBase.

Exact CI-floor local result after using the owner API:

classifier=ScalarSymbolic
SU.shape/generated-map/callable-identity assertions: 3/3 pass

The current-stack owner-API shape assertions pass 2/2. Runic 1.10, typos, and git diff --check pass.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member Author

CI confirmation on head 258f72c:

At this update, 38 checks pass, 66 remain in progress, 2 known base checks fail, and 1 is skipped. No completed branch-relevant check fails.

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.

FullSpecialize: generate device-compatible initialization maps (initializeprobmap/pmap) as RGFs

2 participants