fix: tf set item and add tests

Signed-off-by: anna-charlotte <charlotte.gerhaher@jina.ai>
docarray · 09dd6a9deeffbcddb3b744021437b9af040e2d3e
commit 09dd6a9deeffbcddb3b744021437b9af040e2d3e
diff --git a/docarray/typing/tensor/tensorflow_tensor.py b/docarray/typing/tensor/tensorflow_tensor.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 import tensorflow as tf  # type: ignore
-import tensorflow._api.v2.experimental.numpy as tnp  # type: ignore
 
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
@@ -145,13 +144,11 @@ def __getitem__(self, item):
 
     def __setitem__(self, index, value):
         """Set a slice of this tensor's tf.Tensor"""
-        if tnp.all(tf.math.is_nan(value)):
-            self.tensor = None
-        else:
-            t = self.unwrap()
-            var = tf.Variable(t)
-            var[index].assign(value)
-            self.tensor = tf.constant(var)
+        t = self.unwrap()
+        value = tf.cast(value, dtype=t.dtype)
+        var = tf.Variable(t)
+        var[index].assign(value)
+        self.tensor = tf.constant(var)
 
     def __iter__(self):
         """Iterate over the elements of this tensor's tf.Tensor."""

diff --git a/tests/units/array/test_array_stacked_tf.py b/tests/units/array/test_array_stacked_tf.py
@@ -259,7 +259,6 @@ class MyDoc(BaseDocument):
     ).stack()
 
     assert 'tensor' in da._tensor_columns.keys()
-    assert da.tensor.tensor is None
 
 
 @pytest.mark.tensorflow