Merge pull request #21944 from rxwei/broadcasting-assignment

[TF] [stdlib] Add `.=` for broadcasting assignment.

Author: rxwei

stdlib/public/TensorFlow/Ops.swift

@@ -46,6 +46,7 @@ infix operator .>= : ComparisonPrecedence
 infix operator .> : ComparisonPrecedence
 infix operator .== : ComparisonPrecedence
 infix operator .!= : ComparisonPrecedence
+infix operator .=
 
 // TODO:
 // - Consider explicit broadcasting for elementwise binary ops when
@@ -1440,6 +1441,11 @@ public extension Tensor where Scalar : Numeric {
   func unbroadcast(to shape: TensorShape) -> Tensor {
     return unbroadcast(toShape: Tensor<Int32>(shape.dimensions))
   }
+
+  @inlinable @inline(__always)
+  static func .= (lhs: inout Tensor, rhs: Tensor) {
+    lhs = rhs.broadcast(like: lhs)
+  }
 }
 
 //===----------------------------------------------------------------------===//

test/TensorFlowRuntime/tensor.swift

@@ -466,18 +466,6 @@ TensorTests.testAllBackends("ReshapeTensor") {
   expectEqual([1, 3, 1, 2, 1], result.shape)
 }
 
-// FIXME: This test crashes in dynamic compilation + GPU.
-#if !CUDA
-TensorTests.testAllBackends("BroadcastTensor") {
-  // 1 -> 2 x 3 x 4
-  let one = Tensor<Float>(1)
-  let target = Tensor<Float>(shape: [2, 3, 4], repeating: 0.0)
-  let broadcasted = one.broadcast(like: target)
-  expectEqual([2, 3, 4], broadcasted.shape)
-  expectEqual(Array(repeating: 1, count: 24), broadcasted.scalars)
-}
-#endif // !CUDA
-
 TensorTests.testAllBackends("Unbroadcast1") {
   let x = Tensor<Float>(shape: [2, 3, 4, 5], repeating: 1)
   let y = Tensor<Float>(shape: [4, 5], repeating: 1)

test/TensorFlowRuntime/tensor_api.swift

@@ -30,4 +30,14 @@ TensorNonTPUTests.testAllBackends("SliceUpdate") {
   expectEqual(ShapedArray(shape:[2, 3], repeating: false), t4.array)
 }
 
+TensorNonTPUTests.testAllBackends("BroadcastTensor") {
+  // 1 -> 2 x 3 x 4
+  let one = Tensor<Float>(1)
+  var target = Tensor<Float>(shape: [2, 3, 4], repeating: 0.0)
+  let broadcasted = one.broadcast(like: target)
+  expectEqual(Tensor(shape: [2, 3, 4], repeating: 1), broadcasted)
+  target .= Tensor(shape: [1, 3, 1], repeating: 1)
+  expectEqual(Tensor(shape: [2, 3, 4], repeating: 1), target)
+}
+
 runAllTests()