feat: compile neural operators using Reactant

Project.toml

@@ -1,7 +1,7 @@
 name = "NeuralOperators"
 uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 authors = ["Avik Pal <[email protected]>"]
-version = "0.5.2"
+version = "0.5.3"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -22,10 +22,10 @@ ArgCheck = "2.3"
 ChainRulesCore = "1.24"
 ConcreteStructs = "0.2.3"
 FFTW = "1.8"
-Lux = "1"
-LuxCore = "1"
-LuxLib = "1.2"
-MLDataDevices = "1.2.0"
+Lux = "1.2.1"
+LuxCore = "1.1"
+LuxLib = "1.3.7"
+MLDataDevices = "1.5"
 NNlib = "0.9.21"
 Random = "1.10"
 Static = "1.1.1"

docs/Project.toml

@@ -3,6 +3,7 @@ CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
@@ -11,19 +12,20 @@ NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 CairoMakie = "0.12.11"
 CondaPkg = "0.2.23"
 DataDeps = "0.7.13"
 Documenter = "1.7.0"
-Lux = "1"
+Lux = "1.2.1"
 LuxCUDA = "0.3.3"
 MAT = "0.10.7"
 MLUtils = "0.4.4"
 NeuralOperators = "0.5"
-Optimisers = "0.3.3"
+Optimisers = "0.3.3, 0.4"
 Printf = "1.10"
 PythonCall = "0.9.23"
 Zygote = "0.6.71"

docs/pages.jl

@@ -6,6 +6,7 @@ pages = [
         "NOMAD" => "models/nomad.md"
     ],
     "Tutorials" => [
+        "XLA Compilation" => "tutorials/xla_compilation.md",
         "Burgers Equation" => "tutorials/burgers.md"
     ],
     "API Reference" => "api.md"

docs/src/tutorials/xla_compilation.md

@@ -0,0 +1,78 @@
+# Compiling NeuralOperators.jl using Reactant.jl
+
+```@example xla_compilation
+using NeuralOperators, Lux, Random, Enzyme, Reactant
+
+function sumabs2first(model, ps, st, x)
+    z, _ = model(x, ps, st)
+    return sum(abs2, z)
+end
+
+dev = reactant_device()
+```
+
+## Compiling DeepONet
+
+```@example xla_compilation
+deeponet = DeepONet()
+ps, st = Lux.setup(Random.default_rng(), deeponet) |> dev;
+
+u = rand(Float32, 64, 1024) |> dev;
+y = rand(Float32, 1, 128, 1024) |> dev;
+nothing # hide
+
+deeponet_compiled = @compile deeponet((u, y), ps, st)
+deeponet_compiled((u, y), ps, st)[1]
+```
+
+Computing the gradient of the DeepONet model.
+
+```@example xla_compilation
+function ∇deeponet(model, ps, st, (u, y))
+    dps = Enzyme.make_zero(ps)
+    Enzyme.autodiff(
+        Enzyme.Reverse,
+        sumabs2first,
+        Const(model),
+        Duplicated(ps, dps),
+        Const(st),
+        Const((u, y))
+    )
+    return dps
+end
+
+∇deeponet_compiled = @compile ∇deeponet(deeponet, ps, st, (u, y))
+∇deeponet_compiled(deeponet, ps, st, (u, y))
+```
+
+## Compiling FourierNeuralOperator
+
+```@example xla_compilation
+fno = FourierNeuralOperator()
+ps, st = Lux.setup(Random.default_rng(), fno) |> dev;
+
+x = rand(Float32, 2, 1024, 5) |> dev;
+
+fno_compiled = @compile fno(x, ps, st)
+fno_compiled(x, ps, st)[1]
+```
+
+Computing the gradient of the FourierNeuralOperator model.
+
+```@example xla_compilation
+function ∇fno(model, ps, st, x)
+    dps = Enzyme.make_zero(ps)
+    Enzyme.autodiff(
+        Enzyme.Reverse,
+        sumabs2first,
+        Const(model),
+        Duplicated(ps, dps),
+        Const(st),
+        Const(x)
+    )
+    return dps
+end
+
+∇fno_compiled = @compile ∇fno(fno, ps, st, x)
+∇fno_compiled(fno, ps, st, x)
+```

src/layers.jl

@@ -76,7 +76,7 @@ function operator_conv(x, tform::AbstractTransform, weights)
     x_p = apply_pattern(x_tr, weights)
 
     pad_dims = size(x_t)[1:(end - 2)] .- size(x_p)[1:(end - 2)]
-    x_padded = NNlib.pad_constant(x_p, expand_pad_dims(pad_dims), false;
+    x_padded = NNlib.pad_zeros(x_p, expand_pad_dims(pad_dims);
         dims=ntuple(identity, ndims(x_p) - 2))::typeof(x_p)
 
     return inverse(tform, x_padded, size(x))

test/Project.toml

@@ -30,7 +30,7 @@ LuxCore = "1"
 LuxLib = "1.2"
 LuxTestUtils = "1.1.2"
 MLDataDevices = "1"
-Optimisers = "0.3.3"
+Optimisers = "0.3.3, 0.4"
 Pkg = "1.10"
 Preferences = "1"
 Random = "1.10"