Fix discretization symbolic one hot by bittoby · Pull Request #22048 · keras-team/keras

bittoby · 2026-01-23T18:15:22Z

Fix Discretization layer symbolic tensor shape computation for one_hot mode

Before the fix:

Eager tensor: (2, 3, 4) → (2, 3, 4, 5) ✅

Symbolic tensor: (None, 3, 4) → (None, 3, 4) ❌ (missing one-hot dimension)

After the fix:

Eager tensor: (2, 3, 4) → (2, 3, 4, 5) ✅

Symbolic tensor: (None, 3, 4) → (None, 3, 4, 5) ✅

Solution

Added proper shape computation logic by implementing:

compute_output_shape method: Handles different output modes correctly

Updated compute_output_spec method: Now uses the new shape computation

Verification

The fix resolves the exact issue described in #22044:

### Create layer

layer = keras.layers.Discretization(
    bin_boundaries=[0.0, 1.0, 2.0, 3.0], 
    output_mode='one_hot'
)

Test symbolic tensor (this was broken before)

symbolic_input = keras.Input(shape=(3, 4))
symbolic_output = layer(symbolic_input)
print(symbolic_output.shape)  # Now correctly shows (None, 3, 4, 5)

Closes #22044

google-cla · 2026-01-23T18:15:27Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist · 2026-01-23T18:15:48Z

Summary of Changes

Hello @bittoby, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical issue where the "Discretization" layer in Keras failed to correctly compute output shapes for symbolic tensors when "output_mode" was set to "one_hot". The fix involves implementing a robust "compute_output_shape" method that accurately determines the output dimensions based on the chosen output mode, ensuring proper model building with Keras's Functional and Sequential APIs. This change also includes comprehensive new tests and a reusable utility for verifying symbolic shape consistency.

Highlights

Discretization Layer Output Shape Calculation: Implemented a new "compute_output_shape" method in the "Discretization" layer to accurately determine output shapes for symbolic tensors. For "one_hot" mode, it adds a new dimension representing the number of bins, while for "multi_hot" and "count" modes, it replaces the last input dimension with the number of bins.
Symbolic Tensor Output Specification Update: The "compute_output_spec" method in "Discretization" now correctly utilizes the new "compute_output_shape" method, ensuring proper shape inference for symbolic inputs across all output modes.
Comprehensive Symbolic Shape Testing: Added "test_symbolic_tensor_output_shape" to "discretization_test.py" to validate the output shapes of symbolic tensors for all "Discretization" layer output modes, specifically fixing and testing issue Discretization layer with output_mode='one_hot' produces incorrect output shape for symbolic tensors #22044.
Reusable Test Utility for Shape Consistency: Introduced "assert_symbolic_shape_consistency" in "keras/src/testing/test_utils.py", a utility function designed to verify the consistency between symbolic and eager tensor output shapes for Keras layers.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

The pull request effectively addresses the symbolic tensor shape computation issue for the Discretization layer, particularly for the one_hot output mode. The introduction of the compute_output_shape method and its integration into compute_output_spec correctly infers the output shapes for symbolic inputs. The added test cases in discretization_test.py provide good coverage for various output modes and ensure consistency between eager and symbolic execution. The new assert_symbolic_shape_consistency utility in test_utils.py is a valuable addition for future testing of preprocessing layers.

gemini-code-assist · 2026-01-23T18:17:12Z

keras/src/layers/preprocessing/discretization.py

+            return input_shape
+
+        # Calculate depth (number of bins)
+        depth = len(self.bin_boundaries) + 1 if self.bin_boundaries else self.num_bins


The current logic for calculating depth has a potential issue when self.bin_boundaries is an empty list. In Python, an empty list [] evaluates to False in a boolean context. If self.bin_boundaries is [], the expression len(self.bin_boundaries) + 1 if self.bin_boundaries else self.num_bins will incorrectly fall back to self.num_bins. According to the __init__ validation, if bin_boundaries is explicitly provided (even if empty), num_bins would be None, leading to depth being None and causing a runtime error. The intention is likely to use len(self.bin_boundaries) + 1 if self.bin_boundaries is not None, regardless of whether it's empty or not.

Suggested change

depth = len(self.bin_boundaries) + 1 if self.bin_boundaries else self.num_bins

depth = (len(self.bin_boundaries) + 1) if self.bin_boundaries is not None else self.num_bins

codecov-commenter · 2026-01-23T18:24:38Z

Codecov Report

❌ Patch coverage is 85.71429% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.80%. Comparing base (2ba60cc) to head (65370c4).
⚠️ Report is 10 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/layers/preprocessing/discretization.py	85.71%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##           master   #22048    +/-   ##
========================================
  Coverage   82.80%   82.80%            
========================================
  Files         592      592            
  Lines       62289    62476   +187     
  Branches     9751     9787    +36     
========================================
+ Hits        51576    51734   +158     
- Misses       8190     8211    +21     
- Partials     2523     2531     +8

Flag	Coverage Δ
keras	`82.63% <85.71%> (+<0.01%)`	⬆️
keras-jax	`62.40% <85.71%> (-0.03%)`	⬇️
keras-numpy	`56.50% <85.71%> (-0.07%)`	⬇️
keras-openvino	`37.62% <7.14%> (+0.04%)`	⬆️
keras-tensorflow	`63.66% <85.71%> (-0.01%)`	⬇️
keras-torch	`62.43% <85.71%> (-0.01%)`	⬇️

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.

hertschuh

Thank you for working on this bug!

hertschuh · 2026-01-23T22:59:13Z

keras/src/layers/preprocessing/discretization.py

+            # For multi_hot and count modes, the behavior is more complex
+            # and depends on the backend implementation. For now, we'll use
+            # a simplified approach that works for the most common cases.
+            # This may need refinement based on actual backend behavior.


How is it backend dependent? Can you provide examples?

Replacing the last dimension with "depth" should be correct always.

Yes, you are right. I just thought of TF fastpath. But as you mentioned, output_shape should always be like above.

Please remove this comment. It will cause confusion if somebody ever works on it.

hertschuh · 2026-01-23T23:01:48Z

keras/src/layers/preprocessing/discretization_test.py

+        layer_one_hot = layers.Discretization(
+            bin_boundaries=[0.0, 1.0, 2.0], output_mode="one_hot"
+        )
+        symbolic_input = layers.Input(shape=(3, 4))


Don't use 4 as the value for the last dimension. There are 4 bins, so won't be able to tell if it has actually been replaced with the "depth" value.

Just in general, using unique dimension values is a good approach.

Sure. I will fix it to unique dimension value. Thanks.

hertschuh · 2026-01-23T23:04:22Z

keras/src/layers/preprocessing/discretization_test.py

+        # Shapes should be consistent (ignoring batch dimension)
+        self.assertEqual(eager_output.shape[1:], symbolic_output.shape[1:])


Why do you exclude the batch size in your comparison? The batch size better not change.

You're absolutely right to question this! I just simply thought. The comment "ignoring batch dimension" is misleading and potentially incorrect. Let me compare the two batch sizes. Symbolic should be None while eager should be an integer.

hertschuh · 2026-01-23T23:05:33Z

keras/src/layers/preprocessing/discretization_test.py

+    def test_compute_output_shape_edge_cases(self):
+        """Test edge cases in compute_output_shape to improve coverage."""
+
+        # Test with num_bins instead of bin_boundaries
+        layer_num_bins = layers.Discretization(
+            num_bins=5, output_mode="one_hot"
+        )


Using num_bins hardly counts as and "edge case". Create a dedicated test called test_compute_output_shape_num_bins.

hertschuh · 2026-01-23T23:08:29Z

keras/src/layers/preprocessing/discretization_test.py

+        # Test different output modes
+        modes_and_shapes = [
+            ("int", (None, 3, 4), (None, 3, 4)),  # int mode - no change
+            ("one_hot", (None, 3, 4), (None, 3, 4, 3)),  # one_hot - add dim
+            ("multi_hot", (None, 3, 4), (None, 3, 3)),  # multi_hot - replace
+            ("count", (None, 3, 4), (None, 3, 3)),  # count - replace
+        ]
+
+        for mode, input_shape, expected_shape in modes_and_shapes:
+            layer = layers.Discretization(
+                bin_boundaries=[0.0, 1.0], output_mode=mode
+            )
+            result_shape = layer.compute_output_shape(input_shape)
+            self.assertEqual(result_shape, expected_shape)


As far as I can tell, this covers exactly the same use cases as test_symbolic_tensor_output_shape.

You know you can use parameterized tests instead of a for loop:
https://github.com/keras-team/keras/blob/master/keras/src/layers/preprocessing/discretization_test.py#L41-L93

Make sense.

hertschuh · 2026-01-23T23:09:52Z

keras/src/testing/test_utils.py

+def assert_symbolic_shape_consistency(
+    test_case, layer, input_shape, expected_output_shape=None
+):


This is not used and I don't think we need this. Revert this file.

bittoby · 2026-01-24T05:33:07Z

@hertschuh Thanks for your great advice. Plz review again.

hertschuh · 2026-01-24T20:10:09Z

keras/src/layers/preprocessing/discretization.py

+            # For multi_hot and count modes, the behavior is more complex
+            # and depends on the backend implementation. For now, we'll use
+            # a simplified approach that works for the most common cases.
+            # This may need refinement based on actual backend behavior.


Please remove this comment. It will cause confusion if somebody ever works on it.

hertschuh · 2026-01-24T20:13:30Z

keras/src/layers/preprocessing/discretization_test.py

        )
        self.assertTrue(backend.is_int_dtype(model_call_output.dtype))
+
+    def test_symbolic_tensor_output_shape(self):


Please make this parameterized exactly like test_compute_output_shape_modes below.

hertschuh · 2026-01-24T20:17:56Z

keras/src/layers/preprocessing/discretization_test.py

+        named_product(
+            [


You don't need named_product here since you're specifying the testcase_names. Just do:

@parameterized.named_parameters( [ { "testcase_name": "int_mode", ...

hertschuh · 2026-01-24T20:37:22Z

keras/src/layers/preprocessing/discretization_test.py

+        input_shape = (4, 3)
+
+        for batch_size in batch_sizes:
+            with self.subTest(batch_size=batch_size):


This works, but use a parameterized test instead. This is where named_product helps:

@parameterized.named_parameters(named_product(batch_size=[1, 3, 5, 10])) def test_batch_dimension_consistency(self, batch_size):

hertschuh · 2026-01-24T20:37:47Z

keras/src/layers/preprocessing/discretization_test.py

+        test_cases = [
+            {


Make this a proper parameterized test.

hertschuh · 2026-01-24T20:38:17Z

keras/src/layers/preprocessing/discretization_test.py

+        # Test various num_bins values with one_hot mode
+        num_bins_values = [2, 3, 5, 10, 20, 100]
+        input_shape = (None, 4, 3)
+
+        for num_bins in num_bins_values:


Make this a parameterized test, follow the batch_size pattern for this one.

hertschuh · 2026-01-25T03:48:57Z

keras/src/layers/preprocessing/discretization_test.py

+                    f"Failed for num_bins={num_bins}: expected "
+                    f"{expected_shape}, got {result_shape}",


No need for this, all this information will already be displayed.

hertschuh · 2026-01-25T03:51:47Z

keras/src/layers/preprocessing/discretization_test.py

+        # Test edge case: num_bins=1
+        layer_single_bin = layers.Discretization(
+            num_bins=1, output_mode="one_hot"
+        )
+        result = layer_single_bin.compute_output_shape((None, 3, 2))
+        expected = (None, 3, 2, 1)
+        self.assertEqual(result, expected)


No need for special code for this, just add 1 to the list of num_bins above.

hertschuh · 2026-01-25T03:52:47Z

keras/src/layers/preprocessing/discretization_test.py

+        # Test with different input shapes
+        input_shapes = [
+            (None, 5),  # 2D (batch + 1D data)
+            (None, 3, 4),  # 3D (batch + 2D data)
+            (None, 2, 3, 4),  # 4D (batch + 3D data)
+            (None, 1, 2, 3, 4),  # 5D (batch + 4D data)
+        ]
+
+        layer = layers.Discretization(num_bins=7, output_mode="one_hot")
+
+        for input_shape in input_shapes:
+            with self.subTest(input_shape=input_shape):
+                result_shape = layer.compute_output_shape(input_shape)
+                expected_shape = input_shape + (7,)  # Append depth for one_hot
+
+                self.assertEqual(
+                    result_shape,
+                    expected_shape,
+                    f"Failed for input shape {input_shape}: expected "
+                    f"{expected_shape}, got {result_shape}",
+                )


I don't think this is needed, you got compute_output_shape well covered.

bittoby · 2026-01-25T19:13:02Z

@hertschuh I fixed problems. Plz review again. Thanks.

bittoby · 2026-01-26T23:18:59Z

@hertschuh I already fixed all what you mentioned. Plz review again.

hertschuh

Last nitpicks on the unit tests:

hertschuh · 2026-01-26T23:38:48Z

keras/src/layers/preprocessing/discretization_test.py

+        # Test consistency between eager and symbolic tensors for one_hot mode
+        if output_mode == "one_hot":


Why not for all? There's nothing special about "one_hot", you can do it for all.

Yes, there should be all of the modes. Thanks.

hertschuh · 2026-01-26T23:44:05Z

keras/src/layers/preprocessing/discretization_test.py

+    @parameterized.named_parameters(
+        named_product(num_bins=[1, 2, 3, 5, 10, 20, 100])
+    )
+    def _test_num_bins_one_hot_mode(self, num_bins):
+        """Test various num_bins values with one_hot mode."""
+        input_shape = (None, 4, 3)
+
+        layer = layers.Discretization(num_bins=num_bins, output_mode="one_hot")
+
+        result_shape = layer.compute_output_shape(input_shape)
+        expected_shape = (None, 4, 3, num_bins)
+
+        self.assertEqual(result_shape, expected_shape)


This is already covered by one_hot_mode_4_bins line 434, you can remove.

Thanks, I will fix it.

bittoby · 2026-01-27T12:36:49Z

#hertschuh, thanks to your comments, I found another issue upon multi_hot and count mode and fix the problems. In the previous times, since #22044 issue is about one_hot mode, I only focused on one_hot mode. The issues are about output_shape in multi_hot and count modes and these are slightly different from expectations.

layer = Discretization(
    bin_boundaries=[0.0, 1.0, 2.0, 5.0], output_mode="multi_hot"
)

input_shape = (3, 4)
symbolic_input = layers.Input(shape=input_shape)
symbolic_output = layer(symbolic_input)

Output: 
Eager: Input (2, 3, 4) -> Output torch.Size([2, 4, 5])

Expectation:
Eager: Input (2, 3, 4) -> Output torch.Size([2, 3, 5])

Symbolic:
Input (None, 3, 4) -> Output torch.Size([None, 4, 5])

I fix the problem, which symbolic tensor output_shape is same to original eager tensor output_shape.
Thanks for your kind attention. Plz review again.

hertschuh

Thanks!

bittoby · 2026-01-28T11:11:02Z

hertschuh Thank you.

* Implement logaddexp2 function in keras.ops (keras-team#21691) * [Keras 3 OpenVINO Backend]: Support numpy.sort (keras-team#21687) * [Keras 3 OpenVINO Backend]: Support numpy.median operation (keras-team#21667) * Fix deadlock in `CallbackList`. (keras-team#21701) * [OpenVINO backend] solve randomuniform issue (keras-team#21670) * Bug fixes with variable handling in `LossScaleOptimizer`. (keras-team#21706) * Do not use backend ops in `ProgBar`. (keras-team#21709) * Fix the Doc of the combination relation in func `keras.layers.Normali… * Remove reliance on `__jax_array__` to unwrap variables. (keras-team#21719) * Bump the github-actions group with 6 updates (keras-team#21705) * Add linspace and logspace implementations in OpenVINO NumPy backend (… * Add jvp op (keras-team#21720) * Add unfold op (keras-team#21685) * Add the description that `0` should not in the arg `axes` in `keras.l… * Add daily Python 3.13 CPU-only tests to nightly workflow (keras-team#21566) * Fix histogram op for symbolic inputs (keras-team#21729) * Relax tolerance for svd test (keras-team#21731) * Use jax.enable_x64 in place of jax.experimental.disable_x64 (keras-team#21734) * Refactor variable serialization. (keras-team#21713) * Ensure keras.ops.eye behavior is consistent across backends. (keras-team#21738) * Add `eye` support for OpenVINO backend (keras-team#21739) * Update Torch and Tensorflow versions in cuda requirements files. (keras-team#21… * Implement isreal function in keras.ops (keras-team#21740) * Remove the unused jax `enable_x64`. (keras-team#21737) * Correct implementation for several OpenVINO operations (keras-team#21746) * Sets `is_gptq_calibrated` flag when deserializing GPTQ models (keras-team#21748) * Correct implementation for several OpenVINO operations (keras-team#21752) * Fix the Bug in func `preprocess_input` when `x` in 3D and `data_forma… * Update Torch to 2.9.0 on GPU. (keras-team#21756) * `StringLookup` & `IntegerLookup` now save vocabulary loaded from file… * Implement trapezoid function in keras.ops (keras-team#21757) * Upstream `ReversibleEmbedding` from KerasHub. (keras-team#21753) * Raise exception on batch_size mismatch for stateful RNNs (keras-team#21742) * Propose a method for handling datasets which doesn't explicitly requi… * Use `filter="data"` option of `TarFile.extractall`. (keras-team#21760) * Add Distillation API to Keras (keras-team#21572) * removes unnecessary try-catch blocks and guard conditions (keras-team#21767) * cleanup distillation loss names (keras-team#21766) * Document that `set_backend` requires re-importing keras. (keras-team#21764) * Fix discretization discrepancy (keras-team#21769) * fix sas metrics in jax `fit` (keras-team#21765) * Support for extracting volume patches (keras-team#21759) * Fix negative index handling in MultiHeadAttention attention_axes (keras-team#21… * Make confusion metrics compilable. (keras-team#21775) * Suport keras.op.view() to view the same data bitwise at a new dtype … * Fix: `keras.ops.quantile` works with tf graph execution (keras-team#21782) * Fix typo in Distiller docstring * Add warning to `set_backend` and more detailed example. (keras-team#21787) * Don't fail `Variable.__repr__` if the value cannot be retrieved. (keras-team#21… * Update Keras backend installation instructions * Fix: Support 'jpg' format in keras.utils.save_img() (keras-team#21683) * Fix tf dataset detection logic. (keras-team#21794) * update test after jax.config.jax_vjp3 is enabled (keras-team#21776) * Add keras.ops.array_split for Tensor Parallelism Support (keras-team#21697) * Adding get_device_count function to the distribution_lib (keras-team#21791) * Fix: use raw string for CALIBRATION_TEXT (keras-team#21790) * Add backend compatibility table to documentation (keras-team#21733) * More OpenVINO Operations (keras-team#21774) * Support scalar view for tf backend. (keras-team#21802) * Address bug with convolution using Tensorflow, Numpy, Jax backends (#… * Fix bug with correlate for tensorflow (keras-team#21778) * Pass optional field in a few places to fix None input error. (keras-team#21818) * Fix(backend/torch): Resolve MPS broadcast crash in binary_crossentrop… * Fix broken example indentation in Keras io (keras-team#21807) * Add missing `convert_to_tensor` to `take_along_axis` on JAX. (keras-team#21825) * Added numpy.digitize support for OPENVINO backend (keras-team#21824) * Bump the github-actions group with 4 updates (keras-team#21809) * Fix typo in CONTRIBUTING.md (keras-team#21812) * Fix `Progbar.update` when receiving list, np arrays, and tensors. (#2… * Fix CosineDecay documentation to clarify alpha is a multiplier (keras-team#21827) * Fix noise_shape validation in keras.layers.Dropout (keras-team#21819) * Fix typos in some files (keras-team#21830) * Fix failing sklearn tests following release of pytest 9.0. (keras-team#21843) * Implement empty_like function in keras.ops (keras-team#21840) * Run tests on TPU (keras-team#21425) * Fix typo in variable name 'embeding' to 'embedding' (keras-team#21845) * Fix name_scope_stack AttributeError and IndexError in __exit__ (keras-team#21834) * Update keras3 Softmax mask handling to be more numerically robust. (#… * Support jax2tf in JaxLayer for tf backend (keras-team#21842) * Fix assigning a value to a variable within an autocast scope. (keras-team#21864) * Add note about label noise in CIFAR-10 dataset documentation (keras-team#21855) * Allow None inputs in `Layer.build`. (keras-team#21866) * `standardize_shape` normalizes the dimensions and tuple. (keras-team#21867) * Improve error message when layer/model input validation fails. (keras-team#21869) * Add verbose logging when ModelCheckpoint callback is done saving ... … * [OpenVINO backend] Remove deprecated openvino.runtime import (keras-team#21826) * Fix Torch output_padding constraint for ConvTranspose layers (keras-team#21852) * Support PyDataset in Normalization layer `adapt` methods (keras-team#21817) * Fix test failures when nnx is enabled (keras-team#21875) * Implement ldexp function in keras.ops (keras-team#21863) * Added OrbaxCheckpoint for keras 3.0 for Data centric saving and resto… * Add raise_error option to TerminateOnNaN for immediate termination on… * Fix NNX tests (keras-team#21884) * `keras.utils.set_random_seed` clear the global `SeedGenerator`. (keras-team#21874) * fix tpu test (keras-team#21893) * Model Export to liteRT (keras-team#21674) * Fix: torch layer losses keyword arguments in rematscope (keras-team#21865) * Add label to trigger TPU tests manually. (keras-team#21897) * Support tpu tests allowing tpu precision for matmul (keras-team#21887) * remove log (keras-team#21901) * Introduces layer filtering for quantization and fixes GPTQ dependency… * Replace `np.reshape(x, newshape=y)` with `np.reshape(x, y)`. (keras-team#21899) * Modified Dense layer documentation for use_bias with batch normalizat… * [OpenVINO Backend] Support np.diag (keras-team#20967) * Modify Muon optimizer (keras-team#21885) * Disables implicit GPTQ quantization using dtype_policy setter (keras-team#21895) * Dense: validate units argument (keras-team#21902) * Pin `ai-edge-litert` version to fix CI (keras-team#21912) * Increase JAX GPU tests timeout to 2 hours (keras-team#21915) * Fix TPU tests - for splash attention (keras-team#21891) * Support various filtering functions in OpenVINO (keras-team#21836) * OpenVINO NN Module Functions (keras-team#21803) * fix XLA dynamic shape output of ops.diag (keras-team#21906) * Fix: Remove redundant epsilon in loss mask weight calculation (keras-team#21908) * Implement vander function in keras.ops (keras-team#21882) * Fix Muon optimizer with TensorFlow backend. (keras-team#21924) * OpenVino `device_scope` and data adapters tests (keras-team#21922) * Fix fake quant gradient output shape and use `jax.grad` for tests. (#… * Introduces QuantizationConfig for fine-grained quantization control (… * Extended fix OOM Issue keras-team#21634 on Keras side (keras-team#21755) * Fix ops.tile shape inference issue on TensorFlow backend (keras-team#21860) * Add adaptive pooling (1D, 2D, 3D) support across JAX, NumPy, TensorFl… * More OpenVINO Numpy Operations (keras-team#21925) * Adds Serialization Support for QuantizationConfig based quantized mod… * Refactors AbsMaxQuantizer to accept axis in __call__ (keras-team#21931) * Speed up unit tests on JAX and TensorFlow. (keras-team#21933) * update dev version number (keras-team#21921) * Always use `run_tpu_tests` label to run the TPU tests. (keras-team#21900) * Revert "Always use `run_tpu_tests` label to run the TPU tests. (keras-team#2190… * Forward-fix for JAX API changes (keras-team#21938) * Remove nightly tests with Python 3.13. (keras-team#21943) * Do no always make batch size dynamic during export. (keras-team#21944) * Fix `numpy.mean` with dynamic shape on OpenVino. (keras-team#21947) * Remove NumPy warning with NumPy >= 2. (keras-team#21949) * Always use `run_tpu_tests` label to run the TPU tests. (keras-team#21950) * [OpenVINO backend] Support np.vander, np.trapezoid, np.corrcoef, np.c… * Fixed a bug in _keras_mask (keras-team#21946) * Fix handling of symbolic Tensor in RNN (keras-team#21945) * Add example for arctanh (keras-team#21951) * Fix DoS via malicious HDF5 dataset metadata in KerasFileEditor (keras-team#21880) * Implement nextafter function in keras.ops (keras-team#21960) * fix image.extract_patches strides handling (keras-team#21959) * [OpenVINO backend] Support numpy.flip (keras-team#21963) * Bump the github-actions group with 4 updates (keras-team#21968) * Fix CUDNN flash attention for JAX > 0.6.2. (keras-team#21970) * Skip `PyDataset` tests on TPU. (keras-team#21964) * Add missing `name` to `SeedGenerator.get_config`. (keras-team#21975) * Use `subprocess.run` in `pip_build.py` to escape wheel path. (keras-team#21976) * Update dependencies and `dependabot.yml`. (keras-team#21974) * Use `kokoro:force-run` label for TPU tests too. (keras-team#21956) * Add simple example for keras.layers.Resizing (keras-team#21966) * [OpenVINO backend] Support numpy.diagonal (keras-team#21965) * Bump actions/checkout from 5.0.1 to 6.0.1 in the github-actions group… * Fix ReversibleEmbedding mask error when using reverse=True (keras-team#21961) * Update feature_space.py (keras-team#21935) * Clarify Tracker docstring wording (keras-team#21985) * Remove semi-colon after email in SECURITY.md (keras-team#21993) * Implement cbrt function for OpenVINO backend (keras-team#21987) * Fix config keys for chain depth and num chains (keras-team#21979) * Implement hypot and trace function for OpenVINO backend (keras-team#21991) * Implement ptp function in keras.ops (keras-team#21990) * Orbax Loading and Sharding Support feature (keras-team#21903) * Add usage examples to loss docstrings (keras-team#21989) * Unify extract_patches to support both 2D and 3D patches (keras-team#21980) * Fix ndim to support tf.RaggedTensor by using shape.rank (keras-team#21999) * Implement size and swapaxes function for OpenVINO backend. (keras-team#21995) * Implement kron function for OpenVINO backend (keras-team#22000) * Adds support for AWQ (keras-team#21992) * Trigger TPU tests on kokoro label removal rather than addition. (keras-team#22001) * Document complex dtype limitation in ops.correlate (keras-team#21984) * [OpenVINO backend] Fix and enable numpy.rot90 (keras-team#21967) * Only skip TPU excluded tests on TPU. (keras-team#22008) * Improvements to `JaxLayer` and `FlaxLayer` related to RNG handling an… * Fix typo in contrast adjustment method (keras-team#22012) * Fix typo and improve docstring formatting (keras-team#22017) * Implement nansum function in keras.ops (keras-team#21996) * Fix unreliable Orbax checkpoint detection with custom implementation … * Unpin as many Python packages versions as possible. (keras-team#22023) * Allow `CenterCrop` layer to handle dynamic image sizes. (keras-team#22020) * TPU tests now verify that we can detect TPUs and fails it not. (keras-team#22019) * Refactor ExtractPatches to handle both 2D and 3D (keras-team#22013) * Implement argpartition function for OpenVINO backend (keras-team#22025) * Implement logaddexp2 function for OpenVINO backend (keras-team#22026) * Implement nanmin function in keras.ops (keras-team#22040) * Increase test coverage for IntegerLookup layer (keras-team#22022) * feat: Add documentation examples for image preprocessing augmentation… * Fix: activity regularizer not normalized by batch size (keras-team#22021) * Implement ldexp and select ops for OpenVINO backend (keras-team#22042) * Fix: convert deque to list before tf.transpose in keras.ops.quantile … * Fix timedistributed mask validation (keras-team#22039) * Torch backend: allow explicit device selection and guard DirectML usa… * Implement nanmax function in keras.ops (keras-team#22043) * Add bias support for torch's `dot_product_attention`. (keras-team#22045) * Fix incorrect example in `ops.associative_scan` docstring (keras-team#22051) * Add Batch Renormalisation (keras-team#22047) * Implement round and divide_no_nan ops for OpenVINO backend (keras-team#22052) * Add dynamic shape support for torch backend export (keras-team#22041) * Implement vstack func for OpenVINO backend (keras-team#22059) * Implement ptp function for OpenVINO backend (keras-team#22060) * Implement nanmean function in keras.ops (keras-team#22055) * Do not allow external links in HDF5 files. (keras-team#22057) * Fix discretization symbolic one hot (keras-team#22048) * Implement complete Keras-Orbax checkpoint integration (keras-team#22002) * Increase test coverage for StringLookup preprocessing layer (keras-team#22056) * Set mutable to True by default in nnx_metadata (keras-team#22074) * Adds Asymmetric INT4 Sub-Channel Quantization Support (keras-team#22007) * Allow passing variables to a function with `@custom_gradient`. (keras-team#22069) * Disallow TFSMLayer deserialization in safe_mode to prevent external S… * Remove redundant global seed initialization code. (keras-team#22084) * Add `Muon` to the list of all optimizer classes. (keras-team#22083) * Implement tile function for openvino backend (keras-team#22071) * implement nansum ops for openvino backend (keras-team#22078) * Remove `testing.uses_cpu()` and re-implement for JAX. (keras-team#22087) * benchmarks: add RandomRotation tf.data performance benchmark (keras-team#21986) * Fix arctan2 NaN propagation in OpenVINO backend (keras-team#22064) * Validate positive height and width in image resize (keras-team#22079) * Don't skip some JAX linalg tests on JAX. (keras-team#22091) * Implement nanprod function in keras.ops (keras-team#22089) * Increase test coverage for TextVectorization layer (keras-team#22066) * Bump the github-actions group with 2 updates (keras-team#22093) * fix: pytorch onnx export symbolic test (keras-team#22086) * Improvements to `*_uses_gpu` and `*_uses_tpu`. (keras-team#22088) * Implement cross product operation for OpenVINO backend (keras-team#22096) * Fail fast on invalid convolution output shapes during symbolic execut… * Fix Normalization broadcasting for scalar and multidim mean and varia… * Standardize the way tests are skipped based on backend and accelerato… * Don't call `pythonify_logs` within `get_metrics_result`. (keras-team#22107) * Fix gaussian_blur padding calculation for even kernel sizes (keras-team#22054) * Adjust JAX variable initializer jitting criteria. (keras-team#22116) * Exclude conv transpose tests on TPU. (keras-team#22117) * Remove incorrect but dead code in `BaseOptimizer.stateless_apply`. (#… * Implement tensordot operation for OpenVINO backend (keras-team#22098) * Fix bounding box docstring references (keras-team#22110) * feat: add depth_to_space and space_to_depth ops (keras-team#22112) * Fix sparse reshape test with Numpy 2.4. (keras-team#22141) * Fix vocabulary reload corruption caused by trailing newline handling … * Add support for dynamic dimensions in `ops.slice.compute_output_spec`… * Revamp graph validation in `Function.__init__`. (keras-team#22153) * Fix: draw_bounding_boxes float32 to uint8 conversion (keras-team#22129) * Implement dstack function across all backends (keras-team#22120) * Add exp2 operation to OpenVINO backend (keras-team#22131) * Add trunc operation to OpenVINO backend (keras-team#22134) * Fix: add missing validation for output padding < strides (keras-team#22130) * docs: Add guide on resuming training from weight-only checkpoints (#2… * feat(openvino): upgrade opset to opset15 (keras-team#22159) * Fix order-dependent float16/bfloat16 promotion in cast_to_common_dtyp… * Fix TrackedDict constructor to support iterable (key, value) inputs (… * Implement numpy.gcd using Euclidean algorithm for OpenVINO backend (#… * [Keras 3] Refactor ExportArchive to be a dispatcher for different exp… * [Keras 3] Refactor ExportArchive to be a dispatcher for different exp…

google-ml-butler bot added the size:M label Jan 23, 2026

google-ml-butler bot assigned gbaned Jan 23, 2026

gemini-code-assist bot reviewed Jan 23, 2026

View reviewed changes

bittoby force-pushed the fix-discretization-layer-one-hot branch 2 times, most recently from ec66a66 to e8ac730 Compare January 23, 2026 21:23

Fix discretization symbolic one hot

374aa24

bittoby force-pushed the fix-discretization-layer-one-hot branch from e8ac730 to 374aa24 Compare January 23, 2026 21:45

hertschuh reviewed Jan 23, 2026

View reviewed changes

hertschuh added the stat:awaiting response from contributor label Jan 23, 2026

google-ml-butler bot removed the stat:awaiting response from contributor label Jan 24, 2026

fix: discretization symbolic tensor

33e572c

bittoby force-pushed the fix-discretization-layer-one-hot branch from c1f501b to 33e572c Compare January 24, 2026 05:42

hertschuh reviewed Jan 25, 2026

View reviewed changes

hertschuh added the stat:awaiting response from contributor label Jan 25, 2026

google-ml-butler bot removed the stat:awaiting response from contributor label Jan 25, 2026

fix: discretization symbolic tensor 2

d8cfd9c

bittoby force-pushed the fix-discretization-layer-one-hot branch from 10d42dc to d8cfd9c Compare January 25, 2026 19:36

hertschuh reviewed Jan 26, 2026

View reviewed changes

hertschuh added the stat:awaiting response from contributor label Jan 26, 2026

fix: discretization symbolic all modes

65370c4

google-ml-butler bot removed the stat:awaiting response from contributor label Jan 27, 2026

hertschuh approved these changes Jan 28, 2026

View reviewed changes

google-ml-butler bot added the kokoro:force-run label Jan 28, 2026

google-ml-butler bot added the ready to pull Ready to be merged into the codebase label Jan 28, 2026

kokoro-team removed the kokoro:force-run label Jan 28, 2026

hertschuh merged commit b97f2c7 into keras-team:master Jan 28, 2026
13 checks passed

bittoby deleted the fix-discretization-layer-one-hot branch January 30, 2026 18:02

ChiragSW mentioned this pull request Mar 25, 2026

Discretization one_hot output has incorrect static shape with Keras.Input #22494

Open

	depth = len(self.bin_boundaries) + 1 if self.bin_boundaries else self.num_bins
	depth = (len(self.bin_boundaries) + 1) if self.bin_boundaries is not None else self.num_bins

		# Shapes should be consistent (ignoring batch dimension)
		self.assertEqual(eager_output.shape[1:], symbolic_output.shape[1:])

		f"Failed for num_bins={num_bins}: expected "
		f"{expected_shape}, got {result_shape}",

		# Test consistency between eager and symbolic tensors for one_hot mode
		if output_mode == "one_hot":

Conversation

bittoby commented Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Fix Discretization layer symbolic tensor shape computation for one_hot mode

Before the fix:

After the fix:

Solution

Verification

### Create layer

Test symbolic tensor (this was broken before)

Uh oh!

google-cla bot commented Jan 23, 2026

Uh oh!

gemini-code-assist bot commented Jan 23, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bittoby commented Jan 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bittoby commented Jan 25, 2026

bittoby commented Jan 23, 2026 •

edited

Loading

codecov-commenter commented Jan 23, 2026 •

edited

Loading

bittoby commented Jan 24, 2026 •

edited

Loading