Exposing the experimental _split_transpose JAX scan parameter in Flax.

[copybara-service]

Exposing the experimental _split_transpose JAX scan parameter in Flax.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 3 lines in your changes are missing coverage. Please review.

Project coverage is 60.16%. Comparing base (fe99574) to head (d14c3c1).

Files	Patch %	Lines
flax/core/axes_scan.py	0.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3795      +/-   ##
==========================================
- Coverage   60.17%   60.16%   -0.01%     
==========================================
  Files         101      101              
  Lines       12840    12842       +2     
==========================================
  Hits         7727     7727              
- Misses       5113     5115       +2     

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