Merge pull request #2158 from CliMA/ck/fix_2146

Allow for re-instantiating field-style broadcasted objects
CliMA · Jan 24, 2025 · bbad885 · bbad885
2 parents 75bf60d + 2566ad6
commit bbad885
Show file tree

Hide file tree

Showing 9 changed files with 12 additions and 32 deletions.
diff --git a/benchmarks/scripts/benchmark_utils.jl b/benchmarks/scripts/benchmark_utils.jl
@@ -1,7 +1,7 @@
 # import CUDA
 import ClimaComms
 using BenchmarkTools, Dates
-using ClimaCore: @lazy
+import LazyBroadcast: @lazy
 
 """
     caller_name(@__FILE__)

diff --git a/src/ClimaCore.jl b/src/ClimaCore.jl
@@ -4,9 +4,6 @@ using PkgVersion
 const VERSION = PkgVersion.@Version
 import ClimaComms
 
-# Temporary helpers to work around https://github.com/CliMA/ClimaCore.jl/issues/2146
-include("lazy_broadcast.jl")
-
 include("DebugOnly/DebugOnly.jl")
 include("interface.jl")
 include("devices.jl")

diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl
@@ -249,7 +249,9 @@ function Base.Broadcast.instantiate(sbc::StencilBroadcasted)
         axes = Base.axes(sbc)
     else
         axes = sbc.axes
-        Base.Broadcast.check_broadcast_axes(axes, args...)
+        if axes !== Base.axes(sbc)
+            Base.Broadcast.check_broadcast_axes(axes, args...)
+        end
     end
     Style = AbstractStencilStyle(ClimaComms.device(axes))
     return StencilBroadcasted{Style}(op, args, axes)

diff --git a/src/Operators/spectralelement.jl b/src/Operators/spectralelement.jl
@@ -123,7 +123,9 @@ function Base.Broadcast.instantiate(sbc::SpectralBroadcasted)
         axes = Base.axes(sbc)
     else
         axes = sbc.axes
-        Base.Broadcast.check_broadcast_axes(axes, args...)
+        if axes !== Base.axes(sbc)
+            Base.Broadcast.check_broadcast_axes(axes, args...)
+        end
     end
     op = typeof(op)(axes)
     Style = AbstractSpectralStyle(ClimaComms.device(axes))

diff --git a/src/lazy_broadcast.jl b/src/lazy_broadcast.jl
diff --git a/test/DataLayouts/unit_has_uniform_datalayouts.jl b/test/DataLayouts/unit_has_uniform_datalayouts.jl
@@ -6,7 +6,7 @@ using Test
 using ClimaCore.DataLayouts
 import ClimaCore.Geometry
 import ClimaComms
-import ClimaCore: @lazy
+import LazyBroadcast: @lazy
 using StaticArrays
 import Random
 Random.seed!(1234)
@@ -41,6 +41,5 @@ Random.seed!(1234)
     # This is not a valid broadcast expression,
     # but these two datalayouts can exist in a
     # valid broadcast expression (e.g., interpolation).
-    bc = @lazy @. data_VIJFHᶜ + data_VIJFHᶠ
-    @test DataLayouts.has_uniform_datalayouts(bc)
+    @test_throws DimensionMismatch @lazy @. data_VIJFHᶜ + data_VIJFHᶠ
 end
diff --git a/test/DataLayouts/unit_linear_indexing.jl b/test/DataLayouts/unit_linear_indexing.jl
@@ -15,7 +15,7 @@ using ClimaComms
 ClimaComms.@import_required_backends
 import ClimaCore: Domains, Meshes, Topologies, Spaces, Fields, Operators
 import ClimaCore.DataLayouts: get_struct
-using ClimaCore: @lazy
+import LazyBroadcast: @lazy
 
 Random.seed!(1234)
 

diff --git a/test/DataLayouts/unit_non_extruded_broadcast.jl b/test/DataLayouts/unit_non_extruded_broadcast.jl
@@ -5,7 +5,7 @@ using Revise; include(joinpath("test", "DataLayouts", "unit_non_extruded_broadca
 using Test
 using ClimaComms
 import Base.Broadcast: broadcasted, instantiate
-using ClimaCore: @lazy
+import LazyBroadcast: @lazy
 using ClimaCore.DataLayouts
 using ClimaCore.Geometry
 using ClimaCore: Geometry, Domains, Topologies, Meshes, Spaces, Fields

diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl
@@ -3,7 +3,7 @@ using JET
 import Dates
 import Random: seed!
 import Base.Broadcast: materialize, materialize!
-import ClimaCore: @lazy
+import LazyBroadcast: @lazy
 import BenchmarkTools as BT
 
 import ClimaComms