Heap OOB in nested `tf.map_fn` with `RaggedTensor`s in TensorFlow
Description
TensorFlow is an end-to-end open source platform for machine learning. In affected versions it is possible to nest a tf.map_fn within another tf.map_fn call. However, if the input tensor is a RaggedTensor and there is no function signature provided, code assumes the output is a fully specified tensor and fills output buffer with uninitialized contents from the heap. The t and z outputs should be identical, however this is not the case. The last row of t contains data from the heap which can be used to leak other memory information. The bug lies in the conversion from a Variant tensor to a RaggedTensor. The implementation does not check that all inner shapes match and this results in the additional dimensions. The same implementation can result in data loss, if input tensor is tweaked. We have patched the issue in GitHub commit 4e2565483d0ffcadc719bd44893fb7f609bb5f12. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.
Affected packages
Versions sourced from the GitHub Security Advisory.
| Package | Affected versions | Patched versions |
|---|---|---|
tensorflowPyPI | < 2.3.4 | 2.3.4 |
tensorflowPyPI | >= 2.4.0, < 2.4.3 | 2.4.3 |
tensorflowPyPI | >= 2.5.0, < 2.5.1 | 2.5.1 |
tensorflow-cpuPyPI | < 2.3.4 | 2.3.4 |
tensorflow-cpuPyPI | >= 2.4.0, < 2.4.3 | 2.4.3 |
tensorflow-cpuPyPI | >= 2.5.0, < 2.5.1 | 2.5.1 |
tensorflow-gpuPyPI | < 2.3.4 | 2.3.4 |
tensorflow-gpuPyPI | >= 2.4.0, < 2.4.3 | 2.4.3 |
tensorflow-gpuPyPI | >= 2.5.0, < 2.5.1 | 2.5.1 |
Affected products
1- Range: >= 2.5.0, < 2.5.1
Patches
14e2565483d0fFix bug that could cause map_fn to produce incorrect results (rather than an error)
2 files changed · +39 −0
tensorflow/core/kernels/ragged_tensor_from_variant_op.cc+16 −0 modified@@ -174,7 +174,23 @@ Status NestedStackRaggedTensors( auto output_values_flat = output_ragged->mutable_values()->flat_outer_dims<VALUE_TYPE, 2>(); int values_index = 0; + + TensorShape expected_value_shape = component_values_shape; + expected_value_shape.RemoveDim(0); + for (int i = 0; i < ragged_components.size(); i++) { + // Check that the flat_values tensor shape is compatible. + TensorShape value_shape = ragged_components[i].values().shape(); + value_shape.RemoveDim(0); + if (value_shape != expected_value_shape) { + return errors::InvalidArgument( + "All flat_values must have compatible shapes. Shape at index 0: ", + expected_value_shape, ". Shape at index ", i, ": ", value_shape, + ". If you are using tf.map_fn, then you may need to specify an " + "explicit fn_output_signature with appropriate ragged_rank, and/or " + "convert output tensors to RaggedTensors."); + } + auto component_values_flat = ragged_components[i].values().flat_outer_dims<VALUE_TYPE, 2>(); int num_inner_elements = ragged_components[i].values().NumElements();
tensorflow/python/ops/ragged/ragged_map_fn_op_test.py+23 −0 modified@@ -21,9 +21,11 @@ import numpy as np from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops +from tensorflow.python.ops import map_fn as map_fn_lib from tensorflow.python.ops import math_ops as mo from tensorflow.python.ops import string_ops from tensorflow.python.ops.ragged import ragged_factory_ops @@ -309,6 +311,27 @@ def testMapOnSparseTensor(self): ) self.assertAllEqual(id_t2, [[0, 5], [0, 4]]) + def testRaggedMapWithIncorrectFnOutputSignature(self): + x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]]) + with self.assertRaisesRegex(errors.InvalidArgumentError, + 'All flat_values must have compatible shapes'): + y = map_fn_lib.map_fn(lambda r: map_fn_lib.map_fn(lambda y: r, r), x) + self.evaluate(y) + + def testNestedRaggedMapWithFnOutputSignature(self): + ragged1d = ragged_tensor.RaggedTensorSpec([None], dtypes.int32) + ragged2d = ragged_tensor.RaggedTensorSpec([None, None], dtypes.int32) + + x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]]) + # pylint: disable=g-long-lambda + y = map_fn_lib.map_fn( + lambda r: map_fn_lib.map_fn( + lambda y: r, r, fn_output_signature=ragged1d), + x, + fn_output_signature=ragged2d) + expected = [[[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], [[1]]] + self.assertAllEqual(y, expected) + if __name__ == '__main__': googletest.main()
Vulnerability mechanics
Generated by null/stub on May 9, 2026. Inputs: CWE entries + fix-commit diffs from this CVE's patches. Citations validated against bundle.
References
7- github.com/advisories/GHSA-g8wg-cjwc-xhhpghsaADVISORY
- nvd.nist.gov/vuln/detail/CVE-2021-37679ghsaADVISORY
- github.com/pypa/advisory-database/tree/main/vulns/tensorflow-cpu/PYSEC-2021-592.yamlghsaWEB
- github.com/pypa/advisory-database/tree/main/vulns/tensorflow-gpu/PYSEC-2021-790.yamlghsaWEB
- github.com/pypa/advisory-database/tree/main/vulns/tensorflow/PYSEC-2021-301.yamlghsaWEB
- github.com/tensorflow/tensorflow/commit/4e2565483d0ffcadc719bd44893fb7f609bb5f12ghsax_refsource_MISCWEB
- github.com/tensorflow/tensorflow/security/advisories/GHSA-g8wg-cjwc-xhhpghsax_refsource_CONFIRMWEB
News mentions
0No linked articles in our index yet.