Repository: ivy-llc/ivy

Issue:
general_inner_product



Discussion:

Resolving patch:
diff --git a/ivy/data_classes/array/experimental/linear_algebra.py b/ivy/data_classes/array/experimental/linear_algebra.py
--- a/ivy/data_classes/array/experimental/linear_algebra.py
+++ b/ivy/data_classes/array/experimental/linear_algebra.py
@@ -728,3 +728,59 @@ def dot(
         ivy.array([[-15.28]])
         """
         return ivy.dot(self._data, b, out=out)
+
+    def general_inner_product(
+        self: Union[ivy.Array, ivy.NativeArray],
+        b: Union[ivy.Array, ivy.NativeArray],
+        n_modes: Optional[int] = None,
+        /,
+        *,
+        out: Optional[ivy.Array] = None,
+    ) -> ivy.Array:
+        """
+        ivy.Array instance method variant of ivy.general_inner_product. This method
+        simply wraps the function, and so the docstring for ivy.general_inner_product
+        also applies to this method with minimal changes.
+
+        Parameters
+        ----------
+        self
+            first input tensor.
+        b
+            second input tensor.
+        n_modes
+            int, default is None. If None, the traditional inner product is returned
+            (i.e. a float) otherwise, the product between the `n_modes` last modes of
+            `a` and the `n_modes` first modes of `b` is returned. The resulting tensor's
+            order is `len(a) - n_modes`.
+        out
+            Optional output array. If provided, the output array to store the result.
+
+        Returns
+        -------
+            The inner product of the input arrays.
+
+        Examples
+        --------
+        With :class:`ivy.Array` inputs:
+
+        >>> a = ivy.array([1, 2, 3])
+        >>> b = ivy.array([4, 5, 6])
+        >>> result = a.general_inner_product(b, n_modes=1)
+        >>> print(result)
+        ivy.array(32)
+
+        >>> a = ivy.array([1, 2])
+        >>> b = ivy.array([4, 5])
+        >>> result = a.general_inner_product(b)
+        >>> print(result)
+        ivy.array(14)
+
+        >>> a = ivy.array([[1, 1], [1, 1]])
+        >>> b = ivy.array([[1, 2, 3, 4],[1, 1, 1, 1]])
+        >>> result = a.general_inner_product(b, n_modes=1)
+        >>> print(result)
+        ivy.array([[2, 3, 4, 5],
+            [2, 3, 4, 5]])
+        """
+        return ivy.general_inner_product(self, b, n_modes, out=out)
diff --git a/ivy/data_classes/container/linear_algebra.py b/ivy/data_classes/container/linear_algebra.py
--- a/ivy/data_classes/container/linear_algebra.py
+++ b/ivy/data_classes/container/linear_algebra.py
@@ -3313,3 +3313,106 @@ def vander(
             increasing=increasing,
             out=out,
         )
+
+    @staticmethod
+    def static_general_inner_product(
+        x1: Union[ivy.Container, ivy.Array, ivy.NativeArray],
+        x2: Union[ivy.Container, ivy.Array, ivy.NativeArray],
+        n_modes: Optional[Union[int, ivy.Container]] = None,
+        /,
+        *,
+        key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None,
+        to_apply: Union[bool, ivy.Container] = True,
+        prune_unapplied: Union[bool, ivy.Container] = False,
+        map_sequences: Union[bool, ivy.Container] = False,
+        out: Optional[ivy.Container] = None,
+    ) -> ivy.Container:
+        """
+        ivy.Container static method variant of ivy.general_inner_product. This method
+        simply wraps the function, and so the docstring for ivy.general_inner_product
+        also applies to this method with minimal changes.
+
+        Parameters
+        ----------
+        x1
+            First input container containing input array.
+        x2
+            First input container containing input array.
+        n_modes
+            int, default is None. If None, the traditional inner product is returned
+            (i.e. a float) otherwise, the product between the `n_modes` last modes of
+            `x1` and the `n_modes` first modes of `x2` is returned. The resulting
+            tensor's order is `len(x1) - n_modes`.
+        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
+            Alternate output container in which to place the result.
+            The default is None.
+
+        Returns
+        -------
+        ret
+            Container including the inner product tensor.
+
+        Examples
+        --------
+        >>> x = ivy.Container(
+                a=ivy.reshape(ivy.arange(4), (2, 2)),
+                b=ivy.reshape(ivy.arange(8), (2, 4)),
+            )
+        >>> ivy.Container.general_inner_product(x, 1)
+            {
+                a: ivy.array(6),
+                b: ivy.array(28)
+            }
+        """
+        return ContainerBase.cont_multi_map_in_function(
+            "general_inner_product",
+            x1,
+            x2,
+            n_modes,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )
+
+    def general_inner_product(
+        self: Union[ivy.Container, ivy.Array, ivy.NativeArray],
+        x2: Union[ivy.Container, ivy.Array, ivy.NativeArray],
+        n_modes: Optional[Union[int, ivy.Container]] = None,
+        /,
+        *,
+        key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None,
+        to_apply: Union[bool, ivy.Container] = True,
+        prune_unapplied: Union[bool, ivy.Container] = False,
+        map_sequences: Union[bool, ivy.Container] = False,
+        out: Optional[ivy.Container] = None,
+    ) -> ivy.Container:
+        """
+        ivy.Container instance method variant of ivy.general_inner_product.
+
+        This method simply wraps the function, and so the docstring for
+        ivy.general_inner_product also applies to this method with
+        minimal changes.
+        """
+        return self.static_general_inner_product(
+            self,
+            x2,
+            n_modes,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )
diff --git a/ivy/functional/ivy/experimental/linear_algebra.py b/ivy/functional/ivy/experimental/linear_algebra.py
--- a/ivy/functional/ivy/experimental/linear_algebra.py
+++ b/ivy/functional/ivy/experimental/linear_algebra.py
@@ -1666,3 +1666,90 @@ def dot(
     ivy.array([[-15.28]])
     """
     return current_backend(a, b).dot(a, b, out=out)
+
+
+@handle_exceptions
+@handle_nestable
+@handle_array_like_without_promotion
+@inputs_to_ivy_arrays
+@handle_array_function
+@handle_device_shifting
+def general_inner_product(
+    a: Union[ivy.Array, ivy.NativeArray],
+    b: Union[ivy.Array, ivy.NativeArray],
+    n_modes: Optional[int] = None,
+    /,
+    *,
+    out: Optional[ivy.Array] = None,
+) -> ivy.Array:
+    """
+    Generalised inner products between tensors.
+
+        Takes the inner product between the last (respectively first)
+        `n_modes` of `a` (respectively `b`)
+
+    Parameters
+    ----------
+    a
+        first input tensor.
+    b
+        second input tensor.
+    n_modes
+        int, default is None. If None, the traditional inner product is returned
+        (i.e. a float) otherwise, the product between the `n_modes` last modes of
+        `a` and the `n_modes` first modes of `b` is returned. The resulting tensor's
+        order is `len(a) - n_modes`.
+    out
+        Optional output array. If provided, the output array to store the result.
+
+    Returns
+    -------
+        The inner product of the input arrays.
+
+    Examples
+    --------
+    With :class:`ivy.Array` inputs:
+
+    >>> a = ivy.array([1, 2, 3])
+    >>> b = ivy.array([4, 5, 6])
+    >>> result = ivy.general_inner_product(a, b, n_modes=1)
+    >>> print(result)
+    ivy.array(32)
+
+    >>> a = ivy.array([1, 2])
+    >>> b = ivy.array([4, 5])
+    >>> result = ivy.general_inner_product(a, b)
+    >>> print(result)
+    ivy.array(14)
+
+    >>> a = ivy.array([[1, 1], [1, 1]])
+    >>> b = ivy.array([[1, 2, 3, 4],[1, 1, 1, 1]])
+    >>> result = ivy.general_inner_product(a, b, n_modes=1)
+    >>> print(result)
+    ivy.array([[2, 3, 4, 5],
+       [2, 3, 4, 5]])
+    """
+    shape_a = a.shape
+    shape_b = b.shape
+    if n_modes is None:
+        if shape_a != shape_b:
+            raise ValueError(
+                "Taking a generalised product between two tensors without specifying"
+                " common modes is equivalent to taking inner product.This requires"
+                f" a.shape == b.shape.However, got shapes {a.shape} and {b.shape}"
+            )
+        return ivy.sum(ivy.multiply(a, b), out=out)
+
+    common_modes = shape_a[len(shape_a) - n_modes :]
+    if common_modes != shape_b[:n_modes]:
+        raise ValueError(
+            f"Incorrect shapes for inner product along {n_modes} common modes."
+            f"Shapes {shape_a.shape} and {shape_b.shape}"
+        )
+
+    common_size = int(ivy.prod(common_modes)) if len(common_modes) != 0 else 0
+    output_shape = shape_a[:-n_modes] + shape_b[n_modes:]
+    inner_product = ivy.dot(
+        ivy.reshape(a, (-1, common_size)), ivy.reshape(b, (common_size, -1))
+    )
+    return ivy.reshape(inner_product, output_shape, out=out)


Test patch:
diff --git a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_linalg.py b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_linalg.py
--- a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_linalg.py
+++ b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_linalg.py
@@ -233,6 +233,25 @@ def _generate_eigh_tridiagonal_args(draw):
     return dtype, alpha, beta, eigvals_only, select, select_range, tol
 
 
+@st.composite
+def _generate_general_inner_product_args(draw):
+    dim = draw(st.integers(min_value=1, max_value=3))
+    x_dtype, x = draw(
+        helpers.dtype_and_values(
+            available_dtypes=helpers.get_dtypes("valid"),
+            shape=(dim, dim),
+            min_value=1,
+            max_value=10.0,
+            num_arrays=2,
+            shared_dtype=True,
+            allow_nan=False,
+        )
+    )
+    max_value = dim - 1 if dim > 1 else dim
+    n_modes = draw(st.integers(min_value=1, max_value=max_value) | st.just(None))
+    return x_dtype, x, n_modes
+
+
 # multi_dot
 @st.composite
 def _generate_multi_dot_dtype_and_arrays(draw):
@@ -995,6 +1014,26 @@ def test_eigvals(dtype_x, test_flags, backend_fw, fn_name):
     )
 
 
+@handle_test(
+    fn_tree="functional.ivy.experimental.general_inner_product",
+    data=_generate_general_inner_product_args(),
+)
+def test_general_inner_product(*, data, test_flags, backend_fw, fn_name, on_device):
+    input_dtypes, x, n_modes = data
+    helpers.test_function(
+        backend_to_test=backend_fw,
+        test_flags=test_flags,
+        fn_name=fn_name,
+        on_device=on_device,
+        rtol_=1e-1,
+        atol_=1e-1,
+        input_dtypes=input_dtypes,
+        a=x[0],
+        b=x[1],
+        n_modes=n_modes,
+    )
+
+
 @handle_test(
     fn_tree="functional.ivy.experimental.initialize_tucker",
     data=_initialize_tucker_data(),