fix: tf.function working with set item

ivy/functional/backends/tensorflow/data_type.py

@@ -115,14 +115,16 @@ def broadcast_arrays(
 ) -> List[Union[tf.Tensor, tf.Variable]]:
     if len(arrays) > 1:
         try:
-            desired_shape = tf.broadcast_dynamic_shape(arrays[0].shape, arrays[1].shape)
+            desired_shape = tf.broadcast_dynamic_shape(
+                tf.shape(arrays[0]), tf.shape(arrays[1])
+            )
         except tf.errors.InvalidArgumentError as e:
             raise ivy.utils.exceptions.IvyBroadcastShapeError(e) from e
         if len(arrays) > 2:
             for i in range(2, len(arrays)):
                 try:
                     desired_shape = tf.broadcast_dynamic_shape(
-                        desired_shape, arrays[i].shape
+                        desired_shape, tf.shape(arrays[i])
                     )
                 except tf.errors.InvalidArgumentError as e:
                     raise ivy.utils.exceptions.IvyBroadcastShapeError(e) from e

ivy/functional/backends/tensorflow/general.py

@@ -378,9 +378,9 @@ def scatter_nd(
     )
 
     expected_shape = (
-        list(indices.shape[:-1]) + list(out.shape[indices.shape[-1] :])
+        list(tf.shape(indices)[:-1]) + list(out.shape[tf.shape(indices)[-1] :])
         if ivy.exists(out)
-        else list(indices.shape[:-1]) + list(shape[indices.shape[-1] :])
+        else list(tf.shape(indices)[:-1]) + list(shape[tf.shape(indices)[-1] :])
     )
     updates = _broadcast_to(updates, expected_shape)._data
     if len(updates.shape) == 0:

ivy/functional/backends/tensorflow/manipulation.py

@@ -62,7 +62,7 @@ def expand_dims(
     out: Optional[Union[tf.Tensor, tf.Variable]] = None,
 ) -> Union[tf.Tensor, tf.Variable]:
     try:
-        out_shape = _calculate_out_shape(axis, x.shape)
+        out_shape = _calculate_out_shape(axis, tf.shape(x))
         ret = tf.reshape(x, shape=out_shape)
         return ret
     except (tf.errors.InvalidArgumentError, np.AxisError) as error:

ivy/functional/ivy/general.py

@@ -2887,15 +2887,14 @@ def set_item(
     if ivy.is_array(query) and ivy.is_bool_dtype(query):
         if not len(query.shape):
             query = ivy.tile(query, (x.shape[0],))
-        target_shape = ivy.get_item(x, query).shape
         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 = _broadcast_to(val, target_shape).astype(x.dtype)
+    val = val.astype(x.dtype)
     ret = ivy.scatter_nd(indices, val, reduction="replace", out=x)
     return ret
 
@@ -2977,7 +2976,7 @@ def _parse_query(query, x_shape, scatter=False):
     elif len(array_inds):
         target_shape = (
             [list(query[i].shape) for i in range(0, array_inds[0])]
-            + [list(array_queries[0].shape)]
+            + [list(ivy.shape(array_queries[0], as_array=True))]
             + [[] for _ in range(len(array_inds) - 1)]
             + [list(query[i].shape) for i in range(array_inds[-1] + 1, len(query))]
         )