feat: implement ConvolutionalNeuralOperator (CNO) (closes #122)#134
Open
jitendravjh wants to merge 3 commits into
Open
feat: implement ConvolutionalNeuralOperator (CNO) (closes #122)#134jitendravjh wants to merge 3 commits into
jitendravjh wants to merge 3 commits into
Conversation
…ndexing SciML#125) GridEmbedding built the positional grid using CPU range/meshgrid, then called cat(grid, x) where x may be a CuArray. This caused: ERROR: Scalar indexing is disallowed. Invocation of getindex resulted in scalar indexing of a GPU array. Fix: call Lux.get_device(x)(grid) immediately after building the grid, so the array is moved to the same device as the input before the cat. This is a no-op on CPU and transparently transfers to GPU/Metal/etc. Fixes SciML#125
Implements the Convolutional Neural Operator from: Raonic et al., "Convolutional Neural Operators for robust and accurate learning of PDEs", NeurIPS 2023. https://arxiv.org/abs/2302.01178 Architecture: - Lifting: Conv(1x...x1): in_channels → hidden_channels - CNO blocks: Upsample(:bilinear) → Conv(3x...x3, act, SamePad) → MeanPool - Projection: Conv(1x...x1, act) → Conv(1x...x1): → out_channels Each CNOBlock upsamples by upsample_factor, convolves at higher resolution, then downsamples via MeanPool — ensuring the operator converges to a continuous limit as spatial resolution increases (resolution-invariant). New types exported: - ConvolutionalNeuralOperator (the full model) - CNOBlock (the building block, composable) Closes SciML#122
jitendravjh
force-pushed
the
feat/convolutional-neural-operator
branch
from
d81efeb to
e2ded64
Compare
Member
|
Missing tests |
Author
|
Added |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements CNO from Raonic et al., NeurIPS 2023 (https://arxiv.org/abs/2302.01178).
Each
CNOBlockfollows the paper: Upsample → Conv(3×3, act, SamePad) → MeanPool, ensuring the discrete operator converges to a continuous limit as resolution increases.Architecture: Lifting →
CNOBlock × N→ ProjectionAPI: consistent with
FourierNeuralOperator(modes, in_ch, out_ch, hidden_ch; ...)Tested locally (CPU): 1D, 2D, multi-channel, resolution invariance (spatial=32/64/128 with same weights), finite outputs, doctest shape.
Checklist