Implement ldexp function in keras.ops (#21863)

shashaka · web-flow · commit 74fba843ec2e · 2025-12-01T15:47:37.000-08:00
* Add ldexp initial version

* Add numpy_test for ldexp

* Update code by gemini reveiw

* Update code by review

* Add cast for tensorflow return value

* merge master
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -215,6 +215,7 @@
 from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import kron as kron
 from keras.src.ops.numpy import lcm as lcm
+from keras.src.ops.numpy import ldexp as ldexp
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal
diff --git a/keras/api/_tf_keras/keras/ops/numpy/__init__.py b/keras/api/_tf_keras/keras/ops/numpy/__init__.py
@@ -101,6 +101,7 @@
 from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import kron as kron
 from keras.src.ops.numpy import lcm as lcm
+from keras.src.ops.numpy import ldexp as ldexp
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
@@ -215,6 +215,7 @@
 from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import kron as kron
 from keras.src.ops.numpy import lcm as lcm
+from keras.src.ops.numpy import ldexp as ldexp
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal
diff --git a/keras/api/ops/numpy/__init__.py b/keras/api/ops/numpy/__init__.py
@@ -101,6 +101,7 @@
 from keras.src.ops.numpy import kaiser as kaiser
 from keras.src.ops.numpy import kron as kron
 from keras.src.ops.numpy import lcm as lcm
+from keras.src.ops.numpy import ldexp as ldexp
 from keras.src.ops.numpy import left_shift as left_shift
 from keras.src.ops.numpy import less as less
 from keras.src.ops.numpy import less_equal as less_equal
diff --git a/keras/src/backend/jax/numpy.py b/keras/src/backend/jax/numpy.py
@@ -845,6 +845,19 @@ def lcm(x1, x2):
     return jnp.lcm(x1, x2)
 
 
+def ldexp(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES:
+        raise TypeError(
+            f"ldexp exponent must be an integer type. "
+            f"Received: x2 dtype={x2.dtype}"
+        )
+
+    return jnp.ldexp(x1, x2)
+
+
 def less(x1, x2):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
diff --git a/keras/src/backend/numpy/numpy.py b/keras/src/backend/numpy/numpy.py
@@ -773,6 +773,19 @@ def lcm(x1, x2):
     return np.lcm(x1, x2).astype(dtype)
 
 
+def ldexp(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
+
+    if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES:
+        raise TypeError(
+            f"ldexp exponent must be an integer type. "
+            f"Received: x2 dtype={x2.dtype}"
+        )
+    return np.ldexp(x1, x2).astype(dtype)
+
+
 def less(x1, x2):
     return np.less(x1, x2)
 
diff --git a/keras/src/backend/openvino/excluded_concrete_tests.txt b/keras/src/backend/openvino/excluded_concrete_tests.txt
@@ -29,6 +29,7 @@ NumpyDtypeTest::test_isin
 NumpyDtypeTest::test_isreal
 NumpyDtypeTest::test_kron
 NumpyDtypeTest::test_lcm
+NumpyDtypeTest::test_ldexp
 NumpyDtypeTest::test_logaddexp2
 NumpyDtypeTest::test_matmul_
 NumpyDtypeTest::test_maximum_python_types
@@ -108,6 +109,7 @@ NumpyTwoInputOpsCorrectnessTest::test_inner
 NumpyTwoInputOpsCorrectnessTest::test_isin
 NumpyTwoInputOpsCorrectnessTest::test_kron
 NumpyTwoInputOpsCorrectnessTest::test_lcm
+NumpyTwoInputOpsCorrectnessTest::test_ldexp
 NumpyTwoInputOpsCorrectnessTest::test_quantile
 NumpyTwoInputOpsCorrectnessTest::test_tensordot
 NumpyTwoInputOpsCorrectnessTest::test_vdot
@@ -131,11 +133,13 @@ NumpyTwoInputOpsDynamicShapeTest::test_hypot
 NumpyTwoInputOpsDynamicShapeTest::test_isin
 NumpyTwoInputOpsDynamicShapeTest::test_kron
 NumpyTwoInputOpsDynamicShapeTest::test_lcm
+NumpyTwoInputOpsDynamicShapeTest::test_ldexp
 NumpyTwoInputOpsStaticShapeTest::test_gcd
 NumpyTwoInputOpsStaticShapeTest::test_hypot
 NumpyTwoInputOpsStaticShapeTest::test_isin
 NumpyTwoInputOpsStaticShapeTest::test_kron
 NumpyTwoInputOpsStaticShapeTest::test_lcm
+NumpyTwoInputOpsStaticShapeTest::test_ldexp
 CoreOpsBehaviorTests::test_associative_scan_invalid_arguments
 CoreOpsBehaviorTests::test_scan_invalid_arguments
 CoreOpsCallsTests::test_associative_scan_basic_call
diff --git a/keras/src/backend/openvino/numpy.py b/keras/src/backend/openvino/numpy.py
@@ -1149,6 +1149,10 @@ def lcm(x1, x2):
     raise NotImplementedError("`lcm` is not supported with openvino backend")
 
 
+def ldexp(x1, x2):
+    raise NotImplementedError("`ldexp` is not supported with openvino backend")
+
+
 def less(x1, x2):
     element_type = None
     if isinstance(x1, OpenVINOKerasTensor):
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
@@ -1844,6 +1844,23 @@ def lcm(x1, x2):
     return result
 
 
+def ldexp(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
+
+    if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES:
+        raise TypeError(
+            f"ldexp exponent must be an integer type. "
+            f"Received: x2 dtype={x2.dtype}"
+        )
+
+    x1 = tf.cast(x1, dtype)
+    x2 = tf.cast(x2, x1.dtype)
+    result = x1 * tf.pow(tf.constant(2.0, dtype=x1.dtype), x2)
+    return tf.cast(tf.where(tf.math.is_inf(x1) | (x1 == 0), x1, result), dtype)
+
+
 def less(x1, x2):
     x1 = convert_to_tensor(x1)
     x2 = convert_to_tensor(x2)
diff --git a/keras/src/backend/torch/numpy.py b/keras/src/backend/torch/numpy.py
@@ -975,6 +975,20 @@ def lcm(x1, x2):
     return torch.lcm(x1, x2)
 
 
+def ldexp(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
+
+    if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES:
+        raise TypeError(
+            f"ldexp exponent must be an integer type. "
+            f"Received: x2 dtype={x2.dtype}"
+        )
+
+    return cast(torch.ldexp(x1, x2), dtype)
+
+
 def less(x1, x2):
     x1, x2 = convert_to_tensor(x1), convert_to_tensor(x2)
     return torch.less(x1, x2)
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
@@ -4059,6 +4059,46 @@ def lcm(x1, x2):
     return backend.numpy.lcm(x1, x2)
 
 
+class Ldexp(Operation):
+    def call(self, x1, x2):
+        return backend.numpy.ldexp(x1, x2)
+
+    def compute_output_spec(self, x1, x2):
+        x1_shape = getattr(x1, "shape", [])
+        x2_shape = getattr(x2, "shape", [])
+        output_shape = broadcast_shapes(x1_shape, x2_shape)
+
+        x1_type = backend.standardize_dtype(getattr(x1, "dtype", type(x1)))
+        x2_type = backend.standardize_dtype(getattr(x2, "dtype", type(x2)))
+        dtype = dtypes.result_type(x1_type, x2_type, float)
+        return KerasTensor(output_shape, dtype=dtype)
+
+
+@keras_export(["keras.ops.ldexp", "keras.ops.numpy.ldexp"])
+def ldexp(x1, x2):
+    """Multiply `x1` by 2 raised to the power of `x2`, element-wise.
+
+    This function computes:
+        ldexp(x1, x2) = x1 * 2**x2
+
+    Args:
+        x1: Float input tensor.
+        x2: Integer exponent tensor.
+
+    Returns:
+        Output tensor
+
+    Example:
+    >>> x1 = keras.ops.convert_to_tensor([0.75, 1.5])
+    >>> x2 = keras.ops.convert_to_tensor([1, 2])
+    >>> keras.ops.ldexp(x1, x2)
+    array([1.5, 6. ], dtype=float32)
+    """
+    if any_symbolic_tensors((x1, x2)):
+        return Ldexp().symbolic_call(x1, x2)
+    return backend.numpy.ldexp(x1, x2)
+
+
 class Less(Operation):
     def call(self, x1, x2):
         return backend.numpy.less(x1, x2)
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
@@ -254,6 +254,11 @@ def test_lcm(self):
         y = KerasTensor((2, None))
         self.assertEqual(knp.lcm(x, y).shape, (2, 3))
 
+    def test_ldexp(self):
+        x = KerasTensor((None, 3))
+        y = KerasTensor((1, 3))
+        self.assertEqual(knp.ldexp(x, y).shape, (None, 3))
+
     def test_less(self):
         x = KerasTensor((None, 3))
         y = KerasTensor((2, None))
@@ -837,6 +842,15 @@ def test_lcm(self):
         y = KerasTensor((2, 3))
         self.assertEqual(knp.lcm(x, y).shape, (2, 3))
 
+    def test_ldexp(self):
+        x = KerasTensor((2, 3))
+        y = KerasTensor((2, 3))
+        self.assertEqual(knp.ldexp(x, y).shape, (2, 3))
+
+        x = KerasTensor((2, 3))
+        y = KerasTensor((1, 3))
+        self.assertEqual(knp.ldexp(x, y).shape, (2, 3))
+
     def test_less(self):
         x = KerasTensor((2, 3))
         y = KerasTensor((2, 3))
@@ -3114,6 +3128,12 @@ def test_lcm(self):
         self.assertAllClose(knp.lcm(x, y), np.lcm(x, y))
         self.assertAllClose(knp.Lcm()(x, y), np.lcm(x, y))
 
+    def test_ldexp(self):
+        x = np.array([[1, 2, 3], [3, 2, 1]])
+        y = np.array([[4, 5, 6], [3, 2, 1]])
+        self.assertAllClose(knp.ldexp(x, y), np.ldexp(x, y))
+        self.assertAllClose(knp.Ldexp()(x, y), np.ldexp(x, y))
+
     def test_less(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
         y = np.array([[4, 5, 6], [3, 2, 1]])
@@ -7884,6 +7904,27 @@ def test_lcm(self, dtypes):
             expected_dtype,
         )
 
+    @parameterized.named_parameters(
+        named_product(dtypes=list(itertools.product(ALL_DTYPES, INT_DTYPES)))
+    )
+    def test_ldexp(self, dtypes):
+        import jax.numpy as jnp
+
+        dtype1, dtype2 = dtypes
+        x1 = knp.ones((), dtype=dtype1)
+        x2 = knp.ones((), dtype=dtype2)
+        x1_jax = jnp.ones((), dtype=dtype1)
+        x2_jax = jnp.ones((), dtype=dtype2)
+        expected_dtype = standardize_dtype(jnp.ldexp(x1_jax, x2_jax).dtype)
+
+        self.assertEqual(
+            standardize_dtype(knp.ldexp(x1, x2).dtype), expected_dtype
+        )
+        self.assertEqual(
+            standardize_dtype(knp.Ldexp().symbolic_call(x1, x2).dtype),
+            expected_dtype,
+        )
+
     @parameterized.named_parameters(
         named_product(dtypes=itertools.combinations(ALL_DTYPES, 2))
     )
