Implement trapezoid function in keras.ops (#21757)

shashaka · web-flow · commit 2b8ddf3d0a8f · 2025-10-18T17:12:47.000-07:00
* Add trapezoid for numpy

* Add trapezoid for ops

* Add test cases for initial version

* Add more test cases

* Update openvino side for trapezoid

* Code update by gemini
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -286,6 +286,7 @@
 from keras.src.ops.numpy import tile as tile
 from keras.src.ops.numpy import trace as trace
 from keras.src.ops.numpy import transpose as transpose
+from keras.src.ops.numpy import trapezoid as trapezoid
 from keras.src.ops.numpy import tri as tri
 from keras.src.ops.numpy import tril as tril
 from keras.src.ops.numpy import triu as triu
diff --git a/keras/api/_tf_keras/keras/ops/numpy/__init__.py b/keras/api/_tf_keras/keras/ops/numpy/__init__.py
@@ -172,6 +172,7 @@
 from keras.src.ops.numpy import tile as tile
 from keras.src.ops.numpy import trace as trace
 from keras.src.ops.numpy import transpose as transpose
+from keras.src.ops.numpy import trapezoid as trapezoid
 from keras.src.ops.numpy import tri as tri
 from keras.src.ops.numpy import tril as tril
 from keras.src.ops.numpy import triu as triu
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
@@ -286,6 +286,7 @@
 from keras.src.ops.numpy import tile as tile
 from keras.src.ops.numpy import trace as trace
 from keras.src.ops.numpy import transpose as transpose
+from keras.src.ops.numpy import trapezoid as trapezoid
 from keras.src.ops.numpy import tri as tri
 from keras.src.ops.numpy import tril as tril
 from keras.src.ops.numpy import triu as triu
diff --git a/keras/api/ops/numpy/__init__.py b/keras/api/ops/numpy/__init__.py
@@ -172,6 +172,7 @@
 from keras.src.ops.numpy import tile as tile
 from keras.src.ops.numpy import trace as trace
 from keras.src.ops.numpy import transpose as transpose
+from keras.src.ops.numpy import trapezoid as trapezoid
 from keras.src.ops.numpy import tri as tri
 from keras.src.ops.numpy import tril as tril
 from keras.src.ops.numpy import triu as triu
diff --git a/keras/src/backend/jax/numpy.py b/keras/src/backend/jax/numpy.py
@@ -1361,6 +1361,14 @@ def transpose(x, axes=None):
     return jnp.transpose(x, axes=axes)
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    y = convert_to_tensor(y)
+    if x is not None:
+        x = convert_to_tensor(x)
+    dx = convert_to_tensor(dx)
+    return jnp.trapezoid(y, x, dx=dx, axis=axis)
+
+
 def var(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     # `jnp.var` does not handle low precision (e.g., float16) overflow
diff --git a/keras/src/backend/numpy/numpy.py b/keras/src/backend/numpy/numpy.py
@@ -1336,6 +1336,15 @@ def transpose(x, axes=None):
     return np.transpose(x, axes=axes)
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    y = convert_to_tensor(y)
+    result_dtype = dtypes.result_type(y.dtype, float)
+    if x is not None:
+        x = convert_to_tensor(x)
+    dx = convert_to_tensor(dx)
+    return np.trapezoid(y, x, dx=dx, axis=axis).astype(result_dtype)
+
+
 def var(x, axis=None, keepdims=False):
     axis = standardize_axis_for_numpy(axis)
     x = convert_to_tensor(x)
diff --git a/keras/src/backend/openvino/excluded_concrete_tests.txt b/keras/src/backend/openvino/excluded_concrete_tests.txt
@@ -46,6 +46,7 @@ NumpyDtypeTest::test_swapaxes
 NumpyDtypeTest::test_tensordot_
 NumpyDtypeTest::test_tile
 NumpyDtypeTest::test_trace
+NumpyDtypeTest::test_trapezoid
 NumpyDtypeTest::test_trunc
 NumpyDtypeTest::test_unravel
 NumpyDtypeTest::test_var
@@ -95,6 +96,7 @@ NumpyOneInputOpsCorrectnessTest::test_swapaxes
 NumpyOneInputOpsCorrectnessTest::test_tile
 NumpyOneInputOpsCorrectnessTest::test_trace
 NumpyOneInputOpsCorrectnessTest::test_transpose
+NumpyOneInputOpsCorrectnessTest::test_trapezoid
 NumpyOneInputOpsCorrectnessTest::test_trunc
 NumpyOneInputOpsCorrectnessTest::test_unravel_index
 NumpyOneInputOpsCorrectnessTest::test_var
diff --git a/keras/src/backend/openvino/numpy.py b/keras/src/backend/openvino/numpy.py
@@ -2398,6 +2398,12 @@ def transpose(x, axes=None):
     return OpenVINOKerasTensor(ov_opset.transpose(x, axes).output(0))
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    raise NotImplementedError(
+        "`trapezoid` is not supported with openvino backend"
+    )
+
+
 def var(x, axis=None, keepdims=False):
     x = get_ov_output(x)
     if axis is None:
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
@@ -2984,6 +2984,42 @@ def transpose(x, axes=None):
     return tf.transpose(x, perm=axes)
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    def _move_axis_to_last(tensor, axis):
+        if axis == -1:
+            return tensor
+        rank = tf.rank(tensor)
+        if axis < 0:
+            axis = rank + axis
+        perm = tf.concat(
+            [
+                tf.range(axis, dtype=tf.int32),
+                tf.range(axis + 1, rank, dtype=tf.int32),
+                tf.constant([axis], dtype=tf.int32),
+            ],
+            axis=0,
+        )
+        return tf.transpose(tensor, perm=perm)
+
+    y = convert_to_tensor(y)
+    dtype = dtypes.result_type(y.dtype, float)
+    y = tf.cast(y, dtype)
+
+    if x is None:
+        dx_array = tf.cast(dx, dtype)
+    else:
+        x = convert_to_tensor(x, dtype=dtype)
+        dx_array = diff(x, axis=axis)
+        dx_array = _move_axis_to_last(dx_array, axis)
+
+    y = _move_axis_to_last(y, axis)
+
+    avg_heights = 0.5 * (y[..., 1:] + y[..., :-1])
+    result = tf.reduce_sum(avg_heights * dx_array, axis=-1)
+
+    return result
+
+
 def var(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     compute_dtype = dtypes.result_type(x.dtype, "float32")
diff --git a/keras/src/backend/torch/numpy.py b/keras/src/backend/torch/numpy.py
@@ -1789,6 +1789,18 @@ def transpose(x, axes=None):
     return x.T
 
 
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    y = convert_to_tensor(y)
+    if standardize_dtype(y.dtype) == "bool":
+        y = cast(y, config.floatx())
+    if x is not None:
+        x = convert_to_tensor(x)
+        return torch.trapz(y, x=x, dim=axis)
+    else:
+        dx = convert_to_tensor(dx)
+        return torch.trapz(y, dx=dx, dim=axis)
+
+
 def var(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     compute_dtype = dtypes.result_type(x.dtype, "float32")
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
@@ -7088,6 +7088,48 @@ def transpose(x, axes=None):
     return backend.numpy.transpose(x, axes=axes)
 
 
+class Trapezoid(Operation):
+    def __init__(self, x=None, dx=1.0, axis=-1, *, name=None):
+        super().__init__(name=name)
+        self.x = x
+        self.dx = dx
+        self.axis = axis
+
+    def call(self, y):
+        return backend.numpy.trapezoid(y, x=self.x, dx=self.dx, axis=self.axis)
+
+    def compute_output_spec(self, y):
+        out_shape = list(y.shape)
+        if self.axis is not None and len(out_shape) > 0:
+            out_shape.pop(self.axis % len(out_shape))
+        dtype = backend.result_type(getattr(y, "dtype", type(y)), float)
+        return KerasTensor(tuple(out_shape), dtype=dtype)
+
+
+@keras_export(["keras.ops.trapezoid", "keras.ops.numpy.trapezoid"])
+def trapezoid(y, x=None, dx=1.0, axis=-1):
+    """Integrate along the given axis using the composite trapezoidal rule.
+
+    Args:
+        y: Input tensor.
+        x: Optional tensor specifying sample points corresponding to `y`.
+           If `None`, spacing is assumed to be `dx`.
+        dx: Spacing between sample points when `x` is `None`.
+        axis: Axis along which to integrate. Default is the last axis.
+
+    Returns:
+        The approximate integral of `y` along the given axis.
+
+    Example:
+    >>> y = keras.ops.convert_to_tensor([[1, 2, 3], [4, 5, 6]])
+    >>> keras.ops.trapezoid(y, axis=1)
+    array([ 4., 10.], dtype=float32)
+    """
+    if any_symbolic_tensors((y,)):
+        return Trapezoid(x=x, dx=dx, axis=axis).symbolic_call(y)
+    return backend.numpy.trapezoid(y, x=x, dx=dx, axis=axis)
+
+
 class Mean(Operation):
     def __init__(self, axis=None, keepdims=False, *, name=None):
         super().__init__(name=name)
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
@@ -1133,6 +1133,13 @@ def test_any(self):
         self.assertEqual(knp.any(x, axis=1).shape, (None, 3))
         self.assertEqual(knp.any(x, axis=1, keepdims=True).shape, (None, 1, 3))
 
+    def test_trapezoid(self):
+        x = KerasTensor((None, 3))
+        self.assertEqual(knp.trapezoid(x).shape, (None,))
+
+        x = KerasTensor((None, 3, 3))
+        self.assertEqual(knp.trapezoid(x, axis=1).shape, (None, 3))
+
     def test_var(self):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.var(x).shape, ())
@@ -1865,6 +1872,10 @@ def test_any(self):
         x = KerasTensor((2, 3))
         self.assertEqual(knp.any(x).shape, ())
 
+    def test_trapezoid(self):
+        x = KerasTensor((2, 3))
+        self.assertEqual(knp.trapezoid(x).shape, (2,))
+
     def test_var(self):
         x = KerasTensor((2, 3))
         self.assertEqual(knp.var(x).shape, ())
@@ -3604,6 +3615,19 @@ def test_any(self):
             np.any(x, axis=1, keepdims=True),
         )
 
+    def test_trapezoid(self):
+        y = np.random.random((3, 3, 3))
+        x = np.random.random((3, 3, 3))
+        dx = 2.0
+
+        self.assertAllClose(knp.trapezoid(y), np.trapezoid(y))
+        self.assertAllClose(knp.trapezoid(y, x=x), np.trapezoid(y, x=x))
+        self.assertAllClose(knp.trapezoid(y, dx=dx), np.trapezoid(y, dx=dx))
+        self.assertAllClose(
+            knp.trapezoid(y, x=x, axis=1),
+            np.trapezoid(y, x=x, axis=1),
+        )
+
     def test_var(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         self.assertAllClose(knp.var(x), np.var(x))
@@ -8968,6 +8992,22 @@ def test_trunc(self, dtype):
             expected_dtype,
         )
 
+    @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
+    def test_trapezoid(self, dtype):
+        import jax.numpy as jnp
+
+        x = knp.ones((2,), dtype=dtype)
+        x_jax = jnp.ones((2,), dtype=dtype)
+        expected_dtype = standardize_dtype(jnp.trapezoid(x_jax).dtype)
+
+        self.assertEqual(
+            standardize_dtype(knp.trapezoid(x).dtype), expected_dtype
+        )
+        self.assertEqual(
+            standardize_dtype(knp.Trapezoid().symbolic_call(x).dtype),
+            expected_dtype,
+        )
+
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_var(self, dtype):
         import jax.numpy as jnp
