Implement angle function in keras.ops (#21200)

* Add first version of angle operation on numpy

* Skip test with bfloat16 on numpy

* Remove bfloat16 checking on Angle

* Fix test case for float16 on torch cuda

* exclude openvino test case

* exclude openvino test case

* exclude openvino test case

* Update init files

Author: shashaka

keras/api/_tf_keras/keras/ops/__init__.py

@@ -118,6 +118,7 @@
 from keras.src.ops.numpy import all as all
 from keras.src.ops.numpy import amax as amax
 from keras.src.ops.numpy import amin as amin
+from keras.src.ops.numpy import angle as angle
 from keras.src.ops.numpy import any as any
 from keras.src.ops.numpy import append as append
 from keras.src.ops.numpy import arange as arange

keras/api/_tf_keras/keras/ops/numpy/__init__.py

@@ -10,6 +10,7 @@
 from keras.src.ops.numpy import all as all
 from keras.src.ops.numpy import amax as amax
 from keras.src.ops.numpy import amin as amin
+from keras.src.ops.numpy import angle as angle
 from keras.src.ops.numpy import any as any
 from keras.src.ops.numpy import append as append
 from keras.src.ops.numpy import arange as arange

keras/api/ops/__init__.py

@@ -118,6 +118,7 @@
 from keras.src.ops.numpy import all as all
 from keras.src.ops.numpy import amax as amax
 from keras.src.ops.numpy import amin as amin
+from keras.src.ops.numpy import angle as angle
 from keras.src.ops.numpy import any as any
 from keras.src.ops.numpy import append as append
 from keras.src.ops.numpy import arange as arange

keras/api/ops/numpy/__init__.py

@@ -10,6 +10,7 @@
 from keras.src.ops.numpy import all as all
 from keras.src.ops.numpy import amax as amax
 from keras.src.ops.numpy import amin as amin
+from keras.src.ops.numpy import angle as angle
 from keras.src.ops.numpy import any as any
 from keras.src.ops.numpy import append as append
 from keras.src.ops.numpy import arange as arange

keras/src/backend/jax/numpy.py

@@ -246,6 +246,16 @@ def all(x, axis=None, keepdims=False):
     return jnp.all(x, axis=axis, keepdims=keepdims)
 
 
+def angle(x):
+    x = convert_to_tensor(x)
+    if standardize_dtype(x.dtype) == "int64":
+        dtype = config.floatx()
+    else:
+        dtype = dtypes.result_type(x.dtype, float)
+    x = cast(x, dtype)
+    return jnp.angle(x)
+
+
 def any(x, axis=None, keepdims=False):
     return jnp.any(x, axis=axis, keepdims=keepdims)
 

keras/src/backend/numpy/numpy.py

@@ -138,6 +138,16 @@ def all(x, axis=None, keepdims=False):
     return np.all(x, axis=axis, keepdims=keepdims)
 
 
+def angle(x):
+    x = convert_to_tensor(x)
+    if standardize_dtype(x.dtype) == "int64":
+        dtype = config.floatx()
+    else:
+        dtype = dtypes.result_type(x.dtype, float)
+    x = x.astype(dtype)
+    return np.angle(x)
+
+
 def any(x, axis=None, keepdims=False):
     axis = standardize_axis_for_numpy(axis)
     return np.any(x, axis=axis, keepdims=keepdims)

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -4,6 +4,7 @@ NumpyArrayCreateOpsCorrectnessTest::test_tri
 NumpyDtypeTest::test_absolute_bool
 NumpyDtypeTest::test_add_
 NumpyDtypeTest::test_all
+NumpyDtypeTest::test_angle
 NumpyDtypeTest::test_any
 NumpyDtypeTest::test_argpartition
 NumpyDtypeTest::test_array
@@ -71,6 +72,7 @@ NumpyDtypeTest::test_clip_bool
 NumpyDtypeTest::test_square_bool
 HistogramTest
 NumpyOneInputOpsCorrectnessTest::test_all
+NumpyOneInputOpsCorrectnessTest::test_angle
 NumpyOneInputOpsCorrectnessTest::test_any
 NumpyOneInputOpsCorrectnessTest::test_argpartition
 NumpyOneInputOpsCorrectnessTest::test_array
@@ -150,3 +152,5 @@ NumpyTwoInputOpsCorrectnessTest::test_take_along_axis
 NumpyTwoInputOpsCorrectnessTest::test_tensordot
 NumpyTwoInputOpsCorrectnessTest::test_vdot
 NumpyTwoInputOpsCorrectnessTest::test_where
+NumpyOneInputOpsDynamicShapeTest::test_angle
+NumpyOneInputOpsStaticShapeTest::test_angle

keras/src/backend/tensorflow/numpy.py

@@ -728,6 +728,16 @@ def all(x, axis=None, keepdims=False):
     return tf.reduce_all(x, axis=axis, keepdims=keepdims)
 
 
+def angle(x):
+    x = convert_to_tensor(x)
+    if standardize_dtype(x.dtype) == "int64":
+        dtype = config.floatx()
+    else:
+        dtype = dtypes.result_type(x.dtype, float)
+    x = tf.cast(x, dtype)
+    return tf.math.angle(x)
+
+
 def any(x, axis=None, keepdims=False):
     x = tf.cast(x, "bool")
     return tf.reduce_any(x, axis=axis, keepdims=keepdims)

keras/src/backend/torch/numpy.py

@@ -264,6 +264,17 @@ def all(x, axis=None, keepdims=False):
     return cast(x, "bool")
 
 
+def angle(x):
+    x = convert_to_tensor(x)
+    ori_dtype = standardize_dtype(x.dtype)
+
+    # torch.angle doesn't support float16 with cuda
+    if get_device() != "cpu" and ori_dtype == "float16":
+        x = cast(x, "float32")
+        return cast(torch.angle(x), "float16")
+    return torch.angle(x)
+
+
 def any(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     if axis is None:

keras/src/ops/numpy.py

@@ -328,6 +328,41 @@ def compute_output_spec(self, x):
         )
 
 
+class Angle(Operation):
+    def call(self, x):
+        return backend.numpy.angle(x)
+
+    def compute_output_spec(self, x):
+        dtype = backend.standardize_dtype(getattr(x, "dtype", backend.floatx()))
+        if dtype == "int64":
+            dtype = backend.floatx()
+        else:
+            dtype = dtypes.result_type(dtype, float)
+        return KerasTensor(x.shape, dtype=dtype)
+
+
+@keras_export(["keras.ops.angle", "keras.ops.numpy.angle"])
+def angle(x):
+    """Element-wise angle of a complex tensor.
+
+    Arguments:
+        x: Input tensor. Can be real or complex.
+
+    Returns:
+        Output tensor of same shape as x. containing the angle of each element
+        (in radians).
+
+    Example:
+    >>> x = keras.ops.convert_to_tensor([[1 + 3j, 2 - 5j], [4 - 3j, 3 + 2j]])
+    >>> keras.ops.angle(x)
+    array([[ 1.2490457, -1.19029  ],
+       [-0.6435011,  0.5880026]], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Angle().symbolic_call(x)
+    return backend.numpy.angle(x)
+
+
 @keras_export(["keras.ops.any", "keras.ops.numpy.any"])
 def any(x, axis=None, keepdims=False):
     """Test whether any array element along a given axis evaluates to `True`.