Make inplace_op behavior consistent on CPU/GPU, for empty tensors.
PiperOrigin-RevId: 291990267 Change-Id: I9a9c3ba86381ad1289cecc9335e66c4e271d6894
This commit is contained in:
parent
5467eb43e7
commit
5845168ce3
tensorflow
@ -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);
|
||||
}
|
||||
|
||||
|
@ -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()
|
||||
|
Loading…
Reference in New Issue
Block a user