test: more test fixes

Author: avik-pal

src/Compiler.jl

@@ -6,6 +6,7 @@ import ..Reactant:
     XLA,
     ConcreteRArray,
     TracedRArray,
+    TracedRNumber,
     OrderedIdDict,
     make_tracer,
     TracedToConcrete,
@@ -289,7 +290,7 @@ function compile_mlir!(mod, f, args; optimize=true)
     preserved_args = Tuple{TracedRArray,Int}[]
     results = [MLIR.IR.operand(ret, i) for i in 1:MLIR.IR.noperands(ret)]
     nresults = MLIR.IR.Value[]
-    linear_results2 = TracedRArray[]
+    linear_results2 = Union{TracedRArray,TracedRNumber}[]
     for (i, op) in enumerate(results)
         if !MLIR.IR.is_block_arg(op)
             push!(nresults, op)

src/TracedRNumber.jl

@@ -112,22 +112,45 @@ for (jlop, hloop, hlocomp) in (
     (:(Base.:(<=)), :compare, "LE"),
     (:(Base.:(<)), :compare, "LT"),
 )
-    @eval function $(jlop)(
-        @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::TracedRNumber{T})
-    ) where {T}
-        return TracedRNumber{Bool}(
-            (),
-            MLIR.IR.result(
-                MLIR.Dialects.stablehlo.$(hloop)(
-                    lhs.mlir_data,
-                    rhs.mlir_data;
-                    comparison_direction=MLIR.API.stablehloComparisonDirectionAttrGet(
-                        MLIR.IR.context(), $hlocomp
+    @eval begin
+        function $(jlop)(
+            @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::TracedRNumber{T})
+        ) where {T}
+            return TracedRNumber{Bool}(
+                (),
+                MLIR.IR.result(
+                    MLIR.Dialects.stablehlo.$(hloop)(
+                        lhs.mlir_data,
+                        rhs.mlir_data;
+                        comparison_direction=MLIR.API.stablehloComparisonDirectionAttrGet(
+                            MLIR.IR.context(), $hlocomp
+                        ),
                     ),
+                    1,
                 ),
-                1,
-            ),
-        )
+            )
+        end
+
+        function $(jlop)(
+            @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs)
+        ) where {T}
+            return $(jlop)(lhs, promote_to(lhs, rhs))
+        end
+
+        function $(jlop)(
+            @nospecialize(lhs), @nospecialize(rhs::TracedRNumber{T})
+        ) where {T}
+            return $(jlop)(promote_to(rhs, lhs), rhs)
+        end
+
+        function $(jlop)(
+            @nospecialize(lhs::TracedRNumber{T1}), @nospecialize(rhs::TracedRNumber{T2})
+        ) where {T1,T2}
+            commonTy = TracedRNumber{Base.promote_type(T1, T2)}
+            lhs = promote_to(commonTy, lhs)
+            rhs = promote_to(commonTy, rhs)
+            return $(jlop)(lhs, rhs)
+        end
     end
 end
 
@@ -169,6 +192,9 @@ for (jlop, hloop) in (
     end
 end
 
+# XXX: Enzyme-MLIR doesn't have `abs` adjoint defined
+Base.abs2(x::TracedRNumber{<:Real}) = x^2
+
 struct TypeCast{T<:ReactantPrimitives} <: Function end
 
 (::TypeCast{T})(x::TracedRNumber{T2}) where {T,T2} = promote_to(TracedRNumber{T}, x)