search

EnzymeAD / Reactant.jl

MIT License
54 stars 4 forks source link

Support closures #36

Closed mofeing closed 3 months ago

mofeing commented 3 months ago

This PR adds support for closures. The main changes are to traced_type, so that tracing correctly generates the new type of the traced closure, and to the generated Thunk function so that it correctly prepends the closure as an argument.

It also provides a test.

mofeing commented 3 months ago

mmm none of the checks run the tests?

wsmoses commented 3 months ago

oof we should fix that

wsmoses commented 3 months ago

Alternatively would it make sense to modify the expressions that use args to do an offset.

That’s what’s done for example in the AD rewrite and the batch PR

On Mon, Jul 1, 2024 at 11:23 PM Sergio Sánchez Ramírez < @.***> wrote:

@.**** commented on this pull request.

In src/Reactant.jl https://github.com/EnzymeAD/Reactant.jl/pull/36#discussion_r1661595334:

@@ -865,6 +892,13 @@ end resexpr = create_result(concrete_result_ty, (), result_stores) expr = quote Base.@_inline_meta

  • $(

  • if f is a closure, then prepend the closure into args

  • the closure fields will be correctly extracted from it as the tracer has already passed through it

  • if !(closure_ty <: Nothing)
  • :(args = [thunk.fnwrap, args...])

Sure!

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Reactant.jl/pull/36#discussion_r1661595334, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXBZKW5MSRSZLF2H5F3ZKHJFVAVCNFSM6AAAAABKFWK2QOVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDCNJSGMZDQNZRGA . You are receiving this because your review was requested.Message ID: @.***>

mofeing commented 3 months ago

That's what I started to do, but it soon started to get quirky. The reason is that since we trace Reactant.apply(f, args...) instead of f(args...), then linear_arg_paths and preserved_arg_paths consider that f will be the first argument (e.g. the enclosed ConcreteRArrays in f will be read from args[1] = f).

It can be done, but involves a lot more refactoring and probably changing the way we trace the closures. I did it this way because (1) it has no performance penalty and (2) modifies the least amount of code.

wsmoses commented 3 months ago

Yeah that sounds reasonable to me!

On Mon, Jul 1, 2024 at 11:49 PM Sergio Sánchez Ramírez < @.***> wrote:

That's what I started to do, but it soon started to get quirky. The reason is that since we trace Reactant.apply(f, args...) instead of f(args...), then linear_arg_paths and preserved_arg_paths consider that f will be the first argument (e.g. the enclosed ConcreteRArrays in f will be read from args[1] = f).

It can be done, but involves a lot more refactoring and probably changing the way we trace the closures. I did it this way because (1) it has no performance penalty and (2) modifies the least amount of code.

— Reply to this email directly, view it on GitHub https://github.com/EnzymeAD/Reactant.jl/pull/36#issuecomment-2201285692, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXEQQS3LVZMO7V5Q75LZKHMHVAVCNFSM6AAAAABKFWK2QOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEMBRGI4DKNRZGI . You are receiving this because your review was requested.Message ID: @.***>

mofeing commented 3 months ago

I'm merging this since tests for this feature pass. The failing tests are unrelated to this PR and are the CI's configuration fault.