From eb4ba51e8cbf9a8f763ee375797298936660e471 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Thu, 12 Oct 2023 11:12:58 +0530 Subject: [PATCH 01/24] ivy.index_add issue #26801 --- ivy/functional/backends/jax/manipulation.py | 46 ++++++++++ .../test_core/test_manipulation.py | 92 ++++++++++++++++++- 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/ivy/functional/backends/jax/manipulation.py b/ivy/functional/backends/jax/manipulation.py index 7fe7b070bb8d3..b3c6f89fd2d51 100644 --- a/ivy/functional/backends/jax/manipulation.py +++ b/ivy/functional/backends/jax/manipulation.py @@ -272,3 +272,49 @@ def swapaxes( out: Optional[JaxArray] = None, ) -> JaxArray: return jnp.swapaxes(x, axis0, axis1) + + +def index_add( + x: JaxArray, + index: JaxArray, + axis: int, + value: JaxArray, + /, + *, + name: Optional[str] = None, +) -> JaxArray: + x = jnp.swapaxes(x, axis, 0) + value = jnp.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(jnp.zeros_like(value[0])) + _to_add_cum = value[index[0][1]] + while len(index) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + value[index.pop(1)[1]] + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(jnp.zeros_like(value[0])) + _to_adds = jnp.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = jnp.flatten(_to_adds) + + ret = jnp.add(x, _to_adds) + ret = jnp.swapaxes(ret, axis, 0) + return ret + + +# x1 = jnp.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) +# index1 = jnp.array([0, 2]) +# value1 = jnp.array([[1, 1, 1], [1, 1, 1]]) +# +# print(x1) +# print(index1) +# print(value1) +# +# ret1 = index_add(x1, index1, 0, value1) +# print(ret1) diff --git a/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py b/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py index 8717e429ea643..563337b8528c0 100644 --- a/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py +++ b/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py @@ -14,6 +14,62 @@ # --------------- # +@st.composite +def _arrays_dim_idx_n_dtypes(draw): + 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 = [] + # available_input_types = draw(helpers.get_dtypes("integer")) + available_input_types = ["int32", "int64", "float16", "float32", "float64"] + 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 + + @st.composite def _arrays_idx_n_dtypes(draw): num_dims = draw(st.shared(helpers.ints(min_value=1, max_value=4), key="num_dims")) @@ -394,6 +450,38 @@ def test_flip(*, dtype_value, axis, test_flags, backend_fw, fn_name, on_device): ) +@handle_test( + fn_tree="functional.ivy.index_add", + xs_dtypes_dim_idx=_arrays_dim_idx_n_dtypes(), +) +def test_paddle_index_add( + *, + xs_dtypes_dim_idx, + 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, + backend_to_test=backend_fw, + test_flags=test_flags, + fn_tree=fn_tree, + frontend=frontend, + on_device=on_device, + x=input, + index=indices, + axis=axis, + value=source, + ) + + # permute_dims @handle_test( fn_tree="functional.ivy.permute_dims", @@ -485,7 +573,7 @@ def test_reshape( test_flags, backend_fw, fn_name, - on_device + on_device, ): dtype, value = dtype_value @@ -591,7 +679,7 @@ def test_split( test_flags, backend_fw, fn_name, - on_device + on_device, ): dtype, value = dtype_value if ( From 20c056525ad3d040ea37eca882a79fbdddd4dcbc Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 13 Oct 2023 09:41:39 +0530 Subject: [PATCH 02/24] 3. jax -> exp --- .../backends/jax/experimental/manipulation.py | 34 +++++++ ivy/functional/backends/jax/manipulation.py | 46 ---------- .../test_core/test_manipulation.py | 88 ------------------- 3 files changed, 34 insertions(+), 134 deletions(-) diff --git a/ivy/functional/backends/jax/experimental/manipulation.py b/ivy/functional/backends/jax/experimental/manipulation.py index 43adef47e5dbe..0a6f3daf890cf 100644 --- a/ivy/functional/backends/jax/experimental/manipulation.py +++ b/ivy/functional/backends/jax/experimental/manipulation.py @@ -416,3 +416,37 @@ def fill_diagonal( def trim_zeros(a: JaxArray, /, *, trim: Optional[str] = "bf") -> JaxArray: return jnp.trim_zeros(a, trim=trim) + + +def index_add( + x: JaxArray, + index: JaxArray, + axis: int, + value: JaxArray, + /, + *, + name: Optional[str] = None, +) -> JaxArray: + x = jnp.swapaxes(x, axis, 0) + value = jnp.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(jnp.zeros_like(value[0])) + _to_add_cum = value[index[0][1]] + while len(index) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + value[index.pop(1)[1]] + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(jnp.zeros_like(value[0])) + _to_adds = jnp.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = jnp.flatten(_to_adds) + + ret = jnp.add(x, _to_adds) + ret = jnp.swapaxes(ret, axis, 0) + return ret diff --git a/ivy/functional/backends/jax/manipulation.py b/ivy/functional/backends/jax/manipulation.py index b3c6f89fd2d51..7fe7b070bb8d3 100644 --- a/ivy/functional/backends/jax/manipulation.py +++ b/ivy/functional/backends/jax/manipulation.py @@ -272,49 +272,3 @@ def swapaxes( out: Optional[JaxArray] = None, ) -> JaxArray: return jnp.swapaxes(x, axis0, axis1) - - -def index_add( - x: JaxArray, - index: JaxArray, - axis: int, - value: JaxArray, - /, - *, - name: Optional[str] = None, -) -> JaxArray: - x = jnp.swapaxes(x, axis, 0) - value = jnp.swapaxes(value, axis, 0) - _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) - while index: - _curr_idx = index[0][0] - while len(_to_adds) < _curr_idx: - _to_adds.append(jnp.zeros_like(value[0])) - _to_add_cum = value[index[0][1]] - while len(index) > 1 and (index[0][0] == index[1][0]): - _to_add_cum = _to_add_cum + value[index.pop(1)[1]] - index.pop(0) - _to_adds.append(_to_add_cum) - while len(_to_adds) < x.shape[0]: - _to_adds.append(jnp.zeros_like(value[0])) - _to_adds = jnp.stack(_to_adds) - if len(x.shape) < 2: - # Added this line due to the paddle backend treating scalars as 1-d arrays - _to_adds = jnp.flatten(_to_adds) - - ret = jnp.add(x, _to_adds) - ret = jnp.swapaxes(ret, axis, 0) - return ret - - -# x1 = jnp.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) -# index1 = jnp.array([0, 2]) -# value1 = jnp.array([[1, 1, 1], [1, 1, 1]]) -# -# print(x1) -# print(index1) -# print(value1) -# -# ret1 = index_add(x1, index1, 0, value1) -# print(ret1) diff --git a/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py b/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py index b0b24cc7f4be4..9d3f57a75865b 100644 --- a/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py +++ b/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py @@ -14,62 +14,6 @@ # --------------- # -@st.composite -def _arrays_dim_idx_n_dtypes(draw): - 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 = [] - # available_input_types = draw(helpers.get_dtypes("integer")) - available_input_types = ["int32", "int64", "float16", "float32", "float64"] - 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 - - @st.composite def _arrays_idx_n_dtypes(draw): num_dims = draw(st.shared(helpers.ints(min_value=1, max_value=4), key="num_dims")) @@ -451,38 +395,6 @@ def test_flip(*, dtype_value, axis, test_flags, backend_fw, fn_name, on_device): ) -@handle_test( - fn_tree="functional.ivy.index_add", - xs_dtypes_dim_idx=_arrays_dim_idx_n_dtypes(), -) -def test_paddle_index_add( - *, - xs_dtypes_dim_idx, - 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, - backend_to_test=backend_fw, - test_flags=test_flags, - fn_tree=fn_tree, - frontend=frontend, - on_device=on_device, - x=input, - index=indices, - axis=axis, - value=source, - ) - - # permute_dims @handle_test( fn_tree="functional.ivy.permute_dims", From e49b8bcf0f7573f6c2c510d89a7531d4f185312f Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 13 Oct 2023 09:44:52 +0530 Subject: [PATCH 03/24] 3. jax -> exp --- .../test_ivy/test_functional/test_core/test_manipulation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py b/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py index 9d3f57a75865b..a640051661360 100644 --- a/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py +++ b/ivy_tests/test_ivy/test_functional/test_core/test_manipulation.py @@ -486,7 +486,7 @@ def test_reshape( test_flags, backend_fw, fn_name, - on_device, + on_device ): dtype, value = dtype_value @@ -592,7 +592,7 @@ def test_split( test_flags, backend_fw, fn_name, - on_device, + on_device ): dtype, value = dtype_value if ( From 005e825237d51f38066422488d9d9c4ecca884cd Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 13 Oct 2023 15:17:55 +0530 Subject: [PATCH 04/24] 5. np -> exp --- .../numpy/experimental/manipulation.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/ivy/functional/backends/numpy/experimental/manipulation.py b/ivy/functional/backends/numpy/experimental/manipulation.py index 37adf9bed7f0d..90873f80f07a4 100644 --- a/ivy/functional/backends/numpy/experimental/manipulation.py +++ b/ivy/functional/backends/numpy/experimental/manipulation.py @@ -518,3 +518,37 @@ def put_along_axis( put_along_axis.partial_mixed_handler = lambda *args, mode=None, **kwargs: mode in [ "replace", ] + + +def index_add( + x: np.ndarray, + index: np.ndarray, + axis: int, + value: np.ndarray, + /, + *, + name: Optional[str] = None, +) -> np.ndarray: + x = np.swapaxes(x, axis, 0) + value = np.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(np.zeros_like(value[0])) + _to_add_cum = value[index[0][1]] + while len(index) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + value[index.pop(1)[1]] + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(np.zeros_like(value[0])) + _to_adds = np.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = _to_adds.flatten() + + ret = np.add(x, _to_adds) + ret = np.swapaxes(ret, axis, 0) + return ret From 1e634336a64fa7a91cfad4a01d9ee9ce94fc73ed Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 13 Oct 2023 18:49:15 +0530 Subject: [PATCH 05/24] 6. paddle -> exp --- .../backends/paddle/experimental/manipulation.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/ivy/functional/backends/paddle/experimental/manipulation.py b/ivy/functional/backends/paddle/experimental/manipulation.py index d81ee8852518e..eb43dfffe280c 100644 --- a/ivy/functional/backends/paddle/experimental/manipulation.py +++ b/ivy/functional/backends/paddle/experimental/manipulation.py @@ -742,3 +742,15 @@ def put_along_axis( "sum", "mul", ] + + +def index_add( + x: paddle.Tensor, + index: paddle.Tensor, + axis: int, + value: paddle.Tensor, + /, + *, + name: Optional[str] = None, +) -> paddle.Tensor: + return paddle.index_add(x, index, axis, value) From 649a6505f66081dd7f571450410cbdfb83d848d5 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 13 Oct 2023 19:05:16 +0530 Subject: [PATCH 06/24] 7. tf -> exp --- .../tensorflow/experimental/manipulation.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index 0319733441a97..b4469616f857a 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -431,3 +431,37 @@ def trim_zeros(a: tf.Tensor, /, *, trim: Optional[str] = "bf") -> tf.Tensor: last = tf.minimum(last, tf.cast(tf.shape(a)[0], tf.int64)) return a[first:last] + + +def index_add( + x: tf.Tensor, + index: tf.Tensor, + axis: int, + value: tf.Tensor, + /, + *, + name: Optional[str] = None, +) -> tf.Tensor: + x = tf.experimental.numpy.swapaxes(x, axis, 0) + value = tf.experimental.numpy.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(tf.zeros_like(value[0])) + _to_add_cum = value[index[0][1]] + while len(index) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + value[index.pop(1)[1]] + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(tf.zeros_like(value[0])) + _to_adds = tf.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = tf.flatten(_to_adds) + + ret = tf.add(x, _to_adds) + ret = tf.experimental.numpy.swapaxes(ret, axis, 0) + return ret From 473d31bcffe35f37391501f9f25468aff04f54ee Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 13 Oct 2023 19:10:15 +0530 Subject: [PATCH 07/24] 8. torch -> exp --- .../torch/experimental/manipulation.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/ivy/functional/backends/torch/experimental/manipulation.py b/ivy/functional/backends/torch/experimental/manipulation.py index faaf3f04c6ee5..6096aa21d63af 100644 --- a/ivy/functional/backends/torch/experimental/manipulation.py +++ b/ivy/functional/backends/torch/experimental/manipulation.py @@ -493,3 +493,37 @@ def trim_zeros(a: torch.Tensor, /, *, trim: Optional[str] = "bf") -> torch.Tenso else: last = last - 1 return a[first:last] + + +def index_add( + x: torch.Tensor, + index: torch.Tensor, + axis: int, + value: torch.Tensor, + /, + *, + name: Optional[str] = None, +) -> torch.Tensor: + x = torch.swapaxes(x, axis, 0) + value = torch.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(torch.zeros_like(value[0])) + _to_add_cum = value[index[0][1]] + while len(index) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + value[index.pop(1)[1]] + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(torch.zeros_like(value[0])) + _to_adds = torch.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = torch.flatten(_to_adds) + + ret = torch.add(x, _to_adds) + ret = torch.swapaxes(ret, axis, 0) + return ret From f6e454038f035fc965fe00c90a5b3b4d3a049543 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Sun, 15 Oct 2023 08:24:52 +0530 Subject: [PATCH 08/24] 4. mxnet -> exp --- .../mxnet/experimental/manipulation.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/ivy/functional/backends/mxnet/experimental/manipulation.py b/ivy/functional/backends/mxnet/experimental/manipulation.py index 66b46e813ae47..9987808b61cb0 100644 --- a/ivy/functional/backends/mxnet/experimental/manipulation.py +++ b/ivy/functional/backends/mxnet/experimental/manipulation.py @@ -192,3 +192,39 @@ def concat_from_sequence( out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: raise IvyNotImplementedException() + + +def index_add( + x: Union[(None, mx.ndarray.NDArray)], + index: Union[(None, mx.ndarray.NDArray)], + axis: int, + value: Union[(None, mx.ndarray.NDArray)], + /, + *, + name: Optional[str] = None, +) -> Union[(None, mx.ndarray.NDArray)]: + x = mx.nd.swapaxes(x, axis, 0) + value = mx.nd.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted( + zip(index.asnumpy().tolist(), range(len(index))), key=(lambda i: i[0]) + ) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(mx.nd.zeros_like(value[0])) + _to_add_cum = value[index[0][1]] + while len(index) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + value[index.pop(1)[1]] + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(mx.nd.zeros_like(value[0])) + _to_adds = mx.nd.stack(*_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = mx.nd.flatten(_to_adds) + + ret = mx.nd.add(x, _to_adds) + ret = mx.nd.swapaxes(ret, axis, 0) + return ret From e08573c117028ed9c4e0341c2d7e365e716855b7 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Sun, 15 Oct 2023 15:25:28 +0530 Subject: [PATCH 09/24] 9. ivy -> exp --- .../ivy/experimental/manipulation.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index 644a20391d123..11adee3a52b4b 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -2761,3 +2761,74 @@ def trim_zeros( ), "to_skip": ("inputs_to_ivy_arrays",), } + + +@handle_nestable +@handle_exceptions +@handle_array_like_without_promotion +@inputs_to_ivy_arrays +@handle_array_function +@handle_device +def index_add( + x: ivy.Array, + index: ivy.Array, + axis: int, + value: ivy.Array, + /, + *, + name: Optional[str] = None, +) -> ivy.Array: + """Add the elements of the input tensor with value tensor by selecting the + indices in the order given in index. + + Parameters + ---------- + x : Array + The Destination Array. + index : Array + The 1-D array containing the indices to index. + axis : int + The dimension in which we index. + value : Array + The tensor used to add the elements along the target axis. + name : str, optional + Output array where the output is to be stored. Default value is 'none'. + + Returns + ------- + Array + Same dimention and dtype with x. + + Examples + -------- + >>> x1 = ivy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) + >>> index1 = ivy.array([0, 2]) + >>> value1 = ivy.array([[1, 1, 1], [1, 1, 1]]) + >>> ret1 = ivy.index_add(x1, index1, 0, value1) + array([[2, 2, 2], + [1, 1, 1], + [2, 2, 2]]) + """ + x = ivy.swapaxes(x, axis, 0) + value = ivy.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(ivy.zeros_like(value[0])) + _to_add_cum = ivy.get_item(value, index[0][1]) + while (len(index)) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + ivy.get_item(value, index.pop(1)[1]) + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(ivy.zeros_like(value[0])) + _to_adds = ivy.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = ivy.flatten(_to_adds) + + ret = ivy.add(x, _to_adds) + ret = ivy.swapaxes(ret, axis, 0) + return ret From 23583c1518587fe178db6501dac7481679c99bba Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Mon, 16 Oct 2023 12:10:03 +0530 Subject: [PATCH 10/24] fix --- .../array/experimental/manipulation.py | 65 ++++++++++++++ .../container/experimental/manipulation.py | 64 ++++++++++++++ .../tensorflow/experimental/manipulation.py | 2 +- .../test_core/test_manipulation.py | 86 +++++++++++++++++++ 4 files changed, 216 insertions(+), 1 deletion(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index bee2715195c48..03052cf679b8a 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1396,3 +1396,68 @@ def put_along_axis( changes. """ return ivy.put_along_axis(self._data, indices, values, axis, mode=mode, out=out) + + @handle_view + def index_add( + x: ivy.Array, + index: ivy.Array, + axis: int, + value: ivy.Array, + /, + *, + name: Optional[str] = None, + ) -> ivy.Array: + """Add the elements of the input tensor with value tensor by selecting + the indices in the order given in index. + + Parameters + ---------- + x : Array + The Destination Array. + index : Array + The 1-D array containing the indices to index. + axis : int + The dimension in which we index. + value : Array + The tensor used to add the elements along the target axis. + name : str, optional + Output array where the output is to be stored. Default value is 'none'. + + Returns + ------- + Array + Same dimention and dtype with x. + + Examples + -------- + >>> x1 = ivy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) + >>> index1 = ivy.array([0, 2]) + >>> value1 = ivy.array([[1, 1, 1], [1, 1, 1]]) + >>> ret1 = ivy.index_add(x1, index1, 0, value1) + array([[2, 2, 2], + [1, 1, 1], + [2, 2, 2]]) + """ + x = ivy.swapaxes(x, axis, 0) + value = ivy.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(ivy.zeros_like(value[0])) + _to_add_cum = ivy.get_item(value, index[0][1]) + while (len(index)) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + ivy.get_item(value, index.pop(1)[1]) + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(ivy.zeros_like(value[0])) + _to_adds = ivy.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = ivy.flatten(_to_adds) + + ret = ivy.add(x, _to_adds) + ret = ivy.swapaxes(ret, axis, 0) + return ret diff --git a/ivy/data_classes/container/experimental/manipulation.py b/ivy/data_classes/container/experimental/manipulation.py index 88249d8efd61f..7cc6888fcf0a3 100644 --- a/ivy/data_classes/container/experimental/manipulation.py +++ b/ivy/data_classes/container/experimental/manipulation.py @@ -3865,6 +3865,70 @@ def trim_zeros( """ return self._static_trim_zeros(self, trim=trim) + def index_add( + x: ivy.Array, + index: ivy.Array, + axis: int, + value: ivy.Array, + /, + *, + name: Optional[str] = None, + ) -> ivy.Array: + """Add the elements of the input tensor with value tensor by selecting + the indices in the order given in index. + + Parameters + ---------- + x : Array + The Destination Array. + index : Array + The 1-D array containing the indices to index. + axis : int + The dimension in which we index. + value : Array + The tensor used to add the elements along the target axis. + name : str, optional + Output array where the output is to be stored. Default value is 'none'. + + Returns + ------- + Array + Same dimention and dtype with x. + + Examples + -------- + >>> x1 = ivy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) + >>> index1 = ivy.array([0, 2]) + >>> value1 = ivy.array([[1, 1, 1], [1, 1, 1]]) + >>> ret1 = ivy.index_add(x1, index1, 0, value1) + array([[2, 2, 2], + [1, 1, 1], + [2, 2, 2]]) + """ + x = ivy.swapaxes(x, axis, 0) + value = ivy.swapaxes(value, axis, 0) + _to_adds = [] + index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + while index: + _curr_idx = index[0][0] + while len(_to_adds) < _curr_idx: + _to_adds.append(ivy.zeros_like(value[0])) + _to_add_cum = ivy.get_item(value, index[0][1]) + while (len(index)) > 1 and (index[0][0] == index[1][0]): + _to_add_cum = _to_add_cum + ivy.get_item(value, index.pop(1)[1]) + index.pop(0) + _to_adds.append(_to_add_cum) + while len(_to_adds) < x.shape[0]: + _to_adds.append(ivy.zeros_like(value[0])) + _to_adds = ivy.stack(_to_adds) + if len(x.shape) < 2: + # Added this line due to the paddle backend treating scalars as 1-d arrays + _to_adds = ivy.flatten(_to_adds) + + ret = ivy.add(x, _to_adds) + ret = ivy.swapaxes(ret, axis, 0) + return ret + def concat_from_sequence( self: ivy.Container, diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index b4469616f857a..35f7a9f0bca8a 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -460,7 +460,7 @@ def index_add( _to_adds = tf.stack(_to_adds) if len(x.shape) < 2: # Added this line due to the paddle backend treating scalars as 1-d arrays - _to_adds = tf.flatten(_to_adds) + _to_adds = tf.nest.flatten(_to_adds) ret = tf.add(x, _to_adds) ret = tf.experimental.numpy.swapaxes(ret, axis, 0) diff --git a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py index 67d35c5558197..4b6487aedf793 100644 --- a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py +++ b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py @@ -15,6 +15,61 @@ # --------------- # +@st.composite +def _arrays_dim_idx_n_dtypes(draw): + 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 = [] + available_input_types = ["int32", "int64", "float32", "float64"] + 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 + + @st.composite def _as_strided_helper(draw): dtype, x = draw(helpers.dtype_and_values(min_num_dims=1, max_num_dims=5)) @@ -1028,6 +1083,37 @@ def test_i0(*, dtype_and_x, test_flags, backend_fw, fn_name, on_device): ) +@handle_test( + fn_tree="functional.ivy.experimental.index_add", + xs_dtypes_dim_idx=_arrays_dim_idx_n_dtypes(), +) +def test_index_add( + *, + xs_dtypes_dim_idx, + on_device, + fn_name, + test_flags, + backend_fw, +): + xs, input_dtypes, axis, indices = xs_dtypes_dim_idx + axis = 0 + if xs[0].shape[axis] < xs[1].shape[axis]: + source, x = xs + else: + x, source = xs + helpers.test_function( + input_dtypes=input_dtypes, + backend_to_test=backend_fw, + test_flags=test_flags, + fn_name=fn_name, + on_device=on_device, + x=x, + index=indices, + axis=axis, + value=source, + ) + + @handle_test( fn_tree="functional.ivy.experimental.matricize", data=_matricize_data(), From fd568aaab6d0e0d20f130f4fb879c24f53059854 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Tue, 17 Oct 2023 08:43:46 +0530 Subject: [PATCH 11/24] fix --- ivy/functional/backends/numpy/experimental/manipulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivy/functional/backends/numpy/experimental/manipulation.py b/ivy/functional/backends/numpy/experimental/manipulation.py index 90873f80f07a4..882097b624430 100644 --- a/ivy/functional/backends/numpy/experimental/manipulation.py +++ b/ivy/functional/backends/numpy/experimental/manipulation.py @@ -532,7 +532,7 @@ def index_add( x = np.swapaxes(x, axis, 0) value = np.swapaxes(value, axis, 0) _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + index = sorted(zip(index.tolist(), range(len(index))), key=(lambda i: i[0])) while index: _curr_idx = index[0][0] while len(_to_adds) < _curr_idx: From 1669c257cd2841f00002ac373833e586a81de455 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Tue, 17 Oct 2023 09:38:31 +0530 Subject: [PATCH 12/24] fix --- ivy/functional/backends/jax/experimental/manipulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivy/functional/backends/jax/experimental/manipulation.py b/ivy/functional/backends/jax/experimental/manipulation.py index 0a6f3daf890cf..3254390950649 100644 --- a/ivy/functional/backends/jax/experimental/manipulation.py +++ b/ivy/functional/backends/jax/experimental/manipulation.py @@ -430,7 +430,7 @@ def index_add( x = jnp.swapaxes(x, axis, 0) value = jnp.swapaxes(value, axis, 0) _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + index = sorted(zip(index.tolist(), range(len(index))), key=(lambda i: i[0])) while index: _curr_idx = index[0][0] while len(_to_adds) < _curr_idx: From ab0ca56add76717af03c155da711a3e331b602aa Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Tue, 17 Oct 2023 09:47:38 +0530 Subject: [PATCH 13/24] fix --- ivy/functional/backends/tensorflow/experimental/manipulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index 35f7a9f0bca8a..e0391b1d6e160 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -445,7 +445,7 @@ def index_add( x = tf.experimental.numpy.swapaxes(x, axis, 0) value = tf.experimental.numpy.swapaxes(value, axis, 0) _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + index = sorted(zip(index.numpy().tolist(), range(len(index))), key=(lambda i: i[0])) while index: _curr_idx = index[0][0] while len(_to_adds) < _curr_idx: From afa668cfacf3766e349d1486591b2647d4bce349 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Tue, 17 Oct 2023 09:51:51 +0530 Subject: [PATCH 14/24] fix --- ivy/functional/backends/torch/experimental/manipulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivy/functional/backends/torch/experimental/manipulation.py b/ivy/functional/backends/torch/experimental/manipulation.py index 6096aa21d63af..54b6798fb2a87 100644 --- a/ivy/functional/backends/torch/experimental/manipulation.py +++ b/ivy/functional/backends/torch/experimental/manipulation.py @@ -507,7 +507,7 @@ def index_add( x = torch.swapaxes(x, axis, 0) value = torch.swapaxes(value, axis, 0) _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) + index = sorted(zip(index.tolist(), range(len(index))), key=(lambda i: i[0])) while index: _curr_idx = index[0][0] while len(_to_adds) < _curr_idx: From 831b184b61d3708ce77c46ad14bb3582b7bd1976 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Wed, 18 Oct 2023 11:30:08 +0530 Subject: [PATCH 15/24] fix --- ivy/data_classes/array/experimental/manipulation.py | 5 ++--- ivy/data_classes/container/experimental/manipulation.py | 4 ++-- .../test_experimental/test_core/test_manipulation.py | 1 - 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index 03052cf679b8a..1317d2ad3ac36 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1397,9 +1397,8 @@ def put_along_axis( """ return ivy.put_along_axis(self._data, indices, values, axis, mode=mode, out=out) - @handle_view def index_add( - x: ivy.Array, + self: ivy.Array, index: ivy.Array, axis: int, value: ivy.Array, @@ -1438,7 +1437,7 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - x = ivy.swapaxes(x, axis, 0) + x = ivy.swapaxes(self, axis, 0) value = ivy.swapaxes(value, axis, 0) _to_adds = [] index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) diff --git a/ivy/data_classes/container/experimental/manipulation.py b/ivy/data_classes/container/experimental/manipulation.py index 7cc6888fcf0a3..111a1746ed797 100644 --- a/ivy/data_classes/container/experimental/manipulation.py +++ b/ivy/data_classes/container/experimental/manipulation.py @@ -3866,7 +3866,7 @@ def trim_zeros( return self._static_trim_zeros(self, trim=trim) def index_add( - x: ivy.Array, + self: ivy.Array, index: ivy.Array, axis: int, value: ivy.Array, @@ -3905,7 +3905,7 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - x = ivy.swapaxes(x, axis, 0) + x = ivy.swapaxes(self, axis, 0) value = ivy.swapaxes(value, axis, 0) _to_adds = [] index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) diff --git a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py index 4b6487aedf793..6b6a585ce7635 100644 --- a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py +++ b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_manipulation.py @@ -1096,7 +1096,6 @@ def test_index_add( backend_fw, ): xs, input_dtypes, axis, indices = xs_dtypes_dim_idx - axis = 0 if xs[0].shape[axis] < xs[1].shape[axis]: source, x = xs else: From 369fe246d5e4024a8d6c564fb5670b3c376e5904 Mon Sep 17 00:00:00 2001 From: ivy-branch Date: Wed, 18 Oct 2023 13:03:49 +0000 Subject: [PATCH 16/24] =?UTF-8?q?=F0=9F=A4=96=20Lint=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ivy/data_classes/array/experimental/manipulation.py | 5 +++-- ivy/data_classes/container/experimental/manipulation.py | 5 +++-- ivy/functional/ivy/experimental/manipulation.py | 5 +++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index fcf624ae91bc6..afd6f42a7521a 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1522,8 +1522,9 @@ def index_add( *, name: Optional[str] = None, ) -> ivy.Array: - """Add the elements of the input tensor with value tensor by selecting - the indices in the order given in index. + """ + Add the elements of the input tensor with value tensor by selecting the indices + in the order given in index. Parameters ---------- diff --git a/ivy/data_classes/container/experimental/manipulation.py b/ivy/data_classes/container/experimental/manipulation.py index 54232542c316f..1302a26c8c558 100644 --- a/ivy/data_classes/container/experimental/manipulation.py +++ b/ivy/data_classes/container/experimental/manipulation.py @@ -4146,8 +4146,9 @@ def index_add( *, name: Optional[str] = None, ) -> ivy.Array: - """Add the elements of the input tensor with value tensor by selecting - the indices in the order given in index. + """ + Add the elements of the input tensor with value tensor by selecting the indices + in the order given in index. Parameters ---------- diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index a313b1ceecc7a..f31768317e774 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -2923,8 +2923,9 @@ def index_add( *, name: Optional[str] = None, ) -> ivy.Array: - """Add the elements of the input tensor with value tensor by selecting the - indices in the order given in index. + """ + Add the elements of the input tensor with value tensor by selecting the indices in + the order given in index. Parameters ---------- From 600f8295c1beea104d03bcf7e6ef0d7413e65f77 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Thu, 19 Oct 2023 11:45:24 +0530 Subject: [PATCH 17/24] fix --- .../array/experimental/manipulation.py | 25 +------- .../container/experimental/manipulation.py | 59 +++++++++++-------- .../ivy/experimental/manipulation.py | 24 +------- 3 files changed, 36 insertions(+), 72 deletions(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index afd6f42a7521a..eec60233d2d65 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1521,6 +1521,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ Add the elements of the input tensor with value tensor by selecting the indices @@ -1554,26 +1555,4 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - x = ivy.swapaxes(self, axis, 0) - value = ivy.swapaxes(value, axis, 0) - _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) - while index: - _curr_idx = index[0][0] - while len(_to_adds) < _curr_idx: - _to_adds.append(ivy.zeros_like(value[0])) - _to_add_cum = ivy.get_item(value, index[0][1]) - while (len(index)) > 1 and (index[0][0] == index[1][0]): - _to_add_cum = _to_add_cum + ivy.get_item(value, index.pop(1)[1]) - index.pop(0) - _to_adds.append(_to_add_cum) - while len(_to_adds) < x.shape[0]: - _to_adds.append(ivy.zeros_like(value[0])) - _to_adds = ivy.stack(_to_adds) - if len(x.shape) < 2: - # Added this line due to the paddle backend treating scalars as 1-d arrays - _to_adds = ivy.flatten(_to_adds) - - ret = ivy.add(x, _to_adds) - ret = ivy.swapaxes(ret, axis, 0) - return ret + return ivy.index_add(self._data, index, axis, value, out=out) diff --git a/ivy/data_classes/container/experimental/manipulation.py b/ivy/data_classes/container/experimental/manipulation.py index 1302a26c8c558..e86c4e093b34f 100644 --- a/ivy/data_classes/container/experimental/manipulation.py +++ b/ivy/data_classes/container/experimental/manipulation.py @@ -4137,15 +4137,44 @@ def trim_zeros( """ return self._static_trim_zeros(self, trim=trim) - def index_add( - self: ivy.Array, + @staticmethod + def static_index_add( + x: ivy.Array, index: ivy.Array, axis: int, value: ivy.Array, /, *, name: Optional[str] = None, - ) -> ivy.Array: + 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: + return ContainerBase.cont_multi_map_in_function( + "index_add", + x, + index, + axis, + value, + key_chains=key_chains, + to_apply=to_apply, + prune_unapplied=prune_unapplied, + map_sequences=map_sequences, + out=out, + ) + + def index_add( + self: ivy.Container, + index: ivy.Container, + axis: Union[int, ivy.Container], + value: ivy.Container, + /, + *, + name: Optional[str] = None, + out: Optional[ivy.Container] = None, + ) -> ivy.Container: """ Add the elements of the input tensor with value tensor by selecting the indices in the order given in index. @@ -4178,29 +4207,7 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - x = ivy.swapaxes(self, axis, 0) - value = ivy.swapaxes(value, axis, 0) - _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) - while index: - _curr_idx = index[0][0] - while len(_to_adds) < _curr_idx: - _to_adds.append(ivy.zeros_like(value[0])) - _to_add_cum = ivy.get_item(value, index[0][1]) - while (len(index)) > 1 and (index[0][0] == index[1][0]): - _to_add_cum = _to_add_cum + ivy.get_item(value, index.pop(1)[1]) - index.pop(0) - _to_adds.append(_to_add_cum) - while len(_to_adds) < x.shape[0]: - _to_adds.append(ivy.zeros_like(value[0])) - _to_adds = ivy.stack(_to_adds) - if len(x.shape) < 2: - # Added this line due to the paddle backend treating scalars as 1-d arrays - _to_adds = ivy.flatten(_to_adds) - - ret = ivy.add(x, _to_adds) - ret = ivy.swapaxes(ret, axis, 0) - return ret + return self.static_index_add(self, index, axis, value, out=out) def concat_from_sequence( diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index f31768317e774..284a26f2cd2f1 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -2955,26 +2955,4 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - x = ivy.swapaxes(x, axis, 0) - value = ivy.swapaxes(value, axis, 0) - _to_adds = [] - index = sorted(zip(ivy.to_list(index), range(len(index))), key=(lambda i: i[0])) - while index: - _curr_idx = index[0][0] - while len(_to_adds) < _curr_idx: - _to_adds.append(ivy.zeros_like(value[0])) - _to_add_cum = ivy.get_item(value, index[0][1]) - while (len(index)) > 1 and (index[0][0] == index[1][0]): - _to_add_cum = _to_add_cum + ivy.get_item(value, index.pop(1)[1]) - index.pop(0) - _to_adds.append(_to_add_cum) - while len(_to_adds) < x.shape[0]: - _to_adds.append(ivy.zeros_like(value[0])) - _to_adds = ivy.stack(_to_adds) - if len(x.shape) < 2: - # Added this line due to the paddle backend treating scalars as 1-d arrays - _to_adds = ivy.flatten(_to_adds) - - ret = ivy.add(x, _to_adds) - ret = ivy.swapaxes(ret, axis, 0) - return ret + return ivy.current_backend().index_add(x, index, axis, value) From 6976dba183b5e1d2dbba19d5ae9077a7be8cb0bb Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Sun, 19 Nov 2023 08:28:48 +0530 Subject: [PATCH 18/24] fix --- ivy/data_classes/array/experimental/manipulation.py | 3 +-- ivy/data_classes/container/experimental/manipulation.py | 5 +---- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index 551021a68ea94..a46d2134f3dc8 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1521,7 +1521,6 @@ def index_add( /, *, name: Optional[str] = None, - out: Optional[ivy.Array] = None, ) -> ivy.Array: """ Add the elements of the input tensor with value tensor by selecting the indices @@ -1555,4 +1554,4 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - return ivy.index_add(self._data, index, axis, value, out=out) + return ivy.index_add(self._data, index, axis, value) diff --git a/ivy/data_classes/container/experimental/manipulation.py b/ivy/data_classes/container/experimental/manipulation.py index e86c4e093b34f..1e3e8055c3148 100644 --- a/ivy/data_classes/container/experimental/manipulation.py +++ b/ivy/data_classes/container/experimental/manipulation.py @@ -4150,7 +4150,6 @@ def static_index_add( 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: return ContainerBase.cont_multi_map_in_function( "index_add", @@ -4162,7 +4161,6 @@ def static_index_add( to_apply=to_apply, prune_unapplied=prune_unapplied, map_sequences=map_sequences, - out=out, ) def index_add( @@ -4173,7 +4171,6 @@ def index_add( /, *, name: Optional[str] = None, - out: Optional[ivy.Container] = None, ) -> ivy.Container: """ Add the elements of the input tensor with value tensor by selecting the indices @@ -4207,7 +4204,7 @@ def index_add( [1, 1, 1], [2, 2, 2]]) """ - return self.static_index_add(self, index, axis, value, out=out) + return self.static_index_add(self, index, axis, value) def concat_from_sequence( From dbf7d7339b93342618edbc533286cf2120a3a259 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Thu, 30 Nov 2023 21:55:33 +0530 Subject: [PATCH 19/24] fix --- .../array/experimental/manipulation.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index a46d2134f3dc8..7574677f944f8 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1521,6 +1521,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ Add the elements of the input tensor with value tensor by selecting the indices @@ -1528,7 +1529,7 @@ def index_add( Parameters ---------- - x : Array + self : Array The Destination Array. index : Array The 1-D array containing the indices to index. @@ -1538,6 +1539,8 @@ def index_add( The tensor used to add the elements along the target axis. name : str, optional Output array where the output is to be stored. Default value is 'none'. + out : Array + Output array. Returns ------- @@ -1546,12 +1549,12 @@ def index_add( Examples -------- - >>> x1 = ivy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) - >>> index1 = ivy.array([0, 2]) - >>> value1 = ivy.array([[1, 1, 1], [1, 1, 1]]) - >>> ret1 = ivy.index_add(x1, index1, 0, value1) + >>> x = ivy.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) + >>> idx = ivy.array([0, 2]) + >>> val = ivy.array([[1, 1, 1], [1, 1, 1]]) + >>> ans = ivy.index_add(x, idx, 0, val) array([[2, 2, 2], [1, 1, 1], [2, 2, 2]]) """ - return ivy.index_add(self._data, index, axis, value) + return ivy.index_add(self._data, index, axis, value, out=out) From b0f196ff3e79c2f79fdb24a394cff37af986098c Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Fri, 1 Dec 2023 20:45:07 +0530 Subject: [PATCH 20/24] fix --- ivy/functional/backends/jax/experimental/manipulation.py | 1 + ivy/functional/backends/mxnet/experimental/manipulation.py | 1 + ivy/functional/backends/numpy/experimental/manipulation.py | 1 + ivy/functional/backends/paddle/experimental/manipulation.py | 1 + ivy/functional/backends/tensorflow/experimental/manipulation.py | 1 + ivy/functional/backends/torch/experimental/manipulation.py | 1 + ivy/functional/ivy/experimental/manipulation.py | 1 + 7 files changed, 7 insertions(+) diff --git a/ivy/functional/backends/jax/experimental/manipulation.py b/ivy/functional/backends/jax/experimental/manipulation.py index f9ee09f4d54af..1e35e4b356a71 100644 --- a/ivy/functional/backends/jax/experimental/manipulation.py +++ b/ivy/functional/backends/jax/experimental/manipulation.py @@ -478,6 +478,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[JaxArray] = None, ) -> JaxArray: x = jnp.swapaxes(x, axis, 0) value = jnp.swapaxes(value, axis, 0) diff --git a/ivy/functional/backends/mxnet/experimental/manipulation.py b/ivy/functional/backends/mxnet/experimental/manipulation.py index 3594feeef888b..896cc7084f909 100644 --- a/ivy/functional/backends/mxnet/experimental/manipulation.py +++ b/ivy/functional/backends/mxnet/experimental/manipulation.py @@ -215,6 +215,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[Union[(None, mx.ndarray.NDArray)]] = None, ) -> Union[(None, mx.ndarray.NDArray)]: x = mx.nd.swapaxes(x, axis, 0) value = mx.nd.swapaxes(value, axis, 0) diff --git a/ivy/functional/backends/numpy/experimental/manipulation.py b/ivy/functional/backends/numpy/experimental/manipulation.py index 1eaef0062ae6e..2412bc0eec043 100644 --- a/ivy/functional/backends/numpy/experimental/manipulation.py +++ b/ivy/functional/backends/numpy/experimental/manipulation.py @@ -608,6 +608,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[np.ndarray] = None, ) -> np.ndarray: x = np.swapaxes(x, axis, 0) value = np.swapaxes(value, axis, 0) diff --git a/ivy/functional/backends/paddle/experimental/manipulation.py b/ivy/functional/backends/paddle/experimental/manipulation.py index 0898da2dbfecd..5954d6139cd9f 100644 --- a/ivy/functional/backends/paddle/experimental/manipulation.py +++ b/ivy/functional/backends/paddle/experimental/manipulation.py @@ -915,5 +915,6 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[paddle.Tensor] = None, ) -> paddle.Tensor: return paddle.index_add(x, index, axis, value) diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index 8ad9b1d511d85..bb1edf08bf800 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -571,6 +571,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Union[tf.Tensor, tf.Variable] = None, ) -> tf.Tensor: x = tf.experimental.numpy.swapaxes(x, axis, 0) value = tf.experimental.numpy.swapaxes(value, axis, 0) diff --git a/ivy/functional/backends/torch/experimental/manipulation.py b/ivy/functional/backends/torch/experimental/manipulation.py index ff7336819dfb8..a3f96909fa024 100644 --- a/ivy/functional/backends/torch/experimental/manipulation.py +++ b/ivy/functional/backends/torch/experimental/manipulation.py @@ -649,6 +649,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[torch.Tensor] = None, ) -> torch.Tensor: x = torch.swapaxes(x, axis, 0) value = torch.swapaxes(value, axis, 0) diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index ac18f8f7c81be..38ec19eb07e57 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -2916,6 +2916,7 @@ def index_add( /, *, name: Optional[str] = None, + out: Optional[ivy.Array] = None, ) -> ivy.Array: """ Add the elements of the input tensor with value tensor by selecting the indices in From 4e6aba708844f7e46bacf967098b2f5228958409 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Mon, 4 Dec 2023 10:47:07 +0530 Subject: [PATCH 21/24] fix --- ivy/functional/ivy/experimental/manipulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index 38ec19eb07e57..9cdad5525a4df 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -2905,7 +2905,7 @@ def trim_zeros( @handle_nestable @handle_exceptions @handle_array_like_without_promotion -@inputs_to_ivy_arrays +@to_native_arrays_and_back @handle_array_function @handle_device def index_add( From 5445b29ad21b989d26030a5d09f6b86c94dc6be0 Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Thu, 7 Dec 2023 21:12:53 +0530 Subject: [PATCH 22/24] fix --- .../backends/tensorflow/experimental/manipulation.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index bb1edf08bf800..b2311ed4377b6 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -17,7 +17,7 @@ # local from ivy.func_wrapper import with_unsupported_dtypes -from .. import backend_version +from . import backend_version import ivy from ivy.functional.ivy.experimental.manipulation import _to_tf_padding @@ -596,3 +596,11 @@ def index_add( ret = tf.add(x, _to_adds) ret = tf.experimental.numpy.swapaxes(ret, axis, 0) return ret + + +x1 = tf.constant([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) +index1 = tf.constant([0, 2]) +value1 = tf.constant([[1, 1, 1], [1, 1, 1]]) +ret1 = tf.constant(x1, index1, 0, value1) + +print(x1) From 977a9b32e26696f73af04690579eea783b2f3fbf Mon Sep 17 00:00:00 2001 From: imsoumya18 Date: Sat, 9 Dec 2023 14:22:40 +0530 Subject: [PATCH 23/24] fix --- .../tensorflow/experimental/manipulation.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/ivy/functional/backends/tensorflow/experimental/manipulation.py b/ivy/functional/backends/tensorflow/experimental/manipulation.py index b2311ed4377b6..04ffadd55a6b2 100644 --- a/ivy/functional/backends/tensorflow/experimental/manipulation.py +++ b/ivy/functional/backends/tensorflow/experimental/manipulation.py @@ -564,10 +564,10 @@ def trim_zeros(a: tf.Tensor, /, *, trim: Optional[str] = "bf") -> tf.Tensor: def index_add( - x: tf.Tensor, - index: tf.Tensor, + x: Union[tf.Tensor, tf.Variable], + index: Union[tf.Tensor, tf.Variable], axis: int, - value: tf.Tensor, + value: Union[tf.Tensor, tf.Variable], /, *, name: Optional[str] = None, @@ -596,11 +596,3 @@ def index_add( ret = tf.add(x, _to_adds) ret = tf.experimental.numpy.swapaxes(ret, axis, 0) return ret - - -x1 = tf.constant([[1, 1, 1], [1, 1, 1], [1, 1, 1]]) -index1 = tf.constant([0, 2]) -value1 = tf.constant([[1, 1, 1], [1, 1, 1]]) -ret1 = tf.constant(x1, index1, 0, value1) - -print(x1) From 6a0edb537c4dd1ad73ca7705860a0e770bb847d8 Mon Sep 17 00:00:00 2001 From: ivy-branch Date: Mon, 11 Dec 2023 13:02:03 +0000 Subject: [PATCH 24/24] =?UTF-8?q?=F0=9F=A4=96=20Lint=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ivy/data_classes/array/experimental/manipulation.py | 5 ++--- ivy/data_classes/container/experimental/manipulation.py | 5 ++--- ivy/functional/ivy/experimental/manipulation.py | 5 ++--- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/ivy/data_classes/array/experimental/manipulation.py b/ivy/data_classes/array/experimental/manipulation.py index 50422e30d488e..e91fd44f1c6c4 100644 --- a/ivy/data_classes/array/experimental/manipulation.py +++ b/ivy/data_classes/array/experimental/manipulation.py @@ -1489,9 +1489,8 @@ def index_add( name: Optional[str] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Add the elements of the input tensor with value tensor by selecting the indices - in the order given in index. + """Add the elements of the input tensor with value tensor by selecting + the indices in the order given in index. Parameters ---------- diff --git a/ivy/data_classes/container/experimental/manipulation.py b/ivy/data_classes/container/experimental/manipulation.py index 382184c3b77ab..06bad2e932a22 100644 --- a/ivy/data_classes/container/experimental/manipulation.py +++ b/ivy/data_classes/container/experimental/manipulation.py @@ -4100,9 +4100,8 @@ def index_add( *, name: Optional[str] = None, ) -> ivy.Container: - """ - Add the elements of the input tensor with value tensor by selecting the indices - in the order given in index. + """Add the elements of the input tensor with value tensor by selecting + the indices in the order given in index. Parameters ---------- diff --git a/ivy/functional/ivy/experimental/manipulation.py b/ivy/functional/ivy/experimental/manipulation.py index e405089b3134a..753efc5210a04 100644 --- a/ivy/functional/ivy/experimental/manipulation.py +++ b/ivy/functional/ivy/experimental/manipulation.py @@ -2888,9 +2888,8 @@ def index_add( name: Optional[str] = None, out: Optional[ivy.Array] = None, ) -> ivy.Array: - """ - Add the elements of the input tensor with value tensor by selecting the indices in - the order given in index. + """Add the elements of the input tensor with value tensor by selecting the + indices in the order given in index. Parameters ----------