IGNORE THIS PR #21389
Draft
IGNORE THIS PR #21389
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
divyashreepathihalli
force-pushed
the
nnx-tests
branch
2 times, most recently
from
485f776 to
d8ca752
Compare
divyashreepathihalli
force-pushed
the
nnx-tests
branch
from
f040e7f to
1bdf4e9
Compare
Fix NNX tracing issue for functional
divyashreepathihalli
force-pushed
the
nnx-tests
branch
from
9d5349c to
5500382
Compare
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
integration_tests/import_test.py
Outdated
| "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", ""), |
There was a problem hiding this comment.
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.
integration_tests/numerical_test.py
Outdated
Comment on lines
1
to
147
| # 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() |
There was a problem hiding this comment.
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.
keras/src/backend/config.py
Outdated
Comment on lines
445
to
448
| env_val = os.environ["KERAS_NNX_ENABLED"].lower() | ||
| if env_val == "true": | ||
| _NNX_ENABLED = True | ||
| elif env_val == "false": | ||
| _NNX_ENABLED = False |
There was a problem hiding this comment.
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 |
keras/src/layers/layer.py
Outdated
Comment on lines
1677
to
1679
| self._parent_path = current_path() | ||
| except Exception: | ||
| pass |
There was a problem hiding this comment.
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
force-pushed
the
nnx-tests
branch
from
7493c68 to
46db09b
Compare
divyashreepathihalli
force-pushed
the
nnx-tests
branch
from
8453c4c to
3683dc8
Compare
divyashreepathihalli
force-pushed
the
nnx-tests
branch
from
d7bbcdd to
587bae7
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
just using this for some tests