diff --git a/ext/OceananigansEnzymeExt.jl b/ext/OceananigansEnzymeExt.jl
index e1b7d74a1f..5d03233ac2 100644
--- a/ext/OceananigansEnzymeExt.jl
+++ b/ext/OceananigansEnzymeExt.jl
@@ -7,7 +7,6 @@ using Oceananigans.Utils: contiguousrange
 
 using KernelAbstractions
 
-#isdefined(Base, :get_extension) ? (import EnzymeCore) : (import ..EnzymeCore)
 isdefined(Base, :get_extension) ? (import Enzyme) : (import ..Enzyme)
 
 using Enzyme: EnzymeCore
@@ -17,6 +16,7 @@ EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Base.:(==)), ::Oceananigans.Abstr
 EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.validate_grid), x...) = nothing
 EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.metric_function), x...) = nothing
 EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.flatten_reduced_dimensions), x...) = nothing
+EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.prettytime), x...) = nothing
 EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.Grids.total_size), x...) = nothing
 EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_size_and_offset), x...) = nothing
 @inline EnzymeCore.EnzymeRules.inactive_type(v::Type{Oceananigans.Utils.KernelParameters}) = true
@@ -24,313 +24,6 @@ EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_
 @inline batch(::Val{1}, ::Type{T}) where T = T
 @inline batch(::Val{N}, ::Type{T}) where {T, N} = NTuple{N, T}
 
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  func::EnzymeCore.Const{Type{Field}},
-#                                                  ::Type{<:EnzymeCore.Annotation{RT}},
-#                                                  loc::Union{EnzymeCore.Const{<:Tuple},
-#                                                  EnzymeCore.Duplicated{<:Tuple}},
-#                                                  grid::EnzymeCore.Annotation{<:Oceananigans.Grids.AbstractGrid},
-#                                                  T::EnzymeCore.Const{<:DataType}; kw...) where RT
-# 
-#     primal = if EnzymeCore.EnzymeRules.needs_primal(config)
-#         func.val(loc.val, grid.val, T.val; kw...)
-#     else
-#         nothing
-#     end
-# 
-#     if haskey(kw, :a)
-#         # copy zeroing
-#         kw[:data] = copy(kw[:data])
-#     end
-# 
-#     shadow = if EnzymeCore.EnzymeRules.width(config) == 1
-#         func.val(loc.val, grid.val, T.val; kw...)
-#     else
-#         ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
-#             Base.@_inline_meta
-#             func.val(loc.val, grid.val, T.val; kw...)
-#         end
-#     end
-# 
-#     P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing
-#     B = batch(Val(EnzymeCore.EnzymeRules.width(config)), RT)
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, Nothing}(primal, shadow, nothing)
-# end
-# 
-# #####
-# ##### Field
-# #####
-# 
-# function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1},
-#                                         func::EnzymeCore.Const{Type{Field}},
-#                                         ::RT,
-#                                         tape,
-#                                         loc::Union{EnzymeCore.Const{<:Tuple}, EnzymeCore.Duplicated{<:Tuple}},
-#                                         grid::EnzymeCore.Const{<:Oceananigans.Grids.AbstractGrid},
-#                                         T::EnzymeCore.Const{<:DataType}; kw...) where RT
-#     return (nothing, nothing, nothing)
-# end
-# 
-# function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1},
-#                                         func::EnzymeCore.Const{Type{Field}},
-#                                         ::RT,
-#                                         tape,
-#                                         loc::Union{EnzymeCore.Const{<:Tuple}, EnzymeCore.Duplicated{<:Tuple}},
-#                                         grid::EnzymeCore.Active{<:Oceananigans.Grids.AbstractGrid},
-#                                         T::EnzymeCore.Const{<:DataType}; kw...) where RT
-#     return (nothing, EnzymeCore.make_ero(grid), nothing)
-# end
-
-#####
-##### FunctionField
-#####
-
-# @inline FunctionField(L::Tuple, func, grid) = FunctionField{L[1], L[2], L[3]}(func, grid)
-
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  enzyme_func::Union{EnzymeCore.Const{<:Type{<:FunctionField}}, EnzymeCore.Const{Type{FT2}}},
-#                                                  ::Type{<:EnzymeCore.Annotation{RT}},
-#                                                  function_field_func,
-#                                                  grid;
-#                                                  clock = nothing,
-#                                                  parameters = nothing) where {RT, FT2 <: FunctionField}
-# 
-#     FunctionFieldType = enzyme_func.val     
-# 
-#     primal = if EnzymeCore.EnzymeRules.needs_primal(config)
-#         FunctionFieldType(function_field_func.val, grid.val; clock, parameters)
-#     else
-#         nothing
-#     end
-# 
-#     # function_field_func can be Active, Const (inactive), Duplicated (active but mutable)
-#     function_field_is_active = function_field_func isa Active
-#     # @show function_field_func
-# 
-#     # Support batched differentiation!
-#     config_width = EnzymeCore.EnzymeRules.width(config)
-# 
-#     dactives = if function_field_is_active
-#         if config_width == 1
-#             Ref(EnzymeCore.make_zero(function_field_func.val))
-#         else
-#             ntuple(Val(config_width)) do i
-#                 Base.@_inline_meta
-#                 Ref(EnzymeCore.make_zero(function_field_func.val))
-#             end
-#         end
-#     else
-#         nothing
-#     end
-# 
-#     shadow = if config_width == 1
-#         dfunction_field_func = if function_field_is_active
-#             dactives[]
-#         else
-#             function_field_func.dval
-#         end
-# 
-#         FunctionFieldType(dfunction_field_func, grid.val; clock, parameters)
-#     else
-#   	    ntuple(Val(config_width)) do i
-#   		    Base.@_inline_meta
-# 
-#             dfunction_field_func = if function_field_is_active
-#                 dactives[i][]
-#             else
-#                 function_field_func.dval[i]
-#             end
-# 
-#             FunctionFieldType(dfunction_field_func, grid.val; clock, parameters)
-#   	    end
-#     end
-# 
-#     P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing
-#     B = batch(Val(EnzymeCore.EnzymeRules.width(config)), RT)
-#     D = typeof(dactives)
-# 
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, D}(primal, shadow, dactives)
-# end
-# 
-# function EnzymeCore.EnzymeRules.reverse(config,
-#                                         enzyme_func::Union{EnzymeCore.Const{<:Type{<:FunctionField}}, EnzymeCore.Const{Type{FT2}}},
-#                                         ::RT,
-#                                         tape,
-#                                         function_field_func,
-#                                         grid;
-#                                         clock = nothing,
-#                                         parameters = nothing) where {RT, FT2 <: FunctionField}
-# 
-#     dactives = if function_field_func isa Active
-#         if EnzymeCore.EnzymeRules.width(config) == 1
-#             tape[]
-#         else
-#             ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
-#                 Base.@_inline_meta
-#                 tape[i][]
-#             end
-#         end
-#     else
-#         nothing
-#     end
-# 
-#     # return (dactives, grid (nothing))
-#     return (dactives, nothing)
-# end
-
-#####
-##### launch!
-#####
-
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  func::EnzymeCore.Const{typeof(Oceananigans.Models.flattened_unique_values)},
-#                                                  ::Type{<:EnzymeCore.Annotation{RT}},
-#                                                  a) where RT
-# 
-#     sprimal = if EnzymeCore.EnzymeRules.needs_primal(config) || EnzymeCore.EnzymeRules.needs_shadow(config)
-#         func.val(a.val)
-#     else
-#         nothing
-#     end
-# 
-#     shadow = if EnzymeCore.EnzymeRules.needs_shadow(config)
-#         if EnzymeCore.EnzymeRules.width(config) == 1
-#             (typeof(a) <: Const) ? EnzymeCore.make_zero(sprimal)::RT : func.val(a.dval)
-#         else
-#             ntuple(Val(EnzymeCore.EnzymeRules.width(config))) do i
-#                 Base.@_inline_meta
-#                 (typeof(a) <: Const) ? EnzymeCore.make_zero(sprimal)::RT : func.val(a.dval[i])
-#             end
-#         end
-#     else
-#         nothing
-#     end
-# 
-#     primal = if EnzymeCore.EnzymeRules.needs_primal(config)
-#         sprimal
-#     else
-#         nothing
-#     end
-# 
-#     P = EnzymeCore.EnzymeRules.needs_primal(config) ? RT : Nothing
-#     B = EnzymeCore.EnzymeRules.needs_primal(config) ? batch(Val(EnzymeCore.EnzymeRules.width(config)), RT) : Nothing
-# 
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{P, B, Nothing}(primal, shadow, nothing)
-# end
-# 
-# function EnzymeCore.EnzymeRules.reverse(config,
-#                                         func::EnzymeCore.Const{typeof(Oceananigans.Models.flattened_unique_values)},
-#                                          ::Type{<:EnzymeCore.Annotation{RT}},
-#                                          tape,
-#                                          a) where RT
-# 
-#   return (nothing,)
-# end
-
-#####
-##### launch!
-#####
-
-# function EnzymeCore.EnzymeRules.augmented_primal(config,
-#                                                  func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)},
-#                                                  ::Type{EnzymeCore.Const{Nothing}},
-#                                                  arch,
-#                                                  grid,
-#                                                  workspec,
-#                                                  kernel!,
-#                                                  kernel_args::Vararg{Any,N};
-#                                                  include_right_boundaries = false,
-#                                                  reduced_dimensions = (),
-#                                                  location = nothing,
-#                                                  active_cells_map = nothing,
-#                                                  kwargs...) where N
-# 
-# 
-#     workgroup, worksize = Oceananigans.Utils.work_layout(grid.val, workspec.val;
-#                                                          include_right_boundaries,
-#                                                          reduced_dimensions,
-#                                                          location)
-# 
-#     offset = Oceananigans.Utils.offsets(workspec.val)
-# 
-#     if !isnothing(active_cells_map) 
-#         workgroup, worksize = Oceananigans.Utils.active_cells_work_layout(workgroup, worksize, active_cells_map, grid.val) 
-#         offset = nothing
-#     end
-# 
-#     if worksize != 0
-#       
-#       # We can only launch offset kernels with Static sizes!!!!
-# 
-#       if isnothing(offset)
-#           loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), workgroup, worksize)
-#           dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.dval(Oceananigans.Architectures.device(arch.val), workgroup, worksize)
-#       else
-#           loop! = kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset))) 
-#           dloop! = (typeof(kernel!) <: EnzymeCore.Const) ? nothing : kernel!.val(Oceananigans.Architectures.device(arch.val), KernelAbstractions.StaticSize(workgroup), Oceananigans.Utils.OffsetStaticSize(contiguousrange(worksize, offset)))
-#       end
-# 
-#       @debug "Launching kernel $kernel! with worksize $worksize and offsets $offset from $workspec.val"
-# 
-# 
-#       duploop = (typeof(kernel!) <: EnzymeCore.Const) ? EnzymeCore.Const(loop!) : EnzymeCore.Duplicated(loop!, dloop!)
-# 
-#       config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}()
-#       subtape = EnzymeCore.EnzymeRules.augmented_primal(config2, duploop, EnzymeCore.Const{Nothing}, kernel_args...).tape
-# 
-#       tape = (duploop, subtape)
-#     else
-#       tape = nothing
-#     end
-# 
-#     return EnzymeCore.EnzymeRules.AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, tape)
-# end
-# 
-# @inline arg_elem_type(::Type{T}, ::Val{i}) where {T<:Tuple, i} = eltype(T.parameters[i])
-# 
-# function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1},
-#                                                 func::EnzymeCore.Const{typeof(Oceananigans.Utils.launch!)},
-#                                                  ::Type{EnzymeCore.Const{Nothing}},
-#                                                  tape,
-#                                                  arch,
-#                                                  grid,
-#                                                  workspec,
-#                                                  kernel!,
-#                                                  kernel_args::Vararg{Any,N};
-#                                                  include_right_boundaries = false,
-#                                                  reduced_dimensions = (),
-#                                                  location = nothing,
-#                                                  active_cells_map = nothing,
-#                                                  kwargs...) where N
-# 
-#   subrets = if tape !== nothing
-#     duploop, subtape = tape
-#     config2 = EnzymeCore.EnzymeRules.Config{#=needsprimal=#false, #=needsshadow=#false, #=width=#EnzymeCore.EnzymeRules.width(config), EnzymeCore.EnzymeRules.overwritten(config)[5:end]}()
-#     EnzymeCore.EnzymeRules.reverse(config2, duploop, EnzymeCore.Const{Nothing}, subtape, kernel_args...)
-#   else
-#     res2 = ntuple(Val(N)) do i
-#       Base.@_inline_meta
-#       if kernel_args[i] isa Active
-#         EnzymeCore.make_zero(kernel_args[i].val)
-#       else
-#         nothing
-#       end
-#     end
-#   end
-# 
-#   subrets2 =  ntuple(Val(N)) do i
-#       Base.@_inline_meta
-#       if kernel_args[i] isa Active
-#         subrets[i]::arg_elem_type(typeof(kernel_args), Val(i))
-#       else
-#         nothing
-#       end
-#     end
-# 
-#   return (nothing, nothing, nothing, nothing, subrets2...)
-# 
-# end
-
 #####
 ##### update_model_field_time_series!
 #####
diff --git a/src/BuoyancyModels/seawater_buoyancy.jl b/src/BuoyancyModels/seawater_buoyancy.jl
index ba7fa68390..7472fc5615 100644
--- a/src/BuoyancyModels/seawater_buoyancy.jl
+++ b/src/BuoyancyModels/seawater_buoyancy.jl
@@ -9,15 +9,15 @@ temperature and salinity are active, or of type `FT` if temperature
 or salinity are constant, respectively.
 """
 struct SeawaterBuoyancy{FT, EOS, T, S} <: AbstractBuoyancyModel{EOS}
-             equation_of_state :: EOS
+    equation_of_state :: EOS
     gravitational_acceleration :: FT
-          constant_temperature :: T
-             constant_salinity :: S
+    constant_temperature :: T
+    constant_salinity :: S
 end
 
 required_tracers(::SeawaterBuoyancy) = (:T, :S)
-required_tracers(::SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number}) where {FT, EOS} = (:T,) # active temperature only
-required_tracers(::SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing}) where {FT, EOS} = (:S,) # active salinity only
+required_tracers(::SeawaterBuoyancy{FT, EOS, <:Nothing, <:Number}) where {FT, EOS} = tuple(:T) # active temperature only
+required_tracers(::SeawaterBuoyancy{FT, EOS, <:Number, <:Nothing}) where {FT, EOS} = tuple(:S) # active salinity only
 
 Base.nameof(::Type{SeawaterBuoyancy}) = "SeawaterBuoyancy"
 Base.summary(b::SeawaterBuoyancy) = string(nameof(typeof(b)), " with g=", prettysummary(b.gravitational_acceleration),
diff --git a/src/Fields/field.jl b/src/Fields/field.jl
index ee806db849..c17c120623 100644
--- a/src/Fields/field.jl
+++ b/src/Fields/field.jl
@@ -404,8 +404,8 @@ Base.checkbounds(f::Field, I...) = Base.checkbounds(f.data, I...)
 @propagate_inbounds Base.lastindex(f::Field) = lastindex(f.data)
 @propagate_inbounds Base.lastindex(f::Field, dim) = lastindex(f.data, dim)
 
-Base.fill!(f::Field, val) = fill!(parent(f), val)
-Base.parent(f::Field) = parent(f.data)
+@inline Base.fill!(f::Field, val) = fill!(parent(f), val)
+@inline Base.parent(f::Field) = parent(f.data)
 Adapt.adapt_structure(to, f::Field) = Adapt.adapt(to, f.data)
 
 total_size(f::Field) = total_size(f.grid, location(f), f.indices)
diff --git a/src/Fields/field_tuples.jl b/src/Fields/field_tuples.jl
index debac6aab9..39d3122d53 100644
--- a/src/Fields/field_tuples.jl
+++ b/src/Fields/field_tuples.jl
@@ -55,45 +55,37 @@ Fill halo regions for all `fields`. The algorithm:
 """
 function fill_halo_regions!(maybe_nested_tuple::Union{NamedTuple, Tuple}, args...; kwargs...)
     flattened = flattened_unique_values(maybe_nested_tuple)
-
-    # Sort fields into ReducedField and Field with non-nothing boundary conditions
-    fields_with_bcs = filter(f -> !isnothing(boundary_conditions(f)), flattened)
-    reduced_fields  = filter(f -> f isa ReducedField, fields_with_bcs)
+    return tupled_fill_halo_regions!(flattened, args...; kwargs...)
+end
     
-    for field in reduced_fields
-        fill_halo_regions!(field, args...; kwargs...)
+function tupled_fill_halo_regions!(fields, args...; kwargs...)
+
+    ordinary_fields = Field[]
+    for f in fields
+        if !isnothing(boundary_conditions(f))
+            if f isa ReducedField || !(f isa FullField)
+                # Windowed and reduced fields
+                fill_halo_regions!(f, args...; kwargs...)
+            else
+                push!(ordinary_fields, f)
+            end
+        end
     end
 
-    # MultiRegion fields are considered windowed_fields (indices isa MultiRegionObject))
-    windowed_fields = filter(f -> !(f isa FullField), fields_with_bcs)
-    ordinary_fields = filter(f -> (f isa FullField) && !(f isa ReducedField), fields_with_bcs)
-
-    # Fill halo regions for reduced and windowed fields
-    for field in windowed_fields
-        fill_halo_regions!(field, args...; kwargs...)
-    end
+    ordinary_fields = tuple(ordinary_fields...)
 
-    # Fill the rest
-    if !isempty(ordinary_fields)
+    if !isempty(ordinary_fields) # ie not reduced, and with default_indices
         grid = first(ordinary_fields).grid
-        tupled_fill_halo_regions!(ordinary_fields, grid, args...; kwargs...)
+        fill_halo_regions!(map(data, ordinary_fields),
+                           map(boundary_conditions, ordinary_fields),
+                           default_indices(3),
+                           map(instantiated_location, ordinary_fields),
+                           grid, args...; kwargs...)
     end
 
     return nothing
 end
 
-function tupled_fill_halo_regions!(fields, grid, args...; kwargs...)
-
-    # We cannot group windowed fields together, the indices must be (:, :, :)!
-    indices = default_indices(3)        
-
-    return fill_halo_regions!(map(data, fields),
-                              map(boundary_conditions, fields),
-                              indices,
-                              map(instantiated_location, fields),
-                              grid, args...; kwargs...)
-end
-
 #####
 ##### Tracer names
 #####
diff --git a/src/Fields/set!.jl b/src/Fields/set!.jl
index e4b8a29cb6..a9cff7a0b0 100644
--- a/src/Fields/set!.jl
+++ b/src/Fields/set!.jl
@@ -34,6 +34,11 @@ set!(u::Field, f::Function) = set_to_function!(u, f)
 set!(u::Field, a::Union{Array, CuArray, OffsetArray}) = set_to_array!(u, a)
 set!(u::Field, v::Field) = set_to_field!(u, v)
 
+function set!(u::Field, a::Number)
+    fill!(parent(u), a)
+    return u # return u, not parent(u), for type-stability
+end
+
 function set!(u::Field, v)
     u .= v # fallback
     return u
diff --git a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl
index 56fd8d0049..7d44fe9b44 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl
@@ -32,7 +32,6 @@ on_architecture(to, free_surface::ExplicitFreeSurface) =
 function materialize_free_surface(free_surface::ExplicitFreeSurface{Nothing}, velocities, grid)
     η = free_surface_displacement_field(velocities, free_surface, grid)
     g = convert(eltype(grid), free_surface.gravitational_acceleration)
-
     return ExplicitFreeSurface(η, g)
 end
 
@@ -62,10 +61,16 @@ explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) =
 ##### Kernel
 #####
 
-@kernel function _explicit_ab2_step_free_surface!(η, Δt, χ::FT, Gηⁿ, Gη⁻, Nz) where FT
+@kernel function _explicit_ab2_step_free_surface!(η, Δt, χ, Gηⁿ, Gη⁻, Nz)
     i, j = @index(Global, NTuple)
+    FT = typeof(χ)
+    one_point_five = convert(FT, 1.5)
+    oh_point_five = convert(FT, 0.5)
+    not_euler = χ != convert(FT, -0.5)
 
     @inbounds begin
-        η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1])
+        Gη = (one_point_five + χ) * Gηⁿ[i, j, Nz+1] - (oh_point_five  + χ) * Gη⁻[i, j, Nz+1] * not_euler
+        η[i, j, Nz+1] += Δt * Gη
     end
 end
+
diff --git a/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl
index f216d312c7..a9a2c26006 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl
@@ -33,7 +33,7 @@ T₀[T₀ .< 0.5] .= 0
 
 set!(model, u=u₀, v=v₀, T=T₀)
 
-model.velocities.u
+@model.velocities.u
 
 # output
 
@@ -61,7 +61,8 @@ model.velocities.u
     end
 
     initialize!(model)
-    update_state!(model; compute_tendencies = false)
+    update_state!(model; compute_tendencies=false)
 
     return nothing
 end
+
diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl
index d4d3dec8d8..7d0d854a18 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl
@@ -63,14 +63,13 @@ compute_free_surface_tendency!(::SingleColumnGrid, ::SplitExplicitFreeSurfaceHFS
 
 function update_state!(model::HydrostaticFreeSurfaceModel, grid::SingleColumnGrid, callbacks; compute_tendencies = true)
 
-    fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    tupled_fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
 
     # Compute auxiliaries
     compute_auxiliary_fields!(model.auxiliary_fields)
 
     # Calculate diffusivities
     compute_diffusivities!(model.diffusivity_fields, model.closure, model)
-
     fill_halo_regions!(model.diffusivity_fields, model.clock, fields(model))
 
     for callback in callbacks
diff --git a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl
index 98413d5eb8..e3c224e949 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface_kernels.jl
@@ -168,7 +168,6 @@ end
 end
 
 function initialize_free_surface_state!(state, η, timestepper)
-
     parent(state.U) .= parent(state.U̅)
     parent(state.V) .= parent(state.V̅)
 
@@ -213,11 +212,9 @@ function barotropic_split_explicit_corrector!(u, v, free_surface, grid)
     Hᶠᶜ, Hᶜᶠ   = free_surface.auxiliary.Hᶠᶜ, free_surface.auxiliary.Hᶜᶠ
     arch       = architecture(grid)
 
-
-    # take out "bad" barotropic mode,
+    # Remove incorrect barotropic mode and add correction
     # !!!! reusing U and V for this storage since last timestep doesn't matter
     compute_barotropic_mode!(U, V, grid, u, v)
-    # add in "good" barotropic mode
     launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!,
             u, v, U̅, V̅, U, V, Hᶠᶜ, Hᶜᶠ, grid)
 
@@ -228,14 +225,14 @@ end
 Explicitly step forward η in substeps.
 """
 ab2_step_free_surface!(free_surface::SplitExplicitFreeSurface, model, Δt, χ) =
-    split_explicit_free_surface_step!(free_surface, model, Δt, χ)
+    split_explicit_free_surface_step!(free_surface, model, Δt)
 
 function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities)
     @apply_regionally compute_barotropic_mode!(sefs.state.U̅, sefs.state.V̅, grid, velocities.u, velocities.v)
     fill_halo_regions!((sefs.state.U̅, sefs.state.V̅, sefs.η))
 end
 
-function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurface, model, Δt, χ)
+function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurface, model, Δt)
 
     # Note: free_surface.η.grid != model.grid for DistributedSplitExplicitFreeSurface
     # since halo_size(free_surface.η.grid) != halo_size(model.grid)
@@ -247,16 +244,16 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac
     # Calculate the substepping parameterers
     settings = free_surface.settings
     Nsubsteps = calculate_substeps(settings.substepping, Δt)
-
-    # barotropic time step as fraction of baroclinic step and averaging weights
-    fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps)
+    
+    # Barotropic time step as fraction of baroclinic step and averaging weights
+    fractional_Δt, weights = calculate_adaptive_settings(settings.substepping, Nsubsteps) 
     Nsubsteps = length(weights)
 
-    # barotropic time step in seconds
-    Δτᴮ = fractional_Δt * Δt
-
-    # reset free surface averages
-    @apply_regionally begin
+    # Barotropic time step in time units
+    Δτᴮ = fractional_Δt * Δt  
+    
+    # Reset free surface averages
+    @apply_regionally begin 
         initialize_free_surface_state!(free_surface.state, free_surface.η, settings.timestepper)
 
         # Solve for the free surface at tⁿ⁺¹
@@ -269,8 +266,8 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac
     fields_to_fill = (free_surface.state.U̅, free_surface.state.V̅)
     fill_halo_regions!(fields_to_fill; async = true)
 
-    # Preparing velocities for the barotropic correction
-    @apply_regionally begin
+    # Prepare velocities for the barotropic correction
+    @apply_regionally begin 
         mask_immersed_field!(model.velocities.u)
         mask_immersed_field!(model.velocities.v)
     end
@@ -350,40 +347,55 @@ function iterate_split_explicit!(free_surface, grid, Δτᴮ, weights, ::Val{Nsu
 end
 
 # Calculate RHS for the barotropic time step.
-@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+@kernel function _compute_integrated_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     i, j  = @index(Global, NTuple)
     k_top = grid.Nz + 1
 
-    @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ)
-    @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+    @inbounds begin
+        Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u_tendency_increment(i, j, 1, grid, Gu⁻, Guⁿ, χ)
+        Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v_tendency_increment(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
 
-    for k in 2:grid.Nz
-        @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ)
-        @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+        for k in 2:grid.Nz	
+            Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u_tendency_increment(i, j, k, grid, Gu⁻, Guⁿ, χ)
+            Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v_tendency_increment(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+        end	
     end
 end
 
 # Calculate RHS for the barotropic time step.q
-@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+@kernel function _compute_integrated_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     idx = @index(Global, Linear)
     i, j = active_linear_index_to_tuple(idx, active_cells_map)
     k_top = grid.Nz+1
 
-    @inbounds Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ)
-    @inbounds Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+    @inbounds begin
+        Gᵁ[i, j, k_top-1] = Δzᶠᶜᶜ(i, j, 1, grid) * u_tendency_increment(i, j, 1, grid, Gu⁻, Guⁿ, χ)
+        Gⱽ[i, j, k_top-1] = Δzᶜᶠᶜ(i, j, 1, grid) * v_tendency_increment(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+
+        for k in 2:grid.Nz	
+            Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * u_tendency_increment(i, j, k, grid, Gu⁻, Guⁿ, χ)
+            Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * v_tendency_increment(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+        end	
 
-    for k in 2:grid.Nz
-        @inbounds Gᵁ[i, j, k_top-1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ)
-        @inbounds Gⱽ[i, j, k_top-1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ)
     end
 end
 
-@inline ab2_step_Gu(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT =
-    @inbounds ifelse(peripheral_node(i, j, k, grid, f, c, c), zero(grid), (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ))
-
-@inline ab2_step_Gv(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT =
-    @inbounds ifelse(peripheral_node(i, j, k, grid, c, f, c), zero(grid), (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ))
+@inline function u_tendency_increment(i, j, k, grid, G⁻, Gⁿ, χ)
+    FT = typeof(χ)
+    peripheral = peripheral_node(i, j, k, grid, f, c, c)
+    not_euler = χ != convert(FT, -0.5)
+    ΔGu = @inbounds (convert(FT, 1.5) + χ) * Gⁿ[i, j, k] - (convert(FT, 0.5) + χ) * G⁻[i, j, k] * not_euler
+    return ifelse(peripheral, zero(grid), ΔGu)
+end
 
+@inline function v_tendency_increment(i, j, k, grid, G⁻, Gⁿ, χ)
+    FT = typeof(χ)
+    peripheral = peripheral_node(i, j, k, grid, c, f, c)
+    not_euler = χ != convert(FT, -0.5)
+    ΔGv = @inbounds (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - (convert(FT, 0.5) + χ) * G⁻[i, j, k] * not_euler
+    return ifelse(peripheral, zero(grid), ΔGv)
+end
+                     
 # Setting up the RHS for the barotropic step (tendencies of the barotropic velocity components)
 # This function is called after `calculate_tendency` and before `ab2_step_velocities!`
 function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ)
@@ -399,17 +411,15 @@ function setup_free_surface!(model, free_surface::SplitExplicitFreeSurface, χ)
     @apply_regionally setup_split_explicit_tendency!(auxiliary, model.grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
 
     fields_to_fill = (auxiliary.Gᵁ, auxiliary.Gⱽ)
-    fill_halo_regions!(fields_to_fill; async = true)
+    fill_halo_regions!(fields_to_fill; async=true)
 
     return nothing
 end
 
 @inline function setup_split_explicit_tendency!(auxiliary, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
     active_cells_map = retrieve_surface_active_cells_map(grid)
-
-    launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, auxiliary.Gᵁ, auxiliary.Gⱽ, grid,
-            active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ; active_cells_map)
-
+    kernel_args = (auxiliary.Gᵁ, auxiliary.Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+    launch!(architecture(grid), grid, :xy, _compute_integrated_tendencies!, kernel_args...; active_cells_map)
     return nothing
 end
 
diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl
index cf6b20eade..f8ca935c96 100644
--- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl
+++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl
@@ -8,7 +8,7 @@ using Oceananigans.TurbulenceClosures: compute_diffusivities!
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node
 using Oceananigans.Models: update_model_field_time_series!
 using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!, p_kernel_parameters
-using Oceananigans.Fields: replace_horizontal_vector_halos!
+using Oceananigans.Fields: replace_horizontal_vector_halos!, tupled_fill_halo_regions!
 
 import Oceananigans.Models.NonhydrostaticModels: compute_auxiliaries!
 import Oceananigans.TimeSteppers: update_state!
@@ -38,7 +38,8 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp
     # Update the boundary conditions
     @apply_regionally update_boundary_condition!(fields(model), model)
 
-    fill_halo_regions!(prognostic_fields(model), model.clock, fields(model); async = true)
+    tupled_fill_halo_regions!(prognostic_fields(model), model.clock, fields(model), async=true)
+
     @apply_regionally replace_horizontal_vector_halos!(model.velocities, model.grid)
     @apply_regionally compute_auxiliaries!(model)
 
diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl
index 3b73d85702..3c8e4dbb6b 100644
--- a/src/Simulations/simulation.jl
+++ b/src/Simulations/simulation.jl
@@ -144,7 +144,7 @@ iteration(sim::Simulation) = iteration(sim.model)
 
 Return `sim.model.clock.time` as a prettily formatted string."
 """
-prettytime(sim::Simulation, longform=true) = prettytime(time(sim))
+prettytime(sim::Simulation, longform=true) = prettytime(time(sim); longform)
 
 """
     run_wall_time(sim::Simulation)
diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl
index ed9f55750b..3471132ec3 100644
--- a/src/TimeSteppers/quasi_adams_bashforth_2.jl
+++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl
@@ -79,40 +79,27 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt
     Δt == 0 && @warn "Δt == 0 may cause model blowup!"
 
     # Be paranoid and update state at iteration 0
-    model.clock.iteration == 0 && update_state!(model, callbacks)
-
-    ab2_timestepper = model.timestepper
+    model.clock.iteration == 0 && update_state!(model, callbacks; compute_tendencies=true)
 
-    # Change the default χ if necessary, which occurs if:
+    # Take an euler step if:
     #   * We detect that the time-step size has changed.
     #   * We detect that this is the "first" time-step, which means we
     #     need to take an euler step. Note that model.clock.last_Δt is
     #     initialized as Inf
     #   * The user has passed euler=true to time_step!
     euler = euler || (Δt != model.clock.last_Δt)
-    
+    euler && @debug "Taking a forward Euler step."
+
     # If euler, then set χ = -0.5
     minus_point_five = convert(eltype(model.grid), -0.5)
+    ab2_timestepper = model.timestepper
     χ = ifelse(euler, minus_point_five, ab2_timestepper.χ)
-
-    # Set time-stepper χ (this is used in ab2_step!, but may also be used elsewhere)
     χ₀ = ab2_timestepper.χ # Save initial value
     ab2_timestepper.χ = χ
 
-    # Ensure zeroing out all previous tendency fields to avoid errors in
-    # case G⁻ includes NaNs. See https://github.com/CliMA/Oceananigans.jl/issues/2259
-    if euler
-        @debug "Taking a forward Euler step."
-        for field in ab2_timestepper.G⁻
-            !isnothing(field) && @apply_regionally fill!(field, 0)
-        end
-    end
+    # Full step for tracers, fractional step for velocities.
+    ab2_step!(model, Δt)
 
-    # Be paranoid and update state at iteration 0
-    model.clock.iteration == 0 && update_state!(model, callbacks; compute_tendencies=true)
-    
-    ab2_step!(model, Δt) # full step for tracers, fractional step for velocities.
-    
     tick!(model.clock, Δt)
     model.clock.last_Δt = Δt
     model.clock.last_stage_Δt = Δt # just one stage
@@ -125,7 +112,7 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt
 
     # Return χ to initial value
     ab2_timestepper.χ = χ₀
-    
+
     return nothing
 end
 
@@ -141,7 +128,6 @@ end
 
 """ Generic implementation. """
 function ab2_step!(model, Δt)
-
     grid = model.grid
     arch = architecture(grid)
     model_fields = prognostic_fields(model)
@@ -175,11 +161,16 @@ Time step velocity fields via the 2nd-order quasi Adams-Bashforth method
 @kernel function ab2_step_field!(u, Δt, χ, Gⁿ, G⁻)
     i, j, k = @index(Global, NTuple)
 
-    FT = eltype(χ)
+    FT = typeof(χ)
+    Δt = convert(FT, Δt)
     one_point_five = convert(FT, 1.5)
     oh_point_five  = convert(FT, 0.5)
+    not_euler = χ != convert(FT, -0.5) # use to prevent corruption by leftover NaNs in G⁻
 
-    @inbounds u[i, j, k] += convert(FT, Δt) * ((one_point_five + χ) * Gⁿ[i, j, k] - (oh_point_five + χ) * G⁻[i, j, k])
+    @inbounds begin
+        Gu = (one_point_five + χ) * Gⁿ[i, j, k] - (oh_point_five + χ) * G⁻[i, j, k] * not_euler
+        u[i, j, k] += Δt * Gu
+    end
 end
 
 @kernel ab2_step_field!(::FunctionField, Δt, χ, Gⁿ, G⁻) = nothing
diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl
index 77d0d465fb..acf28a0a87 100644
--- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl
+++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl
@@ -1,5 +1,6 @@
 module TKEBasedVerticalDiffusivities
 
+using Adapt
 using Adapt
 using CUDA
 using KernelAbstractions: @kernel, @index
@@ -18,6 +19,8 @@ using Oceananigans.BoundaryConditions: default_prognostic_bc, DefaultBoundaryCon
 using Oceananigans.BoundaryConditions: BoundaryCondition, FieldBoundaryConditions
 using Oceananigans.BoundaryConditions: DiscreteBoundaryFunction, FluxBoundaryCondition
 using Oceananigans.BuoyancyModels: ∂z_b, top_buoyancy_flux
+using Oceananigans.BuoyancyModels: BuoyancyTracer, SeawaterBuoyancy
+using Oceananigans.BuoyancyModels: TemperatureSeawaterBuoyancy, SalinitySeawaterBuoyancy
 using Oceananigans.Grids: inactive_cell
 
 using Oceananigans.TurbulenceClosures:
@@ -27,7 +30,7 @@ using Oceananigans.TurbulenceClosures:
     VerticallyImplicitTimeDiscretization,
     VerticalFormulation
     
-import Oceananigans.BoundaryConditions: getbc
+import Oceananigans.BoundaryConditions: getbc, fill_halo_regions!
 import Oceananigans.Utils: with_tracers
 import Oceananigans.TurbulenceClosures:
     validate_closure,
@@ -159,8 +162,14 @@ function get_time_step(closure_array::AbstractArray)
     return get_time_step(closure)
 end
 
-include("tke_top_boundary_condition.jl")
+get_top_tracer_bcs(::Nothing, tracers) = NamedTuple()
+get_top_tracer_bcs(::BuoyancyTracer, tracers) = (; b=tracers.b.boundary_conditions.top)
+get_top_tracer_bcs(::SeawaterBuoyancy, tracers) = (T = tracers.T.boundary_conditions.top,
+                                                   S = tracers.S.boundary_conditions.top)
+get_top_tracer_bcs(::TemperatureSeawaterBuoyancy, tracers) = (; T = tracers.T.boundary_conditions.top)
+get_top_tracer_bcs(::SalinitySeawaterBuoyancy, tracers)    = (; S = tracers.S.boundary_conditions.top)
 
+include("tke_top_boundary_condition.jl")
 include("catke_vertical_diffusivity.jl")
 include("catke_mixing_length.jl")
 include("catke_equation.jl")
diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl
index 4847a3802a..9bcf287483 100644
--- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl
+++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl
@@ -1,3 +1,5 @@
+using Oceananigans.BuoyancyModels: BuoyancyTracer, SeawaterBuoyancy
+
 struct CATKEVerticalDiffusivity{TD, CL, FT, DT, TKE} <: AbstractScalarDiffusivity{TD, VerticalFormulation, 2}
     mixing_length :: CL
     turbulent_kinetic_energy_equation :: TKE
@@ -144,6 +146,36 @@ catke_first(catke1::FlavorOfCATKE, catke2::FlavorOfCATKE) = error("Can't have tw
 ##### Diffusivities and diffusivity fields utilities
 #####
 
+struct CATKEDiffusivityFields{K, L, J, T, U, KC, LC}
+    κu :: K
+    κc :: K
+    κe :: K
+    Le :: L
+    Jᵇ :: J
+    previous_compute_time :: T
+    previous_velocities :: U
+    _tupled_tracer_diffusivities :: KC
+    _tupled_implicit_linear_coefficients :: LC
+end
+
+Adapt.adapt_structure(to, catke_diffusivity_fields::CATKEDiffusivityFields) =
+    CATKEDiffusivityFields(adapt(to, catke_diffusivity_fields.κu),
+                           adapt(to, catke_diffusivity_fields.κc),
+                           adapt(to, catke_diffusivity_fields.κe),
+                           adapt(to, catke_diffusivity_fields.Le),
+                           adapt(to, catke_diffusivity_fields.Jᵇ),
+                           previous_compute_time[],
+                           adapt(to, previous_velocities),
+                           adapt(to, _tupled_tracer_diffusivities),
+                           adapt(to, _tupled_implicit_linear_coefficients))
+
+function fill_halo_regions!(catke_diffusivity_fields::CATKEDiffusivityFields, args...; kw...)
+    fill_halo_regions!(catke_diffusivity_fields.κu, args...; kw...)
+    fill_halo_regions!(catke_diffusivity_fields.κc, args...; kw...)
+    fill_halo_regions!(catke_diffusivity_fields.κe, args...; kw...)
+    return nothing
+end
+
 function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE)
 
     default_diffusivity_bcs = (κu = FieldBoundaryConditions(grid, (Center, Center, Face)),
@@ -169,23 +201,22 @@ function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfCATKE)
     _tupled_tracer_diffusivities         = NamedTuple(name => name === :e ? κe : κc          for name in tracer_names)
     _tupled_implicit_linear_coefficients = NamedTuple(name => name === :e ? Le : ZeroField() for name in tracer_names)
 
-    return (; κu, κc, κe, Le, Jᵇ,
-            previous_compute_time, previous_velocities,
-            _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients)
-end
+    return CATKEDiffusivityFields(κu, κc, κe, Le, Jᵇ,
+                                  previous_compute_time, previous_velocities,
+                                  _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients)
+end        
 
 @inline viscosity_location(::FlavorOfCATKE) = (c, c, f)
 @inline diffusivity_location(::FlavorOfCATKE) = (c, c, f)
 
 function compute_diffusivities!(diffusivities, closure::FlavorOfCATKE, model; parameters = :xyz)
-
     arch = model.architecture
     grid = model.grid
     velocities = model.velocities
     tracers = model.tracers
     buoyancy = model.buoyancy
     clock = model.clock
-    top_tracer_bcs = NamedTuple(c => tracers[c].boundary_conditions.top for c in propertynames(tracers))
+    top_tracer_bcs = get_top_tracer_bcs(model.buoyancy.model, tracers)
     Δt = model.clock.time - diffusivities.previous_compute_time[]
     diffusivities.previous_compute_time[] = model.clock.time
 
diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl
index d94035e507..cf8eb52ba5 100644
--- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl
+++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl
@@ -149,6 +149,38 @@ end
 ##### Diffusivities and diffusivity fields utilities
 #####
 
+struct TKEDissipationDiffusivityFields{K, L, U, KC, LC}
+    κu :: K
+    κc :: K
+    κe :: K
+    κϵ :: K
+    Le :: L
+    Lϵ :: L
+    previous_velocities :: U
+    _tupled_tracer_diffusivities :: KC
+    _tupled_implicit_linear_coefficients :: LC
+end
+
+Adapt.adapt_structure(to, tke_dissipation_diffusivity_fields::TKEDissipationDiffusivityFields) =
+    TKEDissipationDiffusivityFields(adapt(to, tke_dissipation_diffusivity_fields.κu),
+                                    adapt(to, tke_dissipation_diffusivity_fields.κc),
+                                    adapt(to, tke_dissipation_diffusivity_fields.κe),
+                                    adapt(to, tke_dissipation_diffusivity_fields.κϵ),
+                                    adapt(to, tke_dissipation_diffusivity_fields.Le),
+                                    adapt(to, tke_dissipation_diffusivity_fields.Lϵ),
+                                    adapt(to, previous_velocities),
+                                    adapt(to, _tupled_tracer_diffusivities),
+                                    adapt(to, _tupled_implicit_linear_coefficients))
+
+function fill_halo_regions!(catke_diffusivity_fields::TKEDissipationDiffusivityFields, args...; kw...)
+    fields_with_halos_to_fill = (tke_dissipation_diffusivity_fields.κu,
+                                 tke_dissipation_diffusivity_fields.κc,
+                                 tke_dissipation_diffusivity_fields.κe,
+                                 tke_dissipation_diffusivity_fields.κϵ)
+
+    return fill_halo_regions!(fields_with_halos_to_fill, args...; kw...)
+end
+
 function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfTD)
 
     default_diffusivity_bcs = (κu = FieldBoundaryConditions(grid, (Center, Center, Face)),
@@ -184,8 +216,10 @@ function DiffusivityFields(grid, tracer_names, bcs, closure::FlavorOfTD)
     _tupled_implicit_linear_coefficients = NamedTuple(name => _tupled_implicit_linear_coefficients[name]
                                                       for name in tracer_names)
 
-    return (; κu, κc, κe, κϵ, Le, Lϵ, previous_velocities,
-            _tupled_tracer_diffusivities, _tupled_implicit_linear_coefficients)
+    return TKEDissipationDiffusivityFields(κu, κc, κe, κϵ, Le, Lϵ,
+                                           previous_velocities,
+                                           _tupled_tracer_diffusivities,
+                                           _tupled_implicit_linear_coefficients)
 end        
 
 @inline viscosity_location(::FlavorOfTD) = (c, c, f)
diff --git a/src/Utils/prettytime.jl b/src/Utils/prettytime.jl
index eb7fa14154..ad1ae172dd 100644
--- a/src/Utils/prettytime.jl
+++ b/src/Utils/prettytime.jl
@@ -19,48 +19,49 @@ minutes, and hours.
 """
 function prettytime(t, longform=true)
     # Modified from: https://github.com/JuliaCI/BenchmarkTools.jl/blob/master/src/trials.jl
-    
-    # Some shortcuts
-    s = longform ? "seconds" : "s" 
-    iszero(t) && return "0 $s"
-    t < 1e-9 && return @sprintf("%.3e %s", t, s) # yah that's small
-
-    t = maybe_int(t)
     value, units = prettytimeunits(t, longform)
 
-    if isinteger(value)
-        return @sprintf("%d %s", value, units)
+    if iszero(value)
+        msg = "0 $units"
+    elseif value < 1e-9 
+        msg = @sprintf("%.3e %s", value, units) # yah that's small
+    elseif isinteger(value)
+        msg = @sprintf("%d %s", value, units)
     else
-        return @sprintf("%.3f %s", value, units)
+        msg = @sprintf("%.3f %s", value, units)
     end
+
+    return msg::String
 end
 
-function prettytimeunits(t, longform=true)
-    t < 1e-9 && return t, "" # just _forget_ picoseconds!
-    t < 1e-6 && return t * 1e9, "ns"
-    t < 1e-3 && return t * 1e6, "μs"
-    t < 1    && return t * 1e3, "ms"
-    if t < minute
+function prettytimeunits(t::T, longform=true) where T
+    if t < 1e-9 # just _forget_ picoseconds!
         value = t
-        !longform && return value, "s"
-        units = value == 1 ? "second" : "seconds"
-        return value, units
+        units = !longform ? "s" : "seconds"
+    elseif t < 1e-6
+        value = t * 1e9
+        units = "ns"
+    elseif t < 1e-3
+        value = t * 1e6
+        units = "μs"
+    elseif t < 1
+        value = t * 1e3
+        units = "ms"
+    elseif t < minute
+        value = t
+        units = !longform ? "s" : value==1 ? "second" : "seconds"
     elseif t < hour
-        value = maybe_int(t / minute)
-        !longform && return value, "m"
-        units = value == 1 ? "minute" : "minutes"
-        return value, units
+        value = t / minute
+        units = !longform ? "m" : value==1 ? "minute" : "minutes"
     elseif t < day
-        value = maybe_int(t / hour)
-        units = value == 1 ? (longform ? "hour" : "hr") :
-                             (longform ? "hours" : "hrs")
-        return value, units
+        value = t / hour
+        units = !longform ? "hr" : value==1 ? "hour" : "hours"
     else
-        value = maybe_int(t / day)
-        !longform && return value, "d"
-        units = value == 1 ? "day" : "days"
-        return value, units
+        value = t / day
+        units = !longform ? "d" : value==1 ? "day" : "days"
     end
+
+    return convert(T, value), units::String
 end
 
 prettytime(dt::AbstractTime) = "$dt"
diff --git a/test/test_enzyme.jl b/test/test_enzyme.jl
index cc86656889..f2e4864ea1 100644
--- a/test/test_enzyme.jl
+++ b/test/test_enzyme.jl
@@ -1,5 +1,7 @@
 include("dependencies_for_runtests.jl")
 
+using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity
+
 # Required presently
 Enzyme.API.looseTypeAnalysis!(true)
 Enzyme.API.maxtypeoffset!(2032)
@@ -7,6 +9,7 @@ Enzyme.API.maxtypeoffset!(2032)
 # OceananigansLogger doesn't work here -- not sure why
 Logging.global_logger(TestLogger())
 
+#=
 f(grid) = CenterField(grid)
 const maximum_diffusivity = 100
 
@@ -106,15 +109,10 @@ function set_initial_condition_via_launch!(model_tracer, amplitude)
 end
 
 @testset "Enzyme + Oceananigans Initialization Broadcast Kernel" begin
-
-    Nx = Ny = 64
-    Nz = 8
-
     x = y = (-π, π)
     z = (-0.5, 0.5)
     topology = (Periodic, Periodic, Bounded)
-
-    grid = RectilinearGrid(size=(Nx, Ny, Nz); x, y, z, topology)
+    grid = RectilinearGrid(size=(64, 64, 8); x, y, z, topology)
     model = HydrostaticFreeSurfaceModel(; grid, tracers=:c)
     model_tracer = model.tracers.c
 
@@ -157,7 +155,6 @@ end
     end
 end
 
-
 @testset "Enzyme for advection and diffusion with various boundary conditions" begin
     Nx = Ny = 64
     Nz = 8
@@ -249,4 +246,175 @@ end
         @test rel_error < tol
     end
 end
+=#
+
+bᵢ(z) = 1e-5 * z
+
+function buoyancy_variance!(model, e_min, Δt=10.0)
+    new_closure = CATKEVerticalDiffusivity(minimum_tke=e_min)
+    model.closure = new_closure
+    model.clock.time = 0
+    model.clock.iteration = 0
+    set!(model, b=bᵢ)
+
+    for n = 1:10
+        time_step!(model, Δt)
+    end
+
+    b = model.tracers.b
+
+    # Another way to compute it
+    Nx, Ny, Nz = size(model.grid)
+    sum_b² = 0.0
+    for k = 1:Nz, j = 1:Ny,  i = 1:Nx
+        sum_b² += b[i, j, k]^2
+    end
+
+    return sum_b²
+end
+
+@testset "Enzyme with CATKEVerticalDiffusivity" begin
+    for arch in archs
+        grid = RectilinearGrid(arch, size=50, z=(-200, 0), topology=(Flat, Flat, Bounded))
+        closure = CATKEVerticalDiffusivity()
+        b_bcs = FieldBoundaryConditions(top=FluxBoundaryCondition(1e-7))
+        tracers = (:b, :e)
+        buoyancy = BuoyancyTracer()
+        model = HydrostaticFreeSurfaceModel(; grid, closure, tracers, buoyancy,
+                                            boundary_conditions=(; b_bcs))
+
+        # Compute derivative by hand
+        e₁, e₂ = 1e-5, 2e-5
+        b²₁ = buoyancy_variance!(model, e₁)
+        b²₂ = buoyancy_variance!(model, e₂)
+        db²_de_fd = (b²₂ - b²₁) / (e₂ - e₁)
+
+        # Now for real
+        e = 1e-5
+        dmodel = Enzyme.make_zero(model)
+
+        @show dmodel
+
+        buoyancy_variance!(dmodel, 0)
+
+        @show dmodel
+
+        db²_de = autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
+                          buoyancy_variance!,
+                          Duplicated(model, dmodel),
+                          Active(e))
+
+        @info """ \n
+            Enzyme computed $db²_de
+            Finite differences computed $db²_de_fd
+        """
+
+        @show db²_de_fd
+        @show db²_de
+
+        #tol = 0.01
+        #rel_error = abs(db²_de[1][3] - db²_de_fd) / abs(db²_de_fd)
+        #@show db²_de, db²_de_fd
+        #@test rel_error < tol
+    end
+end
+
+function set_viscosity!(model, viscosity)
+    new_closure = ScalarDiffusivity(ν=viscosity)
+    names = ()
+    new_closure = with_tracers(names, new_closure)
+    model.closure = new_closure
+    return nothing
+end
+
+function viscous_hydrostatic_turbulence(ν, model, u_init, v_init, Δt, u_truth, v_truth)
+    # Initialize the model
+    model.clock.iteration = 0
+    model.clock.time = 0
+    model.clock.last_Δt = Inf
+    set_viscosity!(model, ν)
+    set!(model, u=u_init, v=v_init, η=0)
+
+    # Step it forward
+    for n = 1:10
+        time_step!(model, Δt)
+    end
+
+    # Compute the sum square error
+    u, v, w = model.velocities
+    Nx, Ny, Nz = size(model.grid)
+    err = 0.0
+    for j = 1:Ny, i = 1:Nx
+        err += @inbounds (u[i, j, 1] - u_truth[i, j, 1])^2 +
+                         (v[i, j, 1] - v_truth[i, j, 1])^2
+    end
+
+    return err::Float64
+end
+
+@testset "Enzyme autodifferentiation of hydrostatic turbulence" begin
+    Random.seed!(123)
+    arch = CPU()
+    Nx = Ny = 32
+    Nz = 1
+    x = y = (0, 2π)
+    z = (0, 1)
+    ν₀ = 1e-2
+
+    grid = RectilinearGrid(arch, size=(Nx, Ny, 1); x, y, z, topology=(Periodic, Periodic, Bounded))
+    closure = ScalarDiffusivity(ν=ν₀)
+    momentum_advection = Centered(order=2)
+
+    g = 4^2
+    c = sqrt(g)
+    free_surface = ExplicitFreeSurface(gravitational_acceleration=g)
+    model = HydrostaticFreeSurfaceModel(; grid, momentum_advection, free_surface, closure)
+
+    ϵ(x, y, z) = 2randn() - 1
+    set!(model, u=ϵ, v=ϵ)
+
+    u_init = deepcopy(model.velocities.u)
+    v_init = deepcopy(model.velocities.v)
+
+    Δx = minimum_xspacing(grid)
+    Δt = 0.01 * Δx / c
+    for n = 1:10
+        time_step!(model, Δt)
+    end
+
+    u_truth = deepcopy(model.velocities.u)
+    v_truth = deepcopy(model.velocities.v)
+    
+    # Use a manual finite difference to compute a gradient
+    Δν = 1e-6
+    ν1 = ν₀ + Δν
+    ν2 = ν1 + Δν
+    e1 = viscous_hydrostatic_turbulence(ν1, model, u_init, v_init, Δt, u_truth, v_truth)
+    e2 = viscous_hydrostatic_turbulence(ν2, model, u_init, v_init, Δt, u_truth, v_truth)
+    ΔeΔν = (e2 - e1) / Δν
+
+    @info "Finite difference computed: $ΔeΔν"
+
+    @info "Now with autodiff..."
+    start_time = time_ns()
+
+    # Use autodiff to compute a gradient
+    dmodel = Enzyme.make_zero(model)
+    dedν = autodiff(set_runtime_activity(Enzyme.Reverse),
+                    viscous_hydrostatic_turbulence,
+                    Active(ν1),
+                    Duplicated(model, dmodel),
+                    Const(u_init),
+                    Const(v_init),
+                    Const(Δt),
+                    Const(u_truth),
+                    Const(v_truth))
+
+    @info "Automatically computed: $dedν."
+    @info "Elapsed time: " * prettytime(1e-9 * (time_ns() - start_time))
+
+    tol = 1e-1
+    rel_error = abs(dedν[1][1] - ΔeΔν) / abs(ΔeΔν)
+    @test rel_error < tol
+end