Test feeding device memory into host-memory expecting op in tf.function

PiperOrigin-RevId: 239020830
This commit is contained in:
Igor Ganichev 2019-03-18 11:05:44 -07:00 committed by TensorFlower Gardener
parent 128add4e08
commit 1a690e6c75

View File

@ -623,6 +623,17 @@ class EagerFunctionTest(xla_test.XLATestCase):
r = f(elems) r = f(elems)
self.assertAllEqual([2., 4., 12., 48., 240., 1440.], self.evaluate(r)) self.assertAllEqual([2., 4., 12., 48., 240., 1440.], self.evaluate(r))
def testFeedDeviceMemoryToOpExpectingHostMemory(self):
@function.defun
def f(dims, value):
return array_ops.fill(dims, value)
with self.test_scope():
x = constant_op.constant([4], dtype=dtypes.int64)
y = f(x, 3)
self.assertAllEqual([3, 3, 3, 3], y)
class ExcessivePaddingTest(xla_test.XLATestCase): class ExcessivePaddingTest(xla_test.XLATestCase):
"""Test that eager execution works with TPU flattened tensors. """Test that eager execution works with TPU flattened tensors.