| # Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Test cases for eager execution using XLA.""" |
| |
| import numpy as np |
| |
| from tensorflow.compiler.tests import xla_test |
| from tensorflow.core.protobuf import config_pb2 |
| from tensorflow.python.eager import backprop |
| from tensorflow.python.eager import context |
| from tensorflow.python.eager import def_function |
| from tensorflow.python.framework import constant_op |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import indexed_slices |
| from tensorflow.python.framework import ops |
| from tensorflow.python.layers import convolutional |
| from tensorflow.python.layers import pooling |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import array_ops_stack # pylint: disable=g-direct-tensorflow-import |
| from tensorflow.python.ops import cond |
| from tensorflow.python.ops import embedding_ops |
| from tensorflow.python.ops import functional_ops |
| from tensorflow.python.ops import gen_random_ops |
| from tensorflow.python.ops import init_ops |
| from tensorflow.python.ops import math_ops |
| from tensorflow.python.ops import nn_ops |
| from tensorflow.python.ops import resource_variable_ops |
| from tensorflow.python.ops import while_loop |
| from tensorflow.python.platform import googletest |
| from tensorflow.python.training import adam |
| |
| |
| class EagerTest(xla_test.XLATestCase): |
| |
| def testBasic(self): |
| with self.test_scope(): |
| three = constant_op.constant(3) |
| five = constant_op.constant(5) |
| product = three * five |
| self.assertAllEqual(15, product) |
| |
| def testGradientTape(self): |
| with self.test_scope(): |
| |
| x = constant_op.constant(1.0) |
| y = constant_op.constant(10.0) |
| with backprop.GradientTape(persistent=True) as tape: |
| tape.watch(x) |
| tape.watch(y) |
| a = x + y + x * y |
| da_dx = tape.gradient(a, x) |
| da_dy = tape.gradient(a, y) |
| |
| self.assertEqual(11.0, da_dx.numpy()) |
| self.assertEqual(2.0, da_dy.numpy()) |
| |
| def testExecuteListOutputLen0(self): |
| with self.test_scope(): |
| empty = constant_op.constant([], dtype=dtypes.float32) |
| result = array_ops_stack.unstack(empty, 0) |
| self.assertTrue(isinstance(result, list)) |
| self.assertEqual(0, len(result)) |
| |
| def testExecuteListOutputLen1(self): |
| with self.test_scope(): |
| split_dim = constant_op.constant(1) |
| value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]]) |
| result = array_ops.split(value, 1, axis=split_dim) |
| self.assertTrue(isinstance(result, list)) |
| self.assertEqual(1, len(result)) |
| self.assertAllEqual([[0, 1, 2], [3, 4, 5]], result[0]) |
| |
| def testExecuteListOutputLen3(self): |
| with self.test_scope(): |
| split_dim = constant_op.constant(1) |
| value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]]) |
| result = array_ops.split(value, 3, axis=split_dim) |
| self.assertTrue(isinstance(result, list)) |
| self.assertEqual(3, len(result)) |
| self.assertAllEqual([[0], [3]], result[0]) |
| self.assertAllEqual([[1], [4]], result[1]) |
| self.assertAllEqual([[2], [5]], result[2]) |
| |
| def testBasicGraph(self): |
| # Run some ops eagerly |
| with self.test_scope(): |
| three = constant_op.constant(3) |
| five = constant_op.constant(5) |
| product = three * five |
| self.assertAllEqual(15, product) |
| |
| # Run some ops graphly |
| with context.graph_mode(), self.session(): |
| with self.test_scope(): |
| three = constant_op.constant(3) |
| five = constant_op.constant(5) |
| product = three * five |
| self.assertAllEqual(15, self.evaluate(product)) |
| |
| def testDegenerateSlices(self): |
| with self.test_scope(): |
| npt = np.arange(1, 19, dtype=np.float32).reshape(3, 2, 3) |
| t = constant_op.constant(npt) |
| # degenerate by offering a forward interval with a negative stride |
| self.assertAllEqual(npt[0:-1:-1, :, :], t[0:-1:-1, :, :]) |
| # degenerate with a reverse interval with a positive stride |
| self.assertAllEqual(npt[-1:0, :, :], t[-1:0, :, :]) |
| # empty interval in every dimension |
| self.assertAllEqual(npt[-1:0, 2:2, 2:3:-1], t[-1:0, 2:2, 2:3:-1]) |
| |
| def testIdentity(self): |
| with self.test_scope(): |
| self.assertAllEqual(2, array_ops.identity(2)) |
| |
| def testRandomOps(self): |
| with self.test_scope(): |
| tensor = gen_random_ops.random_uniform((2, 2), dtypes.float32) |
| row0 = tensor[0].numpy() |
| row1 = tensor[1].numpy() |
| # It should be very unlikely to rng to generate two equal rows. |
| self.assertFalse((row0 == row1).all()) |
| |
| def testIdentityOnVariable(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(True) |
| i = array_ops.identity(v) |
| self.assertAllEqual(True, i.numpy()) |
| |
| def testAssignAddVariable(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(1.0) |
| v.assign_add(2.0) |
| self.assertEqual(3.0, v.numpy()) |
| |
| def testReadAssignRead(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(1.0) |
| val1 = v.read_value() |
| v.assign_add(2.0) |
| val2 = v.read_value() |
| self.assertEqual(1.0, val1.numpy()) |
| self.assertEqual(3.0, val2.numpy()) |
| |
| def testGradient(self): |
| def f(x): |
| return x |
| |
| with self.test_scope(): |
| grad_fn = backprop.gradients_function(f) |
| self.assertAllEqual(2., grad_fn(1., dy=2.)[0]) |
| |
| def testVariableGradient(self): |
| with self.test_scope(): |
| v0 = resource_variable_ops.ResourceVariable(1.0) |
| |
| def f(): |
| x = v0 * v0 |
| return x |
| |
| grads = backprop.implicit_grad(f)() |
| self.assertEqual(2., grads[0][0].numpy()) |
| |
| def testMultipleVariableReads(self): |
| # This test makes sure consecutive variable reads don't copy |
| # the underlying memory. |
| with self.test_scope(): |
| # Create 128MiB variables |
| var = resource_variable_ops.ResourceVariable( |
| array_ops.ones([32, 1024, 1024])) |
| |
| # Read the same variable 100 times. If the underlying tensor |
| # is not copied, this is a trivial operation. If it is copied, |
| # this will eat over 13GB and OOM. |
| values = [] |
| for _ in range(100): |
| values.append(var.value()) |
| |
| # The shape, shape_n, size, and rank are tested here because their |
| # execution kernels (as opposed to compilation only tf2xla kernels) |
| # are distincts from tf2xla kernels. |
| |
| def testShape(self): |
| def const(value): |
| return array_ops.shape( |
| constant_op.constant(value)).numpy() |
| |
| def ones(value): |
| return array_ops.shape( |
| array_ops.ones(value)).numpy() |
| |
| with self.test_scope(): |
| # Shapes of directly constructed tensors |
| self.assertAllEqual([], const(3)) |
| self.assertAllEqual([3], const([1.0, 2.0, 3.0])) |
| self.assertAllEqual([2, 2], const([[1.0, 2.0], [3.0, 4.0]])) |
| self.assertAllEqual([2, 1, 2], const([[[1.0, 2.0]], [[3.0, 4.0]]])) |
| |
| # Shapes of tensors created by op running on device |
| # We make this distinction because directly constructed tensors |
| # are treated differently in a few places that can influence shape: |
| # - they always have on_host_tensor |
| # - they and their shapes can be cached |
| # - they end up on device via a copy, instead of as program output |
| self.assertAllEqual([], ones([])) |
| self.assertAllEqual([3], ones([3])) |
| self.assertAllEqual([2, 2], ones([2, 2])) |
| self.assertAllEqual([2, 1, 2], ones([2, 1, 2])) |
| |
| def testShapeN(self): |
| with self.test_scope(): |
| # Shapes of directly constructed tensors |
| shapes = array_ops.shape_n([ |
| constant_op.constant(1.0), |
| constant_op.constant([1.0, 2.0, 3.0]), |
| constant_op.constant([[1.0, 2.0], [3.0, 4.0]])]) |
| self.assertAllEqual( |
| [[], [3], [2, 2]], |
| [x.numpy().tolist() for x in shapes]) |
| |
| # Shapes of tensors created by op running on device |
| shapes = array_ops.shape_n([ |
| array_ops.ones([]), |
| array_ops.ones([3]), |
| array_ops.ones([2, 2])]) |
| self.assertAllEqual( |
| [[], [3], [2, 2]], |
| [x.numpy().tolist() for x in shapes]) |
| |
| def testSize(self): |
| with self.test_scope(): |
| self.assertEqual( |
| 1, array_ops.size(constant_op.constant(1.0)).numpy()) |
| self.assertEqual( |
| 3, array_ops.size(constant_op.constant([1.0, 2.0, 3.0])).numpy()) |
| self.assertEqual( |
| 4, array_ops.size( |
| constant_op.constant([[1.0, 2.0], [3.0, 4.0]])).numpy()) |
| |
| def testRank(self): |
| with self.test_scope(): |
| self.assertEqual( |
| 0, array_ops.rank(constant_op.constant(1.0)).numpy()) |
| self.assertEqual( |
| 1, array_ops.rank(constant_op.constant([1.0, 2.0, 3.0])).numpy()) |
| self.assertEqual( |
| 2, array_ops.rank( |
| constant_op.constant([[1.0, 2.0], [3.0, 4.0]])).numpy()) |
| |
| def testAdam(self): |
| with self.test_scope(): |
| optimizer = adam.AdamOptimizer(0.1) |
| x = resource_variable_ops.ResourceVariable(10.0) |
| with backprop.GradientTape() as tape: |
| y = x * x |
| dy_dx = tape.gradient(y, x) |
| optimizer.apply_gradients([(dy_dx, x)]) |
| self.assertAlmostEqual(9.9, x.numpy(), places=3) |
| |
| def testAdamSparse(self): |
| with ops.device('/cpu:0'): |
| # Create 2-D embedding for 3 objects on CPU because sparse/sliced updates |
| # are not implemented on TPU. |
| embedding_matrix = resource_variable_ops.ResourceVariable( |
| array_ops.ones([3, 2])) |
| |
| with self.test_scope(): |
| with backprop.GradientTape() as tape: |
| embedding = embedding_ops.embedding_lookup(embedding_matrix, [1]) |
| y = math_ops.reduce_sum(embedding) |
| dy_dx = tape.gradient(y, embedding_matrix) |
| self.assertIsInstance(dy_dx, indexed_slices.IndexedSlices) |
| optimizer = adam.AdamOptimizer(0.1) |
| # The gradient application operations will run on CPU because optimizer |
| # updates are always collocated with the variable. |
| optimizer.apply_gradients([(dy_dx, embedding_matrix)]) |
| |
| # This assign_add will run on CPU because when an input to an |
| # operation is a resource, this operation is placed on the resource's |
| # device by the eager runtime. |
| embedding_matrix.assign_add(array_ops.ones([3, 2])) |
| |
| self.assertAllClose([[2.0, 2.0], |
| [1.9, 1.9], |
| [2.0, 2.0]], embedding_matrix.numpy()) |
| |
| |
| class EagerFunctionTest(xla_test.XLATestCase): |
| |
| def testBasic(self): |
| with self.test_scope(): |
| matmul = def_function.function(math_ops.matmul) |
| t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]]) |
| sq = matmul(t, t, transpose_a=True) |
| self.assertAllEqual(sq.numpy().reshape(-1), [10, 14, 14, 20]) |
| |
| def testConv(self): |
| if 'GPU' in self.device: |
| # TODO(b/32333178) |
| self.skipTest('Current implementation of RandomStandardNormal kernel ' |
| 'is very slow on GPU, and has been denylisted.') |
| with self.test_scope(): |
| data_format = 'channels_last' |
| conv = convolutional.Conv2D( |
| filters=1, kernel_size=2, padding='VALID', |
| data_format=data_format, activation=nn_ops.relu, |
| kernel_initializer=init_ops.ones_initializer(), |
| bias_initializer=init_ops.zeros_initializer()) |
| pool = pooling.MaxPooling2D(2, 2, data_format=data_format) |
| |
| def model(x): |
| x = conv(x) |
| return pool(x) |
| model = def_function.function(model) |
| |
| x = array_ops.ones([1, 4, 4, 1]) |
| y = model(x) |
| self.assertAllEqual(y.numpy(), [[[[4.]]]]) |
| |
| def testReadVariable(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(1.0) |
| |
| @def_function.function |
| def f(): |
| return v.read_value() |
| |
| var = f() |
| self.assertEqual(1.0, var.numpy()) |
| |
| def testResourceVariableNoInlineReadWrite(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(1.0) |
| w = resource_variable_ops.ResourceVariable(0.0) |
| |
| @def_function.function(experimental_attributes={'_noinline': True}) |
| def g(x): |
| w.assign(w.read_value() + x) |
| return v.read_value() + x * w.read_value() |
| |
| @def_function.function(experimental_attributes={'_noinline': True}) |
| def f(): |
| return g(1.0) + g(2.0) + g(3.0) + g(4.0) + g(5.0) |
| |
| # 1 + 1*1 + 1 + 2*3 + 1 + 3*6 + 1 + 4*10 + 1 + 5*15 |
| self.assertEqual(145.0, f().numpy()) |
| self.assertEqual(15.0, w.read_value().numpy()) |
| |
| def testResourceVariableNoInlineReadOnly(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(10.0) |
| |
| @def_function.function(experimental_attributes={'_noinline': True}) |
| def g(): |
| return v.read_value() |
| |
| @def_function.function(experimental_attributes={'_noinline': True}) |
| def f(): |
| return g() + g() + g() + g() + g() |
| |
| self.assertEqual(50.0, f().numpy()) |
| |
| def testResourceVariableNoInlineWriteOnly(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(0.0) |
| |
| @def_function.function(experimental_attributes={'_noinline': True}) |
| def g(x): |
| v.assign(x) |
| |
| @def_function.function(experimental_attributes={'_noinline': True}) |
| def f(): |
| g(1.0) |
| g(2.0) |
| g(3.0) |
| g(4.0) |
| g(5.0) |
| |
| f() |
| self.assertEqual(5.0, v.read_value().numpy()) |
| |
| def testUpdateVariable(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable(1.0) |
| |
| def f(v): |
| v.assign_add(1.0) |
| return v |
| |
| f = def_function.function(f) |
| |
| var = f(v) |
| self.assertEqual(2.0, var.numpy()) |
| |
| def testReturnResourceHandle(self): |
| with self.test_scope(): |
| v = resource_variable_ops.ResourceVariable([[1.0, 2.0], [3.0, 4.0]]) |
| |
| def f(v): |
| return v.handle |
| |
| f = def_function.function(f) |
| handle = f(v) |
| self.assertAllEqual(v.numpy(), |
| resource_variable_ops.read_variable_op( |
| handle, dtypes.float32).numpy()) |
| |
| def testReturnMultipleResourceHandles(self): |
| with self.test_scope(): |
| v1 = resource_variable_ops.ResourceVariable(1.25) |
| v2 = resource_variable_ops.ResourceVariable(2.0) |
| |
| def f(v): |
| return v.handle, 3.0 * v, v2.handle, v + v2 |
| |
| f = def_function.function(f) |
| v1_handle, v1_times_3, v2_handle, variable_sum = f(v1) |
| self.assertAllEqual(v1.numpy(), |
| resource_variable_ops.read_variable_op( |
| v1_handle, dtypes.float32).numpy()) |
| self.assertEqual(3.75, v1_times_3.numpy()) |
| self.assertAllEqual(v2.numpy(), |
| resource_variable_ops.read_variable_op( |
| v2_handle, dtypes.float32).numpy()) |
| self.assertEqual(3.25, variable_sum.numpy()) |
| |
| def testAllArgumentKinds(self): |
| """Test a complex function that takes different argument kinds. |
| |
| tf2xla machinery that translates, compiles, and runs defuns |
| classifies arguments into: compile-time constants, regular tensors, |
| and resources. This test creates a function with a mix of all these |
| kinds. Moreover, the order of function arguments is intentionally mixed up. |
| |
| This also tests the case when the same argument is a compile-time constant |
| as well as used in an operation that normally expects its inputs to be |
| in device memory - addition in this case. |
| """ |
| with self.test_scope(): |
| def foo(c1, r1, v1, c2, v2, r2): |
| # c1 and c2 are compile-time constants |
| # r1 and r2 are regular tensors |
| # v1 and v2 are resource variables |
| a = c1 + r1 |
| b = math_ops.cast(c2, dtypes.float32) + v2 |
| c = array_ops.slice(v1, c1, c2) |
| d = r2 * v2 |
| return a, b, c, d |
| |
| foo = def_function.function(foo) |
| |
| c1 = [0, 0] |
| c2 = array_ops.ones([2], dtype=dtypes.int32) |
| |
| r1 = array_ops.ones([2]) |
| r2 = [[2., 2.], [3., 3.]] |
| |
| v1 = resource_variable_ops.ResourceVariable([[1., 2.], [3., 4.]]) |
| v2 = resource_variable_ops.ResourceVariable([[10., 20.], [30., 40.]]) |
| |
| a, b, c, d = foo(c1, r1, v1, c2, v2, r2) |
| |
| self.assertAllEqual([1, 1], a.numpy()) |
| self.assertAllEqual([[11., 21.], [31., 41.]], b.numpy()) |
| self.assertAllEqual([[1.]], c.numpy()) |
| self.assertAllEqual([[20., 40.], [90., 120.]], d.numpy()) |
| |
| def testDefunInGradientTape(self): |
| with self.test_scope(): |
| v0 = resource_variable_ops.ResourceVariable(5.0) |
| |
| @def_function.function |
| def f(x): |
| x = v0 * v0 * x |
| return x |
| |
| x = constant_op.constant(3.0) |
| with backprop.GradientTape() as tape: |
| y = f(x) |
| dy = tape.gradient(y, v0) |
| |
| self.assertEqual(75, y.numpy()) |
| self.assertEqual(30, dy.numpy()) |
| |
| def testGradientTapeInDefun(self): |
| with self.test_scope(): |
| v0 = resource_variable_ops.ResourceVariable(5.0) |
| |
| @def_function.function |
| def f(): |
| x = constant_op.constant(1.0) |
| with backprop.GradientTape() as tape: |
| y = v0 * x |
| dy = tape.gradient(y, v0) |
| return dy |
| |
| dy = f() |
| self.assertEqual(1.0, dy.numpy()) |
| |
| def testSliceInDefun(self): |
| with self.test_scope(): |
| |
| @def_function.function |
| def f(x, y): |
| return x[0::2, y:, ...] |
| |
| x = array_ops.ones([2, 3, 4], dtype=dtypes.float32) |
| y = array_ops.ones([], dtype=dtypes.int32) |
| with backprop.GradientTape() as tape: |
| tape.watch(x) |
| tape.watch(y) |
| z = f(x, y) |
| dz = tape.gradient(z, x) |
| |
| self.assertAllEqual(np.ones([1, 2, 4]), z.numpy()) |
| self.assertAllEqual((2, 3, 4), dz.shape.as_list()) |
| |
| def testNestedDefun(self): |
| with self.test_scope(): |
| |
| @def_function.function |
| def times_two(x): |
| return 2. * x |
| |
| @def_function.function |
| def two_x_plus_1(x): |
| return times_two(x) + 1. |
| |
| x = constant_op.constant([2., 3., 4.]) |
| y = two_x_plus_1(x) |
| self.assertAllEqual([5., 7., 9.], y.numpy()) |
| |
| def testNestedDefunWithVariable(self): |
| with self.test_scope(): |
| v0 = resource_variable_ops.ResourceVariable(5.0) |
| |
| @def_function.function |
| def g(x): |
| x = v0 * x |
| return x |
| |
| @def_function.function |
| def f(x): |
| x = g(v0 * x) |
| return x |
| |
| x = constant_op.constant(3.0) |
| y = f(x) |
| |
| self.assertEqual(75.0, y.numpy()) |
| |
| def testNestedDefunInGradientTape(self): |
| with self.test_scope(): |
| v0 = resource_variable_ops.ResourceVariable(5.0) |
| |
| @def_function.function |
| def g(x): |
| x = v0 * x |
| return x |
| |
| @def_function.function |
| def f(x): |
| x = g(v0 * x) |
| return x |
| |
| x = constant_op.constant(3.0) |
| with backprop.GradientTape() as tape: |
| y = f(x) |
| dy = tape.gradient(y, v0) |
| |
| self.assertEqual(75, y.numpy()) |
| self.assertEqual(30, dy.numpy()) |
| |
| def testNestedDefunInGradientTapeDifferentVars(self): |
| with self.test_scope(): |
| v0 = resource_variable_ops.ResourceVariable(5.0) |
| v1 = resource_variable_ops.ResourceVariable(3.0) |
| |
| @def_function.function |
| def g(x): |
| x = v1 * x |
| return x |
| |
| @def_function.function |
| def f(x): |
| x = g(v0 * x) |
| return x |
| |
| x = constant_op.constant(3.0) |
| with backprop.GradientTape(persistent=True) as tape: |
| y = f(x) |
| dy_v0 = tape.gradient(y, v0) |
| dy_v1 = tape.gradient(y, v1) |
| |
| self.assertEqual(45, y.numpy()) |
| self.assertEqual(9, dy_v0.numpy()) |
| self.assertEqual(15, dy_v1.numpy()) |
| |
| def testWhileInDefun(self): |
| with self.test_scope(): |
| @def_function.function |
| def f(start): |
| c = lambda x: math_ops.less(x, 13.0) |
| b = lambda x: math_ops.add(x, 1.0) |
| return while_loop.while_loop(c, b, [start]) |
| |
| y = f(constant_op.constant(3.0)) |
| self.assertEqual(13.0, y.numpy()) |
| |
| def testAutoGraphWhileInDefun(self): |
| with self.test_scope(): |
| @def_function.function |
| def f(start): |
| x = start |
| while x < 13.0: |
| x += 1.0 |
| return x |
| |
| y = f(constant_op.constant(3.0)) |
| self.assertEqual(13.0, y.numpy()) |
| |
| def testCondInDefun(self): |
| with self.test_scope(): |
| @def_function.function |
| def f(pred, value): |
| fn1 = lambda: math_ops.add(value, 1.0) |
| fn2 = lambda: math_ops.subtract(value, 1.0) |
| return cond.cond(pred, fn1, fn2) |
| |
| plus_one = f(constant_op.constant(True), constant_op.constant(10.0)) |
| minus_one = f(constant_op.constant(False), constant_op.constant(10.0)) |
| self.assertEqual(11.0, plus_one.numpy()) |
| self.assertEqual(9.0, minus_one.numpy()) |
| |
| def testAutoGraphCondInDefun(self): |
| with self.test_scope(): |
| @def_function.function |
| def f(pred, value): |
| if pred: |
| return value + 1.0 |
| else: |
| return value - 1.0 |
| |
| plus_one = f(constant_op.constant(True), constant_op.constant(10.0)) |
| minus_one = f(constant_op.constant(False), constant_op.constant(10.0)) |
| self.assertEqual(11.0, plus_one.numpy()) |
| self.assertEqual(9.0, minus_one.numpy()) |
| |
| def testScanInDefun(self): |
| with self.test_scope(): |
| elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name='data') |
| v = constant_op.constant(2.0, name='v') |
| |
| @def_function.function |
| def f(y): |
| # pylint: disable=unnecessary-lambda |
| return functional_ops.scan( |
| lambda a, x: math_ops.multiply(a, x), y, initializer=v) |
| # pylint: enable=unnecessary-lambda |
| |
| r = f(elems) |
| self.assertAllEqual([2., 4., 12., 48., 240., 1440.], self.evaluate(r)) |
| |
| def testFeedDeviceMemoryToOpExpectingHostMemory(self): |
| @def_function.function |
| 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) |
| |
| def testRequestNotToCompile(self): |
| with self.test_scope(): |
| def f(x): |
| with ops.device('device:CPU:0'): |
| y = 2.0 * x |
| return x, y |
| |
| wholly_compiled_f = def_function.function(f) |
| op_by_op_f = def_function.function(f, jit_compile=False) |
| |
| x = array_ops.identity([0.0, 2.0], name='data') |
| |
| # When function is wholly compiled, all outputs will be on the |
| # device on which it is run. |
| r_x, r_y = wholly_compiled_f(x) |
| self.assertAllEqual([0.0, 2.0], r_x) |
| self.assertAllEqual([0.0, 4.0], r_y) |
| if context.executing_eagerly(): |
| # backing_device is only available for eager tensors. |
| self.assertRegex(r_x.backing_device, self.device) |
| self.assertRegex(r_y.backing_device, self.device) |
| |
| # When function is executed op-by-op, requested devices will be |
| # respected. |
| r_x, r_y = op_by_op_f(x) |
| self.assertAllEqual([0.0, 2.0], r_x) |
| self.assertAllEqual([0.0, 4.0], r_y) |
| if context.executing_eagerly(): |
| # backing_device is only available for eager tensors. |
| self.assertRegex(r_x.backing_device, self.device) |
| self.assertRegex(r_y.backing_device, 'device:CPU:0') |
| |
| |
| class ExcessivePaddingTest(xla_test.XLATestCase): |
| """Test that eager execution works with TPU flattened tensors. |
| |
| Tensors that would normally be excessively padded when written |
| to TPU memory are reshaped to 1-D flat tensors. |
| |
| This test case verifies that such tensors work with eager execution. |
| |
| The flattening currently only happens on TPU, but tests should work |
| fine with all backends as flattening is transparent. |
| """ |
| |
| def testFromConstant(self): |
| with self.test_scope(): |
| # Create constant of shape [100, 2, 1]. This tensor would be |
| # excessively padded on TPU. |
| tensor = constant_op.constant(100 * [[[10.0], [2.0]]]) |
| # Use reduce_sum since it requires correctly working with |
| # a particular dimension. |
| reduced = math_ops.reduce_sum(tensor, axis=1) |
| self.assertAllEqual(100 * [[12.0]], reduced) |
| |
| def testFromOperation(self): |
| with self.test_scope(): |
| tensor = array_ops.ones([3, 100, 2, 2]) |
| reduced = math_ops.reduce_sum(tensor, axis=[0, 2, 3]) |
| self.assertAllEqual(100 * [12.0], reduced) |
| |
| def testAsFunctionInput(self): |
| with self.test_scope(): |
| |
| @def_function.function |
| def f(x): |
| return math_ops.reduce_sum(x, axis=2) |
| |
| tensor = constant_op.constant(100 * [[[10.0, 2.0]]]) |
| reduced = f(tensor) |
| self.assertAllEqual(100 * [[12.0]], reduced) |
| |
| def testAsFunctionOutput(self): |
| with self.test_scope(): |
| |
| @def_function.function |
| def f(x): |
| return x * constant_op.constant(100 * [[[10.0, 2.0]]]) |
| |
| y = f(3) |
| reduced = math_ops.reduce_sum(y, axis=2) |
| self.assertAllEqual(100 * [[36.0]], reduced) |
| |
| |
| def multiple_tpus(): |
| devices = context.context().devices() |
| return len([d for d in devices if 'device:TPU:' in d]) > 1 |
| |
| |
| class MultiDeviceTest(xla_test.XLATestCase): |
| """Test running TPU computation on more than one core.""" |
| |
| def testBasic(self): |
| if not multiple_tpus(): |
| self.skipTest('MultiDeviceTest requires multiple TPU devices.') |
| |
| # Compute 10 on TPU core 0 |
| with ops.device('device:TPU:0'): |
| two = constant_op.constant(2) |
| five = constant_op.constant(5) |
| ten = two * five |
| self.assertAllEqual(10, ten) |
| |
| # Compute 6 on TPU core 1 |
| with ops.device('device:TPU:1'): |
| two = constant_op.constant(2) |
| three = constant_op.constant(3) |
| six = two * three |
| self.assertAllEqual(6, six) |
| |
| # Copy 10 and 6 to CPU and sum them |
| self.assertAllEqual(16, ten + six) |
| |
| |
| if __name__ == '__main__': |
| ops.enable_eager_execution( |
| config=config_pb2.ConfigProto(log_device_placement=True)) |
| googletest.main() |
