Make inplace_op behavior consistent on CPU/GPU, for empty tensors.

PiperOrigin-RevId: 291990267
Change-Id: I9a9c3ba86381ad1289cecc9335e66c4e271d6894
This commit is contained in:
A. Unique TensorFlower 2020-01-28 12:29:37 -08:00 committed by TensorFlower Gardener
parent 5467eb43e7
commit 5845168ce3
2 changed files with 22 additions and 1 deletions
tensorflow
core/kernels
python/kernel_tests

View File

@ -279,7 +279,10 @@ class InplaceOpBase : public OpKernel {
i.shape().DebugString(), " vs. ", v.shape().DebugString()));
Tensor y = x; // This creates an alias intentionally.
OP_REQUIRES_OK(ctx, DoCompute(ctx, i, v, &y));
// Skip processing if tensors are empty.
if (x.NumElements() > 0 || v.NumElements() > 0) {
OP_REQUIRES_OK(ctx, DoCompute(ctx, i, v, &y));
}
ctx->set_output(0, y);
}

View File

@ -202,6 +202,24 @@ class InplaceOpsTest(test_util.TensorFlowTestCase):
val = inplace_ops.empty((1, 2), dtypes.string, init=False).eval()
self.assertEqual(val.tolist(), [[b"", b""]])
@test_util.run_deprecated_v1
def testInplaceOpOnEmptyTensors(self):
op_fns = [
inplace_ops.inplace_add,
inplace_ops.inplace_sub,
inplace_ops.inplace_update,
]
for dtype in [dtypes.float32, dtypes.int32, dtypes.int64]:
for op_fn in op_fns:
with self.cached_session(use_gpu=True):
x = array_ops.zeros([7, 0], dtype)
y = np.zeros([7, 0], dtype.as_numpy_dtype)
self.assertAllClose(x.eval(), y)
x = op_fn(x, [3], array_ops.ones([1, 0], dtype))
self.assertAllClose(x.eval(), y)
x = op_fn(x, None, array_ops.ones([1, 0], dtype))
self.assertAllClose(x.eval(), y)
if __name__ == "__main__":
test_lib.main()