Keras <> NNX integration

[divyashreepathihalli]

No description provided.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 22.53521% with 55 lines in your changes missing coverage. Please review.

Project coverage is 77.57%. Comparing base (cbb3682) to head (b87c4f9).
Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/core.py	28.57%	40 Missing ⚠️
keras/src/backend/jax/layer.py	0.00%	10 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    54719      +90     
  Branches     8495     8505      +10     
==========================================
- Hits        45108    42448    -2660     
- Misses       7432    10264    +2832     
+ Partials     2089     2007      -82     

Flag	Coverage Δ
keras	77.41% <22.53%> (-4.97%)	⬇️
keras-jax	?
keras-numpy	58.71% <22.53%> (-0.05%)	⬇️
keras-openvino	33.01% <22.53%> (+<0.01%)	⬆️
keras-tensorflow	63.99% <22.53%> (-0.06%)	⬇️
keras-torch	63.65% <22.53%> (-0.06%)	⬇️

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?

[divyashreepathihalli]

yes! it would only work with nnx.jit for now ( They might be working on adding support for jax.jit)