Add VectorizedMap op class. (#21516)

Also: - fix `Map.compute_output_spec` so that it handles nested inputs - test `map` op with nested inputs - added missing `Deg2Rad.compute_output_spec` - added test verifying that all ops implement `compute_output_spec`.

Author: hertschuh

keras/src/ops/core.py

@@ -16,13 +16,16 @@ def call(self, f, xs):
  return backend.core.map(f, xs)
 
  def compute_output_spec(self, f, xs):
- x = xs[0]
- n = xs.shape[0]
+ x = tree.map_structure(lambda t: t[0], xs)
+ n = tree.flatten(xs)[0].shape[0]
  y = backend.compute_output_spec(f, x)
 
- def append_batch_axis(x):
+ def append_batch_axis(t):
  return KerasTensor(
- shape=(n,) + x.shape, dtype=x.dtype, sparse=x.sparse
+ shape=(n,) + t.shape,
+ dtype=t.dtype,
+ sparse=t.sparse,
+ ragged=t.ragged,
  )
 
  y = tree.map_structure(append_batch_axis, y)
@@ -1078,7 +1081,31 @@ def cond(pred, true_fn, false_fn):
  return Cond()(pred, true_fn, false_fn)
 
 
-# TODO: also create an Op subclass VectorizedMap.
+class VectorizedMap(Operation):
+ def __init__(self, function, *, name=None):
+ super().__init__(name=name)
+ self.function = function
+
+ def call(self, elements):
+ return backend.core.vectorized_map(self.function, elements)
+
+ def compute_output_spec(self, elements):
+ x = tree.map_structure(lambda t: t[0], elements)
+ n = tree.flatten(elements)[0].shape[0]
+ y = backend.compute_output_spec(self.function, x)
+
+ def append_batch_axis(t):
+ return KerasTensor(
+ shape=(n,) + t.shape,
+ dtype=t.dtype,
+ sparse=t.sparse,
+ ragged=t.ragged,
+ )
+
+ y = tree.map_structure(append_batch_axis, y)
+ return y
+
+
 @keras_export("keras.ops.vectorized_map")
 def vectorized_map(function, elements):
  """Parallel map of `function` on axis 0 of tensor(s) `elements`.
@@ -1109,6 +1136,8 @@ def vectorized_map(function, elements):
  In this case, `function` is expected to take as input
  a single list of tensor arguments.
  """
+ if any_symbolic_tensors((elements,)):
+ return VectorizedMap(function)(elements)
  return backend.core.vectorized_map(function, elements)
 
 

keras/src/ops/core_test.py

@@ -56,17 +56,24 @@ def test_map(self):
  def f(x):
  return x**2
 
- xs = KerasTensor((None,))
- self.assertEqual(core.map(f, xs).shape, (None,))
+ xs = KerasTensor((None, 5))
+ self.assertEqual(core.map(f, xs).shape, (None, 5))
 
  # Test nested output
  def f2(x):
  return {"a": x**2, "b": x * 10}
 
- xs = KerasTensor((None,))
+ xs = KerasTensor((None, 5))
  ys = core.map(f2, xs)
- self.assertEqual(ys["a"].shape, (None,))
- self.assertEqual(ys["b"].shape, (None,))
+ self.assertEqual(ys["a"].shape, (None, 5))
+ self.assertEqual(ys["b"].shape, (None, 5))
+
+ # Test nested input
+ def f3(x):
+ return x[0] + x[1]
+
+ xs = (KerasTensor((None, 5)), KerasTensor((None, 5)))
+ self.assertEqual(core.map(f3, xs).shape, (None, 5))
 
  def test_saturate_cast(self):
  x = KerasTensor((3, 5, None), dtype="float32")
@@ -125,6 +132,29 @@ def fn(x, y):
  self.assertEqual(result[0].shape, (None,))
  self.assertEqual(result[1].shape, (None,))
 
+ def test_vectorized_map(self):
+ def f(x):
+ return x**2
+
+ xs = KerasTensor((None, 5))
+ self.assertEqual(core.vectorized_map(f, xs).shape, (None, 5))
+
+ # Test nested output
+ def f2(x):
+ return {"a": x**2, "b": x * 10}
+
+ xs = KerasTensor((None, 5))
+ ys = core.vectorized_map(f2, xs)
+ self.assertEqual(ys["a"].shape, (None, 5))
+ self.assertEqual(ys["b"].shape, (None, 5))
+
+ # Test nested input
+ def f3(x):
+ return x[0] + x[1]
+
+ xs = (KerasTensor((None, 5)), KerasTensor((None, 5)))
+ self.assertEqual(core.vectorized_map(f3, xs).shape, (None, 5))
+
  def test_while_loop(self):
  def cond(args):
  return tree.flatten(args)[0] < 10
@@ -203,18 +233,25 @@ def test_map(self):
  def f(x):
  return x**2
 
- xs = KerasTensor((6,))
+ xs = KerasTensor((6, 5))
  ys = core.map(f, xs)
- self.assertEqual(ys.shape, (6,))
+ self.assertEqual(ys.shape, (6, 5))
 
  # Test nested output
  def f2(x):
  return {"a": x**2, "b": x * 10}
 
- xs = KerasTensor((6,))
+ xs = KerasTensor((6, 5))
  ys = core.map(f2, xs)
- self.assertEqual(ys["a"].shape, (6,))
- self.assertEqual(ys["b"].shape, (6,))
+ self.assertEqual(ys["a"].shape, (6, 5))
+ self.assertEqual(ys["b"].shape, (6, 5))
+
+ # Test nested input
+ def f3(x):
+ return x[0] + x[1]
+
+ xs = (KerasTensor((6, 5)), KerasTensor((6, 5)))
+ self.assertEqual(core.map(f3, xs).shape, (6, 5))
 
  def test_saturate_cast(self):
  x = KerasTensor((3, 5, 7), dtype="float32")
@@ -307,6 +344,30 @@ def fn(x, y):
  self.assertEqual(core.switch(index, [fn], x, y)[0].shape, (5,))
  self.assertEqual(core.switch(index, [fn], x, y)[1].shape, (2,))
 
+ def test_vectorized_map(self):
+ def f(x):
+ return x**2
+
+ xs = KerasTensor((6, 5))
+ ys = core.vectorized_map(f, xs)
+ self.assertEqual(ys.shape, (6, 5))
+
+ # Test nested output
+ def f2(x):
+ return {"a": x**2, "b": x * 10}
+
+ xs = KerasTensor((6, 5))
+ ys = core.vectorized_map(f2, xs)
+ self.assertEqual(ys["a"].shape, (6, 5))
+ self.assertEqual(ys["b"].shape, (6, 5))
+
+ # Test nested input
+ def f3(x):
+ return x[0] + x[1]
+
+ xs = (KerasTensor((6, 5)), KerasTensor((6, 5)))
+ self.assertEqual(core.vectorized_map(f3, xs).shape, (6, 5))
+
  def test_while_loop(self):
  def cond(args):
  return tree.flatten(args)[0] < 10

keras/src/ops/numpy.py

@@ -2340,6 +2340,14 @@ class Deg2rad(Operation):
  def call(self, x):
  return backend.numpy.deg2rad(x)
 
+ def compute_output_spec(self, x):
+ dtype = backend.standardize_dtype(x.dtype)
+ if dtype in ["int64", "float64"]:
+ dtype = "float64"
+ elif dtype not in ["bfloat16", "float16"]:
+ dtype = backend.floatx()
+ return KerasTensor(x.shape, dtype)
+
 
 @keras_export(["keras.ops.deg2rad", "keras.ops.numpy.deg2rad"])
 def deg2rad(x):

keras/src/ops/ops_test.py

@@ -182,6 +182,7 @@ def test_class_function_consistency(self, module_name):
  # Check order of parameters.
  if name in (
  "fori_loop",
+ "vectorized_map",
  "while_loop",
  "batch_normalization",
  "dot_product_attention",
@@ -224,6 +225,16 @@ def test_class_function_consistency(self, module_name):
  f"function `{name}` and op class `{op_class.__name__}`",
  )
 
+ # ==== Check compute_output_spec is implement ====
+ # - op class should override Operation's `compute_output_spec`
+ self.assertTrue(
+ hasattr(op_class, "compute_output_spec")
+ and op_class.compute_output_spec
+ is not Operation.compute_output_spec,
+ f"Op class `{op_class.__name__}` should override "
+ "`compute_output_spec`",
+ )
+
  @parameterized.named_parameters(named_product(module_name=OPS_MODULES))
  def test_backend_consistency(self, module_name):
  ops_module = getattr(ops, module_name)