Broadcast with closure

.buildkite/testing.yml

@@ -1,7 +1,7 @@
 steps:
   - group: ":julia: CUDA GPU"
     steps:
-      - label: ":julia: Julia {{matrix.julia}} + CUDA GPU"
+      - label: ":julia: Julia {{matrix.julia}} + {{matrix.testing_group}} + CUDA GPU"
         plugins:
           - JuliaCI/julia#v1:
               version: "{{matrix.julia}}"
@@ -17,12 +17,16 @@ steps:
           cuda: "*"
         env:
           BACKEND_GROUP: "CUDA"
+          LUX_TEST_GROUP: "{{matrix.testing_group}}"
         if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip ci\]/
-        timeout_in_minutes: 240
+        timeout_in_minutes: 35
         matrix:
           setup:
             julia:
               - "1"
+            testing_group:
+              - "distributed"
+              - "!distributed"
 
   - group: ":telescope: Downstream CUDA"
     steps:
@@ -42,7 +46,7 @@ steps:
         env:
           RETESTITEMS_NWORKERS: 2
         if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip downstream\]/ && build.message !~ /\[skip ci\]/ && build.pull_request.labels includes "run downstream test"
-        timeout_in_minutes: 240
+        timeout_in_minutes: 60
         matrix:
           setup:
             repo:
@@ -52,7 +56,7 @@ steps:
 
   - group: ":julia: AMD GPU"
     steps:
-      - label: ":julia: Julia: {{matrix.julia}} + AMD GPU"
+      - label: ":julia: Julia: {{matrix.julia}} + {{matrix.testing_group}} + AMD GPU"
         plugins:
           - JuliaCI/julia#v1:
               version: "{{matrix.julia}}"
@@ -68,16 +72,20 @@ steps:
           JULIA_AMDGPU_HIP_MUST_LOAD: "1"
           JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
           BACKEND_GROUP: "AMDGPU"
+          LUX_TEST_GROUP: "{{matrix.testing_group}}"
         agents:
           queue: "juliagpu"
           rocm: "*"
           rocmgpu: "*"
         if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip ci\]/
-        timeout_in_minutes: 60
+        timeout_in_minutes: 35
         matrix:
           setup:
             julia:
               - "1"
+            testing_group:
+              - "distributed"
+              - "!distributed"
 
   - group: ":telescope: Downstream AMD GPU"
     steps:

ext/LuxForwardDiffExt/utils.jl

@@ -1,7 +1,10 @@
 # Low-Level functions
 @inline function Lux.__partials(::Type{Tag}, x, i) where {Tag}
     x isa ForwardDiff.Dual && return ForwardDiff.partials(Tag, x, i)
-    x isa AbstractArray && return ForwardDiff.partials.(Tag, x, i)
+    if x isa AbstractArray
+        bfn(xᵢ, iᵢ) = ForwardDiff.partials(Tag, xᵢ, iᵢ)
+        return bfn.(x, i)
+    end
     map_fn = @closure(xᵢ->Lux.__partials(Tag, xᵢ, i))
     x isa Tuple && return map(map_fn, x)
     x isa NamedTuple && return NamedTuple{keys(x)}(map(map_fn, values(x)))
@@ -12,8 +15,8 @@ end
 
 @inline function Lux.__dualify(::Type{Tag}, ::Type{T}, x, u) where {Tag, T}
     if x isa AbstractArray
-        return ForwardDiff.Dual{
-            Tag, T, 1}.(x, ForwardDiff.Partials{1, T}.(tuple.(reshape(u, size(x)))))
+        bfn(xᵢ, uᵢ) = ForwardDiff.Dual{Tag, T, 1}(xᵢ, ForwardDiff.Partials{1, T}(uᵢ))
+        return bfn.(x, tuple.(reshape(u, size(x))))
     end
     x isa Tuple && return map((xᵢ, uᵢ) -> Lux.__dualify(Tag, T, xᵢ, uᵢ), x, u)
     x isa NamedTuple &&

src/layers/basic.jl

@@ -247,8 +247,7 @@ function WrappedFunction(f::F) where {F}
     # Not a depwarn but helpful to call this
     Base.depwarn("The current default of `:direct_call` will be replaced with \
                   `:runtime_check` from v0.6). Please make sure that the assumptions of \
-                  this function are correct or specific \
-                  `WrappedFunction{:direct_call}(f)`",
+                  this function are correct or specify `WrappedFunction{:direct_call}(f)`",
         :WrappedFunction)
     return WrappedFunction{:direct_call}(f)
 end

test/helpers/nestedad_tests.jl

@@ -50,8 +50,9 @@
                        (loss_function1, loss_function2, loss_function3, loss_function4)
 
             for loss_fn in loss_fns
-                @test_nowarn loss_fn(model, X, ps, st)
-                @test loss_fn(model, X, ps, st) isa Number
+                l = loss_fn(model, X, ps, st)
+                @test l isa Number
+                @test isfinite(l) && !isnan(l)
 
                 _, ∂x, ∂ps, _ = Zygote.gradient(loss_fn, model, X, ps, st)
 
@@ -130,8 +131,9 @@ end
                        (loss_function1, loss_function2, loss_function3, loss_function4)
 
             for loss_fn in loss_fns
-                @test_nowarn loss_fn(model, X, ps, st)
-                @test loss_fn(model, X, ps, st) isa Number
+                l = loss_fn(model, X, ps, st)
+                @test l isa Number
+                @test isfinite(l) && !isnan(l)
 
                 _, ∂x, ∂ps, _ = Zygote.gradient(loss_fn, model, X, ps, st)
 

test/layers/containers_tests.jl

@@ -4,7 +4,8 @@
 
     @testset "$mode" for (mode, aType, device, ongpu) in MODES
         @testset "zero sum" begin
-            layer = SkipConnection(WrappedFunction(zero), (a, b) -> a .+ b)
+            layer = SkipConnection(
+                WrappedFunction{:direct_call}(Broadcast.BroadcastFunction(zero)), .+)
             display(layer)
             ps, st = Lux.setup(rng, layer) .|> device
             x = randn(rng, 10, 10, 10, 10) |> aType
@@ -36,7 +37,9 @@ end
 
     @testset "$mode" for (mode, aType, device, ongpu) in MODES
         @testset "zero sum" begin
-            layer = Parallel(+, WrappedFunction(zero), NoOpLayer())
+            layer = Parallel(
+                +, WrappedFunction{:direct_call}(Broadcast.BroadcastFunction(zero)),
+                NoOpLayer())
             @test :layer_1 in keys(layer) && :layer_2 in keys(layer)
             display(layer)
             ps, st = Lux.setup(rng, layer) .|> device

test/layers/type_stability_tests.jl

@@ -70,12 +70,13 @@
             model in model_list,
             input in inputs
 
+            model = maybe_rewrite_to_crosscor(mode, model)
             ps, st = Lux.setup(rng, model) |> dev
             x = input |> dev
 
             @inferred model(x, ps, st)
             @inferred loss_function(model, x, ps, st)
-            if mode == "amdgpu" && model isa Conv
+            if mode == "amdgpu" && (model isa Conv || model isa CrossCor)
                 @test_broken @inferred Zygote.gradient(loss_function, model, x, ps, st)
             else
                 @inferred Zygote.gradient(loss_function, model, x, ps, st)

test/runtests.jl

@@ -84,7 +84,7 @@ end
 end
 
 # Distributed Tests
-if ("all" in LUX_TEST_GROUP || "distributed" in LUX_TEST_GROUP) && BACKEND_GROUP != "amdgpu"
+if ("all" in LUX_TEST_GROUP || "distributed" in LUX_TEST_GROUP)
     using MPI
 
     nprocs_str = get(ENV, "JULIA_MPI_TEST_NPROCS", "")