docs: burgers equation example using deeponets (#46)

* refactor: move tutorials to pre-built models

* docs: working burgers example

* fix: minor corrections

* docs: build docs on buildkite

* ci: dont fail CI if codecov fails

* docs: fix codecov badge

Author: avik-pal

.buildkite/documentation.yml

@@ -0,0 +1,29 @@
+steps:
+  - label: ":julia: Documentation"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    command: |
+      julia --project -e '
+        println("--- :julia: Instantiating project")
+        using Pkg
+        Pkg.instantiate()
+        Pkg.activate("docs")
+        Pkg.instantiate()
+        push!(LOAD_PATH, @__DIR__)
+
+        println("+++ :julia: Building documentation")
+        include("docs/make.jl")'
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      DATADEPS_ALWAYS_ACCEPT: true
+    if: build.message !~ /\[skip docs\]/ && !build.pull_request.draft
+    timeout_in_minutes: 1000
+
+env:
+  SECRET_CODECOV_TOKEN: "vn5M+4wSwUFje6fl6UB/Q/rTmLHu3OlCCMgoPOXPQHYpLZTLz2hOHsV44MadAnxw8MsNVxLKZlXBKqP3IydU9gUfV7QUBtnvbUmIvgUHbr+r0bVaIVVhw6cnd0s8/b+561nU483eRJd35bjYDOlO+V5eDxkbdh/0bzLefXNXy5+ALxsBYzsp75Sx/9nuREfRqWwU6S45mne2ZlwCDpZlFvBDXQ2ICKYXpA45MpxhW9RuqfpQdi6sSR6I/HdHkV2cuJO99dqqh8xfUy6vWPC/+HUVrn9ETsrXtayX1MX3McKj869htGICpR8vqd311HTONYVprH2AN1bJqr5MOIZ8Xg==;U2FsdGVkX1+W55pTI7zq+NwYrbK6Cgqe+Gp8wMCmXY+W10aXTB0bS6zshiDYSQ1Y3piT91xFyNhS+9AsajY0yQ=="
+  SECRET_DOCUMENTER_KEY: "hnN1QYai+CvEgtp2M9S9ZxK4xBthWAhx8snqR7MIKk3cqG7qHMsTd49KMSYWLXEsrsMbI2N4OLjNE3SgVZwyfuHop6x3ZqwVWD0OtldK6QvL/c+lmpzYMkIRCqKfXOKm9cQTWrr/eKzEh0uEUEdzCH7vrqLP7ywf5F9BMJbRmHW2f4cE1RKiG5cWVtgyO2OyRg7CxYrWgg0eKba+Vh8dLj1xL1Gs1XpLr7TDYRfSo7VMKYdkvB7Os7hmDN04YwzRFvGQBlroBaiC7KgbI5kZcL33R/whCHUYneSsXozGVyJGGGUbdxCaSslJ3p3Vb3vpnn+eAn0GyriWF/TaKIUlpQ==;U2FsdGVkX1/7Mqz8JnY+1EpErS6yTG4GfBJTNI+3WLtOuhyr+6SWZgIBCoitfIVSqvAws+RXVY09rG5Vsnvperr/lAzuBINsfOdIkSnS4HjfRHrKXGWzwzdTawPw1tn0wJa9h4RNahoBQk7Qs3clEb150AlomRJ//UAOW6UgD/nvr8TASjc5aHPkZodBogtMn93ti7hVRlUCNZAk58zpyXEsJorX9RT4g473y1aJp+CMd+Momo7/eljjwq3uCFyWpBnIwmipal4myUszDnSHuEQHcgh74ghJ1LkBjKcLUrRj06wKM/PfF23/P7h/KZCeryvghPSczKCJHx3wjH10DUuFK4maX9Yli/vg0SBq7/crWxNPXS5GmHYg/NBD2pgbe9qFvSCX4xJvkDl2ixHkNu5wQCCCTixtLfv00/XKWT5nh+PLSapQkY7ouesLFk8adpLAESPlX4pjSbljy0BrcHeFXFclFerq+Ia4StCLehWI4k2o/gfUEx7cnNN5ZKQpcYtQVOKRQriq8ri6dz+znQTmWqIM0xf+ZfWHy1OA3UXi269b0jdIlRQNtAjWhah9WWBaCiPY63/bgW54nIh1TVS1YaHzx2elWtFzHNpeTxIoibmltUivCE9r9Mk+YpfYUHJVUP9DwaLXcCiMZoLPNU4XWAzTejkpKJX/TdVFudLm2WNsCWi7nzjTfWSwlUOjrccLy9JVAx186gMlFhf5rHWeNFydBgEHBUig0FWkASwSPyv5iwD63NoL8UsTCRTNKMDtzn1XuCNxbmxgsbNO5bl+mcCDB/ZJuXexw3Ek04Ur4eZNyhf5qcn2cUzm76Jqi2vHSnY8JS9U9sSCdOd1gHFhxTOmO/OWh4jP5GKtWSWH23T2UOTS95daw1z2+GwLLP3R8chnFUCWvnJ+SUAHKoHYvbiT34gOyhcbnmoYTBCuvy4eghq1pVALUooPHOHzq2WsHAI07y6dUfbq55xaQZYiT++e+UisIID4mamITC5IFAQhkXPSAbVo56VrRCNQ+WnxLhzB6HNceAUdJCvH0IKcC9NqvfhCNioQ0EktfNla4G1BydJwgqexu8Tm4f5dl9e9V65hTn05yRYY2QN4hPCsw9Fr2NOrGfudCp5PEvA9TNuaobXvR5lSHPLL19ImbkCR1FGndrpwzWHR/0xjyS4lfdo3cW6SGVK33WZzOouM2ShGZc5Ma4yJdcPr1g0RsnWMGVX0A2gzJD9EL9RtNLgsDERnRLkEJgvhgVGXI+nmXadWzmn+plWgGEjSCLeZH1a4msFmcco3NFj0c6i5g+6cG09yILwhUTX+iegwm71r71pbPsMBk/hz/0ntVdklqFfCUZOC1n+MGJJmFfHmHAYHXSX5RmJhmvZG59g+oBEKxCh4VaIsSjLbopR4/TxS/uHdq3JJMOHXpgTQDack0qz3iEZFt+tKvKmah15PijdCv/Y8SHlN+txgSpW8+bQpxaMFq2li+/vApT18hBLjgGoMRv+I4ZmEmLUN1C+BzkCBTtErn4REQn4aKQmqlKC0jhQ5jatR5/Olpa1msiNjVyX0l4RBZFmYlpRY0KxKuVyguQtBF2qaxP2Lgrf2lp7tI7bK1z2MMBOxRTKHYs/Jt77mL/oeOOd2263EuZ2ku6K/iaaEY7jOCN7Du60UYVh0ePPorCTICHJ/arKLQ/uhh/d6Lo71bkj72B7CG75KGjZOg3+FRsFI0BLVjqvEJasR3apRCeipt3iMqmWc67wDp6DMXGCb/LlZh+oSfXtTY1VuMK4L4qcPl15Lf9H2xIt7N1YR6riq5R2EAokqQtazfxMfIk8+R25674O2Jksy0HnkIQ/H1/r2vaQYFzB0KVl671RdzHVYpfSMQQUcQtbzJHHsw5OIAPRwGd2RxFwW3GieKPYNS1u01TW+7z7Tq80qCOqewf+L+4ckbgYiiDDux+lzEK+PowtyOnQR1sjw7PE+BCPqF4Uc7b65EItFacseBD3mNVRXFUSG29rehgU4OrBjS7uUjn43QoMfonVCRhiLxydrfgGTdXW0g3qg3NYRsNeX3kP7b8MrzJ1zhPMZCca0YPbLeZ6Ab8lSc/GeJgrJ1m6ollE5yC2YisJMvNyJU525A7pReSk7u77nd1HTcddvngn8glJjcPn5/2aDuPuP4ZSEt67QLUaxVXce83XMY3FqiDngPhbJYSwGb9dlDizu9ph6a/SgZQyVpMBVL76RK9xropTUFDfDf5SOjukv5VWC2klIh8Fc4v9ur6otP5kx0MYV4ndlCoU6CyfEaQIQKXyZIIOWEPxyBvy8GnOSyh+5E82rJ50+PHkU9+iEJLBaZENnHlDJA30ANHSLF33t3CAxlqOvhFC5hx+BWCbK6R92mITBJOmo58+Fcz1e7u6SRLhOrp5rO+t3z6eB7SuVr0j16Tn4CJNGiYVGmhDjA1NfXoI7uEUkLVGuhra0eFzUCH8c+ot1pamb35E0knzaXy7azX5nuRXgt/491L0f9uc20ZqsrP63JtRQO574sbHxM6SnBWcIZSEEhAGRf4busL7DqjfHSFVx8TJH4FO0fCteCveJcxDoc94Wj7J+u5rUQBsRLDQ24XPjSgflqfhyU0lNKXuDWUSJPzqVuUM3fnYCsCdSd4/Y23t5k5dL68wOj3tBEunpGKltzG4zXjopay0VSkpty0EkaQbejvgXWnZv2G/SZlH/dW0fqp4a/M824kuzh9liF34fcV1uReoc+SKXs7iLYHgDXXajWy40HqvQR1NxtD17eIBvfhsAeKM3WR6YqudVkl4DYXU2hGM7L7znex2JAVs1thOtzBD6DbTKtlEciTs9O2B++22XnSmrusweKE68/8P0FN7DKA2oNQCig9J7BMT3pWA2jtoeq5sZU3vgvzlCLELefO4+YUWHY436znIp773oE8+IDdbxwLqKQfRI7zQAUGXNeUNEzAIla51zSARIFocslZfmF8ugx3NDBNqd6cjU+gk4qVAmUZW/Sw7SM+aPbC6hJPrDCfsMYhG7CNO5DZNBSO6Hhh0GSfl/zuB4Lzzg6xJuCfupg/A9DaUMfaIacYcsP7TjqVf07SbiPIXhlrFvuv5d43tkGBQdAQ/dNtpm2R5//u5Fi/mOlELiJBV5aLI2wZn2pGJRuOa6tb8GHSW3eLmv7RlrvBUEXutrfnevzUFu73Ue6jArVLhNQBhq64GhDEMQQvXH8QM/C0aTL8/IhFx80+kESUdnNa1OW7bNrBhOxTJJ4ho2JvenqVP44UQjJKOQGTjD5LSAqy9mkzK73lzRBVTNhpI2MTbKEnfXxN0wCce7cdpzJAUOTYF4fa4A9kkL3dqq/nhvkfzFY7lJbGI0J6hwMa9lWPwhB+wBU5hZu7s9eLDsfMoPE3vti8p8e91SgnPs+QtzPoe72GfBbJe8QQNnGeGW1aI2+qF82yarez1hTFaUPX1dgwSAKJS3/RSy1pMm1EPcM50P5At2ervJRJM5VCKrVJ8rvcFKuzfQfcZ/sqCW4l3E8ojJ7+7MAu4Ilfg0C6ZVY2S+rH52jIJtpTqbtGDaC0Zp7QTEbD5Eu1GkIxUmjpg23VGrNi4SnV+3ZbzOue8qTTstQdwzD2gDgb+zXUSkeGgZYYyEJQpGSjEJOPVi6CByqUtdK+Lfd19WRgSFJMmDfupQFUk7D19/d4K7/dj0NcuH+DQBfPvZl4cJkGqHeBXlK8uPN1A+q1S8D2QA7m4BzHAHz0Sy99KDq4spuFlm5t9oYAielqduZHieY4BugcQ1KcwR6pmc2ayEZvVRwfTNnxyodbOQCg4yh/pV4kX40G+tGnwtimaIVzyeIaZhL+iQ0M1+4KJCbKT2NGnk82hUKyrnBYSTPAG6vjRkg3P6KZj6tj4ABpxnymkZnMrUq2gpx/OlF0nl2CV0iCvj0RqZ5kX3NDVDtW8XmXh1sTATEGltfEclYoKMA/c47adE+rj16CiPwInnI/+JReNSjI14UoNn+27n8yQfxb3LbYkDSYkZcV7aysKPLQunMO0f0n6WUc2nprUIwQSO1Ccdhnio+YJs9JdvICN6JDpXfXM3eM6YGoSFv+XUownwfCW+fc1LMEXb4qjlUv7TIZKgXN5cEokKC2ZPpFJxF1TB2B7M3ib3a0orWuL8RsvexU7fpbUWo0+QI0CmiTqkmtsq1QeC7OCUU6Ocmz0iKUcYzcuLc9gfRyxXE8p7GqamK+okZWA0CJcMGvg0UXZcIaaHvUa9FuYLnjEOwvz1V92GGL0epIR9Vm4N6Aa1HzppbDuyY+donj0KBQkDxhix4+sF79LzdAJq1cvuxWW3km8f9CXYRdJ3oOY2nqwQXsi08K+yE+Zv2YVPd6Z1tJfk7Y+KiyhSIsKxDqIY1SaWwgRRYpi8c52BjTtnaIbYjFI="

.buildkite/pipeline.yml

@@ -19,8 +19,22 @@ steps:
                 agents:
                   queue: "juliagpu"
 
+            - path:
+                - "src/"
+                - "ext/"
+                - "test/"
+                - "docs/"
+                - "Project.toml"
+                - ".buildkite/"
+              config:
+                command: "buildkite-agent pipeline upload .buildkite/documentation.yml"
+                agents:
+                  queue: "juliagpu"
+
   - label: "Triggering Pipelines (Main Branch / Tag)"
     if: build.branch == "main" || build.tag != null
     agents:
       queue: "juliagpu"
-    command: "buildkite-agent pipeline upload .buildkite/testing.yml"
+    command: |
+      buildkite-agent pipeline upload .buildkite/testing.yml
+      buildkite-agent pipeline upload .buildkite/documentation.yml

.github/workflows/CI.yml

@@ -28,7 +28,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - "1"
+          - "1.10"
         os:
           - ubuntu-latest
           - macos-latest
@@ -56,7 +56,7 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           verbose: true
-          fail_ci_if_error: true
+          fail_ci_if_error: false
 
   downgrade:
     if: ${{ !contains(github.event.head_commit.message, '[skip tests]') && github.base_ref == github.event.repository.default_branch }}
@@ -65,7 +65,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        version: ["1"]
+        version: ["1.10"]
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -80,7 +80,7 @@ jobs:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           verbose: true
-          fail_ci_if_error: true
+          fail_ci_if_error: false
 
   invalidations:
     # Only run on PRs to the default branch.
@@ -113,5 +113,3 @@ jobs:
 
 env:
   BACKEND_GROUP: "CPU"
-  RETESTITEMS_NWORKERS: 4
-  RETESTITEMS_NWORKER_THREADS: 2

.github/workflows/Documentation.yml

@@ -3,4 +3,6 @@ Manifest-v*.toml
 .vscode
 wip
 examples
-docs/build
+docs/build
+
+.CondaPkg

.gitignore

@@ -1,7 +1,7 @@
 name = "NeuralOperators"
 uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 authors = ["Avik Pal <[email protected]>"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -11,6 +11,7 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
@@ -24,6 +25,7 @@ FFTW = "1.8"
 Lux = "1"
 LuxCore = "1"
 LuxLib = "1.2"
+MLDataDevices = "1.2.0"
 NNlib = "0.9.21"
 Random = "1.10"
 Static = "1.1.1"

Project.toml

@@ -3,7 +3,7 @@
 [![Join the chat at https://julialang.zulipchat.com #sciml-bridged](https://img.shields.io/static/v1?label=Zulip&message=chat&color=9558b2&labelColor=389826)](https://julialang.zulipchat.com/#narrow/stream/279055-sciml-bridged)
 [![Global Docs](https://img.shields.io/badge/docs-SciML-blue.svg)](https://docs.sciml.ai/NeuralOperators/stable/)
 
-[![codecov](https://codecov.io/gh/SciML/NeuralOperators.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/SciML/NeuralOperators.jl)
+[![codecov](https://codecov.io/gh/SciML/NeuralOperators.jl/branch/main/graph/badge.svg?token=wTS4cxrvB1)](https://codecov.io/gh/SciML/NeuralOperators.jl)
 [![Build Status](https://github.com/SciML/NeuralOperators.jl/workflows/CI/badge.svg)](https://github.com/SciML/NeuralOperators.jl/actions?query=workflow%3ACI)
 [![Build status](https://badge.buildkite.com/dd09599b08f61de1b5c7960aacd5390554c53e3b54f1407ba1.svg?branch=main)](https://buildkite.com/julialang/neuraloperators-dot-jl)
 

README.md

@@ -0,0 +1,2 @@
+[deps]
+gdown = ""

docs/CondaPkg.toml

@@ -1,15 +1,29 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
+DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 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"
+LuxCUDA = "0.3.3"
+MAT = "0.10.7"
+MLUtils = "0.4.4"
 NeuralOperators = "0.5"
 Optimisers = "0.3.3"
+Printf = "1.10"
+PythonCall = "0.9.23"
 Zygote = "0.6.71"

docs/Project.toml

@@ -1,9 +1,12 @@
 pages = [
     "NeuralOperators.jl" => "index.md",
+    "Pre-built Models" => [
+        "FNO" => "models/fno.md",
+        "DeepONet" => "models/deeponet.md",
+        "NOMAD" => "models/nomad.md"
+    ],
     "Tutorials" => [
-        "FNO" => "tutorials/fno.md",
-        "DeepONet" => "tutorials/deeponet.md",
-        "NOMAD" => "tutorials/nomad.md"
+        "Burgers Equation" => "tutorials/burgers.md"
     ],
     "API Reference" => "api.md"
 ]

docs/pages.jl

@@ -14,9 +14,9 @@ Pkg.add("NeuralOperators")
 
 Currently provided operator architectures are :
 
-  - [Fourier Neural Operators (FNOs)](tutorials/fno.md)
-  - [DeepONets](tutorials/deeponet.md)
-  - [Nonlinear Manifold Decoders for Operator Learning (NOMADs)](tutorials/nomad.md)
+  - [Fourier Neural Operators (FNOs)](models/fno.md)
+  - [DeepONets](models/deeponet.md)
+  - [Nonlinear Manifold Decoders for Operator Learning (NOMADs)](models/nomad.md)
 
 ## Reproducibility
 

docs/src/assets/Project.toml

@@ -0,0 +1,131 @@
+# Burgers Equation using DeepONet
+
+## Data Loading
+
+```@example burgers
+using DataDeps, MAT, MLUtils
+using PythonCall, CondaPkg # For `gdown`
+using Printf
+
+const gdown = pyimport("gdown")
+
+register(
+    DataDep(
+    "Burgers",
+    """
+    Burgers' equation dataset from
+    [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator)
+
+    mapping between initial conditions to the solutions at the last point of time \
+    evolution in some function space.
+
+    u(x,0) -> u(x, time_end):
+
+      * `a`: initial conditions u(x,0)
+      * `u`: solutions u(x,t_end)
+    """,
+    "https://drive.google.com/uc?id=16a8od4vidbiNR3WtaBPCSZ0T3moxjhYe",
+    "9cbbe5070556c777b1ba3bacd49da5c36ea8ed138ba51b6ee76a24b971066ecd";
+    fetch_method=(url, local_dir) -> begin
+        pyconvert(String, gdown.download(url, joinpath(local_dir, "Burgers_R10.zip")))
+    end,
+    post_fetch_method=unpack
+)
+)
+
+filepath = joinpath(datadep"Burgers", "burgers_data_R10.mat")
+
+const N = 2048
+const Δsamples = 2^3
+const grid_size = div(2^13, Δsamples)
+const T = Float32
+
+file = matopen(filepath)
+x_data = reshape(T.(collect(read(file, "a")[1:N, 1:Δsamples:end])), N, :, 1)
+y_data = reshape(T.(collect(read(file, "u")[1:N, 1:Δsamples:end])), N, :, 1)
+close(file)
+
+x_data = permutedims(x_data, (2, 1, 3))
+grid = reshape(T.(collect(range(0, 1; length=grid_size)')), :, grid_size, 1)
+```
+
+## Model
+
+```@example burgers
+using Lux, NeuralOperators, Optimisers, Zygote, Random
+using LuxCUDA
+
+const cdev = cpu_device()
+const gdev = gpu_device()
+
+deeponet = DeepONet(;
+    branch=(size(x_data, 1), ntuple(Returns(32), 5)...),
+    trunk=(size(grid, 1), ntuple(Returns(32), 5)...),
+    branch_activation=tanh,
+    trunk_activation=tanh
+)
+ps, st = Lux.setup(Random.default_rng(), deeponet) |> gdev;
+```
+
+## Training
+
+```@example burgers
+x_data_dev = x_data |> gdev
+y_data_dev = y_data |> gdev
+grid_dev = grid |> gdev
+
+function loss_function(model, ps, st, ((v, y), u))
+    û, stₙ = model((v, y), ps, st)
+    return MAELoss()(û, u), stₙ, (;)
+end
+
+function train_model!(model, ps, st, data; epochs=5000)
+    train_state = Training.TrainState(model, ps, st, Adam(0.0001f0))
+
+    for epoch in 1:epochs
+        _, loss, _, train_state = Training.single_train_step!(
+            AutoZygote(), loss_function, data, train_state)
+
+        if epoch % 25 == 1 || epoch == epochs
+            @printf("Epoch %d: loss = %.6e\n", epoch, loss)
+        end
+    end
+
+    return train_state.parameters, train_state.states
+end
+
+ps_trained, st_trained = train_model!(
+    deeponet, ps, st, ((x_data_dev, grid_dev), y_data_dev))
+```
+
+## Plotting
+
+```@example burgers
+using CairoMakie
+
+pred = first(deeponet((x_data_dev, grid_dev), ps_trained, st_trained)) |> cdev
+
+begin
+    fig = Figure(; size=(1024, 1024))
+
+    axs = [Axis(fig[i, j]) for i in 1:4, j in 1:4]
+    for i in 1:4, j in 1:4
+        idx = i + (j - 1) * 4
+        ax = axs[i, j]
+        l1 = lines!(ax, vec(grid), pred[idx, :, 1])
+        l2 = lines!(ax, vec(grid), y_data[idx, :, 1])
+
+        i == 4 && (ax.xlabel = "x")
+        j == 1 && (ax.ylabel = "u(x)")
+
+        if i == 1 && j == 1
+            axislegend(ax, [l1, l2], ["Predictions", "Ground Truth"])
+        end
+    end
+    linkaxes!(axs...)
+
+    fig[0, :] = Label(fig, "Burgers Equation using DeepONet"; tellwidth=false, font=:bold)
+
+    fig
+end
+```

docs/src/index.md

@@ -1,7 +1,7 @@
 module NeuralOperators
 
 using ArgCheck: @argcheck
-using ChainRulesCore: @non_differentiable
+using ChainRulesCore: ChainRulesCore, NoTangent, @non_differentiable
 using ConcreteStructs: @concrete
 using FFTW: FFTW, irfft, rfft
 using Random: Random, AbstractRNG
@@ -10,9 +10,11 @@ using Static: StaticBool, False, True, known, static
 using Lux
 using LuxCore: LuxCore, AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer
 using LuxLib: batched_matmul
-using NNlib: NNlib, batched_adjoint
+using MLDataDevices: AbstractDevice, AbstractGPUDevice
+using NNlib: NNlib
 
 const BoolLike = Union{Bool, StaticBool, Val{true}, Val{false}}
+const CRC = ChainRulesCore
 
 include("utils.jl")