Fix: Validate class_weight keys in model.fit#23222
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces validation for class_weight keys to ensure they are integers and match the target labels, along with a type check for the trainable field during deserialization. The feedback highlights that the trainable validation is out of scope and fragile, potentially breaking backward compatibility. Additionally, the reviewer points out a performance bottleneck caused by duplicate NumPy conversions of target labels and an unhandled ValueError when processing empty class_weight dictionaries in the TensorFlow dataset adapter.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if "trainable" in inner_config and not isinstance( | ||
| inner_config["trainable"], bool | ||
| ): | ||
| raise TypeError( | ||
| "Invalid type for `trainable` field. " | ||
| f"Expected bool. Received: {inner_config['trainable']} " | ||
| f"of type {type(inner_config['trainable'])}." | ||
| ) |
There was a problem hiding this comment.
This validation logic is completely unrelated to the PR's scope, which is to validate class_weight keys in model.fit(). Unrelated changes should be removed from this PR.\n\nFurthermore, this validation is fragile and can break backward compatibility. In Python, bool is a subclass of int, but isinstance(1, bool) is False. If any legacy or third-party serialized model config uses 1 or 0 for trainable (which is common in some JSON serializers), this check will raise a TypeError and prevent loading the model. Please remove this change.
| if y is not None: | ||
| y_numpy = ops.convert_to_numpy(y) | ||
| if len(y_numpy.shape) > 1: | ||
| if y_numpy.shape[-1] != 1: | ||
| y_numpy = np.argmax(y_numpy, axis=-1) | ||
| else: | ||
| y_numpy = np.squeeze(y_numpy, axis=-1) | ||
| y_numpy = np.round(y_numpy).astype("int32") | ||
|
|
||
| unique_classes = set(int(c) for c in np.unique(y_numpy)) | ||
| if unique_classes and class_weight_int: | ||
| intersection = unique_classes.intersection(class_weight_int.keys()) | ||
| if not intersection: | ||
| raise ValueError( | ||
| f"None of the keys in `class_weight` match any class index " | ||
| f"present in the target labels. " | ||
| f"class_weight keys: {list(class_weight_int.keys())}. " | ||
| f"Unique classes in target data: {list(unique_classes)}." | ||
| ) |
There was a problem hiding this comment.
Converting y to a NumPy array and processing its shape is duplicated exactly between validate_class_weight and class_weight_to_sample_weights. Since ops.convert_to_numpy(y) can be an expensive operation (especially when transferring data from GPU/TPU to host memory), doing it twice introduces a significant performance bottleneck.\n\nWe can optimize this by performing the conversion and shape processing once in class_weight_to_sample_weights, and then passing the processed 1D NumPy array of class indices directly to validate_class_weight.
if y is not None:\n unique_classes = set(int(c) for c in np.unique(y))\n if unique_classes and class_weight_int:\n intersection = unique_classes.intersection(class_weight_int.keys())\n if not intersection:\n raise ValueError(\n f\"None of the keys in `class_weight` match any class index \"\n f\"present in the target labels. \"\n f\"class_weight keys: {list(class_weight_int.keys())}. \"\n f\"Unique classes in target data: {list(unique_classes)}.\"\n )| """ | ||
| from keras.src.utils.module_utils import tensorflow as tf | ||
|
|
||
| class_weight = data_adapter_utils.validate_class_weight(class_weight) |
There was a problem hiding this comment.
If class_weight is empty (e.g., {}), validate_class_weight will return {}. In that case, calling max(class_weight.keys()) on line 168 will raise a ValueError: max() arg is an empty sequence.\n\nWe should handle empty class_weight gracefully by returning an identity mapping function.
class_weight = data_adapter_utils.validate_class_weight(class_weight)\n if not class_weight:\n return lambda *data: data| def test_trainable_validation(self): | ||
| """Tests that invalid `trainable` types raise a TypeError.""" | ||
| layer = keras.layers.Dense(units=4) | ||
| config = keras.saving.serialize_keras_object(layer) | ||
|
|
||
| # Test valid boolean | ||
| config["config"]["trainable"] = True | ||
| restored = keras.saving.deserialize_keras_object(config) | ||
| self.assertTrue(restored.trainable) | ||
|
|
||
| # Test invalid type (string) | ||
| config["config"]["trainable"] = "true" | ||
| with self.assertRaisesRegex( | ||
| TypeError, "Invalid type for `trainable` field" | ||
| ): | ||
| keras.saving.deserialize_keras_object(config) | ||
|
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23222 +/- ##
==========================================
- Coverage 84.80% 84.04% -0.76%
==========================================
Files 465 465
Lines 69201 69227 +26
Branches 11379 11388 +9
==========================================
- Hits 58687 58184 -503
- Misses 7577 8116 +539
+ Partials 2937 2927 -10
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
4870071 to
33a6433
Compare
33a6433 to
fd42e66
Compare
|
Hi @gbaned - This PR has been ready for review for 1 day and all CI checks are passing. Summary:
Status: ✅ Ready for review and merge Could you please review and approve when ready? |
This PR introduces validation logic to ensure that keys passed to
class_weightinmodel.fit()are valid class indices present in the target labels, raising a clearValueErrorfor invalid keys instead of silently ignoring them. Fixes #23220.Contributor Agreement
Please review our AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:
Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.