feat: add support for the remaining wrapper types

[avik-pal]

we need this before linear algebra stuff because most of them end up using these wrappers as return types.

fixes #345 #242

Added Wrapper Types

• Symmetric
• UnitLowerTriangular
• LowerTriangular
• UnitUpperTriangular
• UpperTriangular
• Tridiagonal

Updated Linear Algebra Support

• diagm extended for key value pair version and now we use scatter

New Addition to Ops

1. scatter_setindex
2. gather_getindex

Not calling them scatter / gather since they are a significantly restrictive subset of the corresponding stablehlo versions but these are also a very common pattern.

Fixes / Refactor

1. Non-contiguous indexing for both setindex and getindex

[mofeing]

what do you think about adding a section in the developer reference docs? so that wrapper array authors can add support for Reactant and we also remember how it should be done 😅

i actually want to add support for Reactant in my DeltaArrays package

[avik-pal]

what do you think about adding a section in the developer reference docs? so that wrapper array authors can add support for Reactant and we also remember how it should be done 😅

i actually want to add support for Reactant in my DeltaArrays package

That is a good idea, I will add that. 1 caveat of not having "AbstractWrappedArray" in Base Julia is that these new wrapper types won't go through our AnyTracedRArray dispatches. Technically, it is not a unique problem for us, and CUDA has been struggling with this for years now.

xref: JuliaLang/julia#51910

[avik-pal]

How can we introduce a runtime error in the generated MLIR?

[mofeing]

i think you can't? at least not in the ML IR, but the verifier will error when verifying the ops

[avik-pal]

I was thinking of some sort of custom_call which jax is using here https://github.com/openxla/stablehlo/blob/main/docs/tutorials/jax-export.ipynb

stablehlo.custom_call @shape_assertion(%1, %0) {api_version = 2 : i32, error_message = "Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'a'. Using the following polymorphic shapes specifications: args[0].shape = (a, 3, 224, 224). Obtained dimension variables: 'a' = {0} from specification 'a' for dimension args[0].shape[0] (= {0}), . Please see https://jax.readthedocs.io/en/latest/export/shape_poly.html#shape-assertion-errors for more details.", has_side_effect = true} : (tensor<i1>, tensor<i32>) -> ()

For tracedRarray indices we should probably always do a (dynamic_)gather (and we might be able to write an optimization later to transform that into a slice if contiguous.

[avik-pal]

Also we should revisit this at some point, but the current approach means that the generic dispatches convert a structured array into a dense array so we definitely use worse algorithms than are warranted

[avik-pal]

need to add tests

[mofeing]

Also we should revisit this at some point, but the current approach means that the generic dispatches convert a structured array into a dense array so we definitely use worse algorithms than are warranted

yeah, can you leave a note in the docs for this please? IMHO this is ok to leave as default, and we can specialize for performance if we know how to do for the combination of array type and Ops

That is a good idea, I will add that. 1 caveat of not having "AbstractWrappedArray" in Base Julia is that these new wrapper types won't go through our AnyTracedRArray dispatches. Technically, it is not a unique problem for us, and CUDA has been struggling with this for years now.

the way i've solved things like this in Tenet is using Holy traits with a default method that checks the trait (e.g. with LinearAlgebra.diag) but here is gonna be hard because we can't override the default method... maybe now that @wsmoses is readding the method overlay mechanism, we can overlay the default method and do sth like this?

iswrappedtraced(::TracedRArray) = true
iswrappedtraced(::Array) = false
iswrappedtraced(x::Diagonal) = iswrappedtraced(parent(x))
iswrappedtraced(::Type{<:TracedRArray}) = true
iswrappedtraced(::Type{<:Array}) = false
iswrappedtraced(::Type{Diagonal{T,N,A}}) where {T,N,A} = iswrappedtraced(A)

@reactant_override function LinearAlgebra.diag(x::AbstractArray{T,2}, k::Integer=0) where {T}
    if iswrappedtraced(x)
        y = materialize_traced_array(x)
        return diag(y, k)
    else
        # invoke diag(x) on NativeInterpreter
    end
end

function LinearAlgebra.diag(x::TracedRArray{T,2}, k::Integer=0) where {T}
    # here the code we currently have
end

the wrapped array lib dev who would like to add support for Reactant would then just need to correctly implement iswrappedtraced (suggestions welcome for a better name) and materialize_traced_array. furthermore, the dev could implement better methods if they know that this combination of method + array type doesn't need to be materialized (e.g. a Diagonal * Array could be implemented using dot_general with batching dims and no contracting dims).

the only problem i see is whether we can invoke methods in the NativeInterpreter from within a custom interpreter. @wsmoses @vchuravy @gbaraldi is that doable?

---

getting out of scope, i thing that sth Julia could have and would fix a loooot of problems like this are extendable unions. problem with abstract types is that you cannot take an external type and tell it "hey, you're now a 'son' of this abstract type because you fulfill this new interface i've defined". even with multiple inheritance there would still be problems. abstract types block composability.

an extendable union would be similar to abc.ABCMeta in Python, were you can register new subclasses dynamically https://docs.python.org/3.13/library/abc.html#abc.ABCMeta
check out this example where the user creates a new class MyABC and makes the builtin tuple be a subclass of it!

from abc import ABC

class MyABC(ABC):
    pass

MyABC.register(tuple)

assert issubclass(tuple, MyABC)
assert isinstance((), MyABC)

problem with this approach is that if not handled with care, i think it could lead to a lot of invalidation (or maybe no, i don't know much about the internals that deal with this).

[avik-pal]

We can push an is_wrapped_array to ArrayInterface (easier for other packages to use rather than depending on Reactant). Then is_wrapped_tracked_array (EDIT: is_wrapped_traced_array) is simply ancestor(x) isa TracedRArray

[mofeing]

We can push an is_wrapped_array to ArrayInterface (easier for other packages to use rather than depending on Reactant).

super, i like it.

Then is_wrapped_tracked_array is simply ancestor(x) isa TracedRArray

..._traced_...?

the only remaining problem is then how to call the original method in the native interpreter

[avik-pal]

the only remaining problem is then how to call the original method in the native interpreter

Looking at the current implementation, we could possibly do:

struct UseNativeInterpreter end # Open to other names as well

--- a/src/Interpreter.jl
+++ b/src/Interpreter.jl
@@ -43,7 +43,8 @@ function set_reactant_abi(
     # the original call
     if f === Reactant.call_with_reactant
         arginfo2 = ArgInfo(fargs isa Nothing ? nothing : fargs[2:end], argtypes[2:end])
-        return abstract_call(interp, arginfo2::ArgInfo, si, sv, max_methods)
+        ret = abstract_call(interp, arginfo2::ArgInfo, si, sv, max_methods)
+        ret isa UseNativeInterpreter || return ret
     end
 
     return Base.@invoke abstract_call_known(

@reactant_override function LinearAlgebra.diag(x::AbstractArray{T,2}, k::Integer=0) where {T}
    if iswrappedtraced(x)
        y = materialize_traced_array(x)
        return diag(y, k)
    end
    return UseNativeInterpreter()
end

[mofeing]

uuu i like it. i'm not very familiar with set_reactant_abi but it's that what @wsmoses added for solving dynamic dispatch on abs int?

[wsmoses]

if you define the reactant override function as noinline, by default all methods within the override functions will be executed in the native interpreter

[wsmoses]

to re-enter the abstract interp override, do call_with_reactant(f, args...). This same behavior is true of anything in Ops, TracedUtils, or is otherwise in this whitelist:

Reactant.jl/src/utils.jl

Line 92 in 65e9976

function should_rewrite_ft(@nospecialize(ft))

[mofeing]

to re-enter the abstract interp override, do call_with_reactant(f, args...). This same behavior is true of anything in Ops, TracedUtils, or is otherwise in this whitelist:

aha, now i get it

[wsmoses]

@avik-pal is there anything more pending here or good to merge?

[avik-pal]

@avik-pal is there anything more pending here or good to merge?

I will add some tests and then we should be good to merge

[avik-pal]

Let's merge #426 first. We need that to update the dispatches correctly.

[wsmoses]

@avik-pal after this [and perhaps the specialfunctions PR] lets cut a release

[avik-pal]

tests that should be passing are passing.

https://github.com/EnzymeAD/Reactant.jl/actions/runs/12537476469/job/34961608153?pr=369#step:9:823 is probably unrelated to this PR

[wsmoses]

Yeah that one feels like a jax error, let's merge