From 45a46359c15b472609f92c7e07bd88597da1ebcf Mon Sep 17 00:00:00 2001
From: Gregory Wagner <gregory.leclaire.wagner@gmail.com>
Date: Wed, 12 Nov 2025 14:34:05 -0700
Subject: [PATCH 1/4] start adding oceananigans integration tests

---
 test/integration/Oceananigans/Project.toml |  5 +++
 test/integration/Oceananigans/runtests.jl  | 49 ++++++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 test/integration/Oceananigans/Project.toml
 create mode 100644 test/integration/Oceananigans/runtests.jl

diff --git a/test/integration/Oceananigans/Project.toml b/test/integration/Oceananigans/Project.toml
new file mode 100644
index 0000000000..4ab7af1152
--- /dev/null
+++ b/test/integration/Oceananigans/Project.toml
@@ -0,0 +1,5 @@
+[deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/integration/Oceananigans/runtests.jl b/test/integration/Oceananigans/runtests.jl
new file mode 100644
index 0000000000..d73d237222
--- /dev/null
+++ b/test/integration/Oceananigans/runtests.jl
@@ -0,0 +1,49 @@
+using Enzyme
+using GPUArraysCore: @allowscalar
+using Oceananigans
+using Test
+
+function kinetic_energy(model, ν, K, dt=1)
+    vitd = VerticallyImplicitTimeDiscretization()
+    new_closure = ScalarDiffusivity(vitd; ν)
+    tracer_names = keys(model.tracers)
+    new_closure = Oceananigans.TurbulenceClosures.with_tracers(tracer_names, new_closure)
+    model.closure = new_closure
+
+    for n = 1:10
+        time_step!(model, dt)
+    end
+
+    compute!(K)
+
+    return @allowscalar first(K)
+end
+
+@testset "Column model with ScalarDiffusivity" begin
+    ν = 1e-3
+    grid = RectilinearGrid(size=128, z=(-64, 64), topology=(Flat, Flat, Bounded))
+    vitd = VerticallyImplicitTimeDiscretization()
+    closure = ScalarDiffusivity(vitd; ν)
+    m = HydrostaticFreeSurfaceModel(; grid, closure, coriolis=FPlane(f=1e-4),
+        tracers=:b, buoyancy=BuoyancyTracer())
+
+    N² = 1e-6
+    bᵢ(z) = N² * z
+    uᵢ(z) = exp(-z^2 / 10)
+    set!(m, b=bᵢ, u=uᵢ)
+
+    u, v, w = m.velocities
+    ke = (u^2 + v^2) / 2
+    K = Field(Integral(ke))
+
+    dm = Enzyme.make_zero(m)
+    dK = Enzyme.make_zero(K)
+
+    dKdν = autodiff(set_strong_zero(Enzyme.ReverseWithPrimal),
+                    kinetic_energy, Active,
+                    Duplicated(m, dm),
+                    Active(ν),
+                    Duplicated(K, dK))
+
+    @test dKdν[1][1] != 0
+end
\ No newline at end of file

From e728633718deb1e5f2e095ec93f1df744631eb28 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mos=C3=A8=20Giordano?=
 <765740+giordano@users.noreply.github.com>
Date: Wed, 12 Nov 2025 21:42:10 +0000
Subject: [PATCH 2/4] Apply suggestions from code review

---
 test/integration/Oceananigans/Project.toml | 7 +++++++
 test/integration/Oceananigans/runtests.jl  | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/test/integration/Oceananigans/Project.toml b/test/integration/Oceananigans/Project.toml
index 4ab7af1152..adf453c5ca 100644
--- a/test/integration/Oceananigans/Project.toml
+++ b/test/integration/Oceananigans/Project.toml
@@ -3,3 +3,10 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[compat]
+GPUArraysCore = "=0.2.0"
+Oceananigans = "=0.101.0"
+
+[sources]
+Enzyme = { path = "../../.." }
diff --git a/test/integration/Oceananigans/runtests.jl b/test/integration/Oceananigans/runtests.jl
index d73d237222..dab12a362e 100644
--- a/test/integration/Oceananigans/runtests.jl
+++ b/test/integration/Oceananigans/runtests.jl
@@ -46,4 +46,4 @@ end
                     Duplicated(K, dK))
 
     @test dKdν[1][1] != 0
-end
\ No newline at end of file
+end

From d99fe5fb00fa27df6ed289404cb68f6602df404c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mos=C3=A8=20Giordano?=
 <765740+giordano@users.noreply.github.com>
Date: Wed, 12 Nov 2025 21:48:38 +0000
Subject: [PATCH 3/4] Add Oceananigans to integration workflow

---
 .github/workflows/Integration.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/Integration.yml b/.github/workflows/Integration.yml
index 496c6578a2..57c1c20c6d 100644
--- a/.github/workflows/Integration.yml
+++ b/.github/workflows/Integration.yml
@@ -53,6 +53,7 @@ jobs:
           - DynamicExpressions
           - SymbolicRegression
           - Lux
+          - Oceananigans
           - SciML
           - KernelAbstractions
           - MatrixAlgebraKit

From 762da990e041ee6550bf376a0d71adb750fae597 Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Fri, 19 Jun 2026 13:11:59 -0500
Subject: [PATCH 4/4] Update runtests.jl

---
 test/integration/Oceananigans/runtests.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/integration/Oceananigans/runtests.jl b/test/integration/Oceananigans/runtests.jl
index dab12a362e..45f388a40d 100644
--- a/test/integration/Oceananigans/runtests.jl
+++ b/test/integration/Oceananigans/runtests.jl
@@ -16,7 +16,8 @@ function kinetic_energy(model, ν, K, dt=1)
 
     compute!(K)
 
-    return @allowscalar first(K)
+    # TODO return @allowscalar first(K)
+    return first(K)
 end
 
 @testset "Column model with ScalarDiffusivity" begin