feat: more operation coverage

Author: avik-pal

src/TracedRArray.jl

@@ -29,6 +29,10 @@ Base.elsize(::Type{TracedRArray{T,N}}) where {T,N} = sizeof(T)
 # we use it
 Base.convert(T::Type{<:TracedRArray}, x::AbstractArray) = Reactant.promote_to(T, x)
 
+# Base.first is very common usecase for getting first element to get the type
+# inside LinearAlgebra.jl
+Base.first(x::TracedRArray{T,N}) where {T,N} = @allowscalar(x[1])
+
 # Base.complex
 Base.complex(x::TracedRArray{<:Real}) = complex.(x)
 Base.complex(x::TracedRArray{<:Complex}) = x

src/TracedRNumber.jl

@@ -26,13 +26,25 @@ Base.copy(x::TracedRNumber{T}) where {T} = TracedRNumber{T}((), x.mlir_data)
 function Base.eps(::Type{TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, eps(T))
 end
+Base.eps(x::TracedRNumber{T}) where {T} = eps(typeof(x))
 
 function Base.typemin(::Type{TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, typemin(T))
 end
+Base.typemin(x::TracedRNumber{T}) where {T} = typemin(typeof(x))
+
 function Base.typemax(::Type{TracedRNumber{T}}) where {T}
     return Reactant.promote_to(TracedRNumber{T}, typemax(T))
 end
+Base.typemax(x::TracedRNumber{T}) where {T} = typemax(typeof(x))
+
+function Base.nextfloat(x::TracedRNumber{T}) where {T<:AbstractFloat}
+    return @opcall next_after(x, typemax(x))
+end
+
+function Base.prevfloat(x::TracedRNumber{T}) where {T<:AbstractFloat}
+    return @opcall next_after(x, typemin(x))
+end
 
 function Base.rtoldefault(T::Type{<:TracedRNumber})
     return T(Base.rtoldefault(unwrapped_eltype(T)))

src/stdlibs/LinearAlgebra.jl

@@ -5,7 +5,7 @@ using ..Reactant: Reactant, Ops
 using ..Reactant:
     TracedRArray, TracedRNumber, AnyTracedRArray, AnyTracedRMatrix, AnyTracedRVector
 using ..Reactant: call_with_reactant
-using ReactantCore: ReactantCore, materialize_traced_array
+using ReactantCore: ReactantCore, materialize_traced_array, @trace
 using Reactant_jll: Reactant_jll
 
 using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!
@@ -15,7 +15,7 @@ using LinearAlgebra: LinearAlgebra, BLAS
 using LinearAlgebra: Adjoint, Transpose, Factorization, RowMaximum, NoPivot
 using LinearAlgebra: SymTridiagonal, Symmetric, Bidiagonal, Diagonal, Tridiagonal
 using LinearAlgebra: LowerTriangular, UnitLowerTriangular, UpperTriangular
-using LinearAlgebra: diag, diagm, ldiv!
+using LinearAlgebra: diag, diagm, ldiv!, inv, rmul!
 using Libdl: Libdl
 using GPUArraysCore: @allowscalar
 
@@ -694,4 +694,44 @@ function LinearAlgebra.ishermitian(A::AnyTracedRMatrix)
     return all(A .== adjoint(A))
 end
 
+function LinearAlgebra.isbanded(A::AnyTracedRMatrix, kl::Integer, ku::Integer)
+    return LinearAlgebra.istriu(A, kl) && LinearAlgebra.istril(A, ku)
+end
+
+@static if isdefined(LinearAlgebra, :__normalize!)
+    function LinearAlgebra.__normalize!(a::AnyTracedRArray, nrm)
+        # The largest positive floating point number whose inverse is less than infinity
+        δ = inv(prevfloat(typemax(nrm)))
+        @trace if nrm ≥ δ # Safe to multiply with inverse
+            invnrm = inv(nrm)
+            rmul!(a, invnrm)
+        else # scale elements to avoid overflow
+            εδ = eps(one(nrm)) / δ
+            rmul!(a, εδ)
+            rmul!(a, inv(nrm * εδ))
+        end
+        return a
+    end
+end
+
+function LinearAlgebra.rmul!(A::AnyTracedRArray, b::Number)
+    @. A *= b
+    return A
+end
+
+function LinearAlgebra.lmul!(b::Number, A::AnyTracedRArray)
+    @. A = b * A
+    return A
+end
+
+function LinearAlgebra.rdiv!(A::AnyTracedRArray, b::Number)
+    @. A /= b
+    return A
+end
+
+function LinearAlgebra.ldiv!(b::Number, A::AnyTracedRArray)
+    @. A = b \ A
+    return A
+end
+
 end

src/stdlibs/factorization/Cholesky.jl

@@ -8,8 +8,12 @@ end
 Base.size(c::GeneralizedCholesky) = size(c.factors)
 Base.ndims(c::GeneralizedCholesky) = ndims(c.factors)
 
+function overloaded_cholesky(A::AbstractArray, ::NoPivot; check::Bool=false)
+    return overloaded_cholesky(Reactant.promote_to(TracedRArray, A), NoPivot(); check)
+end
+
 function overloaded_cholesky(
-    A::AbstractArray{T,N}, ::NoPivot; check::Bool=false
+    A::AnyTracedRArray{T,N}, ::NoPivot; check::Bool=false
 ) where {T,N}
     # TODO: dont ignore check
     # move the batching dims to the front

src/stdlibs/factorization/SVD.jl

@@ -1 +1,34 @@
+struct GeneralizedSVD{T,Tr,M<:AbstractArray{T},C<:AbstractArray{T}} <: Factorization{T}
+    U::M
+    S::C
+    Vt::M
 
+    function GeneralizedSVD{T,Tr,M,C}(U::M, S::C, Vt::M) where {T,Tr,M,C}
+        @assert ndims(S) == ndims(U) - 1
+        return new{T,Tr,M,C}(U, S, Vt)
+    end
+end
+
+function overloaded_svd(A::AbstractArray; kwargs...)
+    return overloaded_svd(Reactant.promote_to(TracedRArray, A); kwargs...)
+end
+
+function overloaded_svd(
+    A::AnyTracedRArray{T,N}; full::Bool=false, algorithm=nothing
+) where {T,N}
+    # TODO: don't ignore the algorithm kwarg
+    return error("TODO: Not implemented yet")
+end
+
+function overloaded_svd(
+    A::AnyTracedRVector{T}; full::Bool=false, algorithm=nothing
+) where {T}
+    # TODO: don't ignore the algorithm kwarg
+    m = length(A)
+    normA = LinearAlgebra.norm(A)
+
+    return error("TODO: Not implemented yet")
+end
+
+# TODO: compute svdvals without computing the full svd. In principle we should
+#       simple dce the U and Vt inside the compiler itself and simply compute Σ