frexp implementation (#11182)

Co-author: pillarxyz <[email protected]>, nassimberrada <Nassim>

ivy/array/experimental/elementwise.py

@@ -1226,3 +1226,27 @@ def ldexp(
         ivy.array([8.0, 8.0, 6.0])
         """
         return ivy.ldexp(self._data, x2, out=out, **kwargs)
+
+    def frexp(self: ivy.Array, /, *, out: Optional[Tuple[ivy.Array, ivy.Array]] = None, **kwargs) -> ivy.Array:
+        """
+        ivy.Array instance method variant of ivy.frexp. This method simply wraps
+        the function, and so the docstring for ivy.frexp also applies to this
+        method with minimal changes.
+        Parameters
+        ----------
+        self
+            Input array.
+        out
+            Alternate output array in which to place the result.
+            The default is None.
+        Returns
+        -------
+        ret
+            The next representable values of x1 in the direction of x2.
+        Examples
+        --------
+        >>> x = ivy.array([1.0, 2.0, 3.0])
+        >>> x.frexp()
+        ivy.array([[0.5, 0.5, 0.75], [1, 2, 2]])
+        """
+        return ivy.frexp(self._data, out=out, **kwargs)

ivy/container/experimental/elementwise.py

@@ -3448,3 +3448,93 @@ def ldexp(
         }
         """
         return self.static_ldexp(self, x2, out=out)
+
+    @staticmethod
+    def static_frexp(
+        x: Union[ivy.Array, ivy.NativeArray, ivy.Container],
+        /,
+        *,
+        key_chains: Optional[Union[List[str], Dict[str, str]]] = None,
+        to_apply: bool = True,
+        prune_unapplied: bool = False,
+        map_sequences: bool = False,
+        out: Optional[ivy.Container] = None,
+    ) -> ivy.Container:
+        """
+        ivy.Container static method variant of ivy.frexp. This method simply
+        wraps the function, and so the docstring for ivy.frexp also applies to this
+        method with minimal changes.
+
+        Parameters
+        ----------
+        x
+            The container whose arrays should be split into mantissa and exponent.
+        key_chains
+            The key-chains to apply or not apply the method to. Default is ``None``.
+        to_apply
+            If True, the method will be applied to key_chains, otherwise key_chains
+            will be skipped. Default is ``True``.
+        prune_unapplied
+            Whether to prune key_chains for which the function was not applied.
+            Default is ``False``.
+        map_sequences
+            Whether to also map method to sequences (lists, tuples).
+            Default is ``False``.
+        out
+            optional output container, for writing the result to.
+
+        Returns
+        -------
+        ret
+            container including the mantissa and exponent of x.
+
+        Examples
+        --------
+        With one :class:`ivy.Container` input:
+        >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([1, 5, 10]))
+        >>> ivy.Container.static_frexp(x)
+        {
+            a: (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 1, 2])),
+            b: (ivy.array([0.5, 0.625, 0.625]), ivy.array([1, 3, 4]))
+        }
+        """
+        return ContainerBase.cont_multi_map_in_function(
+            "frexp",
+            x,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )
+
+    def frexp(
+        self: ivy.Container,
+        /,
+        *,
+        out: Optional[ivy.Container] = None,
+    ) -> ivy.Container:
+        """ivy.Container instance method variant of ivy.frexp. This method simply
+        wraps the function, and so the docstring for ivy.frexp also applies to this
+        method with minimal changes.
+        Parameters
+        ----------
+        self
+            The container whose arrays should be split into mantissa and exponent.
+        out
+            optional output container, for writing the result to.
+        Returns
+        -------
+        ret
+            container including the mantissa and exponent of x.
+        Examples
+        --------
+        With one :class:`ivy.Container` input:
+        >>> x = ivy.Container(a=ivy.array([1, 2, 3]), b=ivy.array([1, 5, 10]))
+        >>> x.frexp()
+        {
+            a: (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 1, 2])),
+            b: (ivy.array([0.5, 0.625, 0.625]), ivy.array([1, 3, 4]))
+        }
+        """
+        return self.static_frexp(self, out=out)

ivy/functional/backends/jax/experimental/elementwise.py

@@ -539,3 +539,9 @@ def ldexp(
     x1: JaxArray, x2: Union[JaxArray, int], /, *, out: Optional[JaxArray] = None
 ) -> JaxArray:
     return jnp.ldexp(x1, x2)
+
+
+def frexp(
+    x: JaxArray, /, *, out: Optional[Tuple[JaxArray, JaxArray]] = None
+) -> Tuple[JaxArray, JaxArray]:
+    return jnp.frexp(x)

ivy/functional/backends/numpy/experimental/elementwise.py

@@ -448,3 +448,12 @@ def ldexp(
     out: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     return np.ldexp(x1, x2, out=out)
+
+
+def frexp(
+    x: np.ndarray,
+    /,
+    *,
+    out: Optional[Union[Tuple[np.ndarray, np.ndarray], Tuple[None, None]]] = (None, None)
+) -> Tuple[np.ndarray, np.ndarray]:
+    return np.frexp(x, out=out)

ivy/functional/backends/tensorflow/experimental/elementwise.py

@@ -644,3 +644,19 @@ def ldexp(
         x2 = tf.cast(x2, x1.dtype)
         ret = x1 * tf.math.pow(2, x2)
     return ret
+
+
+@with_unsupported_dtypes({"1.11.0 and below": ("unsigned")}, backend_version)
+def frexp(
+    x: Union[tf.Tensor, tf.Variable],
+    /,
+    *,
+    out: Optional[Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]] = None,
+) -> Union[Tuple[tf.Tensor, tf.Tensor], Tuple[tf.Variable, tf.Variable]]:
+    e = tf.math.floor(tf.math.log(tf.math.abs(x)) / tf.cast(tf.math.log(2.), x.dtype))
+    e = tf.cast(e, x.dtype)
+    while tf.reduce_any(tf.abs(x / tf.math.pow(2, e)) >= 1):
+        e += tf.cast(tf.abs(x / tf.math.pow(2, e)) >= 1, e.dtype)
+    m = x / tf.math.pow(2, e)
+    e = tf.cast(e, tf.int32)
+    return m, e

ivy/functional/backends/torch/experimental/elementwise.py

@@ -456,3 +456,12 @@ def ldexp(
     out: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     return torch.ldexp(x1, x2, out=out)
+
+
+def frexp(
+    x: torch.Tensor,
+    /,
+    *,
+    out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return torch.frexp(x, out=out)

ivy/functional/ivy/experimental/elementwise.py

@@ -1451,3 +1451,36 @@ def ldexp(
     ivy.array([1, 4, 12])
     """
     return ivy.current_backend(x1, x2).ldexp(x1, x2, out=out)
+
+
+@to_native_arrays_and_back
+@handle_out_argument
+@handle_nestable
+@handle_exceptions
+@handle_array_like_without_promotion
+def frexp(
+    x: Union[ivy.Array, ivy.NativeArray],
+    /,
+    *,
+    out: Optional[Tuple[ivy.Array, ivy.Array]] = None,
+) -> Tuple[ivy.Array, ivy.Array]:
+    """
+    Decompose the elements of x into mantissa and twos exponent.
+    Parameters
+    ----------
+    x
+        Input array.
+    out
+        optional output array, for writing the result to.
+        It must have a shape that the inputs broadcast to.
+    Returns
+    -------
+    ret
+        A tuple of two arrays, the mantissa and the twos exponent.
+    Examples
+    --------
+    >>> x = ivy.array([1, 2, 3])
+    >>> ivy.frexp(x)
+    (ivy.array([0.5, 0.5, 0.75]), ivy.array([1, 2, 2]))
+    """
+    return ivy.current_backend(x).frexp(x, out=out)

ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_elementwise.py

@@ -1210,3 +1210,38 @@ def test_ldexp(
         x1=x[0],
         x2=x[1],
     )
+
+
+# frexp
+@handle_test(
+    fn_tree="functional.ivy.experimental.frexp",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=["float32", "float64"],
+        num_arrays=1,
+        shared_dtype=True,
+        min_value=-100,
+        max_value=100,
+        min_num_dims=1,
+        max_num_dims=3,
+    ),
+    test_gradients=st.just(False),
+)
+def test_frexp(
+    *,
+    dtype_and_x,
+    test_flags,
+    backend_fw,
+    fn_name,
+    on_device,
+    ground_truth_backend,
+):
+    input_dtype, x = dtype_and_x
+    helpers.test_function(
+        ground_truth_backend=ground_truth_backend,
+        input_dtypes=input_dtype,
+        test_flags=test_flags,
+        fw=backend_fw,
+        fn_name=fn_name,
+        on_device=on_device,
+        x=x[0],
+    )