test: remove Zygote completely

.github/workflows/CI.yml

@@ -26,9 +26,10 @@ jobs:
       matrix:
         version:
           - "1"
-          - "lts"
+          # - "lts"
         os:
           - ubuntu-latest
+          - windows-latest
           - macos-latest
     uses: LuxDL/Lux.jl/.github/workflows/CommonCI.yml@main
     with:
@@ -42,35 +43,3 @@ jobs:
       julia_version: "1.11"
       downgrade_testing: true
       project: "."
-
-  invalidations:
-    # Only run on PRs to the default branch.
-    # In the PR trigger above branches can be specified only explicitly whereas this check should work for master, main, or any other default branch
-    if: github.base_ref == github.event.repository.default_branch
-    runs-on: ubuntu-latest
-    steps:
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: "1"
-      - uses: actions/checkout@v6
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-invalidations@v1
-        id: invs_pr
-
-      - uses: actions/checkout@v6
-        with:
-          ref: ${{ github.event.repository.default_branch }}
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-invalidations@v1
-        id: invs_default
-
-      - name: Report invalidation counts
-        run: |
-          echo "Invalidations on default branch: ${{ steps.invs_default.outputs.total }} (${{ steps.invs_default.outputs.deps }} via deps)" >> $GITHUB_STEP_SUMMARY
-          echo "This branch: ${{ steps.invs_pr.outputs.total }} (${{ steps.invs_pr.outputs.deps }} via deps)" >> $GITHUB_STEP_SUMMARY
-      - name: Check if the PR does increase number of invalidations
-        if: steps.invs_pr.outputs.total > steps.invs_default.outputs.total
-        run: exit 1
-
-env:
-  BACKEND_GROUP: "CPU"

.github/workflows/QualityCheck.yml

@@ -3,14 +3,6 @@ name: Code Quality Check
 on: [pull_request]
 
 jobs:
-  code-style:
-    name: Format Suggestions
-    runs-on: ubuntu-latest
-    steps:
-      - uses: julia-actions/julia-format@v4
-        with:
-          version: "1"
-
   typos-check:
     name: Spell Check with Typos
     runs-on: ubuntu-latest

Project.toml

@@ -1,7 +1,7 @@
 name = "NeuralOperators"
 uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "0.6.2"
+version = "0.6.3"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -16,10 +16,14 @@ WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 [compat]
 AbstractFFTs = "1.5.0"
 ConcreteStructs = "0.2.3"
-Lux = "1.13"
-LuxCore = "1.2"
-LuxLib = "1.8"
-NNlib = "0.9.30"
+Lux = "1.31.3"
+LuxCore = "1.5.1"
+LuxLib = "1.15.1"
+NNlib = "0.9.33"
 Random = "1.10"
-WeightInitializers = "1"
+WeightInitializers = "1.3"
 julia = "1.10"
+
+[workspace]
+projects = ["test", "docs"]
+

docs/Project.toml

@@ -6,6 +6,7 @@ CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Git = "d7ba0133-e1db-5d97-8f8c-041e4b3a1eb2"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
@@ -14,20 +15,21 @@ Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 
-[sources.NeuralOperators]
-path = ".."
+[sources]
+NeuralOperators = {path = ".."}
 
 [compat]
-AlgebraOfGraphics = "0.10.7, 0.11"
+AlgebraOfGraphics = "0.10.7, 0.11, 0.12"
 CSV = "0.10"
 CairoMakie = "0.13, 0.14, 0.15"
 CondaPkg = "0.2.23"
 DataDeps = "0.7.13"
 DataFrames = "1"
 Documenter = "1.7.0"
 Lux = "1"
-MAT = "0.10.7"
+MAT = "0.10.7, 0.11"
 MLUtils = "0.4.4"
 NeuralOperators = "0.6"
 Optimisers = "0.4"

docs/make.jl

@@ -1,7 +1,14 @@
 using Documenter, NeuralOperators
 
-cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml"; force = true)
-cp("./docs/Project.toml", "./docs/src/assets/Project.toml"; force = true)
+if isfile("Manifest.toml")
+    cp("./Manifest.toml", "./docs/src/assets/Manifest.toml"; force = true)
+elseif isfile("../Manifest.toml")
+    cp("../Manifest.toml", "./docs/src/assets/Manifest.toml"; force = true)
+end
+
+if isfile("Project.toml")
+    cp("./Project.toml", "./docs/src/assets/Project.toml"; force = true)
+end
 
 ENV["GKSwstype"] = "100"
 ENV["DATADEPS_ALWAYS_ACCEPT"] = true

docs/src/tutorials/burgers_deeponet.md

@@ -3,9 +3,8 @@
 ## Data Loading
 
 ```@example burgers
-using DataDeps, MAT, MLUtils
+using MAT, MLUtils, Printf, DataDeps
 using PythonCall, CondaPkg # For `gdown`
-using Printf
 
 const gdown = pyimport("gdown")
 

docs/src/tutorials/double_pendulum.md

@@ -3,37 +3,27 @@
 ## Data Loading
 
 ```julia
-using DataDeps, CSV, MLUtils, DataFrames
+using CSV, MLUtils, DataFrames, Git, Scratch
 using Printf
 
-register(
-    DataDep(
-        "DoublePendulumChaotic",
-        """
-        Dataset was generated on the basis of 21 individual runs of a double pendulum.
-        Each of the recorded sequences lasted around 40s and consisted of around 17500 frames.
-
-          * `x_red`: Horizontal pixel coordinate of the red point (the central pivot to the
-            first pendulum)
-          * `y_red`: Vertical pixel coordinate of the red point (the central pivot to the first
-            pendulum)
-          * `x_green`: Horizontal pixel coordinate of the green point (the first pendulum)
-          * `y_green`: Vertical pixel coordinate of the green point (the first pendulum)
-          * `x_blue`: Horizontal pixel coordinate of the blue point (the second pendulum)
-          * `y_blue`: Vertical pixel coordinate of the blue point (the second pendulum)
-
-        Page: https://developer.ibm.com/exchanges/data/all/double-pendulum-chaotic/
-        """,
-        "https://dax-cdn.cdn.appdomain.cloud/dax-double-pendulum-chaotic/2.0.1/double-pendulum-chaotic.tar.gz",
-        "4ca743b4b783094693d313ebedc2e8e53cf29821ee8b20abd99f8fb4c0866f8d";
-        post_fetch_method=unpack,
-    ),
-)
+function get_dataset_path()
+    scratch_path = @get_scratch!("double_pendulum")
+    if !isdir(joinpath(scratch_path, "data", "ibm-double-pendulum"))
+        run(Git.git(
+            [
+                "clone",
+                "https://github.com/erwincoumans/tiny-differentiable-simulator",
+                "$(scratch_path)",
+            ]
+        ))
+    end
+    return joinpath(scratch_path, "data", "ibm-double-pendulum")
+end
 
 function get_data(; i=0, n=-1)
-    data_path = joinpath(datadep"DoublePendulumChaotic", "original", "dpc_dataset_csv")
+    data_path = get_dataset_path()
     df = CSV.read(
-        joinpath(data_path, "$i.csv"),
+        joinpath(data_path, "original", "dpc_dataset_csv", "$i.csv"),
         DataFrame;
         header=[:x_red, :y_red, :x_green, :y_green, :x_blue, :y_blue],
     )
@@ -112,14 +102,9 @@ function train_model!(model, ps, st, trainloader, testloader; epochs=20)
     train_state = Training.TrainState(model, ps, st, AdamW(; eta=3.0f-4, lambda=1.0f-5))
 
     (xtest, ytest) = first(testloader)
-    prediction_loss_compiled = Reactant.with_config(;
-        convolution_precision=PrecisionConfig.HIGH,
-        dot_general_precision=PrecisionConfig.HIGH,
-    ) do
-        @compile prediction_loss(
-            model, xtest, train_state.parameters, train_state.states, ytest
-        )
-    end
+    prediction_loss_compiled = @compile prediction_loss(
+        model, xtest, train_state.parameters, train_state.states, ytest
+    )
 
     for epoch in 1:epochs
         for data in trainloader
@@ -167,15 +152,8 @@ n = 5
 inferenced_data = x_data[:, :, :, 1:1]
 for i in 1:n
     input_data = inferenced_data[:, :, :, i:i] |> xdev
-    prediction = first(
-        Reactant.with_config(;
-            convolution_precision=PrecisionConfig.HIGH,
-            dot_general_precision=PrecisionConfig.HIGH,
-        ) do
-            @jit fno(input_data, ps_trained, st_trained)
-        end,
-    )
-    inferenced_data = cat(inferenced_data, cdev(prediction); dims=4)
+    prediction = @jit fno(input_data, ps_trained, st_trained)
+    inferenced_data = cat(inferenced_data, cdev(prediction[1]); dims=4)
 end
 inferenced_data = cat([inferenced_data[:, :, :, i] for i in 1:n]...; dims=1)[:, :, 1]'
 

docs/src/tutorials/poisson_equation.md

@@ -94,9 +94,9 @@ function predict(model, f_input, x_input, ps, st)
     return vec(pred .* max_u)
 end
 
-compiled_predict_fn = Reactant.with_config(; dot_general_precision=PrecisionConfig.HIGH) do
-    @compile predict(deeponet, f_data[:, 1:1], x_data, ps_trained, st_trained)
-end
+compiled_predict_fn = @compile predict(
+    deeponet, f_data[:, 1:1], x_data, ps_trained, st_trained
+)
 
 # Testing and visualization
 begin

test/Project.toml

@@ -1,39 +1,39 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
-FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
+FastTransforms = "057dd010-8810-581a-b7be-e3fc3b93f78c"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
-Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[sources]
+NeuralOperators = {path = ".."}
+Reactant = {url = "https://github.com/EnzymeAD/Reactant.jl", rev = "ap/fix_complex_fd"}
 
 [compat]
 Aqua = "0.8.7"
 Documenter = "1.5.0"
 Enzyme = "0.13.48"
 ExplicitImports = "1.9.0"
-FFTW = "1.9.0"
-Hwloc = "3.2"
 Lux = "1"
 LuxCore = "1"
 LuxLib = "1.2"
 LuxTestUtils = "1.1.2, 2"
 Optimisers = "0.4"
+ParallelTestRunner = "2.1"
 Random = "1.10"
-ReTestItems = "1.24.0"
-Reactant = "0.2.130"
-Reexport = "1.2.2"
+Reactant = "0.2.203"
 StableRNGs = "1.0.2"
 Test = "1.10"
-Zygote = "0.7"

test/doctests.jl

@@ -0,0 +1,11 @@
+using NeuralOperators, Test, Documenter, FastTransforms
+
+@testset "Doctests: Quality Assurance" begin
+    DocMeta.setdocmeta!(
+        NeuralOperators,
+        :DocTestSetup,
+        :(using Lux, NeuralOperators, Random);
+        recursive = true,
+    )
+    doctest(NeuralOperators; manual = false)
+end