Repository: ivy-llc/ivy

Issue:
index_reduce



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
@@ -1,5 +1,6 @@
 # local
 import ivy
+from ivy.func_wrapper import with_unsupported_dtypes
 from ivy.functional.frontends.torch.func_wrapper import (
     to_ivy_arrays_and_back,
     numpy_to_torch_style_args,
@@ -168,6 +169,105 @@ def index_copy(input, dim, index, source, *, out=None):
     return ivy.swapaxes(res, 0, dim, out=out)
 
 
+@with_unsupported_dtypes(
+    {
+        "2.0.1 and below": (
+            "uint16",
+            "uint32",
+            "uint64",
+            "bfloat16",
+            "complex128",
+            "complex64",
+        )
+    },
+    "torch",
+)
+@to_ivy_arrays_and_back
+def index_reduce(input, dim, index, source, reduce, *, include_self=True, out=None):
+    result = ivy.copy_array(input)
+    counts = (
+        ivy.ones_like(result, dtype=result.dtype)
+        if include_self
+        else ivy.zeros_like(result, dtype=result.dtype)
+    )
+
+    index = index.astype(ivy.int64)
+
+    def init_val(reduce):
+        if reduce == "prod":
+            return 1
+        elif reduce == "amax":
+            return -ivy.inf
+        elif reduce == "amin":
+            return ivy.inf
+        else:
+            return 0
+
+    if not include_self:
+        result[index, ...] = init_val(reduce)
+
+    numel = index.size
+    index_contig = ivy.copy_array(index)
+
+    def update_counts(reduce, counts, dim, input_index):
+        if reduce == "mean":
+            counts_slice = [slice(None)] * counts.ndim
+            counts_slice[dim] = input_index
+            counts[tuple(counts_slice)] += 1
+        return counts
+
+    def update_result(result, reduce, input_data, source_data):
+        if reduce == "prod":
+            return input_data * source_data
+        elif reduce == "amin":
+            return ivy.minimum(input_data, source_data)
+        elif reduce == "amax":
+            return ivy.maximum(input_data, source_data)
+        else:
+            return input_data + source_data
+
+    if result.ndim > 1:
+        for i in range(numel):
+            input_index = index_contig[i]
+            if not (0 <= input_index < result.shape[dim]):
+                raise IndexError("Index out of range in self")
+
+            input_data = ivy.gather(result, [input_index], axis=dim)
+            source_data = ivy.gather(source, [i], axis=dim)
+
+            result_slice = [slice(None)] * result.ndim
+            result_slice[dim] = input_index
+
+            update_data = update_result(result, reduce, input_data, source_data)
+            slide_shape = result[tuple(result_slice)].shape
+            result[tuple(result_slice)] = ivy.reshape(update_data, slide_shape)
+
+            counts = update_counts(reduce, counts, dim, input_index)
+
+    elif result.ndim == 1:
+        for i in range(numel):
+            input_index = index_contig[i]
+            if not (0 <= input_index < result.size):
+                raise IndexError("Index out of range in self")
+
+            input_data = ivy.flatten(result)[input_index]
+            source_data = ivy.flatten(source)[i]
+
+            result[input_index] = update_result(result, reduce, input_data, source_data)
+            counts[input_index] += 1
+
+    if reduce == "mean":
+        if ivy.any(counts == ivy.array(0)):
+            counts[counts == ivy.array(0)] = ivy.array(1)
+        result /= counts
+
+        if not input.is_float_dtype():
+            result = ivy.floor(result)
+            result = result.astype(input.dtype)
+
+    return result
+
+
 @to_ivy_arrays_and_back
 def index_select(input, dim, index, *, out=None):
     return ivy.gather(input, index, axis=dim, out=out)


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
@@ -4,6 +4,7 @@
 
 
 # local
+import ivy
 import ivy_tests.test_ivy.helpers as helpers
 from ivy_tests.test_ivy.helpers import handle_frontend_test
 from ivy_tests.test_ivy.test_functional.test_core.test_manipulation import _get_splits
@@ -95,6 +96,72 @@ def _arrays_dim_idx_n_dtypes(draw):
     return xs, input_dtypes, _dim, _idx
 
 
+@st.composite
+def _arrays_dim_idx_n_dtypes_extend(
+    draw, support_dtypes="numeric", unsupport_dtypes=()
+):
+    num_dims = draw(st.shared(helpers.ints(min_value=1, max_value=4), key="num_dims"))
+    num_arrays = 2
+    common_shape = draw(
+        helpers.lists(
+            x=helpers.ints(min_value=2, max_value=3),
+            min_size=num_dims - 1,
+            max_size=num_dims - 1,
+        )
+    )
+    _dim = draw(helpers.ints(min_value=0, max_value=num_dims - 1))
+    unique_dims = draw(
+        helpers.lists(
+            x=helpers.ints(min_value=2, max_value=3),
+            min_size=num_arrays,
+            max_size=num_arrays,
+        )
+    )
+
+    min_dim = min(unique_dims)
+    max_dim = max(unique_dims)
+    _idx = draw(
+        helpers.array_values(
+            shape=min_dim,
+            dtype="int64",
+            min_value=0,
+            max_value=max_dim,
+            exclude_min=False,
+        )
+    )
+
+    xs = list()
+    available_input_types = draw(helpers.get_dtypes(support_dtypes))
+
+    unstabled_dtypes = ["float16"]
+    available_input_types = [
+        dtype for dtype in available_input_types if dtype not in unstabled_dtypes
+    ]
+    available_input_types = [
+        dtype for dtype in available_input_types if dtype not in unsupport_dtypes
+    ]
+
+    input_dtypes = draw(
+        helpers.array_dtypes(
+            available_dtypes=available_input_types,
+            num_arrays=num_arrays,
+            shared_dtype=True,
+        )
+    )
+    for ud, dt in zip(unique_dims, input_dtypes):
+        x = draw(
+            helpers.array_values(
+                shape=common_shape[:_dim] + [ud] + common_shape[_dim:],
+                dtype=dt,
+                large_abs_safety_factor=2.5,
+                small_abs_safety_factor=2.5,
+                safety_factor_scale="log",
+            )
+        )
+        xs.append(x)
+    return xs, input_dtypes, _dim, _idx
+
+
 # noinspection DuplicatedCode
 @st.composite
 def _arrays_idx_n_dtypes(draw):
@@ -634,6 +701,49 @@ def test_torch_index_copy(
     )
 
 
+# index_reduce
+@handle_frontend_test(
+    fn_tree="torch.index_reduce",
+    xs_dtypes_dim_idx=_arrays_dim_idx_n_dtypes_extend(
+        support_dtypes="numeric",
+        unsupport_dtypes=ivy.function_unsupported_dtypes(
+            ivy.functional.frontends.torch.index_reduce
+        ),
+    ),
+    reduce=st.sampled_from(["prod", "mean", "amin", "amax"]),
+)
+def test_torch_index_reduce(
+    *,
+    xs_dtypes_dim_idx,
+    reduce,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+    backend_fw,
+):
+    xs, input_dtypes, axis, indices = xs_dtypes_dim_idx
+    if xs[0].shape[axis] < xs[1].shape[axis]:
+        source, input = xs
+    else:
+        input, source = xs
+
+    helpers.test_frontend_function(
+        input_dtypes=[input_dtypes[0], "int64", input_dtypes[1]],
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        rtol=1e-03,
+        input=input,
+        dim=axis,
+        index=indices,
+        source=source,
+        reduce=reduce,
+    )
+
+
 # index_select
 @handle_frontend_test(
     fn_tree="torch.index_select",