diff --git a/test/TensorFlowRuntime/tensor.swift b/test/TensorFlowRuntime/tensor.swift
index 82b6f5a7a43b9..192adcac47298 100644
--- a/test/TensorFlowRuntime/tensor.swift
+++ b/test/TensorFlowRuntime/tensor.swift
@@ -719,6 +719,125 @@ TensorTests.testAllBackends("Unbroadcast2") {
               z.array)
 }
 
+TensorTests.testAllBackends("broadcast(toShape:)") {
+  func foo(tensor: Tensor<Float>, shape: Tensor<Int32>) -> Tensor<Float> {
+    tensor.broadcast(toShape: shape)
+  }
+  var inputTensor: Tensor<Float>
+  var expected: Tensor<Float>
+  var pb: (Tensor<Float>) -> Tensor<Float>
+  
+  // [3,] -> [3,3]
+  pb = pullback(at: Tensor([99, 33, 55])) { x in 
+    foo(tensor: x, shape: Tensor([3, 3])) 
+  }
+  
+  // Test 1: same shape as parameter of pullback
+  inputTensor = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
+    [1, 2, 3]]
+  )
+  expected = Tensor([3, 6, 9])
+  expectEqual(expected, pb(inputTensor))
+  
+  // Test 2: different shape than parameter of pullback
+  inputTensor = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
+    [1, 2, 3],
+    [1, 2, 3]]
+  )
+  expected = Tensor([4, 8, 12])
+  expectEqual(expected, pb(inputTensor))
+  
+  // Test 3: same shape as tensor we are differentiating at
+  inputTensor = Tensor([1, 2, 3])
+  expected = Tensor([1, 2, 3])
+  expectEqual(expected, pb(inputTensor))
+
+  // Test 4: extremely padded shape as tensor we are differentiating at
+  inputTensor = Tensor([[[[[[1, 2, 3]]]]]])
+  expected = Tensor([1, 2, 3])
+  expectEqual(expected, pb(inputTensor))
+  
+  // [3,1] -> [3x3]
+  pb = pullback(at: Tensor([[99, 33, 55]])) { x in 
+    foo(tensor: x, shape: Tensor([3, 3])) 
+  }
+
+  // Test 5: same shape as parameter of pullback
+  inputTensor = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
+    [1, 2, 3]]
+  )
+  expected = Tensor([[3, 6, 9]])
+  expectEqual(expected, pb(inputTensor))
+
+  // Test 6: different shape than parameter of pullback
+  inputTensor = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
+    [1, 2, 3],
+    [1, 2, 3]]
+  )
+  expected = Tensor([[4, 8, 12]])
+  expectEqual(expected, pb(inputTensor))
+
+  // Test 7: same shape as tensor we are differentiating at
+  inputTensor = Tensor([[1, 2, 3]])
+  expected = Tensor([[1, 2, 3]])
+  expectEqual(expected, pb(inputTensor))
+
+  // Test 8: extremely padded shape of tensor we are differentiating at
+  inputTensor = Tensor([[[[[[1, 2, 3]]]]]])
+  expected = Tensor([[1, 2, 3]])
+  expectEqual(expected, pb(inputTensor))
+}
+
+TensorADTests.testAllBackends("unbroadcast(toShape:") {
+  func foo(tensor: Tensor<Float>, shape: Tensor<Int32>) -> Tensor<Float> {
+    tensor.unbroadcast(toShape: shape)
+  }
+  var inputTensor: Tensor<Float>
+  var expected: Tensor<Float>
+  var pb: (Tensor<Float>) -> Tensor<Float>
+  
+  // [3,3] -> [1,3]
+  let atTensor: Tensor<Float> = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
+    [1, 2, 3]]
+  )
+  pb = pullback(at: atTensor) { x in 
+    foo(tensor: x, shape: Tensor([1, 3])) 
+  }
+
+  // Test 1: same shape as parameter of pullback
+  inputTensor = Tensor([[1, 2, 3]])
+  expected = atTensor
+  expectEqual(expected, pb(inputTensor))
+  // Test 2: different shape than parameter o
+  f pullback
+  inputTensor = Tensor([2])
+  expected = Tensor([
+    [2, 2, 2], 
+    [2, 2, 2], 
+    [2, 2, 2]]
+  )
+  expectEqual(expected, pb(inputTensor))
+
+  // Test 3: same shape as tensor we are differentiating at
+  inputTensor = Tensor([
+    [8, 1, 3], 
+    [8, 1, 3], 
+    [8, 1, 3]]
+  )
+  expected = inputTensor
+  expectEqual(expected, pb(inputTensor))
+}
+
 // TODO: Merge all rank/shape getter tests into one when we support code motion
 // to avoid sends.
 
diff --git a/test/TensorFlowRuntime/tensor_autodiff_runtime.swift b/test/TensorFlowRuntime/tensor_autodiff_runtime.swift
index 2bf3577bcc1e1..b1b49da672b2d 100644
--- a/test/TensorFlowRuntime/tensor_autodiff_runtime.swift
+++ b/test/TensorFlowRuntime/tensor_autodiff_runtime.swift
@@ -263,130 +263,101 @@ TensorADTests.testAllBackends("broadcast(toShape:)") {
     tensor.broadcast(toShape: shape)
   }
 
-  var inputTensor: Tensor<Float>
-  var expected: Tensor<Float>
-  var pb: (Tensor<Float>) -> Tensor<Float>
-
-  // [3,] -> [3,3]
-  pb = pullback(at: Tensor([99, 33, 55])) { x in
-    foo(tensor: x, shape: Tensor([3, 3]))
+  let pb: (Tensor<Float>) -> Tensor<Float> = pullback(at: Tensor([99, 33, 55])) { x in 
+    foo(tensor: x, shape: Tensor([3, 3])) 
   }
-
-  // Test 1: same shape as parameter of pullback
-  inputTensor = Tensor([
-    [1, 2, 3],
-    [1, 2, 3],
-    [1, 2, 3]]
-  )
-  expected = Tensor([3, 6, 9])
-  expectEqual(expected, pb(inputTensor))
-
-  // Test 2: different shape than parameter of pullback
-  inputTensor = Tensor([
-    [1, 2, 3],
-    [1, 2, 3],
+  let inputTensor: Tensor<Float> = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
     [1, 2, 3],
     [1, 2, 3]]
   )
-  expected = Tensor([4, 8, 12])
-  expectEqual(expected, pb(inputTensor))
-
-  // Test 3: same shape as tensor we are differentiating at
-  inputTensor = Tensor([1, 2, 3])
-  expected = Tensor([1, 2, 3])
+  let expected: Tensor<Float> = Tensor([[4, 8, 12]])
   expectEqual(expected, pb(inputTensor))
+}
 
-  // Test 4: extremely padded shape as tensor we are differentiating at
-  inputTensor = Tensor([[[[[[1, 2, 3]]]]]])
-  expected = Tensor([1, 2, 3])
+TensorADTests.testAllBackends("broadcast(to:") {
+  func foo(tensor: Tensor<Float>, shape: TensorShape) -> Tensor<Float> {
+    tensor.broadcast(to: shape)
+  }
+  let pb: (Tensor<Float>) -> Tensor<Float> = pullback(at: Tensor([99, 33, 55])) { x in 
+    foo(tensor: x, shape: TensorShape([3, 3])) 
+  }
+  let inputTensor: Tensor<Float> = Tensor([1, 2, 3])
+  let expected: Tensor<Float> = Tensor([1, 2, 3])
   expectEqual(expected, pb(inputTensor))
+}
 
-  // [3,1] -> [3x3]
-  pb = pullback(at: Tensor([[99, 33, 55]])) { x in
-    foo(tensor: x, shape: Tensor([3, 3]))
+TensorADTests.testAllBackends("broadcast(like:)") {
+  func foo(tensor: Tensor<Float>, other: Tensor<Double>) -> Tensor<Float> {
+    tensor.broadcast(like: other)
+  }
+  let pb: (Tensor<Float>) -> Tensor<Float> = pullback(at: Tensor([99, 33, 55])) { x in 
+    foo(tensor: x, other: Tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]])) 
   }
+  let inputTensor: Tensor<Float> = Tensor([[[[[[1, 2, 3]]]]]])
+  let expected: Tensor<Float> = Tensor([[1, 2, 3]])
 
-  // Test 5: same shape as parameter of pullback
-  inputTensor = Tensor([
-    [1, 2, 3],
-    [1, 2, 3],
-    [1, 2, 3]]
-  )
-  expected = Tensor([[3, 6, 9]])
   expectEqual(expected, pb(inputTensor))
+}
 
-  // Test 6: different shape than parameter of pullback
-  inputTensor = Tensor([
-    [1, 2, 3],
-    [1, 2, 3],
-    [1, 2, 3],
+TensorADTests.testAllBackends("unbroadcast(toShape:)") {
+  func foo(tensor: Tensor<Float>, shape: Tensor<Int32>) -> Tensor<Float> {
+    tensor.unbroadcast(toShape: shape)
+  }
+  let atTensor: Tensor<Float> = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
     [1, 2, 3]]
   )
-  expected = Tensor([[4, 8, 12]])
-  expectEqual(expected, pb(inputTensor))
-
-  // Test 7: same shape as tensor we are differentiating at
-  inputTensor = Tensor([[1, 2, 3]])
-  expected = Tensor([[1, 2, 3]])
-  expectEqual(expected, pb(inputTensor))
-
-  // Test 8: extremely padded shape of tensor we are differentiating at
-  inputTensor = Tensor([[[[[[1, 2, 3]]]]]])
-  expected = Tensor([[1, 2, 3]])
+  let pb: (Tensor<Float>) -> Tensor<Float> = pullback(at: atTensor) { x in 
+    foo(tensor: x, shape: Tensor([1, 3])) 
+  }
+  let expected = atTensor
+  let inputTensor: Tensor<Float> = Tensor([[1, 2, 3]])
   expectEqual(expected, pb(inputTensor))
 }
 
-TensorADTests.testAllBackends("unbroadcast(toShape:") {
-  func foo(tensor: Tensor<Float>, shape: Tensor<Int32>) -> Tensor<Float> {
-    tensor.unbroadcast(toShape: shape)
+TensorADTests.testAllBackends("unbroadcast(to:") {
+  func foo(tensor: Tensor<Float>, shape: TensorShape) -> Tensor<Float> {
+    tensor.unbroadcast(to: shape)
   }
-
-  var inputTensor: Tensor<Float>
-  var expected: Tensor<Float>
-  var pb: (Tensor<Float>) -> Tensor<Float>
-
-  // [3,3] -> [1,3]
   let atTensor: Tensor<Float> = Tensor([
     [1, 2, 3],
     [1, 2, 3],
     [1, 2, 3]]
   )
-  pb = pullback(at: atTensor) { x in
-    foo(tensor: x, shape: Tensor([1, 3]))
+  let pb: (Tensor<Float>) -> Tensor<Float> = pullback(at: atTensor) { x in 
+    foo(tensor: x, shape: TensorShape([1, 3])) 
   }
-
-  // Test 1: same shape as parameter of pullback
-  inputTensor = Tensor([[1, 2, 3]])
-  expected = atTensor
-  expectEqual(expected, pb(inputTensor))
-
-  // Test 2: different shape than parameter of pullback
-  inputTensor = Tensor([2])
-  expected = Tensor([
-    [2, 2, 2],
-    [2, 2, 2],
+  let inputTensor: Tensor<Float> = Tensor([2])
+  let expected: Tensor<Float> = Tensor([
+    [2, 2, 2], 
+    [2, 2, 2], 
     [2, 2, 2]]
   )
   expectEqual(expected, pb(inputTensor))
+}
 
-  // Test 3: same shape as tensor we are differentiating at
-  inputTensor = Tensor([
-    [8, 1, 3],
-    [8, 1, 3],
+TensorADTests.testAllBackends("unbroadcast(like:") {
+  func foo(tensor: Tensor<Float>, other: Tensor<Double>) -> Tensor<Float> {
+    tensor.unbroadcast(like: other)
+  }
+  let atTensor: Tensor<Float> = Tensor([
+    [1, 2, 3], 
+    [1, 2, 3], 
+    [1, 2, 3]]
+  )
+  let pb: (Tensor<Float>) -> Tensor<Float> = pullback(at: atTensor) { x in 
+    foo(tensor: x, other: Tensor([[1, 2, 3]])) 
+  }
+  let inputTensor: Tensor<Float> = Tensor([
+    [8, 1, 3], 
+    [8, 1, 3], 
     [8, 1, 3]]
   )
-  expected = inputTensor
+  let expected: Tensor<Float> = inputTensor
   expectEqual(expected, pb(inputTensor))
-
-  // TODO
-  // Test 4: extremely padded shape as tensor we are differentiating at
-  // inputTensor = Tensor([
-  //   [[8, 1, 3]],
-  //   [[8, 1, 3]],
-  //   [[8, 1, 3]]]
-  // )
-  // expected = Tensor([1, 2, 3])
-  // expectEqual(expected, pb(inputTensor))
 }
 
 runAllTests()