Tensor.squeezingShape() does nothing for XLA
garymm opened this issue · comments
Gary Miguel commented
Using 0.11 release with Ubuntu 18.04, Cuda 10.1:
import TensorFlow
let eagerT = Tensor<Int32>([[0],[1],[2]], on: Device.defaultTFEager)
print(eagerT.shape)
print(Device.defaultTFEager)
let xlaT = Tensor<Int32>([[0],[1],[2]], on: Device.defaultXLA)
print(xlaT.shape)
print(Device.defaultXLA)
let eagerSqueezed = eagerT.squeezingShape()
print(eagerSqueezed.shape)
let xlaSqueezed = xlaT.squeezingShape()
print(xlaSqueezed.shape)
Output:
[3, 1]
Device(kind: .CPU, ordinal: 0, backend: .TF_EAGER)
[3, 1]
Device(kind: .GPU, ordinal: 0, backend: .XLA)
[3]
[3, 1]
Parker Schuh commented
You can work around this by making the dimensions explicit.