IGNORE THIS PR #21389

divyashreepathihalli · 2025-06-16T23:14:37Z

just using this for some tests

codecov-commenter · 2025-06-16T23:20:38Z

Codecov Report

Attention: Patch coverage is 17.74194% with 204 lines in your changes missing coverage. Please review.

Project coverage is 82.49%. Comparing base (22ab5a3) to head (7493c68).
Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/core.py	6.30%	103 Missing and 1 partial ⚠️
keras/src/models/functional.py	10.52%	49 Missing and 2 partials ⚠️
keras/src/backend/config.py	30.55%	20 Missing and 5 partials ⚠️
keras/src/layers/layer.py	42.85%	5 Missing and 3 partials ⚠️
keras/src/backend/jax/layer.py	20.00%	3 Missing and 1 partial ⚠️
keras/src/backend/common/variables.py	0.00%	2 Missing and 1 partial ⚠️
keras/src/backend/jax/trainer.py	66.66%	2 Missing and 1 partial ⚠️
keras/src/ops/operation.py	25.00%	2 Missing and 1 partial ⚠️
keras/src/backend/jax/__init__.py	50.00%	1 Missing and 1 partial ⚠️
keras/api/_tf_keras/keras/config/__init__.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21389      +/-   ##
==========================================
- Coverage   82.75%   82.49%   -0.27%     
==========================================
  Files         565      565              
  Lines       55346    55571     +225     
  Branches     8634     8691      +57     
==========================================
+ Hits        45804    45842      +38     
- Misses       7440     7614     +174     
- Partials     2102     2115      +13

Flag	Coverage Δ
keras	`82.30% <17.74%> (-0.27%)`	⬇️
keras-jax	`63.13% <17.74%> (-0.21%)`	⬇️
keras-numpy	`58.36% <13.30%> (-0.20%)`	⬇️
keras-openvino	`33.69% <12.50%> (-0.08%)`	⬇️
keras-tensorflow	`63.55% <13.30%> (-0.22%)`	⬇️
keras-torch	`63.17% <13.30%> (-0.21%)`	⬇️

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.

Fix NNX tracing issue for functional

divyashreepathihalli · 2025-07-11T00:38:01Z

/gemini review

gemini-code-assist

Code Review

This pull request integrates Flax NNX support into the JAX backend, involving configuration, core components, layers, and models. Key areas for improvement include test dependency alignment, addressing disabled tests, and enhancing error handling in configuration parsing.

gemini-code-assist · 2025-07-11T00:39:27Z

integration_tests/import_test.py

-    "jax": ("jax[cpu]", ""),
+    # please update the jax version here if jax version is updated in
+    # requirements file
+    "jax": ("jax[cpu]==0.5.1 flax>=0.10.1", ""),


The jax version specified here (0.5.1) differs from the one in requirements.txt (0.5.0). This inconsistency can lead to testing discrepancies. Align the versions to ensure a consistent testing environment.

gemini-code-assist · 2025-07-11T00:39:27Z

integration_tests/numerical_test.py

+#             keras_module.layers.MaxPooling2D(pool_size=(2, 2)),
+#             keras_module.layers.Flatten(),
+#             keras_module.layers.Dense(num_classes, activation="softmax"),
+#         ]
+#     )
+#     return model
+
+
+# def compile_model(model):
+#     model.compile(
+#         loss="categorical_crossentropy",
+#         optimizer="adam",
+#         metrics=["mae", "accuracy"],
+#         jit_compile=False,
+#         run_eagerly=True,
+#     )
+
+
+# def train_model(model, x, y):
+#     return model.fit(
+#         x,
+#         y,
+#         batch_size=BATCH_SIZE,
+#         epochs=EPOCHS,
+#         shuffle=False,
+#         verbose=0,
+#     )
+
+
+# def eval_model(model, x, y):
+#     score = model.evaluate(x, y, verbose=0, batch_size=BATCH_SIZE)
+#     print(score)
+#     return score
+
+
+# def check_history(h1, h2):
+#     for key in h1.history.keys():
+#         print(f"{key}:")
+#         print(h1.history[key])
+#         print(h2.history[key])
+#         np.testing.assert_allclose(
+#             h1.history[key],
+#             h2.history[key],
+#             atol=1e-3,
+#         )
+
+
+# def predict_model(model, x):
+#     return model.predict(x, batch_size=BATCH_SIZE, verbose=0)
+
+
+# def numerical_test():
+#     x_train, y_train = build_mnist_data(NUM_CLASSES)
+#     keras_model = build_keras_model(keras, NUM_CLASSES)
+#     tf_keras_model = build_keras_model(tf_keras, NUM_CLASSES)
+
+#     # Make sure both model have same weights before training
+#     weights = [weight.numpy() for weight in keras_model.weights]
+#     tf_keras_model.set_weights(weights)
+
+#     for kw, kcw in zip(keras_model.weights, tf_keras_model.weights):
+#         np.testing.assert_allclose(kw.numpy(), kcw.numpy())
+
+#     compile_model(keras_model)
+#     compile_model(tf_keras_model)
+
+#     print("Checking training histories:")
+#     keras_history = train_model(keras_model, x_train, y_train)
+#     tf_keras_history = train_model(tf_keras_model, x_train, y_train)
+#     check_history(keras_history, tf_keras_history)
+#     print("Training histories match.")
+#     print()
+
+#     print("Checking trained weights:")
+#     for kw, kcw in zip(keras_model.weights, tf_keras_model.weights):
+#         np.testing.assert_allclose(kw.numpy(), kcw.numpy(), atol=1e-3)
+#     print("Trained weights match.")
+#     print()
+
+#     print("Checking predict:")
+#     outputs1 = predict_model(keras_model, x_train)
+#     outputs2 = predict_model(tf_keras_model, x_train)
+#     np.testing.assert_allclose(outputs1, outputs2, atol=1e-3)
+#     print("Predict results match.")
+#     print()
+
+#     print("Checking evaluate:")
+#     score1 = eval_model(keras_model, x_train, y_train)
+#     score2 = eval_model(tf_keras_model, x_train, y_train)
+#     np.testing.assert_allclose(score1, score2, atol=1e-3)
+#     print("Evaluate results match.")
+
+
+# if __name__ == "__main__":
+#     if keras.backend.backend() == "openvino":
+#         # this test requires trainable backend
+#         sys.exit(0)
+#     keras.utils.set_random_seed(1337)
+#     tf_keras.utils.set_random_seed(1337)
+#     numerical_test()


This test file is fully commented out, which removes its test coverage. If these tests are no longer needed, remove the file. If temporarily disabled, add a comment explaining the reason and a tracking issue for re-enabling.

gemini-code-assist · 2025-07-11T00:39:27Z

keras/src/backend/config.py

+    env_val = os.environ["KERAS_NNX_ENABLED"].lower()
+    if env_val == "true":
+        _NNX_ENABLED = True
+    elif env_val == "false":
+        _NNX_ENABLED = False


The parsing logic for KERAS_NNX_ENABLED only handles "true" and "false", ignoring other boolean representations. Enhance it to handle common boolean values like "1" and "0" for better usability.

Suggested change

env_val = os.environ["KERAS_NNX_ENABLED"].lower()

if env_val == "true":

_NNX_ENABLED = True

elif env_val == "false":

_NNX_ENABLED = False

if "KERAS_NNX_ENABLED" in os.environ:

env_val = os.environ["KERAS_NNX_ENABLED"].lower()

if env_val in ("true", "1"):

_NNX_ENABLED = True

elif env_val in ("false", "0"):

_NNX_ENABLED = False

gemini-code-assist · 2025-07-11T00:39:27Z

keras/src/layers/layer.py

+                    self._parent_path = current_path()
+                except Exception:
+                    pass


The try...except Exception: pass block can hide issues when getting current_path(). Consider catching a specific exception or logging the exception for debugging.

divyashreepathihalli added 30 commits May 5, 2025 23:03

_valu

8e1c008

update variables

7159709

add nnx.jit

e378cfb

revert changes to JaxLayer

8701fc7

Merge branch 'master' into nnx

b599727

format fix

4f7b3b8

make variables subclass nnx.Variable

0234e27

more tweaks

b87c4f9

update init

91b9a73

refactor jax Variable class

aee1789

code reformat

141487f

more cleanup

5ccc31e

Merge branch 'keras-team:master' into nnx

4e35416

update flax version

dd9c77d

update flax version

48983c6

fix jax error

b22f9ef

update Variables implementation

4dbffa6

fix import

627e581

add a test

f58ef60

needs updates in operation

c2b73b7

remove __new__ from JaxLayer

a662f5e

update base optimizers

396f973

code reformat+ model saving tests

30e971d

add __hash__

968d804

update variable value updates

b99571a

sync value properly

ed0bc00

update flag based routing between nnx and jax

460e0e2

clean up

34f27e9

fix circular import error

427ff82

fix is nnx call enabled flag

c4ee191

divyashreepathihalli marked this pull request as draft June 17, 2025 05:54

Merge branch 'keras-team:master' into nnx

94ea1a2

divyashreepathihalli marked this pull request as ready for review June 17, 2025 18:50

divyashreepathihalli marked this pull request as draft June 17, 2025 18:50

divyashreepathihalli added the kokoro:force-run label Jun 17, 2025

kokoro-team removed the kokoro:force-run label Jun 17, 2025

try to fix memory error

f0b10ef

divyashreepathihalli force-pushed the nnx-tests branch 3 times, most recently from 64b99f0 to f0b10ef Compare June 20, 2025 02:34

divyashreepathihalli added 5 commits June 20, 2025 02:53

revert memory fix

d18dd33

fix test

51aa455

fix test

a02f410

fix test

adca8da

revert version back

d8ca752

divyashreepathihalli force-pushed the nnx-tests branch 2 times, most recently from 485f776 to d8ca752 Compare June 24, 2025 02:39

divyashreepathihalli force-pushed the nnx-tests branch from f040e7f to 1bdf4e9 Compare July 8, 2025 00:36

divyashreepathihalli added 4 commits July 8, 2025 00:48

Merge remote-tracking branch 'origin/master' into nnx

95e67e0

Update functional.py

a5741be

Update core.py

02d607b

Merge pull request #1 from divyashreepathihalli/nnx-fix-fn

5500382

Fix NNX tracing issue for functional

divyashreepathihalli force-pushed the nnx-tests branch from 9d5349c to 5500382 Compare July 9, 2025 04:35

divyashreepathihalli added 4 commits July 9, 2025 04:38

fix integration test#

9cf8fc5

etg

2b4a4b3

fix integration test

fda0c24

jhv

7493c68

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

IGNORE THIS PR #21389

IGNORE THIS PR #21389

Uh oh!

divyashreepathihalli commented Jun 16, 2025

Uh oh!

codecov-commenter commented Jun 16, 2025 •

edited

Loading

Uh oh!

divyashreepathihalli commented Jul 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Uh oh!

gemini-code-assist bot Jul 11, 2025

Uh oh!

gemini-code-assist bot Jul 11, 2025

Uh oh!

gemini-code-assist bot Jul 11, 2025

Uh oh!

Uh oh!

IGNORE THIS PR #21389

Are you sure you want to change the base?

IGNORE THIS PR #21389

Uh oh!

Conversation

divyashreepathihalli commented Jun 16, 2025

Uh oh!

codecov-commenter commented Jun 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

divyashreepathihalli commented Jul 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov-commenter commented Jun 16, 2025 •

edited

Loading