tensorflow / swift

Swift for TensorFlow

Home Page:https://tensorflow.org/swift

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Tensor.squeezingShape() does nothing for XLA

garymm opened this issue · comments

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]

You can work around this by making the dimensions explicit.