Merge pull request #129414 from Nishuuzz:hunt-2026-10-21

PiperOrigin-RevId: 1005364077

tensorflow/python/kernel_tests/control_flow/functional_ops_test.py

@@ -73,6 +73,26 @@ class FunctionalOpsTest(test.TestCase):
         initializer=10)
     self.assertAllEqual(880, self.evaluate(r))
 
+  @test_util.run_in_graph_and_eager_modes
+  def testScalarElemsRejected(self):
+    # These take the number of iterations from elems.shape[0], which gives
+    # IndexError out of tensor_shape for a scalar. map_fn already reports it,
+    # so they should say the same thing.
+    elems = constant_op.constant(5, name="data")
+    vec = constant_op.constant([1, 2], name="vec")
+    for op in (functional_ops.foldl, functional_ops.foldr, functional_ops.scan):
+      with self.assertRaisesRegex(ValueError, "must be 1\\+ dimensional"):
+        op(lambda a, x: a + x, elems)
+      # a scalar anywhere in elems is a problem, not just the first one
+      with self.assertRaisesRegex(ValueError, "must be 1\\+ dimensional"):
+        op(lambda a, x: a + x[0] + x[1], (vec, elems), initializer=0)
+      with self.assertRaisesRegex(
+          ValueError, "elems must contain at least one Tensor"
+      ):
+        op(lambda a, x: a + x, [])
+    with self.assertRaisesRegex(ValueError, "must be 1\\+ dimensional"):
+      functional_ops.foldl(lambda a, x: a + x, elems, initializer=0)
+
   @test_util.run_in_graph_and_eager_modes
   def testFoldl_SingleInputMultiOutput(self):
     elems = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])

tensorflow/python/ops/functional_ops.py

@@ -38,6 +38,32 @@ from tensorflow.python.util import nest
 from tensorflow.python.util.tf_export import tf_export
 
 
+def _check_elems_not_scalar(elems_flat):
+  """Rejects elements whose known rank is below 1.
+
+  foldl, foldr and scan all take the number of iterations from
+  `elems_flat[0].shape[0]`, which raises IndexError out of tensor_shape for a
+  scalar instead of saying what is wrong. They then unstack every element, so
+  a scalar anywhere in `elems` is a problem, not just the first one. map_fn
+  reports this for its first element, so use the same message.
+
+  Args:
+    elems_flat: the flattened, converted `elems`.
+
+  Raises:
+    ValueError: if `elems_flat` is empty or any element has a statically known
+      rank below 1.
+  """
+  if not elems_flat:
+    raise ValueError("elems must contain at least one Tensor.")
+  for elem in elems_flat:
+    shape = elem.shape
+    if shape.ndims is not None and shape.ndims < 1:
+      raise ValueError(
+          "Elements in elems must be 1+ dimensional Tensors, not scalars"
+      )
+
+
 # TODO(yuanbyu, mrry): Handle stride to support sliding windows.
 @tf_export(v1=["foldl"])
 @dispatch.add_dispatch_support
@@ -124,6 +150,7 @@ def foldl(fn,
     elems_flat = [
         ops.convert_to_tensor(elem, name="elem") for elem in nest.flatten(elems)
     ]
+    _check_elems_not_scalar(elems_flat)
     n = (
         tensor_shape.dimension_value(elems_flat[0].shape[0]) or
         array_ops.shape(elems_flat[0])[0])
@@ -320,6 +347,7 @@ def foldr(fn,
     elems_flat = [
         ops.convert_to_tensor(elem, name="elem") for elem in nest.flatten(elems)
     ]
+    _check_elems_not_scalar(elems_flat)
     n = (
         tensor_shape.dimension_value(elems_flat[0].shape[0]) or
         array_ops.shape(elems_flat[0])[0])
@@ -579,6 +607,7 @@ def scan(fn,
     elems_flat = [
         ops.convert_to_tensor(elem, name="elem") for elem in elems_flat
     ]
+    _check_elems_not_scalar(elems_flat)
 
     # Convert elems to tensor array. n may be known statically.
     n = tensor_shape.dimension_value(elems_flat[0].shape[0])