Tagging zero scalar tensors. (#1081)

* Tagging scalar zero tensors.

* Update Sources/TensorFlow/Core/Tensor.swift

Co-authored-by: Dan Zheng <[email protected]>

* Attempt to workaround SR-13263: debug info generation crash.

Co-authored-by: Brad Larson <[email protected]>
Co-authored-by: Brad Larson <[email protected]>

Author: dan-zheng

Sources/TensorFlow/Core/Tensor.swift

@@ -35,6 +35,17 @@ public struct Tensor<Scalar: TensorFlowScalar> {
   /// - Note: `handle` is public to allow user defined ops, but should not normally be used.
   public let handle: TensorHandle<Scalar>
 
+  /// An internal marker to identify scalar zero tensors, for use in optimizations.
+  @usableFromInline
+  internal var _isScalarZero = false
+
+  /// An internal workaround for SR-13263: debug info generation crash.
+  @usableFromInline
+  class SR13263Workaround {}
+
+  /// An internal workaround for SR-13263: debug info generation crash.
+  internal var _sr13263Workaround: SR13263Workaround?
+  
   @inlinable
   public init(handle: TensorHandle<Scalar>) {
     self.handle = handle
@@ -669,6 +680,7 @@ extension Tensor: AdditiveArithmetic where Scalar: Numeric {
       if _DeviceThreadLocalState.local.isReducedPrecision {
         zero = zero.toReducedPrecision
       }
+      zero._isScalarZero = true
       return zero
     }
   #else
@@ -681,15 +693,23 @@ extension Tensor: AdditiveArithmetic where Scalar: Numeric {
   @inlinable
   @differentiable(where Scalar: TensorFlowFloatingPoint)
   public static func + (lhs: Tensor, rhs: Tensor) -> Tensor {
-    _Raw.addV2(lhs, rhs)
+    if lhs._isScalarZero {
+      return rhs
+    } else if rhs._isScalarZero {
+      return lhs
+    }
+    return _Raw.addV2(lhs, rhs)
   }
 
   /// Subtracts one tensor from another and produces their difference.
   /// - Note: `-` supports broadcasting.
   @inlinable
   @differentiable(where Scalar: TensorFlowFloatingPoint)
   public static func - (lhs: Tensor, rhs: Tensor) -> Tensor {
-    _Raw.sub(lhs, rhs)
+    if rhs._isScalarZero {
+      return lhs
+    }
+    return _Raw.sub(lhs, rhs)
   }
 }