fix: tf backend get_item and add set_item (#28823)

Co-authored-by: ivy-dev-bot <[email protected]>

binaries.json

@@ -5,11 +5,9 @@
  {
  "utils": [
  "C",
- "CC",
  "CD",
  "CI",
  "CL",
- "CM",
  "CV",
  "CX",
  "D",
@@ -43,7 +41,6 @@
  "III",
  "IIL",
  "IIM",
- "IIV",
  "IIX",
  "IL",
  "ILC",
@@ -62,7 +59,6 @@
  "IMV",
  "IMX",
  "IV",
- "IVC",
  "IVD",
  "IVI",
  "IVL",
@@ -103,7 +99,13 @@
  "VCV",
  "VCX",
  "VD",
+ "VDC",
+ "VDD",
  "VDI",
+ "VDL",
+ "VDM",
+ "VDV",
+ "VDX",
  "VI",
  "VIC",
  "VID",
@@ -149,4 +151,4 @@
  }
  ]
  }
-}
+}

ivy/compiler/compiler.py

@@ -2,15 +2,15 @@
 
 
 def clear_graph_cache():
- """Clears the graph cache which gets populated if `graph_caching` is set to
- `True` in `ivy.trace_graph`, `ivy.transpile` or `ivy.unify`. Use this to
+ """Clears the graph cache which gets populated if `graph_caching` is set
+ to `True` in `ivy.trace_graph`, `ivy.transpile` or `ivy.unify`. Use this to
  reset or clear the graph cache if needed.
 
  Examples
  --------
  >>> import ivy
- >>> ivy.clear_graph_cache()
- """
+ >>> ivy.clear_graph_cache()"""
+
  from ._compiler import clear_graph_cache as _clear_graph_cache
 
  return _clear_graph_cache()
@@ -55,8 +55,8 @@ def graph_transpile(
 
  Returns
  -------
- Either a transpiled Graph or a non-initialized LazyGraph.
- """
+ Either a transpiled Graph or a non-initialized LazyGraph."""
+
  from ._compiler import graph_transpile as _graph_transpile
 
  return _graph_transpile(
@@ -96,7 +96,6 @@ def source_to_source(
  e.g. (source="torch_frontend", target="ivy") or (source="torch_frontend", target="tensorflow") etc.
 
  Args:
- ----
  object: The object (class/function) to be translated.
  source (str, optional): The source framework. Defaults to 'torch'.
  target (str, optional): The target framework. Defaults to 'tensorflow'.
@@ -107,9 +106,8 @@ def source_to_source(
  the old implementation. Defaults to 'True'.
 
  Returns:
- -------
- The translated object.
- """
+ The translated object."""
+
  from ._compiler import source_to_source as _source_to_source
 
  return _source_to_source(
@@ -140,8 +138,7 @@ def trace_graph(
  params_v=None,
  v=None
 ):
- """Takes `fn` and traces it into a more efficient composition of backend
- operations.
+ """Takes `fn` and traces it into a more efficient composition of backend operations.
 
  Parameters
  ----------
@@ -211,8 +208,8 @@ def trace_graph(
  >>> start = time.time()
  >>> graph(x)
  >>> print(time.time() - start)
- 0.0001785755157470703
- """
+ 0.0001785755157470703"""
+
  from ._compiler import trace_graph as _trace_graph
 
  return _trace_graph(
@@ -252,7 +249,6 @@ def transpile(
  e.g. (source="torch_frontend", target="ivy") or (source="torch_frontend", target="tensorflow") etc.
 
  Args:
- ----
  object: The object (class/function) to be translated.
  source (str, optional): The source framework. Defaults to 'torch'.
  target (str, optional): The target framework. Defaults to 'tensorflow'.
@@ -263,9 +259,8 @@ def transpile(
  the old implementation. Defaults to 'True'.
 
  Returns:
- -------
- The translated object.
- """
+ The translated object."""
+
  from ._compiler import transpile as _transpile
 
  return _transpile(

ivy/functional/backends/tensorflow/general.py

@@ -61,7 +61,51 @@ def get_item(
 ) -> Union[tf.Tensor, tf.Variable]:
  if ivy.is_array(query) and ivy.is_bool_dtype(query) and not len(query.shape):
  return tf.expand_dims(x, 0)
- return x[query]
+ if isinstance(query, (tf.Tensor, tf.Variable)):
+ if query.dtype == tf.bool:
+ return tf.boolean_mask(x, query, axis=0)
+ else:
+ query = tf.cast(query, tf.int64)
+ return tf.gather(x, query, axis=0)
+ else:
+ if any([isinstance(q, slice) for q in query]):
+ # convert any lists/tuples within the query to slices
+ query = tuple([slice(*q) if isinstance(q, (list, tuple)) else q for q in query])
+ # for slices and other basic indexing, use __getitem__
+ return x[query]
+
+
+def set_item(
+ x: Union[tf.Tensor, tf.Variable],
+ query: Union[tf.Tensor, tf.Variable, Tuple],
+ val: Union[tf.Tensor, tf.Variable],
+ /,
+ *,
+ copy: Optional[bool] = False,
+) -> Union[tf.Tensor, tf.Variable]:
+ # TODO: we should re-write this at some point so it's compatible with tf.function (don't use numpy as an intermediary)
+ # when doing this, be sure to check the performance of the function on large tensors, compared to this implementation
+
+ if tf.is_tensor(x):
+ x = x.numpy()
+ if tf.is_tensor(val):
+ val = val.numpy()
+
+ if isinstance(query, (tf.Tensor, tf.Variable)):
+ query = query.numpy()
+ elif isinstance(query, tuple):
+ query = tuple(
+ q.numpy() if isinstance(q, (tf.Tensor, tf.Variable)) else q
+ for q in query
+ )
+
+ x[query] = val
+
+ if isinstance(x, tf.Variable) and not copy:
+ x.assign(x)
+ return x
+ else:
+ return tf.Variable(x) if isinstance(x, tf.Variable) else tf.convert_to_tensor(x)
 
 
 def to_numpy(x: Union[tf.Tensor, tf.Variable], /, *, copy: bool = True) -> np.ndarray:

ivy/functional/ivy/general.py

@@ -2879,25 +2879,7 @@ def set_item(
  ivy.array([[ 0, -1, 20],
  [10, 10, 10]])
  """
- if copy:
- x = ivy.copy_array(x)
- if not ivy.is_array(val):
- val = ivy.array(val)
- if 0 in x.shape or 0 in val.shape:
- return x
- if ivy.is_array(query) and ivy.is_bool_dtype(query):
- if not len(query.shape):
- query = ivy.tile(query, (x.shape[0],))
- indices = ivy.nonzero(query, as_tuple=False)
- else:
- indices, target_shape, _ = _parse_query(
- query, ivy.shape(x, as_array=True), scatter=True
- )
- if indices is None:
- return x
- val = val.astype(x.dtype)
- ret = ivy.scatter_nd(indices, val, reduction="replace", out=x)
- return ret
+ return current_backend(x).set_item(x, query, val, copy=copy)
 
 
 set_item.mixed_backend_wrappers = {

ivy/utils/decorator_utils.py

@@ -183,8 +183,6 @@ def handle_get_item(fn):
  def wrapper(inp, query, **kwargs):
  try:
  res = inp.__getitem__(query)
- except IndexError:
- raise
  except Exception:
  res = fn(inp, query, **kwargs)
  return res