Repository: ivy-llc/ivy

Issue:
diagonal_scatter



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/torch/indexing_slicing_joining_mutating_ops.py b/ivy/functional/frontends/torch/indexing_slicing_joining_mutating_ops.py
--- a/ivy/functional/frontends/torch/indexing_slicing_joining_mutating_ops.py
+++ b/ivy/functional/frontends/torch/indexing_slicing_joining_mutating_ops.py
@@ -6,6 +6,7 @@
     numpy_to_torch_style_args,
     to_ivy_shape,
 )
+from ivy.func_wrapper import with_unsupported_dtypes
 
 
 @to_ivy_arrays_and_back
@@ -70,6 +71,35 @@ def conj(input):
     return ivy.conj(input)
 
 
+# diagonal_scatter
+@with_unsupported_dtypes(
+    {
+        "2.0.1 and below": (
+            "bfloat16",
+            "float16",
+        )
+    },
+    "torch",
+)
+@to_ivy_arrays_and_back
+def diagonal_scatter(input, src, offset=0, dim1=0, dim2=1):
+    input = ivy.copy_array(input)
+    input_shape = input.shape
+    indices = ivy.arange(0, input.size)
+    diagonal_indices = ivy.diagonal(
+        indices.reshape(input.shape), offset=offset, axis1=dim1, axis2=dim2
+    )
+    if not (src.shape == diagonal_indices.shape):
+        raise ivy.utils.exceptions.IvyException(
+            "src must have shape equal to specified diagonal of input. src size ="
+            f" {src.shape}, diagonal size = {diagonal_indices.shape}"
+        )
+    input = input.reshape((-1,))
+    input[diagonal_indices.reshape((-1,))] = src.reshape((-1,))
+    input = input.reshape(input_shape)
+    return input
+
+
 @to_ivy_arrays_and_back
 def dsplit(input, indices_or_sections, /):
     if isinstance(indices_or_sections, (list, tuple, ivy.Array)):


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_indexing_slicing_joining_mutating_ops.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_indexing_slicing_joining_mutating_ops.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_indexing_slicing_joining_mutating_ops.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_indexing_slicing_joining_mutating_ops.py
@@ -1,4 +1,6 @@
 # global
+import random
+
 from hypothesis import strategies as st
 import math
 
@@ -6,6 +8,7 @@
 # local
 import ivy
 import ivy_tests.test_ivy.helpers as helpers
+import ivy_tests.test_ivy.helpers.globals as test_globals
 from ivy_tests.test_ivy.helpers import handle_frontend_test
 from ivy_tests.test_ivy.test_functional.test_core.test_manipulation import _get_splits
 from ivy_tests.test_ivy.test_functional.test_core.test_searching import (
@@ -218,6 +221,63 @@ def _chunk_helper(draw):
     return dtype, x, axis, chunks
 
 
+# diagonal_scatter
+@st.composite
+def _diag_x_y_offset_axes(draw):
+    currentshape = random.randint(2, 4)
+
+    if test_globals.CURRENT_BACKEND == "paddle":
+        currentshape = 2
+
+    offset = draw(
+        helpers.ints(min_value=-(currentshape - 1), max_value=currentshape - 1)
+    )
+    available_input_types = draw(helpers.get_dtypes("float"))
+
+    available_input_types = helpers.array_dtypes(available_dtypes=available_input_types)
+
+    dtype, x = draw(
+        helpers.dtype_and_values(
+            min_num_dims=currentshape,
+            max_num_dims=currentshape,
+            min_dim_size=currentshape,
+            max_dim_size=currentshape,
+            num_arrays=1,
+            available_dtypes=available_input_types,
+        ),
+    )
+
+    diagonal_shape = draw(
+        helpers.get_shape(
+            min_num_dims=currentshape - 1,
+            max_num_dims=currentshape - 1,
+            min_dim_size=currentshape,
+            max_dim_size=currentshape,
+        ),
+    )
+    diagonal_shape = diagonal_shape[:-1] + (diagonal_shape[-1] - abs(offset),)
+    y = draw(
+        helpers.array_values(
+            shape=diagonal_shape,
+            dtype=available_input_types,
+            exclude_min=False,
+        )
+    )
+
+    prohibited_pairs = {(2, -1), (-2, 1), (1, -2), (-1, 2)}
+
+    axes = draw(
+        st.lists(
+            helpers.ints(min_value=-2, max_value=1), min_size=2, max_size=2, unique=True
+        ).filter(
+            lambda axes: (axes[0] % 2 != axes[1] % 2)
+            and tuple(axes) not in prohibited_pairs,
+        )
+    )
+
+    return dtype, x, y, offset, axes
+
+
 @st.composite
 def _dtype_input_dim_start_length(draw):
     _shape = draw(helpers.get_shape(min_num_dims=1, min_dim_size=1))
@@ -495,6 +555,35 @@ def test_torch_conj(
     )
 
 
+@handle_frontend_test(
+    fn_tree="torch.diagonal_scatter", dtype_and_values=_diag_x_y_offset_axes()
+)
+def test_torch_diagonal_scatter(
+    *,
+    dtype_and_values,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+    backend_fw,
+):
+    input_dtype, value, src, offset, axes = dtype_and_values
+
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        input=value[0],
+        src=src,
+        offset=offset,
+        dim1=axes[0],
+        dim2=axes[1],
+    )
+
+
 # dsplit
 @handle_frontend_test(
     fn_tree="torch.dsplit",