From 7e9e7b7e9a35e8b85b84ad2acd8afdad535d4857 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 12:04:31 +0530
Subject: [PATCH 01/12] chore: add compat entries

---
 .github/workflows/CompatHelper.yml |  2 +-
 test/Project.toml                  | 21 +++++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml
index ee0b4b6f74..b90370aa1a 100644
--- a/.github/workflows/CompatHelper.yml
+++ b/.github/workflows/CompatHelper.yml
@@ -38,7 +38,7 @@ jobs:
       - name: "Run CompatHelper"
         run: |
           import CompatHelper
-          CompatHelper.main()
+          CompatHelper.main(; subdirs=[".", "test", "lib/ReactantCore"])
         shell: julia --color=yes {0}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/test/Project.toml b/test/Project.toml
index c9526e3624..700daf90ae 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -19,3 +19,24 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[compat]
+ArrayInterface = "7.17.1"
+BenchmarkTools = "1.5.0"
+Enzyme = "0.13.20"
+FFTW = "1.8"
+Flux = "0.15"
+Functors = "0.4, 0.5"
+InteractiveUtils = "1.10"
+LinearAlgebra = "1.10"
+Lux = "1.4.1"
+LuxLib = "1.3"
+MLUtils = "0.4.4"
+NNlib = "0.9.26"
+OneHotArrays = "0.2.6"
+Optimisers = "0.4"
+Random = "1.10"
+SafeTestsets = "0.1"
+SpecialFunctions = "2.4"
+Statistics = "1.10"
+Test = "1.10"

From 554e6e6ab3efb878e68b03c33c91b43e3f26b5b5 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 12:16:49 +0530
Subject: [PATCH 02/12] feat: add copyto! for ConcreteRArray

---
 src/ConcreteRArray.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 46efa14a3f..4dbf677acf 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -301,3 +301,8 @@ function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteR
     aux = copyto!(similar(Array{ElType}, axes(bc)), bc)
     return ConcreteRArray(aux)
 end
+
+function Base.copyto!(dest::ConcreteRArray, src::ConcreteRArray)
+    dest.data = src.data
+    return dest
+end

From 2f76da5878ec169ea5c99cd035486452fa3c1387 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 12:25:28 +0530
Subject: [PATCH 03/12] feat: `@jit` compile ConcreteRArray broadcasting

---
 src/ConcreteRArray.jl | 13 ++-----------
 src/TracedRNumber.jl  |  2 ++
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 4dbf677acf..b308a27340 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -289,17 +289,8 @@ end
 
 # TODO replace this copy for `setindex!` maybe? how to copy data to already existing buffer? (i.e. `copyto!`)
 function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteRArray}})
-    ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
-    if !Base.isconcretetype(ElType)
-        throw(
-            ErrorException(
-                "`copy` on `ConcreteRArray` for non-concrete eltype is not implemented"
-            ),
-        )
-    end
-
-    aux = copyto!(similar(Array{ElType}, axes(bc)), bc)
-    return ConcreteRArray(aux)
+    fn = Reactant.compile(Broadcast.BroadcastFunction(bc.f), (bc.args...,))
+    return fn(bc.args...)
 end
 
 function Base.copyto!(dest::ConcreteRArray, src::ConcreteRArray)
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
index 12c5123b75..d3a6f3b990 100644
--- a/src/TracedRNumber.jl
+++ b/src/TracedRNumber.jl
@@ -240,6 +240,8 @@ for (jlop, hloop) in (
     (:(Base.FastMath.exp_fast), :exponential),
     (:(Base.log), :log),
     (:(Base.sqrt), :sqrt),
+    (:(Base.ceil), :ceil),
+    (:(Base.floor), :floor),
 )
     @eval function $(jlop)(@nospecialize(lhs::TracedRNumber{T})) where {T}
         OutTy = $(hloop === :abs) ? real(T) : T

From 136997161eaf4ae14e1db82b62df776f6a51f7e2 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 12:44:22 +0530
Subject: [PATCH 04/12] fix: manually zero out the lower triangular and upper
 triangular values

---
 src/Ops.jl | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/Ops.jl b/src/Ops.jl
index 2148cb5ebd..3da16de546 100644
--- a/src/Ops.jl
+++ b/src/Ops.jl
@@ -13,6 +13,8 @@ using ..Reactant:
     mlir_type,
     mlir_stacktrace
 
+using LinearAlgebra: triu!, tril!
+
 struct Token
     mlir_data::MLIR.IR.Value
 end
@@ -485,13 +487,20 @@ function cholesky(
     lower::Bool=false,
     location=mlir_stacktrace("cholesky", @__FILE__, @__LINE__),
 ) where {T,N}
-    lower = MLIR.IR.Attribute(lower)
     res = MLIR.IR.result(
         stablehlo.cholesky(
-            x.mlir_data; result=mlir_type(TracedRArray{T,N}, size(x)), lower, location
+            x.mlir_data;
+            result=mlir_type(TracedRArray{T,N}, size(x)),
+            lower=MLIR.IR.Attribute(lower),
+            location,
         ),
     )
-    return TracedRArray{T,N}((), res, size(x))
+    res = TracedRArray{T,N}((), res, size(x))
+
+    # See https://github.com/EnzymeAD/Reactant.jl/issues/338 for why we need to do this
+    lower ? tril!(res) : triu!(res)
+
+    return res
 end
 
 function clamp(

From 96b4a8dae4c9c742235df7b43435c8343f8e165b Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 12:59:50 +0530
Subject: [PATCH 05/12] fix: only do it in tests

---
 src/ConcreteRArray.jl |  2 ++
 src/Ops.jl            |  9 +--------
 src/Reactant.jl       | 12 ++++++++++++
 src/TracedRArray.jl   |  8 --------
 test/ops.jl           | 13 ++++++++-----
 5 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index b308a27340..f32e7d33c4 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -8,6 +8,8 @@ mutable struct ConcreteRArray{T,N} <: RArray{T,N}
     shape::NTuple{N,Int}
 end
 
+const WrappedConcreteRArray{T,N} = WrappedArray{T,N,ConcreteRArray,ConcreteRArray{T,N}}
+
 mutable struct ConcreteRNumber{T} <: RNumber{T}
     data::XLA.AsyncBuffer
 end
diff --git a/src/Ops.jl b/src/Ops.jl
index 3da16de546..50c643760c 100644
--- a/src/Ops.jl
+++ b/src/Ops.jl
@@ -13,8 +13,6 @@ using ..Reactant:
     mlir_type,
     mlir_stacktrace
 
-using LinearAlgebra: triu!, tril!
-
 struct Token
     mlir_data::MLIR.IR.Value
 end
@@ -495,12 +493,7 @@ function cholesky(
             location,
         ),
     )
-    res = TracedRArray{T,N}((), res, size(x))
-
-    # See https://github.com/EnzymeAD/Reactant.jl/issues/338 for why we need to do this
-    lower ? tril!(res) : triu!(res)
-
-    return res
+    return TracedRArray{T,N}((), res, size(x))
 end
 
 function clamp(
diff --git a/src/Reactant.jl b/src/Reactant.jl
index 0b73d3d96e..348a541f5d 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -89,6 +89,18 @@ function Enzyme.make_zero(
     return res
 end
 
+function ancestor(x::AbstractArray)
+    p_x = parent(x)
+    p_x === x && return x
+    return ancestor(p_x)
+end
+
+function get_ancestor_indices(x::AbstractArray, indices...)
+    p_x = parent(x)
+    p_x === x && return indices
+    return get_ancestor_indices(p_x, Base.reindex(parentindices(x), indices)...)
+end
+
 include("mlir/MLIR.jl")
 include("XLA.jl")
 include("Interpreter.jl")
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index 97a29e56f9..0dd4b7261c 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -72,14 +72,6 @@ function set_mlir_data!(x::AnyTracedRArray, data)
     return x
 end
 
-ancestor(x::TracedRArray) = x
-ancestor(x::WrappedTracedRArray) = ancestor(parent(x))
-
-get_ancestor_indices(::TracedRArray, indices...) = indices
-function get_ancestor_indices(x::WrappedTracedRArray, indices...)
-    return get_ancestor_indices(parent(x), Base.reindex(parentindices(x), indices)...)
-end
-
 function Base.getindex(
     a::TracedRArray{T,N}, index::Vararg{Union{Int,TracedRNumber{Int}},N}
 ) where {T,N}
diff --git a/test/ops.jl b/test/ops.jl
index eb4981b80e..e178a660f9 100644
--- a/test/ops.jl
+++ b/test/ops.jl
@@ -82,14 +82,16 @@ end
 end
 
 @testset "cholesky" begin
-    g(x) = Ops.cholesky(x; lower=true)
+    g1(x) = triu(Ops.cholesky(x))
+    g2(x) = tril(Ops.cholesky(x; lower=true))
+
     x = ConcreteRArray([
         10.0 2.0 3.0
         2.0 5.0 6.0
         3.0 6.0 9.0
     ])
-    @test cholesky(Array(x)).U ≈ @jit Ops.cholesky(x)
-    @test transpose(cholesky(Array(x)).U) ≈ @jit g(x)
+    @test cholesky(Array(x)).U ≈ @jit g1(x)
+    @test transpose(cholesky(Array(x)).U) ≈ @jit g2(x)
 
     x = ConcreteRArray(
         [
@@ -98,8 +100,9 @@ end
             3.0+4.0im 3.0+2.0im 9.0+0.0im
         ],
     )
-    @test cholesky(Array(x)).U ≈ @jit Ops.cholesky(x)
-    @test adjoint(cholesky(Array(x)).U) ≈ @jit g(x)
+
+    @test cholesky(Array(x)).U ≈ @jit g1(x)
+    @test adjoint(cholesky(Array(x)).U) ≈ @jit g2(x)
 end
 
 @testset "clamp" begin

From 739d830f1010e07b31806efb084790e1a03eef46 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 13:11:20 +0530
Subject: [PATCH 06/12] feat: compile mapreduce for ConcreteRArray

---
 src/ConcreteRArray.jl | 12 ++++++++++++
 test/ops.jl           |  5 +++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index f32e7d33c4..37def1342f 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -299,3 +299,15 @@ function Base.copyto!(dest::ConcreteRArray, src::ConcreteRArray)
     dest.data = src.data
     return dest
 end
+
+function Base.mapreduce(
+    @nospecialize(f),
+    @nospecialize(op),
+    @nospecialize(A::ConcreteRArray{T,N});
+    dims=:,
+    init=nothing,
+) where {T,N}
+    mapreduce_fn(x) = Base.mapreduce(f, op, x; dims, init)
+    fn = Reactant.compile(mapreduce_fn, (A,))
+    return fn(A)
+end
diff --git a/test/ops.jl b/test/ops.jl
index e178a660f9..a1064e4030 100644
--- a/test/ops.jl
+++ b/test/ops.jl
@@ -213,13 +213,14 @@ end
     ]
         # NOTE `LinearAlgebra.dot` is not equal to `sum(a .* b)` on complex numbers due to conjugation
         @test sum(a .* b) ≈ @jit f1(a, b)
-        @test kron(reshape(a, length(a), 1), reshape(b, 1, length(b))) ≈ @jit fouter(a, b)
+        @test kron(reshape(Array(a), length(a), 1), reshape(Array(b), 1, length(b))) ≈
+            @jit fouter(a, b)
         @test a .* b ≈ @jit fouter_batch1(a, b)
     end
 
     a = ConcreteRArray([1 2; 3 4])
     b = ConcreteRArray([5 6; -7 -8])
-    @test a' * b == @jit f1(a, b)
+    @test Array(a)' * Array(b) == @jit f1(a, b)
 end
 
 @testset "einsum" begin

From 4d93782fc7e3406c69a5d40842ed54ca200fa960 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 13:12:59 +0530
Subject: [PATCH 07/12] test: manual array conversion

---
 test/ops.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/ops.jl b/test/ops.jl
index a1064e4030..1d24f4cec0 100644
--- a/test/ops.jl
+++ b/test/ops.jl
@@ -243,7 +243,7 @@ end
         x = reshape(a, (2, 2))
         y = reshape(b, (2, 2))
         @test x .* y ≈ @jit f3(x, y)
-        @test x * y ≈ @jit f4(x, y)
+        @test Array(x) * Array(y) ≈ @jit f4(x, y)
     end
 end
 

From 9934d6dff1e96a7efcdc0813aaca2b242a755e5d Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 13:13:53 +0530
Subject: [PATCH 08/12] revert: change in Ops.cholesky

---
 src/Ops.jl | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/Ops.jl b/src/Ops.jl
index 50c643760c..2148cb5ebd 100644
--- a/src/Ops.jl
+++ b/src/Ops.jl
@@ -485,12 +485,10 @@ function cholesky(
     lower::Bool=false,
     location=mlir_stacktrace("cholesky", @__FILE__, @__LINE__),
 ) where {T,N}
+    lower = MLIR.IR.Attribute(lower)
     res = MLIR.IR.result(
         stablehlo.cholesky(
-            x.mlir_data;
-            result=mlir_type(TracedRArray{T,N}, size(x)),
-            lower=MLIR.IR.Attribute(lower),
-            location,
+            x.mlir_data; result=mlir_type(TracedRArray{T,N}, size(x)), lower, location
         ),
     )
     return TracedRArray{T,N}((), res, size(x))

From bfabd262291292f11254f01be77a57ef04dcfa22 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 7 Dec 2024 13:15:15 +0530
Subject: [PATCH 09/12] revert: remove unnecessary changes

---
 src/ConcreteRArray.jl |  2 --
 src/Reactant.jl       | 12 ------------
 src/TracedRArray.jl   |  8 ++++++++
 3 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 37def1342f..b466f66d45 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -8,8 +8,6 @@ mutable struct ConcreteRArray{T,N} <: RArray{T,N}
     shape::NTuple{N,Int}
 end
 
-const WrappedConcreteRArray{T,N} = WrappedArray{T,N,ConcreteRArray,ConcreteRArray{T,N}}
-
 mutable struct ConcreteRNumber{T} <: RNumber{T}
     data::XLA.AsyncBuffer
 end
diff --git a/src/Reactant.jl b/src/Reactant.jl
index 348a541f5d..0b73d3d96e 100644
--- a/src/Reactant.jl
+++ b/src/Reactant.jl
@@ -89,18 +89,6 @@ function Enzyme.make_zero(
     return res
 end
 
-function ancestor(x::AbstractArray)
-    p_x = parent(x)
-    p_x === x && return x
-    return ancestor(p_x)
-end
-
-function get_ancestor_indices(x::AbstractArray, indices...)
-    p_x = parent(x)
-    p_x === x && return indices
-    return get_ancestor_indices(p_x, Base.reindex(parentindices(x), indices)...)
-end
-
 include("mlir/MLIR.jl")
 include("XLA.jl")
 include("Interpreter.jl")
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
index 0dd4b7261c..97a29e56f9 100644
--- a/src/TracedRArray.jl
+++ b/src/TracedRArray.jl
@@ -72,6 +72,14 @@ function set_mlir_data!(x::AnyTracedRArray, data)
     return x
 end
 
+ancestor(x::TracedRArray) = x
+ancestor(x::WrappedTracedRArray) = ancestor(parent(x))
+
+get_ancestor_indices(::TracedRArray, indices...) = indices
+function get_ancestor_indices(x::WrappedTracedRArray, indices...)
+    return get_ancestor_indices(parent(x), Base.reindex(parentindices(x), indices)...)
+end
+
 function Base.getindex(
     a::TracedRArray{T,N}, index::Vararg{Union{Int,TracedRNumber{Int}},N}
 ) where {T,N}

From 57a319236f950a90bbc2c8ebfa0e4b51dbb38109 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 8 Dec 2024 21:27:08 +0530
Subject: [PATCH 10/12] fix: only compile non-CPU broadcasting

---
 Project.toml          |  4 ++--
 src/ConcreteRArray.jl | 21 +++++++++++++++++++++
 test/Project.toml     |  8 ++++----
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index f0762b756a..db90c4658b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -39,10 +39,10 @@ ArrayInterface = "7.10"
 CEnum = "0.4, 0.5"
 Downloads = "1.6"
 Enzyme = "0.13.21"
-EnzymeCore = "0.8.6, 0.8.7, 0.8.8"
+EnzymeCore = "0.8.8"
 GPUArraysCore = "0.1.6, 0.2"
 LinearAlgebra = "1.10"
-NNlib = "0.9.24"
+NNlib = "0.9.26"
 OrderedCollections = "1"
 Preferences = "1.4"
 ReactantCore = "0.1.2"
diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index b466f66d45..9c4a1130c1 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -289,6 +289,27 @@ end
 
 # TODO replace this copy for `setindex!` maybe? how to copy data to already existing buffer? (i.e. `copyto!`)
 function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteRArray}})
+    foreach(bc.args) do x
+        x isa ConcreteRArray && XLA.await(x.data)
+    end
+
+    all_on_cpu = all(bc.args) do x
+        x isa ConcreteRArray && return XLA.BufferOnCPU(x.data.buffer)
+        return true
+    end
+    if all_on_cpu
+        ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
+        if !Base.isconcretetype(ElType)
+            throw(
+                ErrorException(
+                    "`copy` on `ConcreteRArray` for non-concrete eltype is not implemented"
+                ),
+            )
+        end
+        aux = copyto!(similar(Array{ElType}, axes(bc)), bc)
+        return ConcreteRArray(aux)
+    end
+
     fn = Reactant.compile(Broadcast.BroadcastFunction(bc.f), (bc.args...,))
     return fn(bc.args...)
 end
diff --git a/test/Project.toml b/test/Project.toml
index 700daf90ae..4b50a487fc 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -21,12 +21,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ArrayInterface = "7.17.1"
-BenchmarkTools = "1.5.0"
-Enzyme = "0.13.20"
+ArrayInterface = "7.10"
+BenchmarkTools = "1.5"
+Enzyme = "0.13.21"
 FFTW = "1.8"
 Flux = "0.15"
-Functors = "0.4, 0.5"
+Functors = "0.5"
 InteractiveUtils = "1.10"
 LinearAlgebra = "1.10"
 Lux = "1.4.1"

From 5a13b569e766f7c2d75707a98f89c77e7865855f Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Mon, 9 Dec 2024 09:13:11 +0530
Subject: [PATCH 11/12] fix: address reviewer comments

---
 src/ConcreteRArray.jl | 26 +++++++++++++++++---------
 test/ops.jl           |  2 ++
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 9c4a1130c1..7cda47fc71 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -215,7 +215,7 @@ function Base.getindex(a::ConcreteRArray{T}, args::Vararg{Int,N}) where {T,N}
     end
 
     XLA.await(a.data)
-    if XLA.BufferOnCPU(a.data.buffer)
+    if buffer_on_cpu(a)
         buf = a.data.buffer
         GC.@preserve buf begin
             ptr = Base.unsafe_convert(Ptr{T}, XLA.UnsafeBufferPointer(buf))
@@ -246,7 +246,7 @@ function Base.setindex!(a::ConcreteRArray{T}, v, args::Vararg{Int,N}) where {T,N
     end
 
     XLA.await(a.data)
-    if XLA.BufferOnCPU(a.data.buffer)
+    if buffer_on_cpu(a)
         buf = a.data.buffer
         GC.@preserve buf begin
             ptr = Base.unsafe_convert(Ptr{T}, XLA.UnsafeBufferPointer(buf))
@@ -289,14 +289,11 @@ end
 
 # TODO replace this copy for `setindex!` maybe? how to copy data to already existing buffer? (i.e. `copyto!`)
 function Base.copy(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{ConcreteRArray}})
-    foreach(bc.args) do x
+    for x in bc.args
         x isa ConcreteRArray && XLA.await(x.data)
     end
 
-    all_on_cpu = all(bc.args) do x
-        x isa ConcreteRArray && return XLA.BufferOnCPU(x.data.buffer)
-        return true
-    end
+    all_on_cpu = all(buffer_on_cpu, bc.args)
     if all_on_cpu
         ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
         if !Base.isconcretetype(ElType)
@@ -326,7 +323,18 @@ function Base.mapreduce(
     dims=:,
     init=nothing,
 ) where {T,N}
-    mapreduce_fn(x) = Base.mapreduce(f, op, x; dims, init)
-    fn = Reactant.compile(mapreduce_fn, (A,))
+    fn = Reactant.compile(CallMapReduce(f, op, dims, init), (A,))
     return fn(A)
 end
+
+struct CallMapReduce{Fn, Op, Dims, Init}
+    f::Fn
+    op::Op
+    dims::Dims
+    init::Init
+end
+
+(f::CallMapReduce)(A) = Base.mapreduce(f.f, f.op, A; f.dims, f.init)
+
+buffer_on_cpu(::Any) = true
+buffer_on_cpu(x::ConcreteRArray) = XLA.BufferOnCPU(x.data.buffer)
diff --git a/test/ops.jl b/test/ops.jl
index 1d24f4cec0..0437b2723f 100644
--- a/test/ops.jl
+++ b/test/ops.jl
@@ -82,6 +82,8 @@ end
 end
 
 @testset "cholesky" begin
+    # cholesky in stablehlo for the other triangle is implementation defined.
+    # See https://github.com/EnzymeAD/Reactant.jl/issues/338 for more details.
     g1(x) = triu(Ops.cholesky(x))
     g2(x) = tril(Ops.cholesky(x; lower=true))
 

From 937916c19e02dee346c9153b0d7dbb73fdc5690c Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Mon, 9 Dec 2024 10:10:48 +0530
Subject: [PATCH 12/12] chore: apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 src/ConcreteRArray.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
index 7cda47fc71..42c454b31a 100644
--- a/src/ConcreteRArray.jl
+++ b/src/ConcreteRArray.jl
@@ -327,7 +327,7 @@ function Base.mapreduce(
     return fn(A)
 end
 
-struct CallMapReduce{Fn, Op, Dims, Init}
+struct CallMapReduce{Fn,Op,Dims,Init}
     f::Fn
     op::Op
     dims::Dims