fix tensorflow tensor implementation with norm fixes and device preservation

Author: darshil929

nncf/tensor/functions/numeric.py

@@ -105,6 +105,16 @@ def abs(a: Tensor) -> Tensor:
     """
 
 
+@tensor_dispatcher
+def neg(a: Tensor) -> Tensor:
+    """
+    Numerical negative, element-wise.
+
+    :param a: The input tensor.
+    :return: A tensor containing the negative value of each element in a.
+    """
+
+
 @tensor_dispatcher
 def astype(a: Tensor, dtype: TensorDataType) -> Tensor:
     """
@@ -493,6 +503,28 @@ def sum(a: Tensor, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims:
     """
 
 
+@tensor_dispatcher
+def add(x1: Tensor, x2: Union[Tensor, float]) -> Tensor:
+    """
+    Add two tensors element-wise.
+
+    :param x1: The first input tensor.
+    :param x2: The second input tensor or number.
+    :return: The sum of x1 and x2, element-wise.
+    """
+
+
+@tensor_dispatcher
+def subtract(x1: Tensor, x2: Union[Tensor, float]) -> Tensor:
+    """
+    Subtract two tensors element-wise.
+
+    :param x1: The first input tensor.
+    :param x2: The second input tensor or number.
+    :return: The result of x1 - x2, element-wise.
+    """
+
+
 @tensor_dispatcher
 def multiply(x1: Tensor, x2: Union[Tensor, float]) -> Tensor:
     """

nncf/tensor/functions/tf_io.py

@@ -18,12 +18,26 @@
 
 from nncf.tensor import TensorDeviceType
 from nncf.tensor.functions import io as io
+from nncf.tensor.functions.tf_numeric import DEVICE_MAP
 
 
 def load_file(file_path: Path, *, device: Optional[TensorDeviceType] = None) -> Dict[str, tf.Tensor]:
-    return tf_load_file(file_path)
+    loaded_tensors = tf_load_file(file_path)
+
+    if device is not None:
+        device_str = DEVICE_MAP[device]
+        with tf.device(device_str):
+            loaded_tensors = {k: tf.identity(v) for k, v in loaded_tensors.items()}
+
+    return loaded_tensors
 
 
 @io.save_file.register
 def _(data: Dict[str, tf.Tensor], file_path: Path) -> None:
+    if file_path.is_symlink():
+        from nncf.errors import ValidationError
+
+        error_msg = "Cannot save tensor to a symbolic link"
+        raise ValidationError(error_msg)
+
     return tf_save_file(data, file_path)

nncf/tensor/functions/tf_linalg.py

@@ -24,9 +24,14 @@ def _(
     axis: Optional[Union[int, Tuple[int, ...]]] = None,
     keepdims: bool = False,
 ) -> tf.Tensor:
-    if ord is None:
-        ord = "euclidean"
     rank = tf.rank(a)
+
+    if ord is None:
+        if axis is None and rank == 2:
+            ord = "fro"
+        else:
+            ord = 2
+
     if rank == 2 and axis is None:
         axis = (0, 1)
 
@@ -49,41 +54,75 @@ def _(
             if rank != 2:
                 error_msg = "ord=-1 is only supported for 2D tensors"
                 raise ValueError(error_msg)
-            return tf.reduce_min(tf.reduce_sum(tf.abs(a), axis=axis[0]), keepdims=keepdims)
+            result = tf.reduce_min(tf.reduce_sum(tf.abs(a), axis=axis[0]), keepdims=keepdims)
+            if keepdims:
+                result = tf.reshape(result, [1, 1])
+            return result
 
         if ord == 1 and isinstance(axis, tuple) and len(axis) != 1:
             if rank != 2:
                 error_msg = "ord=1 is only supported for 2D tensors"
                 raise ValueError(error_msg)
-            return tf.reduce_max(tf.reduce_sum(tf.abs(a), axis=axis[0]), keepdims=keepdims)
+            result = tf.reduce_max(tf.reduce_sum(tf.abs(a), axis=axis[0]), keepdims=keepdims)
+            if keepdims:
+                result = tf.reshape(result, [1, 1])
+            return result
 
         if ord == -2 and isinstance(axis, tuple) and len(axis) != 1:
             if rank != 2:
                 error_msg = "ord=-2 is only supported for 2D tensors"
                 raise ValueError(error_msg)
             s = tf.linalg.svd(a, compute_uv=False)
-            return tf.reduce_min(s, axis=-1)
+            result = tf.reduce_min(s, axis=-1)
+            if keepdims:
+                result = tf.reshape(result, [1, 1])
+            return result
 
         if ord == 2 and isinstance(axis, tuple) and len(axis) != 1:
             if rank != 2:
                 error_msg = "ord=2 is only supported for 2D tensors"
                 raise ValueError(error_msg)
             s = tf.linalg.svd(a, compute_uv=False)
-            return tf.reduce_max(s, axis=-1)
+            result = tf.reduce_max(s, axis=-1)
+            if keepdims:
+                result = tf.reshape(result, [1, 1])
+            return result
 
         if ord == float("inf") and isinstance(axis, tuple) and len(axis) != 1:
             if rank != 2:
                 error_msg = "ord=inf is only supported for 2D tensors"
                 raise ValueError(error_msg)
-            return tf.reduce_max(tf.reduce_sum(tf.abs(a), axis=axis[1]), keepdims=keepdims)
+            result = tf.reduce_max(tf.reduce_sum(tf.abs(a), axis=axis[1]), keepdims=keepdims)
+            if keepdims:
+                result = tf.reshape(result, [1, 1])
+            return result
 
         if ord == -float("inf") and isinstance(axis, tuple) and len(axis) != 1:
             if rank != 2:
                 error_msg = "ord=-inf is only supported for 2D tensors"
                 raise ValueError(error_msg)
-            return tf.reduce_min(tf.reduce_sum(tf.abs(a), axis=axis[1]), keepdims=keepdims)
+            result = tf.reduce_min(tf.reduce_sum(tf.abs(a), axis=axis[1]), keepdims=keepdims)
+            if keepdims:
+                result = tf.reshape(result, [1, 1])
+            return result
 
-        return tf.linalg.norm(a, ord=ord, axis=axis, keepdims=keepdims)
+        try:
+            return tf.linalg.norm(a, ord=ord, axis=axis, keepdims=keepdims)
+        except (TypeError, ValueError):
+            if axis is not None:
+                if ord == 2:
+                    squared = tf.square(a)
+                    sum_squares = tf.reduce_sum(squared, axis=axis, keepdims=keepdims)
+                    return tf.sqrt(sum_squares)
+                elif ord == 1:
+                    return tf.reduce_sum(tf.abs(a), axis=axis, keepdims=keepdims)
+                elif ord == float("inf"):
+                    return tf.reduce_max(tf.abs(a), axis=axis, keepdims=keepdims)
+                elif ord == -float("inf"):
+                    return tf.reduce_min(tf.abs(a), axis=axis, keepdims=keepdims)
+
+            error_msg = f"Unsupported combination of ord={ord} and axis={axis}"
+            raise ValueError(error_msg)
 
 
 @linalg.cholesky.register