Keras <> NNX integration

[divyashreepathihalli]

No description provided.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 17.24138% with 48 lines in your changes missing coverage. Please review.

Project coverage is 77.57%. Comparing base (cbb3682) to head (4f7b3b8).

Files with missing lines	Patch %	Lines
keras/src/backend/jax/core.py	22.22%	35 Missing ⚠️
keras/src/backend/jax/layer.py	0.00%	8 Missing ⚠️
keras/src/backend/jax/trainer.py	0.00%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21252      +/-   ##
==========================================
- Coverage   82.57%   77.57%   -5.00%     
==========================================
  Files         564      564              
  Lines       54629    54679      +50     
  Branches     8495     8504       +9     
==========================================
- Hits        45108    42419    -2689     
- Misses       7432    10254    +2822     
+ Partials     2089     2006      -83     

Flag	Coverage Δ
keras	77.42% <17.24%> (-4.97%)	⬇️
keras-jax	?
keras-numpy	58.72% <17.24%> (-0.04%)	⬇️
keras-openvino	32.99% <17.24%> (-0.02%)	⬇️
keras-tensorflow	64.00% <17.24%> (-0.05%)	⬇️
keras-torch	63.66% <17.24%> (-0.05%)	⬇️

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.

[fchollet]

Would the integration prevent the use of jax.jit with Keras layers?