WIP: fix RandAugment behavior in tf graph mode

[gregstarr]

fix for #21169

RandAugment changes:

• samples random execution order
• samples transform params for all sub layers
• uses fori and switch to execute the sub layers according to the sampled order

this is passing the tests in all three backends. However for the test_rand_augment_model test, jax runs very slowly, way slower than the other backends, by a factor of 10 - 50. I have no experience with jax, can you suggest anything that might help?

thanks

[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

All modified and coverable lines are covered by tests ✅

Project coverage is 73.11%. Comparing base (6d52164) to head (57ed25a).
Report is 9 commits behind head on master.

❗ There is a different number of reports uploaded between BASE (6d52164) and HEAD (57ed25a). Click for more details.

HEAD has 4 uploads less than BASE

Flag	BASE (6d52164)	HEAD (57ed25a)
keras	5	3
keras-numpy	1	0
keras-jax	1	0

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21185      +/-   ##
==========================================
- Coverage   82.69%   73.11%   -9.58%     
==========================================
  Files         564      564              
  Lines       54226    54404     +178     
  Branches     8425     8448      +23     
==========================================
- Hits        44840    39776    -5064     
- Misses       7310    12711    +5401     
+ Partials     2076     1917     -159     

Flag	Coverage Δ
keras	72.96% <100.00%> (-9.54%)	⬇️
keras-jax	?
keras-numpy	?
keras-openvino	32.94% <7.69%> (-0.04%)	⬇️
keras-tensorflow	64.28% <100.00%> (-0.02%)	⬇️
keras-torch	64.00% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ 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.

[shashaka]

In my opinion, some exception handling should be added for the NumPy backend in your test case, particularly due to the use of the fit method.

[gregstarr]

In my opinion, some exception handling should be added for the NumPy backend in your test case, particularly due to the use of the fit method.

yes I agree, the test should be skipped for numpy backend.

Any ideas about the layer implementation? it seems to have problems for jax backend

[shashaka]

I’m not entirely sure about that. I think a step-by-step breakdown is needed to identify what’s causing the slowdown in the JAX backend.

[gregstarr]

I’m not entirely sure about that. I think a step-by-step breakdown is needed to identify what’s causing the slowdown in the JAX backend.

Ok I will take a look when I get a chance.