Add Reactant

[wsmoses]

It's not quite an AD engine in its own right, but given how it traces things and can make for much better performance/compatibility, and can be used with an AD engine, it seems like a natural fit.

[gdalle]

Why not but I don't understand what the package does or how it would fit within DI, so you'll have to hold my hand on this one ^^

[wsmoses]

Yeah so basically, Reactant exports a compile utility which takes a julia function and arguments, and compiles it to a nice/fast version of it. It does require that all data one cares about to be passed in using ConcreteRArrays rather than Ararys [and structs of those/etc]. We do have a tracer utility which will take an arg and auotmatically changes arrays to ConcreteRArrays for you though.

This can obviously compile code with an autodiff call on the inside, such as Enzyme. See here: https://github.com/EnzymeAD/Reactant.jl/blob/292dc03593ceb1a7a1f022fd7d3289bd69b000b5/test/basic.jl#L81. Enzyme usage should end up with quite good performance since the fancy optimizations that Reactant applies will be able to interoperate with AD.

[wsmoses]

We also now have a brief tldr in the readme: https://github.com/EnzymeAD/Reactant.jl/tree/main

[gdalle]

So how would you see this in the context of DI? As a variant of the Enzyme backend which Reactant-compiles the function during preparation?

[wsmoses]

I think the other way around. Reactant can take another AD backend (perhaps with enzyme as default)

…

On Fri, May 17, 2024 at 2:34 AM Guillaume Dalle ***@***.***> wrote: So how would you see this in the context of DI? As a variant of the Enzyme backend which Reactant-compiles the function during preparation? — Reply to this email directly, view it on GitHub <#265 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXH6GJ5F3VCRQQWRCYLZCXFMDAVCNFSM6AAAAABH3ROHEKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMJXGE2DMMZYGM> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[gdalle]

And do what with it? Compile and convert everything before and after use?

[wsmoses]

Yup

[gdalle]

I would welcome a PR with an example!

[wsmoses]

I'm not quite sure how to make an example PR, but here's an example from our tests: https://github.com/EnzymeAD/Reactant.jl/blob/59d9304948cc7e28acdd4351db5e069d62a4f4ec/test/basic.jl#L92

function sumcos(x)
    return sum(cos.(x))
end

function resgrad_ip(x)
    dx = Enzyme.make_zero(x)
    res = Enzyme.autodiff(ReverseWithPrimal, sumcos, Active, Duplicated(x, dx))
    return (res, dx)
end

 c = Reactant.ConcreteRArray(ones(3, 2))
    f = Reactant.compile(resgrad_ip, (c,))
    orig, r = f(c)

    @test orig[2] ≈ sum(cos.(ones(3, 2)))
    @test r ≈ -sin.(ones(3, 2))

Essentially you have to convert all inputs to ConcreteRarrays, then compile your desired function (which in this case will call an autodiff tool).

Happy to hop on a call or help however else!

This would be really great to get in and would give me a lot more use cases for DI.

Also fun fact the Reactant backend wouldn't suffer from any potential closure issues (since it would trace closures out anyways)

[gdalle]

I'm open to testing it but I can't do much until it is registered

[wsmoses]

oh it's been registered for a while now!

[gdalle]

My bad, I only checked the GitHub repo and the releases aren't tagged there

[gdalle]

I wonder how to add this to DI with as little code duplication as possible. I see several options:

• A backend wrapper ReactantBackend(AutoSomething)
• A function wrapper ReactantFunction(f)
• An option in the preparation prepare_operator(f, backend, x; reactant=true)

The first one seems easier to just put in a Reactant.jl extension