From dfb64bafe6ab123216888ab3af38fa058dade49a Mon Sep 17 00:00:00 2001 From: abhijeet117 Date: Sun, 23 Aug 2026 23:53:33 +0530 Subject: [PATCH 01/28] Support axis=None in experimental.numpy.concatenate NumPy flattens every input before concatenating when axis is None, but the value was passed straight to array_ops.concat and failed with a cryptic conversion error. Flatten the inputs and concatenate along axis 0 instead. --- tensorflow/python/ops/numpy_ops/np_math_ops.py | 4 ++++ tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 9 +++++++++ 2 files changed, 13 insertions(+) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index f35cb18eba0650..a4677f26608787 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -1389,6 +1389,10 @@ def concatenate(arys, axis=0): # pylint: disable=missing-function-docstring ) dtype = np_utils.result_type(*arys) arys = [np_array_ops.array(array, dtype=dtype) for array in arys] + if axis is None: + # NumPy flattens every input before concatenating when axis is None. + arys = [array_ops.reshape(array, [-1]) for array in arys] + axis = 0 return array_ops.concat(arys, axis) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index f6c50d1e02babe..d92d3a83fc3527 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -596,6 +596,15 @@ def testIsInf(self): self.assertFalse(np_math_ops.isneginf(x1)) self.assertFalse(np_math_ops.isneginf(x2)) + def testConcatenateAxisNone(self): + a = np_array_ops.array([1, 2]) + b = np_array_ops.array([[3], [4]]) + self.assertAllEqual( + np_math_ops.concatenate([a, b], axis=None), [1, 2, 3, 4]) + self.assertAllEqual( + np_math_ops.concatenate(np_array_ops.array([[5, 6]]), axis=None), + [5, 6]) + if __name__ == '__main__': tensor.enable_tensor_equality() ops.enable_eager_execution() From a8e930324369d05fec7a23cab33c765bb291ba76 Mon Sep 17 00:00:00 2001 From: abhijeet117 Date: Tue, 25 Aug 2026 14:44:44 +0530 Subject: [PATCH 02/28] Skip reshaping arrays that are already flat in concatenate Reshaping a 1-D array to [-1] is a no-op, so skip the extra op dispatch for inputs that are already flat. --- tensorflow/python/ops/numpy_ops/np_math_ops.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index a4677f26608787..8f0742aa6a6ffe 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -1391,7 +1391,11 @@ def concatenate(arys, axis=0): # pylint: disable=missing-function-docstring arys = [np_array_ops.array(array, dtype=dtype) for array in arys] if axis is None: # NumPy flattens every input before concatenating when axis is None. - arys = [array_ops.reshape(array, [-1]) for array in arys] + # Reshaping an already flat array is a no-op, so skip the op dispatch. + arys = [ + array if array.shape.ndims == 1 else array_ops.reshape(array, [-1]) + for array in arys + ] axis = 0 return array_ops.concat(arys, axis) From fcfa07673c6c33a67b38baa109cfff13ae678a23 Mon Sep 17 00:00:00 2001 From: Alexander Grund Date: Tue, 25 Aug 2026 14:21:39 +0200 Subject: [PATCH 03/28] Fix SYSTEM pybind11 & flatbuffers The `TF_SYSTEMLIBS` for those 2 broke: - The LICENSE filegroup should not have an extension. - Duplicate `xla` in project reference - Missing config-setting referenced by other targets. --- third_party/flatbuffers/BUILD.system | 2 +- third_party/systemlibs/pybind11.BUILD | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/third_party/flatbuffers/BUILD.system b/third_party/flatbuffers/BUILD.system index 8fe4d7a590719f..b1d63b4ca0fd77 100644 --- a/third_party/flatbuffers/BUILD.system +++ b/third_party/flatbuffers/BUILD.system @@ -1,7 +1,7 @@ licenses(["notice"]) # Apache 2.0 filegroup( - name = "LICENSE.txt", + name = "LICENSE", visibility = ["//visibility:public"], ) diff --git a/third_party/systemlibs/pybind11.BUILD b/third_party/systemlibs/pybind11.BUILD index 711fcf13b5d863..e90af0bb5879d2 100644 --- a/third_party/systemlibs/pybind11.BUILD +++ b/third_party/systemlibs/pybind11.BUILD @@ -18,6 +18,13 @@ package(default_visibility = ["//visibility:public"]) cc_library( name = "pybind11", deps = [ - "@xla@xla//third_party/python_runtime:headers", + "@xla//third_party/python_runtime:headers", ], ) + +# Needed by pybind11_bazel. +config_setting( + name = "msvc_compiler", + flag_values = {"@bazel_tools//tools/cpp:compiler": "msvc-cl"}, + visibility = ["//visibility:public"], +) From 19ea445df2cd9399f4ebee4ddaabccf32ebbfcc9 Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Wed, 26 Aug 2026 18:43:41 +0530 Subject: [PATCH 04/28] feat: Add Advanced Quantization Dense layer --- tensorflow/python/keras/layers/BUILD | 18 +++ tensorflow/python/keras/layers/__init__.py | 1 + tensorflow/python/keras/layers/quantized.py | 128 ++++++++++++++++++ .../python/keras/layers/quantized_test.py | 34 +++++ 4 files changed, 181 insertions(+) create mode 100644 tensorflow/python/keras/layers/quantized.py create mode 100644 tensorflow/python/keras/layers/quantized_test.py diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD index 83925764d3fbb0..3b4ad8af90e8a5 100644 --- a/tensorflow/python/keras/layers/BUILD +++ b/tensorflow/python/keras/layers/BUILD @@ -54,6 +54,7 @@ py_library( ":convolutional", ":convolutional_recurrent", ":core", + ":quantized", ":dense_attention", ":embeddings", ":merge", @@ -131,6 +132,23 @@ py_library( ], ) +py_library( + name = "quantized", + srcs = ["quantized.py"], + srcs_version = "PY3", + strict_deps = False, + deps = [ + "//tensorflow/python/keras:activations", + "//tensorflow/python/keras:backend", + "//tensorflow/python/keras:base_layer", + "//tensorflow/python/keras:constraints", + "//tensorflow/python/keras:regularizers", + "//tensorflow/python/keras/engine:input_spec", + "//tensorflow/python/keras/initializers", + "//tensorflow/python/ops:math_ops", + ], +) + py_library( name = "core", srcs = ["core.py"], diff --git a/tensorflow/python/keras/layers/__init__.py b/tensorflow/python/keras/layers/__init__.py index 889e9d181fb3af..984ee58afbd836 100644 --- a/tensorflow/python/keras/layers/__init__.py +++ b/tensorflow/python/keras/layers/__init__.py @@ -77,6 +77,7 @@ from tensorflow.python.keras.layers.core import RepeatVector from tensorflow.python.keras.layers.core import Lambda from tensorflow.python.keras.layers.core import Dense +from tensorflow.python.keras.layers.quantized import QuantizedDense from tensorflow.python.keras.layers.core import ActivityRegularization # Dense Attention layers. diff --git a/tensorflow/python/keras/layers/quantized.py b/tensorflow/python/keras/layers/quantized.py new file mode 100644 index 00000000000000..beff91287292d0 --- /dev/null +++ b/tensorflow/python/keras/layers/quantized.py @@ -0,0 +1,128 @@ +import tensorflow as tf +from tensorflow.python.keras.engine.base_layer import Layer +from tensorflow.python.keras import initializers +from tensorflow.python.keras import regularizers +from tensorflow.python.keras import constraints +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import math_ops +from tensorflow.python.keras.engine.input_spec import InputSpec +from tensorflow.python.ops import nn + +class QuantizedDense(Layer): + """A densely-connected layer with weight quantization. + + This layer acts like a standard Dense layer but simulates + 4-bit or 8-bit weight quantization using fake quantization nodes. + """ + + def __init__(self, + units, + bits=8, + activation=None, + use_bias=True, + kernel_initializer='glorot_uniform', + bias_initializer='zeros', + kernel_regularizer=None, + bias_regularizer=None, + activity_regularizer=None, + kernel_constraint=None, + bias_constraint=None, + **kwargs): + super(QuantizedDense, self).__init__( + activity_regularizer=activity_regularizer, **kwargs) + self.units = int(units) + self.bits = int(bits) + if self.bits not in [4, 8]: + raise ValueError('Only 4-bit and 8-bit quantization are supported.') + self.activation = tf.keras.activations.get(activation) + self.use_bias = use_bias + self.kernel_initializer = initializers.get(kernel_initializer) + self.bias_initializer = initializers.get(bias_initializer) + self.kernel_regularizer = regularizers.get(kernel_regularizer) + self.bias_regularizer = regularizers.get(bias_regularizer) + self.kernel_constraint = constraints.get(kernel_constraint) + self.bias_constraint = constraints.get(bias_constraint) + self.input_spec = InputSpec(min_ndim=2) + + def build(self, input_shape): + input_shape = tf.TensorShape(input_shape) + last_dim = tf.compat.dimension_value(input_shape[-1]) + if last_dim is None: + raise ValueError('The last dimension of the inputs to `QuantizedDense` ' + 'should be defined. Found `None`.') + self.input_spec = InputSpec(min_ndim=2, axes={-1: last_dim}) + self.kernel = self.add_weight( + 'kernel', + shape=[last_dim, self.units], + initializer=self.kernel_initializer, + regularizer=self.kernel_regularizer, + constraint=self.kernel_constraint, + dtype=self.dtype, + trainable=True) + if self.use_bias: + self.bias = self.add_weight( + 'bias', + shape=[self.units,], + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + dtype=self.dtype, + trainable=True) + else: + self.bias = None + self.built = True + + def call(self, inputs): + # Apply Fake Quantization to weights + min_val = tf.math.reduce_min(self.kernel) + max_val = tf.math.reduce_max(self.kernel) + + quant_bits = self.bits + + quantized_kernel = tf.quantization.fake_quant_with_min_max_vars( + self.kernel, + min_val, + max_val, + num_bits=quant_bits, + narrow_range=True) + + rank = inputs.shape.rank + if rank == 2 or rank is None: + if isinstance(inputs, tf.SparseTensor): + raise NotImplementedError("Sparse inputs not supported.") + outputs = tf.matmul(a=inputs, b=quantized_kernel) + else: + outputs = tf.tensordot(inputs, quantized_kernel, [[rank - 1], [0]]) + + if self.use_bias: + outputs = tf.nn.bias_add(outputs, self.bias) + + if self.activation is not None: + outputs = self.activation(outputs) + return outputs + + def compute_output_shape(self, input_shape): + input_shape = tf.TensorShape(input_shape) + input_shape = input_shape.with_rank_at_least(2) + if tf.compat.dimension_value(input_shape[-1]) is None: + raise ValueError( + 'The innermost dimension of input_shape must be defined, but saw: %s' + % input_shape) + return input_shape[:-1].concatenate(self.units) + + def get_config(self): + config = super(QuantizedDense, self).get_config() + config.update({ + 'units': self.units, + 'bits': self.bits, + 'activation': tf.keras.activations.serialize(self.activation), + 'use_bias': self.use_bias, + 'kernel_initializer': initializers.serialize(self.kernel_initializer), + 'bias_initializer': initializers.serialize(self.bias_initializer), + 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), + 'bias_regularizer': regularizers.serialize(self.bias_regularizer), + 'activity_regularizer': regularizers.serialize(self.activity_regularizer), + 'kernel_constraint': constraints.serialize(self.kernel_constraint), + 'bias_constraint': constraints.serialize(self.bias_constraint) + }) + return config diff --git a/tensorflow/python/keras/layers/quantized_test.py b/tensorflow/python/keras/layers/quantized_test.py new file mode 100644 index 00000000000000..85d35cae00f6cd --- /dev/null +++ b/tensorflow/python/keras/layers/quantized_test.py @@ -0,0 +1,34 @@ +import numpy as np +import tensorflow as tf +from tensorflow.python.platform import test +from tensorflow.python.keras.layers.quantized import QuantizedDense + +class QuantizedDenseTest(test.TestCase): + + def test_quantized_dense_basic(self): + inputs = tf.random.uniform((32, 128)) + + # Test 8-bit quantization + layer_8bit = QuantizedDense(64, bits=8) + out_8bit = layer_8bit(inputs) + + self.assertEqual(out_8bit.shape, (32, 64)) + self.assertEqual(layer_8bit.kernel.shape, (128, 64)) + self.assertEqual(layer_8bit.bias.shape, (64,)) + + def test_quantized_dense_4bit(self): + inputs = tf.random.uniform((16, 32)) + + # Test 4-bit quantization + layer_4bit = QuantizedDense(16, bits=4, use_bias=False) + out_4bit = layer_4bit(inputs) + + self.assertEqual(out_4bit.shape, (16, 16)) + self.assertIsNone(layer_4bit.bias) + + def test_invalid_bits(self): + with self.assertRaisesRegex(ValueError, 'Only 4-bit and 8-bit quantization'): + QuantizedDense(32, bits=16) + +if __name__ == '__main__': + test.main() From d09cbd9fa882ff98dc6ed6b3be9fa8e49d456b1d Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Wed, 26 Aug 2026 19:00:45 +0530 Subject: [PATCH 05/28] fix: address PR feedback for quantized dense layer --- tensorflow/python/keras/layers/quantized.py | 16 +++++++++++----- tensorflow/python/keras/layers/quantized_test.py | 5 +++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/keras/layers/quantized.py b/tensorflow/python/keras/layers/quantized.py index beff91287292d0..cba74939f1d1f2 100644 --- a/tensorflow/python/keras/layers/quantized.py +++ b/tensorflow/python/keras/layers/quantized.py @@ -3,6 +3,7 @@ from tensorflow.python.keras import initializers from tensorflow.python.keras import regularizers from tensorflow.python.keras import constraints +from tensorflow.python.keras import activations from tensorflow.python.framework import dtypes from tensorflow.python.ops import math_ops from tensorflow.python.keras.engine.input_spec import InputSpec @@ -34,7 +35,7 @@ def __init__(self, self.bits = int(bits) if self.bits not in [4, 8]: raise ValueError('Only 4-bit and 8-bit quantization are supported.') - self.activation = tf.keras.activations.get(activation) + self.activation = activations.get(activation) self.use_bias = use_bias self.kernel_initializer = initializers.get(kernel_initializer) self.bias_initializer = initializers.get(bias_initializer) @@ -74,17 +75,22 @@ def build(self, input_shape): def call(self, inputs): # Apply Fake Quantization to weights - min_val = tf.math.reduce_min(self.kernel) - max_val = tf.math.reduce_max(self.kernel) + # fake_quant_with_min_max_vars requires float32 inputs. + kernel = tf.cast(self.kernel, dtypes.float32) + min_val = tf.math.reduce_min(kernel) + max_val = tf.math.reduce_max(kernel) + # Ensure min_val and max_val are not equal to avoid division by zero or NaN gradients + max_val = tf.math.maximum(max_val, min_val + 1e-5) quant_bits = self.bits quantized_kernel = tf.quantization.fake_quant_with_min_max_vars( - self.kernel, + kernel, min_val, max_val, num_bits=quant_bits, narrow_range=True) + quantized_kernel = tf.cast(quantized_kernel, self.dtype) rank = inputs.shape.rank if rank == 2 or rank is None: @@ -115,7 +121,7 @@ def get_config(self): config.update({ 'units': self.units, 'bits': self.bits, - 'activation': tf.keras.activations.serialize(self.activation), + 'activation': activations.serialize(self.activation), 'use_bias': self.use_bias, 'kernel_initializer': initializers.serialize(self.kernel_initializer), 'bias_initializer': initializers.serialize(self.bias_initializer), diff --git a/tensorflow/python/keras/layers/quantized_test.py b/tensorflow/python/keras/layers/quantized_test.py index 85d35cae00f6cd..6c5c1e5b8a4825 100644 --- a/tensorflow/python/keras/layers/quantized_test.py +++ b/tensorflow/python/keras/layers/quantized_test.py @@ -5,6 +5,11 @@ class QuantizedDenseTest(test.TestCase): + def setUp(self): + super(QuantizedDenseTest, self).setUp() + tf.random.set_seed(0) + np.random.seed(0) + def test_quantized_dense_basic(self): inputs = tf.random.uniform((32, 128)) From 1595ef53b6244ce65ce86d4e7f91ea059fc9e972 Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Thu, 27 Aug 2026 18:49:12 +0530 Subject: [PATCH 06/28] Address review comments: fix architecture, tests, and formatting --- tensorflow/python/keras/layers/BUILD | 19 ++- tensorflow/python/keras/layers/quantized.py | 129 ++++++++++-------- .../python/keras/layers/quantized_test.py | 35 +++-- 3 files changed, 115 insertions(+), 68 deletions(-) diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD index 3b4ad8af90e8a5..1cc69ed3d93bc4 100644 --- a/tensorflow/python/keras/layers/BUILD +++ b/tensorflow/python/keras/layers/BUILD @@ -136,16 +136,31 @@ py_library( name = "quantized", srcs = ["quantized.py"], srcs_version = "PY3", - strict_deps = False, deps = [ + "//tensorflow/python/framework:dtypes", + "//tensorflow/python/framework:tensor_shape", "//tensorflow/python/keras:activations", - "//tensorflow/python/keras:backend", "//tensorflow/python/keras:base_layer", "//tensorflow/python/keras:constraints", "//tensorflow/python/keras:regularizers", "//tensorflow/python/keras/engine:input_spec", "//tensorflow/python/keras/initializers", "//tensorflow/python/ops:math_ops", + "//tensorflow/python/ops:nn", + "//tensorflow/python/ops:quantized_ops", + ], +) + +tf_py_test( + name = "quantized_test", + srcs = ["quantized_test.py"], + python_version = "PY3", + srcs_version = "PY3", + deps = [ + ":quantized", + "//tensorflow/python/framework:test_lib", + "//tensorflow/python/platform:client_testlib", + "//third_party/py/numpy", ], ) diff --git a/tensorflow/python/keras/layers/quantized.py b/tensorflow/python/keras/layers/quantized.py index cba74939f1d1f2..bedca40ff4ab3a 100644 --- a/tensorflow/python/keras/layers/quantized.py +++ b/tensorflow/python/keras/layers/quantized.py @@ -1,28 +1,48 @@ -import tensorflow as tf -from tensorflow.python.keras.engine.base_layer import Layer +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Quantized Dense layer.""" + +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework import ops +from tensorflow.python.keras import activations +from tensorflow.python.keras import constraints from tensorflow.python.keras import initializers from tensorflow.python.keras import regularizers -from tensorflow.python.keras import constraints -from tensorflow.python.keras import activations -from tensorflow.python.framework import dtypes -from tensorflow.python.ops import math_ops +from tensorflow.python.keras.engine.base_layer import Layer from tensorflow.python.keras.engine.input_spec import InputSpec +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn +from tensorflow.python.ops import nn_ops + class QuantizedDense(Layer): """A densely-connected layer with weight quantization. - + This layer acts like a standard Dense layer but simulates 4-bit or 8-bit weight quantization using fake quantization nodes. """ - + def __init__(self, units, bits=8, activation=None, use_bias=True, - kernel_initializer='glorot_uniform', - bias_initializer='zeros', + kernel_initializer="glorot_uniform", + bias_initializer="zeros", kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, @@ -34,7 +54,7 @@ def __init__(self, self.units = int(units) self.bits = int(bits) if self.bits not in [4, 8]: - raise ValueError('Only 4-bit and 8-bit quantization are supported.') + raise ValueError("Only 4-bit and 8-bit quantization are supported.") self.activation = activations.get(activation) self.use_bias = use_bias self.kernel_initializer = initializers.get(kernel_initializer) @@ -46,14 +66,14 @@ def __init__(self, self.input_spec = InputSpec(min_ndim=2) def build(self, input_shape): - input_shape = tf.TensorShape(input_shape) - last_dim = tf.compat.dimension_value(input_shape[-1]) + input_shape = tensor_shape.TensorShape(input_shape) + last_dim = tensor_shape.dimension_value(input_shape[-1]) if last_dim is None: - raise ValueError('The last dimension of the inputs to `QuantizedDense` ' - 'should be defined. Found `None`.') + raise ValueError("The last dimension of the inputs to `QuantizedDense` " + "should be defined. Found `None`.") self.input_spec = InputSpec(min_ndim=2, axes={-1: last_dim}) self.kernel = self.add_weight( - 'kernel', + "kernel", shape=[last_dim, self.units], initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, @@ -62,7 +82,7 @@ def build(self, input_shape): trainable=True) if self.use_bias: self.bias = self.add_weight( - 'bias', + "bias", shape=[self.units,], initializer=self.bias_initializer, regularizer=self.bias_regularizer, @@ -74,61 +94,54 @@ def build(self, input_shape): self.built = True def call(self, inputs): - # Apply Fake Quantization to weights - # fake_quant_with_min_max_vars requires float32 inputs. - kernel = tf.cast(self.kernel, dtypes.float32) - min_val = tf.math.reduce_min(kernel) - max_val = tf.math.reduce_max(kernel) - # Ensure min_val and max_val are not equal to avoid division by zero or NaN gradients - max_val = tf.math.maximum(max_val, min_val + 1e-5) - - quant_bits = self.bits - - quantized_kernel = tf.quantization.fake_quant_with_min_max_vars( - kernel, - min_val, - max_val, - num_bits=quant_bits, + kernel = math_ops.cast(self.kernel, dtypes.float32) + min_val = math_ops.reduce_min(kernel) + max_val = math_ops.reduce_max(kernel) + max_val = math_ops.maximum(max_val, min_val + 1e-5) + + quantized_kernel = array_ops.fake_quant_with_min_max_vars( + kernel, + min_val, + max_val, + num_bits=self.bits, narrow_range=True) - quantized_kernel = tf.cast(quantized_kernel, self.dtype) - + quantized_kernel = math_ops.cast(quantized_kernel, self.dtype) + rank = inputs.shape.rank - if rank == 2 or rank is None: - if isinstance(inputs, tf.SparseTensor): - raise NotImplementedError("Sparse inputs not supported.") - outputs = tf.matmul(a=inputs, b=quantized_kernel) + if rank is not None and rank <= 2: + outputs = math_ops.matmul(a=inputs, b=quantized_kernel) else: - outputs = tf.tensordot(inputs, quantized_kernel, [[rank - 1], [0]]) - + outputs = math_ops.tensordot(inputs, quantized_kernel, [[rank - 1 if rank else -1], [0]]) + if self.use_bias: - outputs = tf.nn.bias_add(outputs, self.bias) - + outputs = nn_ops.bias_add(outputs, self.bias) + if self.activation is not None: outputs = self.activation(outputs) return outputs - + def compute_output_shape(self, input_shape): - input_shape = tf.TensorShape(input_shape) + input_shape = tensor_shape.TensorShape(input_shape) input_shape = input_shape.with_rank_at_least(2) - if tf.compat.dimension_value(input_shape[-1]) is None: + if tensor_shape.dimension_value(input_shape[-1]) is None: raise ValueError( - 'The innermost dimension of input_shape must be defined, but saw: %s' + "The innermost dimension of input_shape must be defined, but saw: %s" % input_shape) return input_shape[:-1].concatenate(self.units) - + def get_config(self): config = super(QuantizedDense, self).get_config() config.update({ - 'units': self.units, - 'bits': self.bits, - 'activation': activations.serialize(self.activation), - 'use_bias': self.use_bias, - 'kernel_initializer': initializers.serialize(self.kernel_initializer), - 'bias_initializer': initializers.serialize(self.bias_initializer), - 'kernel_regularizer': regularizers.serialize(self.kernel_regularizer), - 'bias_regularizer': regularizers.serialize(self.bias_regularizer), - 'activity_regularizer': regularizers.serialize(self.activity_regularizer), - 'kernel_constraint': constraints.serialize(self.kernel_constraint), - 'bias_constraint': constraints.serialize(self.bias_constraint) + "units": self.units, + "bits": self.bits, + "activation": activations.serialize(self.activation), + "use_bias": self.use_bias, + "kernel_initializer": initializers.serialize(self.kernel_initializer), + "bias_initializer": initializers.serialize(self.bias_initializer), + "kernel_regularizer": regularizers.serialize(self.kernel_regularizer), + "bias_regularizer": regularizers.serialize(self.bias_regularizer), + "activity_regularizer": regularizers.serialize(self.activity_regularizer), + "kernel_constraint": constraints.serialize(self.kernel_constraint), + "bias_constraint": constraints.serialize(self.bias_constraint) }) return config diff --git a/tensorflow/python/keras/layers/quantized_test.py b/tensorflow/python/keras/layers/quantized_test.py index 6c5c1e5b8a4825..187b64c4b23718 100644 --- a/tensorflow/python/keras/layers/quantized_test.py +++ b/tensorflow/python/keras/layers/quantized_test.py @@ -1,7 +1,25 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for quantized Dense layer.""" + import numpy as np import tensorflow as tf -from tensorflow.python.platform import test + from tensorflow.python.keras.layers.quantized import QuantizedDense +from tensorflow.python.platform import test + class QuantizedDenseTest(test.TestCase): @@ -12,28 +30,29 @@ def setUp(self): def test_quantized_dense_basic(self): inputs = tf.random.uniform((32, 128)) - + # Test 8-bit quantization layer_8bit = QuantizedDense(64, bits=8) out_8bit = layer_8bit(inputs) - + self.assertEqual(out_8bit.shape, (32, 64)) self.assertEqual(layer_8bit.kernel.shape, (128, 64)) self.assertEqual(layer_8bit.bias.shape, (64,)) def test_quantized_dense_4bit(self): inputs = tf.random.uniform((16, 32)) - + # Test 4-bit quantization layer_4bit = QuantizedDense(16, bits=4, use_bias=False) out_4bit = layer_4bit(inputs) - + self.assertEqual(out_4bit.shape, (16, 16)) self.assertIsNone(layer_4bit.bias) - + def test_invalid_bits(self): - with self.assertRaisesRegex(ValueError, 'Only 4-bit and 8-bit quantization'): + with self.assertRaisesRegex(ValueError, + "Only 4-bit and 8-bit quantization"): QuantizedDense(32, bits=16) -if __name__ == '__main__': +if __name__ == "__main__": test.main() From ba1b7f7ad0940722483e62ecc9f390f2851855cf Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Thu, 27 Aug 2026 19:02:58 +0530 Subject: [PATCH 07/28] Fix pylint indentation to 2 spaces and long lines --- tensorflow/python/keras/layers/quantized.py | 221 +++++++++--------- .../python/keras/layers/quantized_test.py | 48 ++-- 2 files changed, 138 insertions(+), 131 deletions(-) diff --git a/tensorflow/python/keras/layers/quantized.py b/tensorflow/python/keras/layers/quantized.py index bedca40ff4ab3a..9bb2e3f42dd9e8 100644 --- a/tensorflow/python/keras/layers/quantized.py +++ b/tensorflow/python/keras/layers/quantized.py @@ -16,7 +16,6 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import tensor_shape -from tensorflow.python.framework import ops from tensorflow.python.keras import activations from tensorflow.python.keras import constraints from tensorflow.python.keras import initializers @@ -25,123 +24,131 @@ from tensorflow.python.keras.engine.input_spec import InputSpec from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn from tensorflow.python.ops import nn_ops class QuantizedDense(Layer): - """A densely-connected layer with weight quantization. + """A densely-connected layer with weight quantization. - This layer acts like a standard Dense layer but simulates - 4-bit or 8-bit weight quantization using fake quantization nodes. - """ + This layer acts like a standard Dense layer but simulates + 4-bit or 8-bit weight quantization using fake quantization nodes. + """ - def __init__(self, - units, - bits=8, - activation=None, - use_bias=True, - kernel_initializer="glorot_uniform", - bias_initializer="zeros", - kernel_regularizer=None, - bias_regularizer=None, - activity_regularizer=None, - kernel_constraint=None, - bias_constraint=None, - **kwargs): - super(QuantizedDense, self).__init__( - activity_regularizer=activity_regularizer, **kwargs) - self.units = int(units) - self.bits = int(bits) - if self.bits not in [4, 8]: - raise ValueError("Only 4-bit and 8-bit quantization are supported.") - self.activation = activations.get(activation) - self.use_bias = use_bias - self.kernel_initializer = initializers.get(kernel_initializer) - self.bias_initializer = initializers.get(bias_initializer) - self.kernel_regularizer = regularizers.get(kernel_regularizer) - self.bias_regularizer = regularizers.get(bias_regularizer) - self.kernel_constraint = constraints.get(kernel_constraint) - self.bias_constraint = constraints.get(bias_constraint) - self.input_spec = InputSpec(min_ndim=2) + def __init__(self, + units, + bits=8, + activation=None, + use_bias=True, + kernel_initializer="glorot_uniform", + bias_initializer="zeros", + kernel_regularizer=None, + bias_regularizer=None, + activity_regularizer=None, + kernel_constraint=None, + bias_constraint=None, + **kwargs): + super(QuantizedDense, self).__init__( + activity_regularizer=activity_regularizer, **kwargs) + self.units = int(units) + self.bits = int(bits) + if self.bits not in [4, 8]: + raise ValueError("Only 4-bit and 8-bit quantization are supported.") + self.activation = activations.get(activation) + self.use_bias = use_bias + self.kernel_initializer = initializers.get(kernel_initializer) + self.bias_initializer = initializers.get(bias_initializer) + self.kernel_regularizer = regularizers.get(kernel_regularizer) + self.bias_regularizer = regularizers.get(bias_regularizer) + self.kernel_constraint = constraints.get(kernel_constraint) + self.bias_constraint = constraints.get(bias_constraint) + self.input_spec = InputSpec(min_ndim=2) - def build(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - last_dim = tensor_shape.dimension_value(input_shape[-1]) - if last_dim is None: - raise ValueError("The last dimension of the inputs to `QuantizedDense` " - "should be defined. Found `None`.") - self.input_spec = InputSpec(min_ndim=2, axes={-1: last_dim}) - self.kernel = self.add_weight( - "kernel", - shape=[last_dim, self.units], - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - constraint=self.kernel_constraint, - dtype=self.dtype, - trainable=True) - if self.use_bias: - self.bias = self.add_weight( - "bias", - shape=[self.units,], - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - dtype=self.dtype, - trainable=True) - else: - self.bias = None - self.built = True + def build(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + last_dim = tensor_shape.dimension_value(input_shape[-1]) + if last_dim is None: + raise ValueError( + "The last dimension of the inputs to `QuantizedDense` " + "should be defined. Found `None`.") + self.input_spec = InputSpec(min_ndim=2, axes={-1: last_dim}) + self.kernel = self.add_weight( + "kernel", + shape=[last_dim, self.units], + initializer=self.kernel_initializer, + regularizer=self.kernel_regularizer, + constraint=self.kernel_constraint, + dtype=self.dtype, + trainable=True) + if self.use_bias: + self.bias = self.add_weight( + "bias", + shape=[self.units,], + initializer=self.bias_initializer, + regularizer=self.bias_regularizer, + constraint=self.bias_constraint, + dtype=self.dtype, + trainable=True) + else: + self.bias = None + self.built = True - def call(self, inputs): - kernel = math_ops.cast(self.kernel, dtypes.float32) - min_val = math_ops.reduce_min(kernel) - max_val = math_ops.reduce_max(kernel) - max_val = math_ops.maximum(max_val, min_val + 1e-5) + def call(self, inputs): + kernel = math_ops.cast(self.kernel, dtypes.float32) + min_val = math_ops.reduce_min(kernel) + max_val = math_ops.reduce_max(kernel) + max_val = math_ops.maximum(max_val, min_val + 1e-5) - quantized_kernel = array_ops.fake_quant_with_min_max_vars( - kernel, - min_val, - max_val, - num_bits=self.bits, - narrow_range=True) - quantized_kernel = math_ops.cast(quantized_kernel, self.dtype) + quantized_kernel = array_ops.fake_quant_with_min_max_vars( + kernel, + min_val, + max_val, + num_bits=self.bits, + narrow_range=True) + quantized_kernel = math_ops.cast(quantized_kernel, self.dtype) - rank = inputs.shape.rank - if rank is not None and rank <= 2: - outputs = math_ops.matmul(a=inputs, b=quantized_kernel) - else: - outputs = math_ops.tensordot(inputs, quantized_kernel, [[rank - 1 if rank else -1], [0]]) + rank = inputs.shape.rank + if rank is not None and rank <= 2: + outputs = math_ops.matmul(a=inputs, b=quantized_kernel) + else: + axes = [[rank - 1 if rank else -1], [0]] + outputs = math_ops.tensordot(inputs, quantized_kernel, axes) - if self.use_bias: - outputs = nn_ops.bias_add(outputs, self.bias) + if self.use_bias: + outputs = nn_ops.bias_add(outputs, self.bias) - if self.activation is not None: - outputs = self.activation(outputs) - return outputs + if self.activation is not None: + outputs = self.activation(outputs) + return outputs - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - input_shape = input_shape.with_rank_at_least(2) - if tensor_shape.dimension_value(input_shape[-1]) is None: - raise ValueError( - "The innermost dimension of input_shape must be defined, but saw: %s" - % input_shape) - return input_shape[:-1].concatenate(self.units) + def compute_output_shape(self, input_shape): + input_shape = tensor_shape.TensorShape(input_shape) + input_shape = input_shape.with_rank_at_least(2) + if tensor_shape.dimension_value(input_shape[-1]) is None: + raise ValueError( + "The innermost dimension of input_shape must be defined, " + "but saw: %s" % input_shape) + return input_shape[:-1].concatenate(self.units) - def get_config(self): - config = super(QuantizedDense, self).get_config() - config.update({ - "units": self.units, - "bits": self.bits, - "activation": activations.serialize(self.activation), - "use_bias": self.use_bias, - "kernel_initializer": initializers.serialize(self.kernel_initializer), - "bias_initializer": initializers.serialize(self.bias_initializer), - "kernel_regularizer": regularizers.serialize(self.kernel_regularizer), - "bias_regularizer": regularizers.serialize(self.bias_regularizer), - "activity_regularizer": regularizers.serialize(self.activity_regularizer), - "kernel_constraint": constraints.serialize(self.kernel_constraint), - "bias_constraint": constraints.serialize(self.bias_constraint) - }) - return config + def get_config(self): + config = super(QuantizedDense, self).get_config() + config.update({ + "units": self.units, + "bits": self.bits, + "activation": activations.serialize(self.activation), + "use_bias": self.use_bias, + "kernel_initializer": initializers.serialize( + self.kernel_initializer), + "bias_initializer": initializers.serialize( + self.bias_initializer), + "kernel_regularizer": regularizers.serialize( + self.kernel_regularizer), + "bias_regularizer": regularizers.serialize( + self.bias_regularizer), + "activity_regularizer": regularizers.serialize( + self.activity_regularizer), + "kernel_constraint": constraints.serialize( + self.kernel_constraint), + "bias_constraint": constraints.serialize( + self.bias_constraint) + }) + return config diff --git a/tensorflow/python/keras/layers/quantized_test.py b/tensorflow/python/keras/layers/quantized_test.py index 187b64c4b23718..7b2dd80407dee6 100644 --- a/tensorflow/python/keras/layers/quantized_test.py +++ b/tensorflow/python/keras/layers/quantized_test.py @@ -23,36 +23,36 @@ class QuantizedDenseTest(test.TestCase): - def setUp(self): - super(QuantizedDenseTest, self).setUp() - tf.random.set_seed(0) - np.random.seed(0) + def setUp(self): + super(QuantizedDenseTest, self).setUp() + tf.random.set_seed(0) + np.random.seed(0) - def test_quantized_dense_basic(self): - inputs = tf.random.uniform((32, 128)) + def test_quantized_dense_basic(self): + inputs = tf.random.uniform((32, 128)) - # Test 8-bit quantization - layer_8bit = QuantizedDense(64, bits=8) - out_8bit = layer_8bit(inputs) + # Test 8-bit quantization + layer_8bit = QuantizedDense(64, bits=8) + out_8bit = layer_8bit(inputs) - self.assertEqual(out_8bit.shape, (32, 64)) - self.assertEqual(layer_8bit.kernel.shape, (128, 64)) - self.assertEqual(layer_8bit.bias.shape, (64,)) + self.assertEqual(out_8bit.shape, (32, 64)) + self.assertEqual(layer_8bit.kernel.shape, (128, 64)) + self.assertEqual(layer_8bit.bias.shape, (64,)) - def test_quantized_dense_4bit(self): - inputs = tf.random.uniform((16, 32)) + def test_quantized_dense_4bit(self): + inputs = tf.random.uniform((16, 32)) - # Test 4-bit quantization - layer_4bit = QuantizedDense(16, bits=4, use_bias=False) - out_4bit = layer_4bit(inputs) + # Test 4-bit quantization + layer_4bit = QuantizedDense(16, bits=4, use_bias=False) + out_4bit = layer_4bit(inputs) - self.assertEqual(out_4bit.shape, (16, 16)) - self.assertIsNone(layer_4bit.bias) + self.assertEqual(out_4bit.shape, (16, 16)) + self.assertIsNone(layer_4bit.bias) - def test_invalid_bits(self): - with self.assertRaisesRegex(ValueError, - "Only 4-bit and 8-bit quantization"): - QuantizedDense(32, bits=16) + def test_invalid_bits(self): + with self.assertRaisesRegex(ValueError, + "Only 4-bit and 8-bit quantization"): + QuantizedDense(32, bits=16) if __name__ == "__main__": - test.main() + test.main() From c29985a6c5669fde1271f13ad7f3c745b469149d Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Thu, 27 Aug 2026 19:43:09 +0530 Subject: [PATCH 08/28] Address latest review comments: fix BUILD targets and modular imports --- tensorflow/python/keras/layers/BUILD | 4 +++- tensorflow/python/keras/layers/quantized_test.py | 13 ++++++------- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD index 1cc69ed3d93bc4..1ef5de313ac7d8 100644 --- a/tensorflow/python/keras/layers/BUILD +++ b/tensorflow/python/keras/layers/BUILD @@ -17,6 +17,7 @@ # Contains the Keras layers (internal TensorFlow version). load("@xla//third_party/rules_python/python:defs.bzl", "py_library") +load("//tensorflow:tensorflow.default.bzl", "tf_py_test") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -158,9 +159,10 @@ tf_py_test( srcs_version = "PY3", deps = [ ":quantized", + "//tensorflow/python/framework:random_seed", "//tensorflow/python/framework:test_lib", + "//tensorflow/python/ops:random_ops", "//tensorflow/python/platform:client_testlib", - "//third_party/py/numpy", ], ) diff --git a/tensorflow/python/keras/layers/quantized_test.py b/tensorflow/python/keras/layers/quantized_test.py index 7b2dd80407dee6..7a7b04680bb653 100644 --- a/tensorflow/python/keras/layers/quantized_test.py +++ b/tensorflow/python/keras/layers/quantized_test.py @@ -14,10 +14,9 @@ # ============================================================================== """Tests for quantized Dense layer.""" -import numpy as np -import tensorflow as tf - +from tensorflow.python.framework import random_seed from tensorflow.python.keras.layers.quantized import QuantizedDense +from tensorflow.python.ops import random_ops from tensorflow.python.platform import test @@ -25,11 +24,10 @@ class QuantizedDenseTest(test.TestCase): def setUp(self): super(QuantizedDenseTest, self).setUp() - tf.random.set_seed(0) - np.random.seed(0) + random_seed.set_random_seed(0) def test_quantized_dense_basic(self): - inputs = tf.random.uniform((32, 128)) + inputs = random_ops.random_uniform((32, 128)) # Test 8-bit quantization layer_8bit = QuantizedDense(64, bits=8) @@ -40,7 +38,7 @@ def test_quantized_dense_basic(self): self.assertEqual(layer_8bit.bias.shape, (64,)) def test_quantized_dense_4bit(self): - inputs = tf.random.uniform((16, 32)) + inputs = random_ops.random_uniform((16, 32)) # Test 4-bit quantization layer_4bit = QuantizedDense(16, bits=4, use_bias=False) @@ -54,5 +52,6 @@ def test_invalid_bits(self): "Only 4-bit and 8-bit quantization"): QuantizedDense(32, bits=16) + if __name__ == "__main__": test.main() From 2a764da0ec4fbc2289d40fcea2d2278076dcc724 Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Thu, 27 Aug 2026 20:13:54 +0530 Subject: [PATCH 09/28] Fix BUILD error: replace quantized_ops with array_ops and nn_ops in deps --- tensorflow/python/keras/layers/BUILD | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD index 1ef5de313ac7d8..eeec755ae556c7 100644 --- a/tensorflow/python/keras/layers/BUILD +++ b/tensorflow/python/keras/layers/BUILD @@ -148,7 +148,8 @@ py_library( "//tensorflow/python/keras/initializers", "//tensorflow/python/ops:math_ops", "//tensorflow/python/ops:nn", - "//tensorflow/python/ops:quantized_ops", + "//tensorflow/python/ops:array_ops", + "//tensorflow/python/ops:nn_ops", ], ) From d4a40a9379991a8ea27406cc34b5387274c89d79 Mon Sep 17 00:00:00 2001 From: Alexander Grund Date: Fri, 28 Aug 2026 09:41:13 +0200 Subject: [PATCH 10/28] Move `BUILD.system` files to `systemlibs` as `.BUILD` files Allow them to be recognized as build extensions. --- third_party/flatbuffers/workspace.bzl | 2 +- third_party/icu/workspace.bzl | 1 + third_party/jpeg/workspace.bzl | 2 +- .../{flatbuffers/BUILD.system => systemlibs/flatbuffers.BUILD} | 0 third_party/{icu/BUILD.system => systemlibs/icu.BUILD} | 0 third_party/{jpeg/BUILD.system => systemlibs/jpeg.BUILD} | 0 6 files changed, 3 insertions(+), 2 deletions(-) rename third_party/{flatbuffers/BUILD.system => systemlibs/flatbuffers.BUILD} (100%) rename third_party/{icu/BUILD.system => systemlibs/icu.BUILD} (100%) rename third_party/{jpeg/BUILD.system => systemlibs/jpeg.BUILD} (100%) diff --git a/third_party/flatbuffers/workspace.bzl b/third_party/flatbuffers/workspace.bzl index aa0ba40e8aec9d..89aa1bdc59ce74 100644 --- a/third_party/flatbuffers/workspace.bzl +++ b/third_party/flatbuffers/workspace.bzl @@ -14,7 +14,7 @@ def repo(): sha256 = _FLATBUFFERS_SHA256, urls = tf_mirror_urls("https://github.com/google/flatbuffers/archive/v%s.tar.gz" % _FLATBUFFERS_VERSION), build_file = "//third_party/flatbuffers:flatbuffers.BUILD", - system_build_file = "//third_party/flatbuffers:BUILD.system", + system_build_file = "//third_party/systemlibs:flatbuffers.BUILD", link_files = { "//third_party/flatbuffers:build_defs.bzl": "build_defs.bzl", }, diff --git a/third_party/icu/workspace.bzl b/third_party/icu/workspace.bzl index 3773cc43ddaeb6..800dfaede26576 100644 --- a/third_party/icu/workspace.bzl +++ b/third_party/icu/workspace.bzl @@ -11,6 +11,7 @@ def repo(): sha256 = "588e431f77327c39031ffbb8843c0e3bc122c211374485fa87dc5f3faff24061", urls = tf_mirror_urls("https://github.com/unicode-org/icu/releases/download/release-77-1/icu4c-77_1-src.tgz"), build_file = "//third_party/icu:icu.BUILD", + system_build_file = "//third_party/systemlibs/icu.BUILD", patch_file = ["//third_party/icu:udata.patch"], patch_cmds = [ "rm -f source/common/BUILD.bazel", diff --git a/third_party/jpeg/workspace.bzl b/third_party/jpeg/workspace.bzl index 631cc933bc60d9..75457d9cd1c2b2 100644 --- a/third_party/jpeg/workspace.bzl +++ b/third_party/jpeg/workspace.bzl @@ -9,5 +9,5 @@ def repo(): sha256 = "a78b05c0d8427a90eb5b4eb08af25309770c8379592bb0b8a863373128e6143f", strip_prefix = "libjpeg-turbo-2.1.4", build_file = "//third_party/jpeg:jpeg.BUILD", - system_build_file = "//third_party/jpeg:BUILD.system", + system_build_file = "//third_party/systemlibs/jpeg.BUILD", ) diff --git a/third_party/flatbuffers/BUILD.system b/third_party/systemlibs/flatbuffers.BUILD similarity index 100% rename from third_party/flatbuffers/BUILD.system rename to third_party/systemlibs/flatbuffers.BUILD diff --git a/third_party/icu/BUILD.system b/third_party/systemlibs/icu.BUILD similarity index 100% rename from third_party/icu/BUILD.system rename to third_party/systemlibs/icu.BUILD diff --git a/third_party/jpeg/BUILD.system b/third_party/systemlibs/jpeg.BUILD similarity index 100% rename from third_party/jpeg/BUILD.system rename to third_party/systemlibs/jpeg.BUILD From ab2091fbef6f2844143d7cd472553c0637a010b8 Mon Sep 17 00:00:00 2001 From: Alexander Grund Date: Sat, 29 Aug 2026 14:07:56 +0200 Subject: [PATCH 11/28] Fix wrong syntax --- third_party/icu/workspace.bzl | 2 +- third_party/jpeg/workspace.bzl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/icu/workspace.bzl b/third_party/icu/workspace.bzl index 800dfaede26576..6fb281e548c391 100644 --- a/third_party/icu/workspace.bzl +++ b/third_party/icu/workspace.bzl @@ -11,7 +11,7 @@ def repo(): sha256 = "588e431f77327c39031ffbb8843c0e3bc122c211374485fa87dc5f3faff24061", urls = tf_mirror_urls("https://github.com/unicode-org/icu/releases/download/release-77-1/icu4c-77_1-src.tgz"), build_file = "//third_party/icu:icu.BUILD", - system_build_file = "//third_party/systemlibs/icu.BUILD", + system_build_file = "//third_party/systemlibs:icu.BUILD", patch_file = ["//third_party/icu:udata.patch"], patch_cmds = [ "rm -f source/common/BUILD.bazel", diff --git a/third_party/jpeg/workspace.bzl b/third_party/jpeg/workspace.bzl index 75457d9cd1c2b2..579d95ba4fed32 100644 --- a/third_party/jpeg/workspace.bzl +++ b/third_party/jpeg/workspace.bzl @@ -9,5 +9,5 @@ def repo(): sha256 = "a78b05c0d8427a90eb5b4eb08af25309770c8379592bb0b8a863373128e6143f", strip_prefix = "libjpeg-turbo-2.1.4", build_file = "//third_party/jpeg:jpeg.BUILD", - system_build_file = "//third_party/systemlibs/jpeg.BUILD", + system_build_file = "//third_party/systemlibs:jpeg.BUILD", ) From 30a12786ab44139bcc26af38d52c91c5032fec3b Mon Sep 17 00:00:00 2001 From: adi-IL Date: Mon, 31 Aug 2026 02:00:17 +0530 Subject: [PATCH 12/28] [XLA:tf2xla] Support half, bfloat16, and integer dtypes in RangeOp --- tensorflow/compiler/tests/ternary_ops_test.py | 25 ++++---- .../compiler/tf2xla/kernels/sequence_ops.cc | 61 +++++++++++++++---- 2 files changed, 61 insertions(+), 25 deletions(-) diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py index 101ca75f8b68be..63629688263dca 100644 --- a/tensorflow/compiler/tests/ternary_ops_test.py +++ b/tensorflow/compiler/tests/ternary_ops_test.py @@ -60,18 +60,19 @@ def testLinspace(self, start, end, num): self.assertEqual(result[0], expected[0]) def testRange(self): - self._testTernary( - math_ops.range, - np.int32(1), - np.int32(2), - np.int32(1), - expected=np.array([1], dtype=np.int32)) - self._testTernary( - math_ops.range, - np.int32(1), - np.int32(7), - np.int32(2), - expected=np.array([1, 3, 5], dtype=np.int32)) + for dtype in self.int_types | self.float_types: + self._testTernary( + math_ops.range, + dtype(1), + dtype(2), + dtype(1), + expected=np.array([1], dtype=dtype)) + self._testTernary( + math_ops.range, + dtype(1), + dtype(7), + dtype(2), + expected=np.array([1, 3, 5], dtype=dtype)) def testSelect(self): for dtype in self.numeric_types: diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc index d24d1688d188a6..f22932e15b45c1 100644 --- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc @@ -48,10 +48,10 @@ absl::StatusOr CreateRangeTensor( T limit = limit_literal.Get({}); T delta = delta_literal.Get({}); - if (delta == 0) { + if (delta == static_cast(0)) { return errors::InvalidArgument("Requires delta != 0: ", delta); } - if (delta > 0) { + if (delta > static_cast(0)) { if (start > limit) { return errors::InvalidArgument( "Requires start <= limit when delta > 0: ", start, "/", limit); @@ -62,13 +62,21 @@ absl::StatusOr CreateRangeTensor( "Requires start >= limit when delta < 0: ", start, "/", limit); } } - int64_t size = - (std::is_integral::value - ? static_cast( - limit == start - ? 0 - : (std::abs(limit - start) - 1) / std::abs(delta) + 1) - : std::ceil(std::abs((limit - start) / delta))); + int64_t size; + if constexpr (std::is_integral::value) { + int64_t start_i = static_cast(start); + int64_t limit_i = static_cast(limit); + int64_t delta_i = static_cast(delta); + size = (limit_i == start_i + ? 0 + : (std::abs(limit_i - start_i) - 1) / std::abs(delta_i) + 1); + } else { + double start_f = static_cast(start); + double limit_f = static_cast(limit); + double delta_f = static_cast(delta); + size = static_cast( + std::ceil(std::abs((limit_f - start_f) / delta_f))); + } return xla::ConstantR0(builder, start) + xla::ConstantR0(builder, delta) * @@ -86,13 +94,13 @@ class RangeOp : public XlaOpKernel { const TensorShape delta_in_shape = ctx->InputShape(2); OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(start_in_shape), errors::InvalidArgument("start must be a scalar, not shape ", - start_in_shape.DebugString())); + start_in_shape.DebugString())); OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(limit_in_shape), errors::InvalidArgument("limit must be a scalar, not shape ", - limit_in_shape.DebugString())); + limit_in_shape.DebugString())); OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(delta_in_shape), errors::InvalidArgument("delta must be a scalar, not shape ", - delta_in_shape.DebugString())); + delta_in_shape.DebugString())); xla::Literal start, limit, delta; OP_REQUIRES_OK(ctx, ctx->ConstantInput( 0, &start, xla::ValueInferenceMode::kLowerBound)); @@ -103,6 +111,13 @@ class RangeOp : public XlaOpKernel { DataType type = input_type(0); absl::StatusOr output; switch (type) { + case DT_INT8: + output = CreateRangeTensor(start, limit, delta, ctx->builder()); + break; + case DT_INT16: + output = + CreateRangeTensor(start, limit, delta, ctx->builder()); + break; case DT_INT32: output = CreateRangeTensor(start, limit, delta, ctx->builder()); @@ -111,6 +126,26 @@ class RangeOp : public XlaOpKernel { output = CreateRangeTensor(start, limit, delta, ctx->builder()); break; + case DT_UINT16: + output = + CreateRangeTensor(start, limit, delta, ctx->builder()); + break; + case DT_UINT32: + output = + CreateRangeTensor(start, limit, delta, ctx->builder()); + break; + case DT_UINT64: + output = + CreateRangeTensor(start, limit, delta, ctx->builder()); + break; + case DT_HALF: + output = + CreateRangeTensor(start, limit, delta, ctx->builder()); + break; + case DT_BFLOAT16: + output = + CreateRangeTensor(start, limit, delta, ctx->builder()); + break; case DT_FLOAT: output = CreateRangeTensor(start, limit, delta, ctx->builder()); break; @@ -133,7 +168,7 @@ class RangeOp : public XlaOpKernel { xla::XlaOp delta = ctx->Input(2); xla::XlaOp limit = ctx->Input(1); xla::XlaOp start = ctx->Input(0); - if (type == DT_INT32 || type == DT_INT64) { + if (DataTypeIsInteger(type)) { auto dynamic_size = (xla::Abs(limit - start) + xla::Abs(delta) - xla::One(ctx->builder(), ctx->input_xla_type(0))) / xla::Abs(delta); From 170787f4548121e0311325d3c1d29ba0cbd5688f Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Mon, 31 Aug 2026 06:50:20 +0530 Subject: [PATCH 13/28] Enhance accelerator detection and JSON output Refactor accelerator detection and JSON summary generation in tf_env_collect.sh. Improve temporary file handling and streamline pip list checks. --- tools/tf_env_collect.sh | 79 ++++++++++++++++++++++++++++++++--------- 1 file changed, 62 insertions(+), 17 deletions(-) diff --git a/tools/tf_env_collect.sh b/tools/tf_env_collect.sh index 390b4b6fd18f1e..7d09b97887f3fc 100755 --- a/tools/tf_env_collect.sh +++ b/tools/tf_env_collect.sh @@ -18,8 +18,10 @@ set -u # Check for undefined variables # Track temporary files so they are removed on exit, including on interrupt. LOADED_LIBS_FILE="" +ACCEL_FLAGS_FILE="" cleanup() { [ -n "${LOADED_LIBS_FILE:-}" ] && rm -f "$LOADED_LIBS_FILE" + [ -n "${ACCEL_FLAGS_FILE:-}" ] && rm -f "$ACCEL_FLAGS_FILE" } trap cleanup EXIT INT TERM @@ -79,6 +81,14 @@ case "${OUTPUT_FILE##*/}" in *) JSON_FILE="${OUTPUT_FILE}.json" ;; esac +# Only pay for a temp file when the JSON summary is actually requested; it's +# used to smuggle a couple of accelerator-detection booleans out of the +# report-generation subshell below so we don't have to re-probe nvidia-smi / +# rocm-smi / tensorflow-metal a second time just for the JSON section. +if [ "$EMIT_JSON" -eq 1 ]; then + ACCEL_FLAGS_FILE="$(mktemp 2>/dev/null || mktemp -t tfenv)" +fi + echo "Collecting system information..." PYTHON_BIN_PATH="$(command -v python || command -v python3 || die "Cannot find Python binary")" @@ -93,11 +103,15 @@ have_cmd() { run_cmd() { # Run a command if it exists, otherwise note that it is missing instead of - # erroring out. Captures stderr so the report stays readable. + # erroring out. Captures stderr so the report stays readable. Propagates + # the real exit status of the command (or 127 if it wasn't found) so + # callers can check success without re-invoking the command. if have_cmd "$1"; then "$@" 2>&1 + return $? else echo "$1 not found" + return 127 fi } @@ -125,8 +139,10 @@ pip_run() { TF_PKG_PATTERN='^(tensorflow|tf-nightly|tensorflow-cpu|tensorflow-gpu|tensorflow-rocm|tensorflow-macos|tensorflow-metal|intel-tensorflow)\b' HEADER_WIDTH=68 -# Create a string of HEADER_WIDTH "=" characters -HEADER=$(printf "%*s" "$HEADER_WIDTH" "" | sed 's/ /=/g') +# Build a string of HEADER_WIDTH "=" characters using shell builtins only +# (printf -v + parameter expansion), avoiding a fork+pipe through sed. +printf -v HEADER '%*s' "$HEADER_WIDTH" '' +HEADER=${HEADER// /=} print_header () { # This function simply prints the header with even spacing, @@ -212,8 +228,12 @@ EOF echo "Not found" fi + # Fetch "pip list" once and reuse it below for the TensorFlow package + # conflict check, instead of shelling out to pip a second time. + PIP_LIST_OUTPUT="$(pip_run list 2>&1)" + print_header 'check pips' - pip_run list 2>&1 | grep -E 'proto|numpy|keras|tensorflow|tf_nightly|tf-nightly' + grep -E 'proto|numpy|keras|tensorflow|tf_nightly|tf-nightly' <<<"$PIP_LIST_OUTPUT" print_header 'check for virtualenv' @@ -232,12 +252,12 @@ EOF print_header 'tensorflow package conflicts' # Multiple TensorFlow distributions in the same environment frequently cause # confusing import errors; surface them so triage can spot the conflict. - TF_PKGS="$(pip_run list 2>/dev/null | grep -iE "$TF_PKG_PATTERN" || true)" + TF_PKGS="$(grep -iE "$TF_PKG_PATTERN" <<<"$PIP_LIST_OUTPUT" || true)" if [ -z "$TF_PKGS" ]; then echo "No TensorFlow packages found via pip." else echo "$TF_PKGS" - TF_COUNT="$(echo "$TF_PKGS" | grep -icE "$TF_PKG_PATTERN")" + TF_COUNT="$(grep -icE "$TF_PKG_PATTERN" <<<"$TF_PKGS")" if [ "$TF_COUNT" -gt 1 ]; then echo "WARNING: multiple TensorFlow distributions detected; this can cause import conflicts." fi @@ -315,16 +335,17 @@ EOF print_header 'build / hermetic accelerator config' # Surface the environment variables that control modern (hermetic) CUDA and - # ROCm builds. See .bazelrc for how these are consumed. + # ROCm builds. See .bazelrc for how these are consumed. Uses indirect + # parameter expansion instead of eval - one less string re-parse per + # variable, and no eval footgun. for var in CC CXX \ TF_NEED_CUDA TF_NEED_ROCM TF_CUDA_VERSION TF_CUDNN_VERSION \ HERMETIC_CUDA_VERSION HERMETIC_CUDNN_VERSION \ CUDA_HOME CUDA_PATH CUDA_TOOLKIT_PATH \ ROCM_PATH HIP_PATH \ XLA_FLAGS TF_XLA_FLAGS TPU_NAME; do - eval "marker=\${$var+set} val=\"\$$var\"" - if [ "${marker:-}" = "set" ]; then - echo "$var=$val" + if [ -n "${!var+set}" ]; then + echo "$var=${!var}" else echo "$var is unset" fi @@ -332,6 +353,10 @@ EOF print_header 'accelerator: nvidia gpu' run_cmd nvidia-smi + NVIDIA_STATUS=$? + if [ -n "$ACCEL_FLAGS_FILE" ] && [ "$NVIDIA_STATUS" -eq 0 ]; then + echo "HAS_NVIDIA=1" >> "$ACCEL_FLAGS_FILE" + fi print_header 'cuda libs' # Find cudart/cudnn files @@ -343,6 +368,10 @@ EOF print_header 'accelerator: amd / rocm gpu' run_cmd rocm-smi + ROCM_STATUS=$? + if [ -n "$ACCEL_FLAGS_FILE" ] && [ "$ROCM_STATUS" -eq 0 ]; then + echo "HAS_ROCM=1" >> "$ACCEL_FLAGS_FILE" + fi if [ "$VERBOSE" -eq 1 ]; then print_header 'rocminfo' run_cmd rocminfo @@ -355,9 +384,16 @@ EOF # tensorflow-metal is the PluggableDevice that enables GPU acceleration on # Apple Silicon; report whether it is installed and the GPU chipset. if [ "$(uname -s)" = "Darwin" ]; then - if pip_run show tensorflow-metal >/dev/null 2>&1; then + # Single "pip show" call, reused both for the existence check and for + # the Name/Version detail line below (previously called twice). + TF_METAL_INFO="$(pip_run show tensorflow-metal 2>&1)" + TF_METAL_STATUS=$? + if [ "$TF_METAL_STATUS" -eq 0 ]; then echo "tensorflow-metal installed:" - pip_run show tensorflow-metal 2>&1 | grep -iE '^(Name|Version):' + grep -iE '^(Name|Version):' <<<"$TF_METAL_INFO" + if [ -n "$ACCEL_FLAGS_FILE" ]; then + echo "HAS_METAL=1" >> "$ACCEL_FLAGS_FILE" + fi else echo "tensorflow-metal not installed" fi @@ -396,14 +432,23 @@ EOF # Optional machine-readable JSON summary # ---------------------------------------------------------------------------- if [ "$EMIT_JSON" -eq 1 ]; then - # Detect accelerators in the shell and hand the booleans to Python, which - # assembles a structured, easy-to-parse summary of the key facts. + # nvidia-smi / rocm-smi / tensorflow-metal detection already happened once + # above (inside the report-generation block); read the results back in + # from ACCEL_FLAGS_FILE instead of re-invoking those (potentially slow) + # tools a second time. HAS_NVIDIA=0 - if have_cmd nvidia-smi && nvidia-smi >/dev/null 2>&1; then HAS_NVIDIA=1; fi HAS_ROCM=0 - if have_cmd rocm-smi && rocm-smi >/dev/null 2>&1; then HAS_ROCM=1; fi HAS_METAL=0 - if pip_run show tensorflow-metal >/dev/null 2>&1; then HAS_METAL=1; fi + if [ -n "$ACCEL_FLAGS_FILE" ] && [ -s "$ACCEL_FLAGS_FILE" ]; then + # shellcheck disable=SC1090 + . "$ACCEL_FLAGS_FILE" + fi + # The block above only probes tensorflow-metal on Darwin hosts (matching + # the text report). On any other platform, fall back to a direct check so + # the JSON summary's accelerator info stays complete. + if [ "$HAS_METAL" -ne 1 ] && [ "$(uname -s)" != "Darwin" ]; then + if pip_run show tensorflow-metal >/dev/null 2>&1; then HAS_METAL=1; fi + fi TFENV_HAS_NVIDIA="$HAS_NVIDIA" \ TFENV_HAS_ROCM="$HAS_ROCM" \ From b4ac550b146d4ae92c8ca03030230ccdbab450a0 Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Sat, 29 Aug 2026 12:37:11 +0530 Subject: [PATCH 14/28] Refactor: Move QuantizedDense from Keras to core TF ops --- tensorflow/python/keras/layers/BUILD | 36 ---- tensorflow/python/keras/layers/__init__.py | 1 - tensorflow/python/keras/layers/quantized.py | 154 ------------------ tensorflow/python/ops/BUILD | 30 ++++ tensorflow/python/ops/quantized_dense.py | 98 +++++++++++ .../quantized_dense_test.py} | 9 +- 6 files changed, 133 insertions(+), 195 deletions(-) delete mode 100644 tensorflow/python/keras/layers/quantized.py create mode 100644 tensorflow/python/ops/quantized_dense.py rename tensorflow/python/{keras/layers/quantized_test.py => ops/quantized_dense_test.py} (88%) diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD index eeec755ae556c7..83925764d3fbb0 100644 --- a/tensorflow/python/keras/layers/BUILD +++ b/tensorflow/python/keras/layers/BUILD @@ -17,7 +17,6 @@ # Contains the Keras layers (internal TensorFlow version). load("@xla//third_party/rules_python/python:defs.bzl", "py_library") -load("//tensorflow:tensorflow.default.bzl", "tf_py_test") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -55,7 +54,6 @@ py_library( ":convolutional", ":convolutional_recurrent", ":core", - ":quantized", ":dense_attention", ":embeddings", ":merge", @@ -133,40 +131,6 @@ py_library( ], ) -py_library( - name = "quantized", - srcs = ["quantized.py"], - srcs_version = "PY3", - deps = [ - "//tensorflow/python/framework:dtypes", - "//tensorflow/python/framework:tensor_shape", - "//tensorflow/python/keras:activations", - "//tensorflow/python/keras:base_layer", - "//tensorflow/python/keras:constraints", - "//tensorflow/python/keras:regularizers", - "//tensorflow/python/keras/engine:input_spec", - "//tensorflow/python/keras/initializers", - "//tensorflow/python/ops:math_ops", - "//tensorflow/python/ops:nn", - "//tensorflow/python/ops:array_ops", - "//tensorflow/python/ops:nn_ops", - ], -) - -tf_py_test( - name = "quantized_test", - srcs = ["quantized_test.py"], - python_version = "PY3", - srcs_version = "PY3", - deps = [ - ":quantized", - "//tensorflow/python/framework:random_seed", - "//tensorflow/python/framework:test_lib", - "//tensorflow/python/ops:random_ops", - "//tensorflow/python/platform:client_testlib", - ], -) - py_library( name = "core", srcs = ["core.py"], diff --git a/tensorflow/python/keras/layers/__init__.py b/tensorflow/python/keras/layers/__init__.py index 984ee58afbd836..889e9d181fb3af 100644 --- a/tensorflow/python/keras/layers/__init__.py +++ b/tensorflow/python/keras/layers/__init__.py @@ -77,7 +77,6 @@ from tensorflow.python.keras.layers.core import RepeatVector from tensorflow.python.keras.layers.core import Lambda from tensorflow.python.keras.layers.core import Dense -from tensorflow.python.keras.layers.quantized import QuantizedDense from tensorflow.python.keras.layers.core import ActivityRegularization # Dense Attention layers. diff --git a/tensorflow/python/keras/layers/quantized.py b/tensorflow/python/keras/layers/quantized.py deleted file mode 100644 index 9bb2e3f42dd9e8..00000000000000 --- a/tensorflow/python/keras/layers/quantized.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright 2026 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Quantized Dense layer.""" - -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import tensor_shape -from tensorflow.python.keras import activations -from tensorflow.python.keras import constraints -from tensorflow.python.keras import initializers -from tensorflow.python.keras import regularizers -from tensorflow.python.keras.engine.base_layer import Layer -from tensorflow.python.keras.engine.input_spec import InputSpec -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn_ops - - -class QuantizedDense(Layer): - """A densely-connected layer with weight quantization. - - This layer acts like a standard Dense layer but simulates - 4-bit or 8-bit weight quantization using fake quantization nodes. - """ - - def __init__(self, - units, - bits=8, - activation=None, - use_bias=True, - kernel_initializer="glorot_uniform", - bias_initializer="zeros", - kernel_regularizer=None, - bias_regularizer=None, - activity_regularizer=None, - kernel_constraint=None, - bias_constraint=None, - **kwargs): - super(QuantizedDense, self).__init__( - activity_regularizer=activity_regularizer, **kwargs) - self.units = int(units) - self.bits = int(bits) - if self.bits not in [4, 8]: - raise ValueError("Only 4-bit and 8-bit quantization are supported.") - self.activation = activations.get(activation) - self.use_bias = use_bias - self.kernel_initializer = initializers.get(kernel_initializer) - self.bias_initializer = initializers.get(bias_initializer) - self.kernel_regularizer = regularizers.get(kernel_regularizer) - self.bias_regularizer = regularizers.get(bias_regularizer) - self.kernel_constraint = constraints.get(kernel_constraint) - self.bias_constraint = constraints.get(bias_constraint) - self.input_spec = InputSpec(min_ndim=2) - - def build(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - last_dim = tensor_shape.dimension_value(input_shape[-1]) - if last_dim is None: - raise ValueError( - "The last dimension of the inputs to `QuantizedDense` " - "should be defined. Found `None`.") - self.input_spec = InputSpec(min_ndim=2, axes={-1: last_dim}) - self.kernel = self.add_weight( - "kernel", - shape=[last_dim, self.units], - initializer=self.kernel_initializer, - regularizer=self.kernel_regularizer, - constraint=self.kernel_constraint, - dtype=self.dtype, - trainable=True) - if self.use_bias: - self.bias = self.add_weight( - "bias", - shape=[self.units,], - initializer=self.bias_initializer, - regularizer=self.bias_regularizer, - constraint=self.bias_constraint, - dtype=self.dtype, - trainable=True) - else: - self.bias = None - self.built = True - - def call(self, inputs): - kernel = math_ops.cast(self.kernel, dtypes.float32) - min_val = math_ops.reduce_min(kernel) - max_val = math_ops.reduce_max(kernel) - max_val = math_ops.maximum(max_val, min_val + 1e-5) - - quantized_kernel = array_ops.fake_quant_with_min_max_vars( - kernel, - min_val, - max_val, - num_bits=self.bits, - narrow_range=True) - quantized_kernel = math_ops.cast(quantized_kernel, self.dtype) - - rank = inputs.shape.rank - if rank is not None and rank <= 2: - outputs = math_ops.matmul(a=inputs, b=quantized_kernel) - else: - axes = [[rank - 1 if rank else -1], [0]] - outputs = math_ops.tensordot(inputs, quantized_kernel, axes) - - if self.use_bias: - outputs = nn_ops.bias_add(outputs, self.bias) - - if self.activation is not None: - outputs = self.activation(outputs) - return outputs - - def compute_output_shape(self, input_shape): - input_shape = tensor_shape.TensorShape(input_shape) - input_shape = input_shape.with_rank_at_least(2) - if tensor_shape.dimension_value(input_shape[-1]) is None: - raise ValueError( - "The innermost dimension of input_shape must be defined, " - "but saw: %s" % input_shape) - return input_shape[:-1].concatenate(self.units) - - def get_config(self): - config = super(QuantizedDense, self).get_config() - config.update({ - "units": self.units, - "bits": self.bits, - "activation": activations.serialize(self.activation), - "use_bias": self.use_bias, - "kernel_initializer": initializers.serialize( - self.kernel_initializer), - "bias_initializer": initializers.serialize( - self.bias_initializer), - "kernel_regularizer": regularizers.serialize( - self.kernel_regularizer), - "bias_regularizer": regularizers.serialize( - self.bias_regularizer), - "activity_regularizer": regularizers.serialize( - self.activity_regularizer), - "kernel_constraint": constraints.serialize( - self.kernel_constraint), - "bias_constraint": constraints.serialize( - self.bias_constraint) - }) - return config diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 34817b3c0bd601..bdc8af4a1fc928 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -4827,3 +4827,33 @@ py_test( "//tensorflow/python/platform:client_testlib", ], ) + +py_library( + name = "quantized_dense", + srcs = ["quantized_dense.py"], + srcs_version = "PY3", + deps = [ + ":array_ops", + ":math_ops", + ":nn_ops", + ":random_ops", + ":variables", + "//tensorflow/python/framework:dtypes", + "//tensorflow/python/framework:ops", + "//tensorflow/python/module", + ], +) + +tf_py_strict_test( + name = "quantized_dense_test", + size = "small", + srcs = ["quantized_dense_test.py"], + python_version = "PY3", + deps = [ + ":quantized_dense", + ":random_ops", + "//tensorflow/python/framework:random_seed", + "//tensorflow/python/platform:client_testlib", + ], +) + diff --git a/tensorflow/python/ops/quantized_dense.py b/tensorflow/python/ops/quantized_dense.py new file mode 100644 index 00000000000000..5d4ca6ee37c3e8 --- /dev/null +++ b/tensorflow/python/ops/quantized_dense.py @@ -0,0 +1,98 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Quantized Dense module.""" + +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.module import module +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import random_ops +from tensorflow.python.ops import variables + + +class QuantizedDense(module.Module): + """A densely-connected layer with weight quantization. + + This module acts like a standard Dense layer but simulates + 4-bit or 8-bit weight quantization using fake quantization nodes. + """ + + def __init__(self, units, bits=8, use_bias=True, name=None): + super(QuantizedDense, self).__init__(name=name) + self.units = int(units) + self.bits = int(bits) + if self.bits not in [4, 8]: + raise ValueError("Only 4-bit and 8-bit quantization are supported.") + self.use_bias = use_bias + self.kernel = None + self.bias = None + + def __call__(self, inputs): + inputs = ops.convert_to_tensor(inputs) + if self.kernel is None: + last_dim = inputs.shape[-1] + if last_dim is None: + raise ValueError( + "The last dimension of the inputs to `QuantizedDense` should be" + " defined." + ) + # Initialize weights with glorot uniform + limit = math_ops.sqrt(6.0 / (last_dim + self.units)) + self.kernel = variables.Variable( + initial_value=random_ops.random_uniform( + [last_dim, self.units], + minval=-limit, + maxval=limit, + dtype=inputs.dtype, + ), + name="kernel", + trainable=True, + ) + if self.use_bias: + self.bias = variables.Variable( + initial_value=array_ops.zeros( + [ + self.units, + ], + dtype=inputs.dtype, + ), + name="bias", + trainable=True, + ) + + kernel = math_ops.cast(self.kernel, dtypes.float32) + min_val = math_ops.reduce_min(kernel) + max_val = math_ops.reduce_max(kernel) + max_val = math_ops.maximum(max_val, min_val + 1e-5) + + quantized_kernel = array_ops.fake_quant_with_min_max_vars( + kernel, min_val, max_val, num_bits=self.bits, narrow_range=True + ) + quantized_kernel = math_ops.cast(quantized_kernel, inputs.dtype) + + rank = inputs.shape.rank + if rank is not None and rank <= 2: + outputs = math_ops.matmul(a=inputs, b=quantized_kernel) + else: + outputs = math_ops.tensordot( + inputs, quantized_kernel, [[rank - 1 if rank else -1], [0]] + ) + + if self.use_bias: + outputs = nn_ops.bias_add(outputs, self.bias) + + return outputs diff --git a/tensorflow/python/keras/layers/quantized_test.py b/tensorflow/python/ops/quantized_dense_test.py similarity index 88% rename from tensorflow/python/keras/layers/quantized_test.py rename to tensorflow/python/ops/quantized_dense_test.py index 7a7b04680bb653..8734cc3aae55cc 100644 --- a/tensorflow/python/keras/layers/quantized_test.py +++ b/tensorflow/python/ops/quantized_dense_test.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for quantized Dense layer.""" +"""Tests for quantized Dense module.""" from tensorflow.python.framework import random_seed -from tensorflow.python.keras.layers.quantized import QuantizedDense from tensorflow.python.ops import random_ops +from tensorflow.python.ops.quantized_dense import QuantizedDense from tensorflow.python.platform import test @@ -48,8 +48,9 @@ def test_quantized_dense_4bit(self): self.assertIsNone(layer_4bit.bias) def test_invalid_bits(self): - with self.assertRaisesRegex(ValueError, - "Only 4-bit and 8-bit quantization"): + with self.assertRaisesRegex( + ValueError, "Only 4-bit and 8-bit quantization" + ): QuantizedDense(32, bits=16) From 32fe50eb5943bd8f0bcf5dfe41d9b13c06683d35 Mon Sep 17 00:00:00 2001 From: adi-IL Date: Mon, 31 Aug 2026 18:21:21 +0530 Subject: [PATCH 15/28] Exclude uint8 from XLA Range ternary tests --- tensorflow/compiler/tests/ternary_ops_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py index 63629688263dca..473f6a7812aa99 100644 --- a/tensorflow/compiler/tests/ternary_ops_test.py +++ b/tensorflow/compiler/tests/ternary_ops_test.py @@ -60,7 +60,7 @@ def testLinspace(self, start, end, num): self.assertEqual(result[0], expected[0]) def testRange(self): - for dtype in self.int_types | self.float_types: + for dtype in (self.int_types | self.float_types) - {np.uint8}: self._testTernary( math_ops.range, dtype(1), From b4d16ef9a94c9c7c49a9374a94a43c0c5e6c0443 Mon Sep 17 00:00:00 2001 From: AshiteshSingh Date: Mon, 31 Aug 2026 18:55:48 +0530 Subject: [PATCH 16/28] fix: cast last_dim to int to fix TypeError --- tensorflow/python/ops/quantized_dense.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/python/ops/quantized_dense.py b/tensorflow/python/ops/quantized_dense.py index 5d4ca6ee37c3e8..890e94fe65d872 100644 --- a/tensorflow/python/ops/quantized_dense.py +++ b/tensorflow/python/ops/quantized_dense.py @@ -50,6 +50,7 @@ def __call__(self, inputs): "The last dimension of the inputs to `QuantizedDense` should be" " defined." ) + last_dim = int(last_dim) # Initialize weights with glorot uniform limit = math_ops.sqrt(6.0 / (last_dim + self.units)) self.kernel = variables.Variable( From a584bd38508aa130ab6f3f3c0582b23dada7e329 Mon Sep 17 00:00:00 2001 From: Song Date: Wed, 29 Apr 2026 17:57:10 +0000 Subject: [PATCH 17/28] Fix MapUnstageNoKey crash on out-of-range index (#112757) In StagingMap::popitem(), the local key tensor was passed to copy_or_move_tensors() before being assigned the actual key from the map iterator. When an out-of-range index is supplied, the downstream check_index() helper formats its InvalidArgument message via key.scalar()() on the still-empty (0-element) key tensor. Tensor::scalar() then invokes CheckIsAlignedAndSingleElement() which CHECK_EQ(1, NumElements()) fails (1 vs. 0) and aborts the process with SIGABRT (exit 134), turning a recoverable input error into a fatal CHECK. Move the assignment of *key = it->first above the call to copy_or_move_tensors() so that check_index() always sees a fully constructed scalar key when formatting its error message. The existing bounds-check and InvalidArgument semantics are unchanged. Fixes #112757 --- tensorflow/core/kernels/map_stage_op.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/map_stage_op.cc b/tensorflow/core/kernels/map_stage_op.cc index e4c674646cb890..cff40deb0b4e9d 100644 --- a/tensorflow/core/kernels/map_stage_op.cc +++ b/tensorflow/core/kernels/map_stage_op.cc @@ -448,11 +448,11 @@ class StagingMap : public ResourceBase { auto it = map_.begin(); + *key = it->first; + TF_RETURN_IF_ERROR( copy_or_move_tensors(&it->second, *key, *indices, tuple)); - *key = it->first; - // Remove entry if all the values have been consumed if (!std::any_of( it->second.begin(), it->second.end(), From d604357bb65b684736f2c784d9db0c7964bafb8a Mon Sep 17 00:00:00 2001 From: Song Date: Wed, 29 Apr 2026 18:14:36 +0000 Subject: [PATCH 18/28] Add regression tests for MapUnstageNoKey out-of-range index (#112757) Adds testMapUnstageNoKeyOutOfRangeIndex and testOrderedMapUnstageNoKeyOutOfRangeIndex to MapStageTest. Each test stages one float32 value at index 0 with int64 key 1, then calls the corresponding *MapUnstageNoKey op with index 1 and asserts that the call raises tf.errors.InvalidArgumentError matching "out of bounds" instead of aborting with a fatal CHECK in Tensor::CheckIsAlignedAndSingleElement. The ordered variant covers OrderedMapUnstageNoKeyOp, which shares StagingMap::popitem() with the unordered MapUnstageNoKeyOp through the StagingMap template, so both variants exhibit the same out-of-range-index bug and need explicit coverage. Without the popitem() reorder fix in the previous commit, both tests crash the Python process with SIGABRT (exit 134) and the tensor.cc:904 "Check failed: 1 == NumElements() (1 vs. 0)" signature. With the fix, copy_or_move_tensors() sees a 1-element scalar key and check_index() returns the expected InvalidArgumentError. --- .../data_structures/map_stage_op_test.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py index 832fce050ca3d6..553736b5541c83 100644 --- a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py +++ b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py @@ -655,6 +655,69 @@ def testNonScalarKeyMapUnStage(self): ) self.evaluate(v) + def testMapUnstageNoKeyOutOfRangeIndex(self): + # MapUnstageNoKey with an out-of-range index after MapStage must surface + # a normal InvalidArgumentError from check_index(), not a fatal CHECK + # inside Tensor::CheckIsAlignedAndSingleElement. The CHECK can fire if + # popitem() reaches copy_or_move_tensors() with an empty local key + # tensor, because check_index() formats its error using + # key.scalar()() and Tensor::scalar() requires a 1-element + # tensor. + data_flow_ops.gen_data_flow_ops.map_stage( + key=constant_op.constant([1], dtype=dtypes.int64), + indices=constant_op.constant([0], dtype=dtypes.int32), + values=[constant_op.constant([1.0], dtype=dtypes.float32)], + dtypes=[dtypes.float32], + capacity=10, + memory_limit=0, + container='', + shared_name='test_map_unstage_no_key_oob', + name=None, + ) + with self.assertRaisesRegex( + errors.InvalidArgumentError, 'out of bounds' + ): + result = data_flow_ops.gen_data_flow_ops.map_unstage_no_key( + indices=[1], + dtypes=[dtypes.int64, dtypes.float32], + capacity=10, + memory_limit=0, + container='', + shared_name='test_map_unstage_no_key_oob', + name=None, + ) + self.evaluate(result) + + def testOrderedMapUnstageNoKeyOutOfRangeIndex(self): + # Parallel coverage for the ordered variant. OrderedMapUnstageNoKey + # shares StagingMap::popitem() with MapUnstageNoKey via the + # StagingMap template, so the same out-of-range-index + # path must surface InvalidArgumentError rather than abort. + data_flow_ops.gen_data_flow_ops.ordered_map_stage( + key=constant_op.constant([1], dtype=dtypes.int64), + indices=constant_op.constant([0], dtype=dtypes.int32), + values=[constant_op.constant([1.0], dtype=dtypes.float32)], + dtypes=[dtypes.float32], + capacity=10, + memory_limit=0, + container='', + shared_name='test_ordered_map_unstage_no_key_oob', + name=None, + ) + with self.assertRaisesRegex( + errors.InvalidArgumentError, 'out of bounds' + ): + result = data_flow_ops.gen_data_flow_ops.ordered_map_unstage_no_key( + indices=[1], + dtypes=[dtypes.int64, dtypes.float32], + capacity=10, + memory_limit=0, + container='', + shared_name='test_ordered_map_unstage_no_key_oob', + name=None, + ) + self.evaluate(result) + if __name__ == '__main__': test.main() From d2ec841a5877b5ce711cd0e083c33adf04209659 Mon Sep 17 00:00:00 2001 From: Song Date: Wed, 26 Aug 2026 16:40:26 -0700 Subject: [PATCH 19/28] Evaluate MapStage before unstaging in OOB tests to avoid graph-mode hang. Change-Id: Idf4224600b1fcdc29bd708f9b89d3ac4ab14b72d --- .../kernel_tests/data_structures/map_stage_op_test.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py index 553736b5541c83..8b1fb880ed9973 100644 --- a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py +++ b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py @@ -663,7 +663,7 @@ def testMapUnstageNoKeyOutOfRangeIndex(self): # tensor, because check_index() formats its error using # key.scalar()() and Tensor::scalar() requires a 1-element # tensor. - data_flow_ops.gen_data_flow_ops.map_stage( + stage_op = data_flow_ops.gen_data_flow_ops.map_stage( key=constant_op.constant([1], dtype=dtypes.int64), indices=constant_op.constant([0], dtype=dtypes.int32), values=[constant_op.constant([1.0], dtype=dtypes.float32)], @@ -674,6 +674,7 @@ def testMapUnstageNoKeyOutOfRangeIndex(self): shared_name='test_map_unstage_no_key_oob', name=None, ) + self.evaluate(stage_op) with self.assertRaisesRegex( errors.InvalidArgumentError, 'out of bounds' ): @@ -693,7 +694,7 @@ def testOrderedMapUnstageNoKeyOutOfRangeIndex(self): # shares StagingMap::popitem() with MapUnstageNoKey via the # StagingMap template, so the same out-of-range-index # path must surface InvalidArgumentError rather than abort. - data_flow_ops.gen_data_flow_ops.ordered_map_stage( + stage_op = data_flow_ops.gen_data_flow_ops.ordered_map_stage( key=constant_op.constant([1], dtype=dtypes.int64), indices=constant_op.constant([0], dtype=dtypes.int32), values=[constant_op.constant([1.0], dtype=dtypes.float32)], @@ -704,6 +705,7 @@ def testOrderedMapUnstageNoKeyOutOfRangeIndex(self): shared_name='test_ordered_map_unstage_no_key_oob', name=None, ) + self.evaluate(stage_op) with self.assertRaisesRegex( errors.InvalidArgumentError, 'out of bounds' ): From 66306d4bc3eb5ddc94eef37c00a4522a9eb13d4d Mon Sep 17 00:00:00 2001 From: adi-IL Date: Tue, 1 Sep 2026 06:10:02 +0530 Subject: [PATCH 20/28] Support integer dtypes in range dtype_hierarchy --- tensorflow/python/ops/math_ops.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py index c9676bf6fabcf0..eb4dc0a28abe64 100644 --- a/tensorflow/python/ops/math_ops.py +++ b/tensorflow/python/ops/math_ops.py @@ -2097,8 +2097,13 @@ def range(start, limit=None, delta=1, dtype=None, name="range"): # pylint: disa # infer dtype if not explicitly provided if dtype is None: dtype_hierarchy = [ + dtypes.int8, + dtypes.int16, dtypes.int32, dtypes.int64, + dtypes.uint16, + dtypes.uint32, + dtypes.uint64, dtypes.float16, dtypes.bfloat16, dtypes.float32, From 9fd77a8096870e4c3c0879eaabcf8f0c6d449eff Mon Sep 17 00:00:00 2001 From: Tori Baker Date: Wed, 2 Sep 2026 05:27:56 -0700 Subject: [PATCH 21/28] Disable cuDNN fusion for F64 data types and fix tests. cudnn doesn't accept f64 data types, so we should reject it in IsCudnnSupportedFusion. This was discovered in the test from DoNotExecuteGemmFusionWithCuDnnWhenNotSupported when enabling gemm fusion V2. I believe this test was never testing what it was intending to. Previously in V1, gemm fusion didn't support f64, so it created a fusion around the dot and another around the negate. Since cudnn rejected the _loop_ fusion around the negate, it returned "No supported configs", but this was unrelated to the dot/f64 fusion it was trying to test. I have updated the HLO to be a premade fusion to ensure it's testing the correct fusion. Then I updated IsCudnnSupportedFusion to successfully reject it (and added relevant test there). PiperOrigin-RevId: 975071652 --- .../xla/xla/backends/gpu/autotuner/cudnn.cc | 5 ++++ .../xla/backends/gpu/autotuner/cudnn_test.cc | 27 +++++++++++++++++++ .../xla/xla/backends/gpu/codegen/BUILD | 7 ++--- .../xla/backends/gpu/codegen/cudnn_test.cc | 18 ++++++++----- 4 files changed, 45 insertions(+), 12 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc index 1f1d6ff1a06ab3..4c0577acedec25 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc @@ -153,6 +153,11 @@ bool IsSupportedCudnnFusion(const HloInstruction& instr, return true; } + if (hero->shape().element_type() == PrimitiveType::F64) { + VLOG(1) << "cuDNN GEMM fusion does not support F64."; + return false; + } + stream_executor::CudaComputeCapability compute_capability = target_config.device_description.cuda_compute_capability(); if ((compute_capability.IsAtLeastAmpere() && diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc index ecba9d7af5786c..f6ab20f5dc74b5 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc @@ -47,7 +47,9 @@ namespace gpu { using CudnnBackendConfig = stream_executor::dnn::AlgorithmProto; +using ::absl_testing::IsOkAndHolds; using ::testing::Gt; +using ::testing::IsEmpty; using ::testing::SizeIs; using ::tsl::proto_testing::EqualsProto; @@ -104,6 +106,22 @@ absl::string_view kTritonGemmFusionHlo = R"hlo( backend_config={"fusion_backend_config": {kind: "__triton_gemm"}} })hlo"; +absl::string_view kF64GemmFusionHlo = R"hlo( + fusion1 { + p0 = f64[3,28,32] parameter(0) + p1 = f64[3,28,32] parameter(1) + ROOT d = f64[3,32,32] dot(p0, p1), + lhs_batch_dims={0}, rhs_batch_dims={0}, + lhs_contracting_dims={1}, rhs_contracting_dims={1} + } + + e { + p0 = f64[3,28,32] parameter(0) + p1 = f64[3,28,32] parameter(1) + ROOT _ = f64[3,32,32] fusion(p0, p1), kind=kCustom, calls=fusion1, + backend_config={"fusion_backend_config": {kind: "__triton_gemm"}} + })hlo"; + absl::string_view kScaledDotGemmFusionHlo = R"hlo( block_scaled_dot { lhs = f8e4m3fn[256,128] parameter(0) @@ -215,6 +233,15 @@ TEST_F(CudnnBackendTest, GetSupportedConfigsFromTritonGemmFusion) { EXPECT_THAT(configs, absl_testing::IsOkAndHolds(SizeIs(Gt(0)))); } +TEST_F(CudnnBackendTest, GetSupportedConfigsFromF64GemmFusionReturnsEmpty) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(kF64GemmFusionHlo)); + absl::StatusOr>> configs = + backend_->GetSupportedConfigs( + (*hlo_module->entry_computation()->root_instruction())); + EXPECT_THAT(configs, IsOkAndHolds(IsEmpty())); +} + TEST_F(CudnnBackendTest, GetSupportedConfigsFromScaledDotGemmFusion) { se::CudaComputeCapability cc = stream_executor_->GetDeviceDescription().cuda_compute_capability(); diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD index 94c732c4dc10e1..0877e5e0921bf2 100644 --- a/third_party/xla/xla/backends/gpu/codegen/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/BUILD @@ -94,16 +94,13 @@ xla_test( "//xla/service:pattern_matcher", "//xla/service/gpu:cudnn_support_utils", "//xla/service/gpu:ir_emission_utils", - "//xla/service/gpu:stream_executor_util", "//xla/stream_executor:device_description", - "//xla/stream_executor:dnn", "//xla/stream_executor:platform_manager", + "//xla/stream_executor:semantic_version", "//xla/stream_executor:stream_executor_h", "//xla/stream_executor/cuda:cuda_compute_capability", - "//xla/tests:hlo_pjrt_interpreter_reference_mixin", + "//xla/tests:hlo_interpreter_reference_mixin", "//xla/tsl/platform:env", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:status_macros", diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc index 216528dd11e2f9..8f47dd1e52192c 100644 --- a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc @@ -44,18 +44,15 @@ limitations under the License. #include "xla/service/dump.h" #include "xla/service/gpu/cudnn_support_utils.h" #include "xla/service/gpu/ir_emission_utils.h" -#include "xla/service/gpu/stream_executor_util.h" #include "xla/service/hlo_module_config.h" #include "xla/service/pattern_matcher.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" -#include "xla/stream_executor/dnn.h" #include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/semantic_version.h" #include "xla/stream_executor/stream_executor.h" -#include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" +#include "xla/tests/hlo_interpreter_reference_mixin.h" #include "xla/tsl/platform/env.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" #include "xla/xla.pb.h" #include "xla/xla_data.pb.h" @@ -1407,13 +1404,20 @@ TEST_F(CuDnnFusionRewriteTest, // With other backends disabled, compilation must fail. ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(R"( -e { +triton_gemm_dot { p0 = f64[20,40,64] parameter(0) p0n = f64[20,40,64] negate(p0) p1 = f64[20,80,64] parameter(1) - r = f64[20,40,80] dot(p0n, p1), + ROOT r = f64[20,40,80] dot(p0n, p1), lhs_batch_dims={0}, rhs_batch_dims={0}, lhs_contracting_dims={2}, rhs_contracting_dims={2} +} + +e { + p0 = f64[20,40,64] parameter(0) + p1 = f64[20,80,64] parameter(1) + ROOT fusion = f64[20,40,80] fusion(p0, p1), kind=kCustom, calls=triton_gemm_dot, + backend_config={"fusion_backend_config": {kind: "__triton_gemm"}} })")); auto status = CreateExecutable(std::move(module), /*run_hlo_passes=*/true).status(); From 9f72d450ded3fdc487b661d5434834daa86d1769 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 2 Sep 2026 06:09:40 -0700 Subject: [PATCH 22/28] Upgrade rules_cc to 0.2.20 and bazel_skylib to 1.9.0 in XLA - Bump rules_cc to 0.2.20 in MODULE.bazel and workspace3.bzl. - Bump bazel_skylib to 1.9.0 in workspace3.bzl to support the 'scope' attribute on bool_flag used by rules_cc 0.2.20. - Remove rules_cc_protobuf.patch from rules_cc in workspace3.bzl since rules_cc 0.2.20 no longer contains cc_proto_library in cc/defs.bzl. - Call compatibility_proxy_repo() early in WORKSPACE to define @cc_compatibility_proxy in WORKSPACE mode. - Revert compiler flag workaround from third_party/llvm/build.patch now that rules_cc 0.2.20 natively provides @rules_cc//cc/compiler:compiler. PiperOrigin-RevId: 975088002 --- third_party/xla/MODULE.bazel | 2 +- third_party/xla/WORKSPACE | 4 ++ third_party/xla/third_party/llvm/build.patch | 68 -------------------- third_party/xla/workspace1.bzl | 3 +- third_party/xla/workspace3.bzl | 15 ++--- 5 files changed, 14 insertions(+), 78 deletions(-) diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 64b3d29e143192..6de764de21d23a 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -23,7 +23,7 @@ bazel_dep(name = "pybind11_abseil", version = "202402.0") bazel_dep(name = "pybind11_bazel", version = "3.0.0") bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7") bazel_dep(name = "re2", version = "2025-11-05.bcr.1", repo_name = "com_googlesource_code_re2") -bazel_dep(name = "rules_cc", version = "0.2.18") +bazel_dep(name = "rules_cc", version = "0.2.20") bazel_dep(name = "rules_java", version = "8.16.1") bazel_dep(name = "rules_license", version = "1.0.0") bazel_dep(name = "rules_python", version = "2.2.0") diff --git a/third_party/xla/WORKSPACE b/third_party/xla/WORKSPACE index 37e4edee2a579f..14cf2eb79476c2 100644 --- a/third_party/xla/WORKSPACE +++ b/third_party/xla/WORKSPACE @@ -35,6 +35,10 @@ load("@bazel_features//:deps.bzl", "bazel_features_deps") bazel_features_deps() +load("@rules_cc//cc:extensions.bzl", "compatibility_proxy_repo") + +compatibility_proxy_repo() + # Initialize hermetic C++ load("@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl", "cc_toolchain_deps") diff --git a/third_party/xla/third_party/llvm/build.patch b/third_party/xla/third_party/llvm/build.patch index 641bdb95bc262c..032406b4e16483 100644 --- a/third_party/xla/third_party/llvm/build.patch +++ b/third_party/xla/third_party/llvm/build.patch @@ -15,74 +15,6 @@ diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -317,19 +317,19 @@ - config_setting( - name = "is_windows_clang_mingw", - constraint_values = ["@platforms//os:windows"], -- flag_values = {"@rules_cc//cc/compiler:compiler": "clang"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"}, - ) - - config_setting( - name = "is_windows_clang_cl", - constraint_values = ["@platforms//os:windows"], -- flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"}, - ) - - config_setting( - name = "is_windows_msvc", - constraint_values = ["@platforms//os:windows"], -- flag_values = {"@rules_cc//cc/compiler:compiler": "msvc-cl"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "msvc-cl"}, - ) - - config_setting( -@@ -338,7 +338,7 @@ - "@platforms//cpu:aarch64", - "@platforms//os:windows", - ], -- flag_values = {"@rules_cc//cc/compiler:compiler": "clang"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"}, - ) - - config_setting( -@@ -347,7 +347,7 @@ - "@platforms//cpu:aarch64", - "@platforms//os:windows", - ], -- flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"}, - ) - - config_setting( -@@ -356,7 +356,7 @@ - "@platforms//cpu:aarch64", - "@platforms//os:windows", - ], -- flag_values = {"@rules_cc//cc/compiler:compiler": "msvc-cl"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "msvc-cl"}, - ) - - config_setting( -@@ -365,7 +365,7 @@ - "@platforms//cpu:x86_64", - "@platforms//os:windows", - ], -- flag_values = {"@rules_cc//cc/compiler:compiler": "clang"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"}, - ) - - config_setting( -@@ -374,7 +374,7 @@ - "@platforms//cpu:x86_64", - "@platforms//os:windows", - ], -- flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"}, -+ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"}, - ) - - BLAKE3_x86_64_ASM_SOURCE_PATTERNS = [ @@ -430,7 +430,8 @@ "@platforms//cpu:aarch64": [ "lib/Support/BLAKE3/blake3_neon.c", diff --git a/third_party/xla/workspace1.bzl b/third_party/xla/workspace1.bzl index 54b16631be4f72..ffcaf9ddb768c7 100644 --- a/third_party/xla/workspace1.bzl +++ b/third_party/xla/workspace1.bzl @@ -27,7 +27,8 @@ def workspace(): llvm_setup(name = "llvm-project") native.register_toolchains("@local_config_python//:py_toolchain") rules_pkg_dependencies() - compatibility_proxy_repo() + if "cc_compatibility_proxy" not in native.existing_rules(): + compatibility_proxy_repo() tf_http_archive( name = "bazel_toolchains", diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl index 79a81c7a644602..a52a078a5f3840 100644 --- a/third_party/xla/workspace3.bzl +++ b/third_party/xla/workspace3.bzl @@ -32,9 +32,9 @@ def workspace(): # https://github.com/bazelbuild/bazel-skylib/releases tf_http_archive( name = "bazel_skylib", - sha256 = "bc283cdfcd526a52c3201279cda4bc298652efa898b10b4db0837dc51652756f", + sha256 = "3b5b49006181f5f8ff626ef8ddceaa95e9bb8ad294f7b5d7b11ea9f7ddaf8c59", urls = tf_mirror_urls( - "https://github.com/bazelbuild/bazel-skylib/releases/download/1.7.1/bazel-skylib-1.7.1.tar.gz", + "https://github.com/bazelbuild/bazel-skylib/releases/download/1.9.0/bazel-skylib-1.9.0.tar.gz", ), ) @@ -63,12 +63,11 @@ def workspace(): tf_http_archive( name = "rules_cc", - urls = tf_mirror_urls("https://github.com/bazelbuild/rules_cc/releases/download/0.2.0/rules_cc-0.2.0.tar.gz"), - strip_prefix = "rules_cc-0.2.0", - sha256 = "ae244f400218f4a12ee81658ff246c0be5cb02c5ca2de5519ed505a6795431e9", - patch_file = [ - "@xla//third_party/py:rules_cc_protobuf.patch", - ], + sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b", + strip_prefix = "rules_cc-0.2.20", + urls = tf_mirror_urls( + "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz", + ), ) # Toolchains for ML projects hermetic builds. From 5ff1f844eefda986b0e1f33fd7995a070b39c20e Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 09:44:25 +0300 Subject: [PATCH 23/28] Validate strides/rates in Dilation2D and widen effective-filter arithmetic `ParseAttributes` in dilation_ops.cc checked only that `strides` and `rates` have 4 entries and that entries 0 and 3 equal 1. It never checked that the spatial entries are positive, so zero and negative values reached the kernel. `ParseSizes` then computes, in `int`: const int filter_rows_eff = filter_rows + (filter_rows - 1) * (rate_rows - 1); A negative rate makes the effective filter size negative directly, and a large negative one overflows. `GetWindowedOutputSizeVerbose` rejects `stride <= 0` and a negative output size, but not a negative filter size: with dilation_rate == 1 the effective filter size is passed straight through, so `output_size = (input - negative + stride) / stride` is larger than the input and positive, and the guard passes. Observed on 2.22.0-dev (input 2x4x4x2, filter 2x2x2, strides 1): rates=[1,0,0,1] -> accepted, output 2x4x4x2 rates=[1,-1,-1,1] -> accepted, output 2x5x5x2 rates=[1,-100,-100,1]-> accepted, output 2x104x104x2 (larger than the input) rates=[1,INT32_MIN,INT32_MIN,1] -> aborts the process: F tensor_shape.cc:204] Check failed: InitDims(dim_sizes) is OK (INVALID_ARGUMENT: Encountered overflow when multiplying ...) @ tensorflow::TensorShapeBase<>::TensorShapeBase() @ tensorflow::DilationOp<>::Compute() Because the CHECK is inside TensorShape's constructor, callers cannot catch it. This mirrors what 210e2943a0f7 ("Fix integer overflow DoS in ExtractImagePatches and ExtractVolumePatches") did for the sibling ops: that change routed attribute parsing through `ParseAttributeVec4`, which enforces `(*attr)[1] >= 1 && (*attr)[2] >= 1`, and widened the shape arithmetic to int64_t. Dilation2D was not covered. With the same inputs above, ExtractImagePatches already returns OutOfRangeError for every non-positive rate. `ParseAttributes` is shared by Dilation2D, Dilation2DBackpropInput and Dilation2DBackpropFilter, so all three are covered. --- tensorflow/core/kernels/dilation_ops.cc | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/kernels/dilation_ops.cc b/tensorflow/core/kernels/dilation_ops.cc index 8919bead32fca5..ef466fde72a13e 100644 --- a/tensorflow/core/kernels/dilation_ops.cc +++ b/tensorflow/core/kernels/dilation_ops.cc @@ -53,6 +53,9 @@ void ParseAttributes(OpKernelConstruction* context, OP_REQUIRES(context, (*strides)[0] == 1 && (*strides)[3] == 1, absl::UnimplementedError( "Stride is only supported across spatial dimensions.")); + OP_REQUIRES( + context, (*strides)[1] >= 1 && (*strides)[2] >= 1, + absl::OutOfRangeError("Strides in the spatial dimensions must be >= 1.")); OP_REQUIRES_OK(context, context->GetAttr("rates", rates)); OP_REQUIRES(context, rates->size() == 4, @@ -61,6 +64,9 @@ void ParseAttributes(OpKernelConstruction* context, OP_REQUIRES(context, (*rates)[0] == 1 && (*rates)[3] == 1, absl::UnimplementedError( "Rate is only supported across spatial dimensions.")); + OP_REQUIRES( + context, (*rates)[1] >= 1 && (*rates)[2] >= 1, + absl::OutOfRangeError("Rates in the spatial dimensions must be >= 1.")); OP_REQUIRES_OK(context, context->GetAttr("padding", padding)); } @@ -103,10 +109,12 @@ void ParseSizes(OpKernelContext* context, const std::vector& strides, // Effective filter size, after introducing rate - 1 zeros between each // non-zero filter element. - const int filter_rows_eff = - filter_rows + (filter_rows - 1) * (*rate_rows - 1); - const int filter_cols_eff = - filter_cols + (filter_cols - 1) * (*rate_cols - 1); + const int64_t filter_rows_eff = + static_cast(filter_rows) + + static_cast(filter_rows - 1) * (*rate_rows - 1); + const int64_t filter_cols_eff = + static_cast(filter_cols) + + static_cast(filter_cols - 1) * (*rate_cols - 1); OP_REQUIRES_OK(context, GetWindowedOutputSize( input_rows, filter_rows_eff, /*dilation_rate=*/1, From dcbb856dab241c8f0f1b8547d73e3a1bcbf36809 Mon Sep 17 00:00:00 2001 From: endorphin13 Date: Wed, 26 Aug 2026 10:06:47 +0300 Subject: [PATCH 24/28] Use InvalidArgumentError for non-positive strides/rates Matches the error type used by the other attribute validation in this file. --- tensorflow/core/kernels/dilation_ops.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/dilation_ops.cc b/tensorflow/core/kernels/dilation_ops.cc index ef466fde72a13e..99a90efa62c024 100644 --- a/tensorflow/core/kernels/dilation_ops.cc +++ b/tensorflow/core/kernels/dilation_ops.cc @@ -55,7 +55,8 @@ void ParseAttributes(OpKernelConstruction* context, "Stride is only supported across spatial dimensions.")); OP_REQUIRES( context, (*strides)[1] >= 1 && (*strides)[2] >= 1, - absl::OutOfRangeError("Strides in the spatial dimensions must be >= 1.")); + absl::InvalidArgumentError( + "Strides in the spatial dimensions must be >= 1.")); OP_REQUIRES_OK(context, context->GetAttr("rates", rates)); OP_REQUIRES(context, rates->size() == 4, @@ -66,7 +67,8 @@ void ParseAttributes(OpKernelConstruction* context, "Rate is only supported across spatial dimensions.")); OP_REQUIRES( context, (*rates)[1] >= 1 && (*rates)[2] >= 1, - absl::OutOfRangeError("Rates in the spatial dimensions must be >= 1.")); + absl::InvalidArgumentError( + "Rates in the spatial dimensions must be >= 1.")); OP_REQUIRES_OK(context, context->GetAttr("padding", padding)); } From 84ac1e068f19940f737d41d87d93066b468428b2 Mon Sep 17 00:00:00 2001 From: DEEVEN SERU Date: Wed, 2 Sep 2026 20:45:53 +0530 Subject: [PATCH 25/28] Fix unguarded vector resize in IteratorRandomAccessCache::Get This adds bounds checking to cache resizes and updates the cache truncation warning. --- tensorflow/core/kernels/data/BUILD | 1 + .../core/kernels/data/cache_dataset_ops.cc | 33 +++++++++---- .../kernels/data/cache_dataset_ops_test.cc | 47 +++++++++++++++++-- 3 files changed, 68 insertions(+), 13 deletions(-) diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD index 3c5b008c1a74fb..7d1973e434b200 100644 --- a/tensorflow/core/kernels/data/BUILD +++ b/tensorflow/core/kernels/data/BUILD @@ -95,6 +95,7 @@ tf_cc_test( deps = [ ":cache_dataset_ops", ":iterator_ops", + ":range_dataset_op", ":tensor_slice_dataset_op", "//tensorflow/core:framework", "//tensorflow/core:lib", diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc index 9bfcbe60a7ffb1..db799198e253e4 100644 --- a/tensorflow/core/kernels/data/cache_dataset_ops.cc +++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc @@ -75,9 +75,15 @@ constexpr char kCacheDataset[] = "CacheDataset"; constexpr char kIncompleteCacheErrorMessage[] = "The calling iterator did not fully read the dataset being cached. In " "order to avoid unexpected truncation of the dataset, the partially cached " - "contents of the dataset will be discarded. This can happen if you have " - "an input pipeline similar to `dataset.cache().take(k).repeat()`. You " - "should use `dataset.take(k).cache().repeat()` instead."; + "contents of the dataset will be discarded. This can happen if you have " + "an input pipeline similar to `dataset.cache().take(k).repeat()`, or if " + "downstream operations drop elements (e.g. `batch(drop_remainder=True)`). " + "You should use `dataset.take(k).cache().repeat()` instead, or ensure the " + "dataset size is a multiple of the batch size before caching. Another " + "common workaround is to place the `.cache()` operation after the " + "operation that drops elements (like `.batch(...)`), if caching the " + "transformed data is acceptable."; +constexpr size_t kMaxItems = 10000000; // 10 million } // namespace class DatasetRandomAccessCache { @@ -89,15 +95,15 @@ class DatasetRandomAccessCache { // out_tensors with the element at that index. absl::Status Get(OpKernelContext* ctx, int64_t index, std::vector* out_tensors) { + if (index < 0) { + return absl::InvalidArgumentError( + absl::StrCat("Expected index >= 0; Received index: ", index)); + } if (!iter_resource_) { TF_ASSIGN_OR_RETURN(iter_resource_, GetIteratorResourceFromDataset(ctx, input_)); TF_RETURN_IF_ERROR(iter_resource_->SetIteratorFromDataset(ctx, input_)); } - if (index < 0) { - return absl::InvalidArgumentError( - absl::StrCat("Expected index >= 0; Received index: ", index)); - } if (index >= static_cast(cache_.size())) { TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx)); } @@ -158,6 +164,12 @@ class IteratorRandomAccessCache { absl::StrCat("Element position must be non-negative; Received: ", element_position)); } + + if (static_cast(element_position) == std::numeric_limits::max() || + static_cast(element_position) >= cache_.max_size()) { + return absl::InvalidArgumentError( + absl::StrCat("Element position too large or invalid.")); + } if (element_position < static_cast(cache_.size()) && !cache_[element_position].empty()) { @@ -165,6 +177,12 @@ class IteratorRandomAccessCache { return absl::OkStatus(); } + if (element_position >= kMaxItems) { + return absl::InvalidArgumentError(absl::StrCat( + "Requested element_position ", element_position, + " exceeds the maximum allowed cache size of ", kMaxItems)); + } + TF_RETURN_IF_ERROR(input_->Get(ctx, element_position, out_tensors)); if (element_position >= static_cast(cache_.size())) { cache_.resize(element_position + 1); @@ -721,7 +739,6 @@ class CacheDatasetOp::FileDatasetBase : public DatasetBase { Env* const env_; const size_t num_tensors_; const size_t tensor_index_padding_size_; - static constexpr size_t kMaxItems = 10000000; // 10 million const size_t item_index_padding_size_; }; // FileDatasetBase diff --git a/tensorflow/core/kernels/data/cache_dataset_ops_test.cc b/tensorflow/core/kernels/data/cache_dataset_ops_test.cc index ba6dada3d704e6..799734d7a588c8 100644 --- a/tensorflow/core/kernels/data/cache_dataset_ops_test.cc +++ b/tensorflow/core/kernels/data/cache_dataset_ops_test.cc @@ -375,14 +375,51 @@ INSTANTIATE_TEST_CASE_P(CacheDatasetOpTest, ParameterizedIteratorSaveAndRestoreTest, ::testing::ValuesIn(IteratorSaveAndRestoreTestCases())); -TEST_F(CacheDatasetOpTest, NegativeIndexTest) { - auto params = CacheDatasetParams3(); +TEST_F(CacheDatasetOpTest, NegativeIndexEarlyRejection) { + auto range_dataset_params = RangeDatasetParams(0, 20000000, 1); + auto params = + CacheDatasetParams(range_dataset_params, + /*filename=*/"", + /*output_dtypes=*/{DT_INT64}, + /*output_shapes=*/{PartialTensorShape({})}, kNodeName); TF_ASSERT_OK(Initialize(params)); std::vector out_tensors; absl::Status status = - dataset_->Get(AnyContext(iterator_ctx_.get()), -1, &out_tensors); - EXPECT_TRUE(status.code() == absl::StatusCode::kOutOfRange); - EXPECT_EQ(status.message(), "Index out of range [0, 3):-1"); + dataset_->Get(AnyContext(iterator_ctx_.get()), -1LL, &out_tensors); + EXPECT_TRUE(status.code() == absl::StatusCode::kInvalidArgument || + status.code() == absl::StatusCode::kOutOfRange); +} + +TEST_F(CacheDatasetOpTest, LargeIndexTest) { + auto range_dataset_params = RangeDatasetParams(0, 20000000, 1); + auto params = + CacheDatasetParams(range_dataset_params, + /*filename=*/"", + /*output_dtypes=*/{DT_INT64}, + /*output_shapes=*/{PartialTensorShape({})}, kNodeName); + TF_ASSERT_OK(Initialize(params)); + std::vector out_tensors; + int64_t huge_index = std::numeric_limits::max(); + absl::Status status = + dataset_->Get(AnyContext(iterator_ctx_.get()), huge_index, &out_tensors); + EXPECT_TRUE(status.code() == absl::StatusCode::kInvalidArgument || + status.code() == absl::StatusCode::kOutOfRange); +} + +TEST_F(CacheDatasetOpTest, BadAllocCrashTest) { + auto range_dataset_params = RangeDatasetParams(0, 20000000, 1); + auto params = + CacheDatasetParams(range_dataset_params, + /*filename=*/"", + /*output_dtypes=*/{DT_INT64}, + /*output_shapes=*/{PartialTensorShape({})}, kNodeName); + TF_ASSERT_OK(Initialize(params)); + std::vector out_tensors; + absl::Status status = + dataset_->Get(AnyContext(iterator_ctx_.get()), 15000000, &out_tensors); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_TRUE(absl::StrContains(status.message(), + "exceeds the maximum allowed cache size")); } } // namespace From 36a0e001b825227b3113163384fb8d4b31a802c2 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Wed, 2 Sep 2026 09:06:48 -0700 Subject: [PATCH 26/28] Reverts a30ea91b11ca72297b9f4f00b2a710e84e55cbef PiperOrigin-RevId: 975162224 --- .../xla/third_party/stablehlo/temporary.patch | 571 ------------------ 1 file changed, 571 deletions(-) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index e1b5fa813d250c..ebb9dda1528d7c 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -53,295 +53,6 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`. If `x` is a value or placeholder, this function is a shortcut for -diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp ---- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp -+++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp -@@ -20,6 +20,7 @@ - #include - - #include "llvm/ADT/DenseSet.h" -+#include "llvm/ADT/STLExtras.h" - #include "llvm/ADT/SmallVector.h" - #include "llvm/ADT/StringRef.h" - #include "mlir/IR/Attributes.h" -@@ -31,57 +32,53 @@ - namespace mlir { - namespace stablehlo { - --static SmallVector> -+namespace { -+ -+struct ReindexedAxes { -+ SmallVector splitAxisSizes; -+ SmallVector groupedAxisIndices; -+}; -+ -+// Generates replica groups from the reshaped mesh axis sizes and the indices of -+// the communication axes using a Reshape-Transpose permutation. -+SmallVector> - flattenedReplicaGroupsFromTransposePermutation( -- const SmallVector& meshAxisNames, -- const SmallVector& commAxisNames, -- const llvm::DenseSet& commAxisSet, -- const SmallVector& axisSizes, -- const SmallVector& deviceIds, int64_t totalDevices) { -- // Reshape and Transpose equivalence bridging XLA TileAssignment behavior. -+ ArrayRef axisSizes, ArrayRef groupedAxisIndices, -+ ArrayRef deviceIds, int64_t totalDevices) { -+ llvm::DenseSet groupedAxisSet(groupedAxisIndices.begin(), -+ groupedAxisIndices.end()); - SmallVector transposeAxes; - // Non-grouped axes first -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (!commAxisSet.count(meshAxisNames[i])) { -+ for (size_t i = 0; i < axisSizes.size(); ++i) { -+ if (!groupedAxisSet.count(i)) { - transposeAxes.push_back(i); - } - } -- // Grouped axes -- for (const auto& name : commAxisNames) { -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (meshAxisNames[i] == name) { -- transposeAxes.push_back(i); -- break; -- } -- } -- } -- -- SmallVector transposedSizes(meshAxisNames.size()); -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -+ // Grouped axes in the specified order -+ for (int64_t idx : groupedAxisIndices) { -+ transposeAxes.push_back(idx); -+ } -+ -+ SmallVector transposedSizes(axisSizes.size()); -+ for (size_t i = 0; i < axisSizes.size(); ++i) { - transposedSizes[i] = axisSizes[transposeAxes[i]]; - } - -- // Compute strides for original shape -- SmallVector originalStrides(meshAxisNames.size(), 1); -- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { -+ // Compute strides for reshaped shape -+ SmallVector originalStrides(axisSizes.size(), 1); -+ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { - originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1]; - } - - // Compute strides for transposed shape -- SmallVector transposedStrides(meshAxisNames.size(), 1); -- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { -+ SmallVector transposedStrides(axisSizes.size(), 1); -+ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { - transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1]; - } - -- // Generate chunks - int64_t numDevicesPerGroup = 1; -- for (auto name : commAxisNames) { -- for (size_t i = 0; i < meshAxisNames.size(); ++i) { -- if (meshAxisNames[i] == name) { -- numDevicesPerGroup *= axisSizes[i]; -- break; -- } -- } -+ for (int64_t idx : groupedAxisIndices) { -+ numDevicesPerGroup *= axisSizes[idx]; - } - int64_t numGroups = totalDevices / numDevicesPerGroup; - -@@ -93,7 +90,7 @@ - for (int64_t j = 0; j < numDevicesPerGroup; ++j) { - int64_t linearTransposeIdx = i * numDevicesPerGroup + j; - int64_t originalIndex = 0; -- for (size_t k = 0; k < meshAxisNames.size(); ++k) { -+ for (size_t k = 0; k < axisSizes.size(); ++k) { - int64_t coord = - (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; - originalIndex += coord * originalStrides[transposeAxes[k]]; -@@ -102,9 +99,128 @@ - } - groups.push_back(std::move(group)); - } -- - return groups; - } -+ -+// Splits mesh axes based on sub-axis references and computes the corresponding -+// indices for the communication axes. -+FailureOr computeReindexedAxes(ArrayRef axesInMesh, -+ ArrayAttr commAxes, -+ Location loc) { -+ ReindexedAxes result; -+ -+ // Validate commAxes and verify that all mesh axes exist and have valid sizes. -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::dyn_cast(attr); -+ if (!shloAxisRef) { -+ return emitError(loc) << "expected AxisRefAttr in comm_axes"; -+ } -+ StringRef axisName = shloAxisRef.getName(); -+ bool found = false; -+ for (auto meshAxis : axesInMesh) { -+ if (meshAxis.getName() == axisName) { -+ found = true; -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ int64_t preSize = subAxisInfo.getPreSize(); -+ int64_t size = subAxisInfo.getSize(); -+ if (preSize < 1 || size < 1) { -+ return emitError(loc) -+ << "sub-axis pre_size and size must be at least 1"; -+ } -+ int64_t nextPreSize = preSize * size; -+ if (nextPreSize > meshAxis.getSize() || -+ meshAxis.getSize() % nextPreSize != 0) { -+ return emitError(loc) -+ << "sub-axis (pre_size * size) must divide mesh axis size"; -+ } -+ } -+ break; -+ } -+ } -+ if (!found) { -+ return emitError(loc) -+ << "axis '" << axisName << "' not found in mesh definition"; -+ } -+ } -+ -+ // Split each mesh axis according to the referenced subaxes. -+ struct SplitDim { -+ StringRef axisName; -+ int64_t preSize; -+ int64_t size; -+ int64_t dimIndex; -+ }; -+ SmallVector splitDims; -+ -+ for (auto meshAxis : axesInMesh) { -+ StringRef axisName = meshAxis.getName(); -+ int64_t axisSize = meshAxis.getSize(); -+ -+ SmallVector preSizes = {1, axisSize}; -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::cast(attr); -+ if (shloAxisRef.getName() == axisName) { -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ preSizes.push_back(subAxisInfo.getPreSize()); -+ preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize()); -+ } -+ } -+ } -+ -+ llvm::sort(preSizes); -+ preSizes.erase(llvm::unique(preSizes), preSizes.end()); -+ -+ for (size_t j = 0; j < preSizes.size() - 1; ++j) { -+ int64_t segPreSize = preSizes[j]; -+ int64_t segSize = preSizes[j + 1] / segPreSize; -+ int64_t dimIdx = result.splitAxisSizes.size(); -+ result.splitAxisSizes.push_back(segSize); -+ splitDims.push_back({axisName, segPreSize, segSize, dimIdx}); -+ } -+ } -+ -+ // Map each communication axis to its corresponding split dimension. -+ llvm::DenseSet groupedSet; -+ for (auto attr : commAxes) { -+ auto shloAxisRef = llvm::cast(attr); -+ StringRef axisName = shloAxisRef.getName(); -+ int64_t reqPreSize = 1; -+ int64_t reqSize = 0; -+ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { -+ reqPreSize = subAxisInfo.getPreSize(); -+ reqSize = subAxisInfo.getSize(); -+ } else { -+ for (auto meshAxis : axesInMesh) { -+ if (meshAxis.getName() == axisName) { -+ reqSize = meshAxis.getSize(); -+ break; -+ } -+ } -+ } -+ -+ bool matched = false; -+ for (const auto& splitDim : splitDims) { -+ if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize && -+ splitDim.size == reqSize) { -+ if (!groupedSet.insert(splitDim.dimIndex).second) { -+ return emitError(loc) -+ << "Duplicate or overlapping communication axis: " << axisName; -+ } -+ result.groupedAxisIndices.push_back(splitDim.dimIndex); -+ matched = true; -+ break; -+ } -+ } -+ if (!matched) { -+ return emitError(loc) << "Invalid or overlapping communication axis on '" -+ << axisName << "'"; -+ } -+ } -+ -+ return result; -+} -+ -+} // namespace - - FailureOr>> flattenReplicaGroupMeshAxes( - Attribute meshAttr, ArrayAttr commAxes, Location loc) { -@@ -120,34 +236,13 @@ - if (!mesh) - return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); - -- auto axesInMesh = mesh.getAxes(); -- -- // Identify which axes are communication axes. -- llvm::SmallVector commAxisNames; -- llvm::DenseSet commAxisSet; -- for (auto attr : commAxes) { -- auto shloAxisRef = llvm::dyn_cast(attr); -- if (!shloAxisRef) { -- return emitError(loc) << "expected AxisRefAttr in comm_axes"; -- } -- if (shloAxisRef.getSubAxisInfo()) { -- return emitError(loc) << "Subaxes are not supported in " -- "flattenReplicaGroupMeshAxes"; -- } -- commAxisNames.push_back(shloAxisRef.getName()); -- commAxisSet.insert(shloAxisRef.getName()); -- } -- -- // Calculate total devices and axis sizes -+ FailureOr reindexedAxes = -+ computeReindexedAxes(mesh.getAxes(), commAxes, loc); -+ if (failed(reindexedAxes)) return failure(); - - int64_t totalDevices = 1; -- SmallVector axisSizes; -- SmallVector meshAxisNames; -- for (auto meshAxis : axesInMesh) { -- auto typedMeshAxis = llvm::cast(meshAxis); -- axisSizes.push_back(typedMeshAxis.getSize()); -- meshAxisNames.push_back(typedMeshAxis.getName()); -- totalDevices *= typedMeshAxis.getSize(); -+ for (auto meshAxis : mesh.getAxes()) { -+ totalDevices *= llvm::cast(meshAxis).getSize(); - } - - SmallVector deviceIds; -@@ -160,8 +255,8 @@ - } - - return flattenedReplicaGroupsFromTransposePermutation( -- meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds, -- totalDevices); -+ reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices, -+ deviceIds, totalDevices); - } - - } // namespace stablehlo diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp --- stablehlo/stablehlo/dialect/Serialization.cpp +++ stablehlo/stablehlo/dialect/Serialization.cpp @@ -776,196 +487,6 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st // CHECK-LABEL: func.func @ragged_dot_mode_3( // CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x5xf32>, // CHECK-SAME: %[[ARG1:.*]]: tensor<2x5x7xf32>, -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir ---- stablehlo/stablehlo/tests/interpret/all_gather.mlir -+++ stablehlo/stablehlo/tests/interpret/all_gather.mlir -@@ -133,3 +133,41 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> { -+ %result = "stablehlo.all_gather"(%operand) { -+ all_gather_dim = 0 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<1xi64>) -> tensor<2xi64> -+ return %result : tensor<2xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[0]> : tensor<1xi64> -+ %p1 = stablehlo.constant dense<[1]> : tensor<1xi64> -+ %p2 = stablehlo.constant dense<[2]> : tensor<1xi64> -+ %p3 = stablehlo.constant dense<[3]> : tensor<1xi64> -+ %p4 = stablehlo.constant dense<[4]> : tensor<1xi64> -+ %p5 = stablehlo.constant dense<[5]> : tensor<1xi64> -+ %p6 = stablehlo.constant dense<[6]> : tensor<1xi64> -+ %p7 = stablehlo.constant dense<[7]> : tensor<1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather], -+ [@all_gather], [@all_gather], [@all_gather], [@all_gather]] -+ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> -+ (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, -+ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -+ check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64> -+ check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64> -+ check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64> -+ check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64> -+ check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64> -+ check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64> -+ check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64> -+ check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir ---- stablehlo/stablehlo/tests/interpret/all_reduce.mlir -+++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir -@@ -135,3 +135,45 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> { -+ %result = "stablehlo.all_reduce"(%operand) ({ -+ ^bb0(%arg0: tensor, %arg1: tensor): -+ %0 = stablehlo.add %arg0, %arg1 : tensor -+ stablehlo.return %0 : tensor -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]>, -+ channel_handle = #stablehlo.channel_handle -+ } : (tensor<1xi64>) -> tensor<1xi64> -+ return %result : tensor<1xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[10]> : tensor<1xi64> -+ %p1 = stablehlo.constant dense<[20]> : tensor<1xi64> -+ %p2 = stablehlo.constant dense<[30]> : tensor<1xi64> -+ %p3 = stablehlo.constant dense<[40]> : tensor<1xi64> -+ %p4 = stablehlo.constant dense<[50]> : tensor<1xi64> -+ %p5 = stablehlo.constant dense<[60]> : tensor<1xi64> -+ %p6 = stablehlo.constant dense<[70]> : tensor<1xi64> -+ %p7 = stablehlo.constant dense<[80]> : tensor<1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce], -+ [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]] -+ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> -+ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -+ check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64> -+ check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64> -+ check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64> -+ check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64> -+ check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64> -+ check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64> -+ check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64> -+ check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir ---- stablehlo/stablehlo/tests/interpret/all_to_all.mlir -+++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir -@@ -172,3 +172,43 @@ - func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32> - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> { -+ %result = "stablehlo.all_to_all"(%operand) { -+ split_dimension = 0 : i64, -+ concat_dimension = 1 : i64, -+ split_count = 2 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<2x1xi64>) -> tensor<1x2xi64> -+ return %result : tensor<1x2xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64> -+ %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64> -+ %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64> -+ %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64> -+ %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64> -+ %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64> -+ %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64> -+ %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all], -+ [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]] -+ } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, -+ tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) -> -+ (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, -+ tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>) -+ check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64> -+ check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64> -+ func.return -+ } -+} -diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir ---- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir -+++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir -@@ -90,3 +90,45 @@ - func.return - } - } -+ -+// ----- -+ -+module @mesh_axes_subaxis { -+ func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> { -+ %result = "stablehlo.reduce_scatter"(%operand) ({ -+ ^bb0(%arg0: tensor, %arg1: tensor): -+ %0 = stablehlo.add %arg0, %arg1 : tensor -+ stablehlo.return %0 : tensor -+ }) { -+ scatter_dimension = 0 : i64, -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<2xi64>) -> tensor<1xi64> -+ return %result : tensor<1xi64> -+ } -+ func.func @main() { -+ %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64> -+ %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64> -+ %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64> -+ %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64> -+ %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64> -+ %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64> -+ %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64> -+ %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64> -+ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { -+ programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], -+ [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]] -+ } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, -+ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> -+ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, -+ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -+ check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64> -+ check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64> -+ check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64> -+ check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64> -+ check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64> -+ check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64> -+ check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64> -+ check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64> -+ func.return -+ } -+} diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stablehlo/tests/ops_broadcasting.mlir --- stablehlo/stablehlo/tests/ops_broadcasting.mlir +++ stablehlo/stablehlo/tests/ops_broadcasting.mlir @@ -984,98 +505,6 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stableh + return %0 : tensor<3x4x5xf64> +} + -diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir ---- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir -+++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir -@@ -5,7 +5,7 @@ - - // CHECK-LABEL: @all_reduce_rgv3 - func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -19,7 +19,7 @@ - - // CHECK-LABEL: @all_gather_rgv3 - func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_gather"(%arg0) { - all_gather_dim = 0 : i64, - replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -@@ -30,7 +30,7 @@ - - // CHECK-LABEL: @all_to_all_rgv3 - func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_to_all"(%arg0) { - concat_dimension = 0 : i64, - split_dimension = 0 : i64, -@@ -44,7 +44,7 @@ - - // CHECK-LABEL: @all_reduce_sdy_mesh - func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -59,7 +59,7 @@ - - // CHECK-LABEL: @all_reduce_sdy_mesh_dev - func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> { -- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> - %0 = "stablehlo.all_reduce"(%arg0) ({ - ^bb0(%arg1: tensor, %arg2: tensor): - %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -@@ -84,4 +84,43 @@ - } : (tensor<4xf32>) -> tensor<4xf32> - return %0 : tensor<4xf32> - } -+ -+ // CHECK-LABEL: @all_reduce_subaxis -+ func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_subaxis_order_1 -+ func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } -+ -+ // CHECK-LABEL: @all_reduce_subaxis_order_2 -+ func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> { -+ // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64> -+ %0 = "stablehlo.all_reduce"(%arg0) ({ -+ ^bb0(%arg1: tensor, %arg2: tensor): -+ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor -+ "stablehlo.return"(%1) : (tensor) -> () -+ }) { -+ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> -+ } : (tensor<4xf32>) -> tensor<4xf32> -+ return %0 : tensor<4xf32> -+ } - } diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir --- stablehlo/stablehlo/tests/verify_convolution.mlir +++ stablehlo/stablehlo/tests/verify_convolution.mlir From 0655d0a0197391422cde1c98c1849097d8346603 Mon Sep 17 00:00:00 2001 From: Penporn Koanantakool Date: Wed, 2 Sep 2026 09:52:30 -0700 Subject: [PATCH 27/28] [XLA:Build] Fix toolchain resolution for AArch64. Expose the `llvm_linux_aarch64` repository from `@rules_ml_toolchain` in `MODULE.bazel`. XLA registers Linux AArch64 toolchains: - `@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64` - `@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64_cuda` These toolchains depend on `@llvm_linux_aarch64`. Under Bzlmod, repositories generated by module extensions must be explicitly imported via `use_repo`. Otherwise, the toolchain resolution will fail. PiperOrigin-RevId: 975182846 --- third_party/xla/MODULE.bazel | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 6de764de21d23a..ae3b0122f1d23f 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -281,7 +281,7 @@ nvshmem_redist = use_extension("@rules_ml_toolchain//extensions:nvshmem_redist.b use_repo(nvshmem_redist, "nvidia_nvshmem") toolchain_ext = use_extension("@rules_ml_toolchain//extensions:toolchain.bzl", "toolchain_ext") -use_repo(toolchain_ext, "llvm18_linux_x86_64", "llvm_linux_x86_64") +use_repo(toolchain_ext, "llvm18_linux_x86_64", "llvm_linux_aarch64", "llvm_linux_x86_64") register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64") From f05b1abd9e540930bbaa75bf6fa6c778997ed4a2 Mon Sep 17 00:00:00 2001 From: Eric Yang Date: Wed, 2 Sep 2026 09:53:28 -0700 Subject: [PATCH 28/28] Internal changes only PiperOrigin-RevId: 975183363 --- tensorflow/core/profiler/utils/BUILD | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD index 621adf043ee11e..bcc08265beddf0 100644 --- a/tensorflow/core/profiler/utils/BUILD +++ b/tensorflow/core/profiler/utils/BUILD @@ -261,7 +261,6 @@ cc_library( hdrs = ["hlo_module_utils.h"], visibility = internal_visibility([ "//third_party/odml/model_explorer/backend/adapters/hlo:__pkg__", - "//tensorflow/compiler/mlir/lite/experimental/google/tooling/hlo_adapter:__pkg__", ]), deps = [ "@org_xprof//xprof/utils:hlo_module_utils",