feat: batched_jacobian for Reactant

[avik-pal]

• Implementation
• Testing
• Upstream
  • feat: reverse mode batched AD EnzymeAD/Reactant.jl#1740

[gemini-code-assist]

Summary of Changes

Hello @avik-pal, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the automatic differentiation capabilities by integrating batched_jacobian functionality for the Reactant.jl backend, utilizing Enzyme.jl. The changes encompass the core implementation for both forward and reverse differentiation modes, alongside a strategic refactoring of Enzyme-related utility functions into a shared Utils module. This not only streamlines the AD API but also improves the configuration and management of the ReactantBackend for various Enzyme settings, ultimately providing more robust and flexible gradient computations.

Highlights

• Batched Jacobian for Reactant: Introduced a new batched_jacobian implementation specifically for the Reactant.jl automatic differentiation backend, leveraging Enzyme.jl for efficient computation in both forward and reverse modes.
• Utility Function Refactoring: Refactored Enzyme.jl-specific utility functions, normalize_autoenzyme_mode and annotate_enzyme_function, into a centralized Utils module. This improves code organization and reusability across different extensions.
• Enhanced Reactant Backend Configuration: The ReactantBackend struct has been updated to store the AutoEnzyme instance directly, allowing for more precise configuration and integration of Enzyme's differentiation modes.
• Generalized AD API: The batched_jacobian API has been generalized to simplify the addition of new automatic differentiation backends, now explicitly including AutoEnzyme for Reactant.jl.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for batched_jacobian with the Reactant.jl backend. The implementation for both forward and reverse modes using Enzyme is added in a new file. The changes also include refactoring of some Enzyme-related utility functions into a central utils.jl module and cleaning up the batched_jacobian API definitions. The overall changes are well-structured and improve the codebase. I've found one critical issue where a function was renamed/moved, but one of its usages was not updated, which will lead to a runtime error.

[codecov]

Codecov Report

❌ Patch coverage is 12.19512% with 36 lines in your changes missing coverage. Please review.
✅ Project coverage is 16.78%. Comparing base (c254958) to head (ef7f2f5).

Files with missing lines	Patch %	Lines
ext/LuxReactantExt/batched_jacobian.jl	0.00%	30 Missing ⚠️
src/utils.jl	0.00%	4 Missing ⚠️
ext/LuxEnzymeExt/autodiff.jl	0.00%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (c254958) and HEAD (ef7f2f5). Click for more details.

HEAD has 44 uploads less than BASE

Flag	BASE (c254958)	HEAD (ef7f2f5)
	52	8

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1507       +/-   ##
===========================================
- Coverage   82.79%   16.78%   -66.01%     
===========================================
  Files         168      156       -12     
  Lines        6939     6784      -155     
===========================================
- Hits         5745     1139     -4606     
- Misses       1194     5645     +4451     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Benchmark Results (Julia v1.11)

Time benchmarks

	main	ef7f2f5...	main / ef7f2f5...
basics/MHA	5.76 ± 2.9 μs	5.79 ± 2.9 μs	0.995 ± 0.7
basics/MHA (first run)	4.46 ± 0.5 μs	4.45 ± 0.55 μs	1 ± 0.17
basics/MHA reactant	0.059 ± 0.0051 ms	0.0564 ± 0.0039 ms	1.05 ± 0.12
basics/MHA reactant (comp + run)	0.134 ± 0.0023 s	0.133 ± 0.0023 s	1 ± 0.025
basics/conv	12.8 ± 19 μs	12.3 ± 16 μs	1.05 ± 2.1
basics/conv (first run)	10.4 ± 0.6 μs	10.5 ± 0.61 μs	0.992 ± 0.081
basics/conv reactant	0.0497 ± 0.0025 ms	0.0565 ± 0.0044 ms	0.88 ± 0.081
basics/conv reactant (comp + run)	0.0874 ± 0.0047 s	0.0879 ± 0.004 s	0.995 ± 0.07
basics/dense	0.181 ± 0.011 μs	0.18 ± 0.001 μs	1.01 ± 0.061
basics/dense (first run)	0.21 ± 0.001 μs	0.201 ± 0.01 μs	1.04 ± 0.052
basics/dense reactant	0.0465 ± 0.0023 ms	0.051 ± 0.0047 ms	0.912 ± 0.096
basics/dense reactant (comp + run)	0.0688 ± 0.0024 s	0.069 ± 0.0025 s	0.998 ± 0.05
time_to_load	0.889 ± 0.0022 s	0.894 ± 0.0035 s	0.995 ± 0.0046

Memory benchmarks

	main	ef7f2f5...	main / ef7f2f5...
basics/MHA	0.087 k allocs: 6.05 kB	0.087 k allocs: 6.05 kB	1
basics/MHA (first run)	0.087 k allocs: 6.05 kB	0.087 k allocs: 6.05 kB	1
basics/MHA reactant	19 allocs: 0.578 kB	19 allocs: 0.578 kB	1
basics/MHA reactant (comp + run)	17.8 k allocs: 1.34 MB	17.8 k allocs: 1.34 MB	1
basics/conv	0.038 k allocs: 5.12 kB	0.038 k allocs: 5.12 kB	1
basics/conv (first run)	0.038 k allocs: 5.12 kB	0.038 k allocs: 5.12 kB	1
basics/conv reactant	15 allocs: 0.438 kB	15 allocs: 0.438 kB	1
basics/conv reactant (comp + run)	6.13 k allocs: 0.791 MB	6.13 k allocs: 0.791 MB	1
basics/dense	2 allocs: 0.109 kB	2 allocs: 0.109 kB	1
basics/dense (first run)	2 allocs: 0.109 kB	2 allocs: 0.109 kB	1
basics/dense reactant	15 allocs: 0.422 kB	15 allocs: 0.422 kB	1
basics/dense reactant (comp + run)	5.87 k allocs: 0.773 MB	5.87 k allocs: 0.773 MB	1
time_to_load	0.159 k allocs: 11.2 kB	0.159 k allocs: 11.2 kB	1