def testMultipleWhileLoops(self):
x = constant_op.constant(2.)
ret1 = while_loop_v2(lambda v: v < 4., lambda v: v * v, [x]) # x**2
ret2 = while_loop_v2(lambda v: v < 16., lambda v: v * v, ret1) # x**4
grad = gradients_impl.gradients(ret2, [x]) # 4x**3
grad_grad = gradients_impl.gradients(grad, [x]) # 12x**2
with self.cached_session() as sess:
self.assertSequenceEqual(sess.run(grad), [32.])
self.assertSequenceEqual(sess.run(grad_grad), [48.])
def testNestedWhileAndTensorArray(self):
n = constant_op.constant(3.0)
def Body(row, ta, n):
def InnerBody(row, col, ta, n):
# Note: row and col are 1-based.
ta = ta.write(
math_ops.cast(n * (row - 1.) + col - 1., dtypes.int32), row * col)
return row, col + 1., ta, n
# TODO(b/118457764): Remove n from loop_vars from both loops once fixed.
ta = while_loop_v2(
lambda _, col, _1, n: col <= n,
InnerBody, [row, constant_op.constant(1.), ta, n],
return_same_structure=False)[2]
return row + 1., ta, n
ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=9)
ta = while_loop_v2(
lambda row, _, _1: row <= n,
Body, [constant_op.constant(1.), ta, n],
return_same_structure=False)[1]
output = array_ops.reshape(ta.stack(), [3, 3])
self.assertAllEqual(
self.evaluate(output), [[1., 2., 3.], [2., 4., 6.], [3., 6., 9.]])
def fnWithLoop(): # pylint: disable=invalid-name
with backprop.GradientTape() as tape:
_, x = while_loop_v2(
lambda i, _: i < 2,
lambda i, x: (i + 1, x * v),
[0, 2.])
return tape.gradient(x, v)
def Fn():
x = constant_op.constant(2.)
with backprop.GradientTape() as tape:
tape.watch(x)
ret1 = while_loop_v2(
lambda v: v < 4.,
lambda v: v * v, [x],
return_same_structure=False,
name="while_1") # x**2
ret2 = while_loop_v2(
lambda v: v < 16.,
lambda v: v * v, [x],
return_same_structure=False,
name="while_2") # x**4
loss = ret1 + ret2
return tape.gradient(loss, x)
def testMultipleLoopVars(self):
x = constant_op.constant(5.)
y = constant_op.constant(3.)
# x = 5.
# y = 3.
# while x < 45.:
# x = x * y
# y = x + y
ret = while_loop_v2(lambda v, _: v < 45.,
lambda v, w: (v * w, v + w), [x, y],
return_same_structure=False)
# ret = [y*x**2 + x*y**2, x*y + x + y]
gradx_0 = gradients_impl.gradients(ret[0], [x]) # [2*x*y + y**2]
gradx_1 = gradients_impl.gradients(ret[1], [x]) # [y + 1]
gradx_2 = gradients_impl.gradients(ret,
[x]) # [2*x*y + y**2 + 2*y + 1]
grady_0 = gradients_impl.gradients(ret[0], [y]) # [2*x*y + x**2]
grady_1 = gradients_impl.gradients(ret[1], [y]) # [x + 1]
grady_2 = gradients_impl.gradients(ret, [y]) # [2*x*y + x**2 + x + 1]
with self.cached_session() as sess:
self.assertSequenceEqual(self.evaluate(ret), [120., 23.])
self.assertSequenceEqual(self.evaluate(gradx_0), [39.])
self.assertSequenceEqual(self.evaluate(gradx_1), [4.])
self.assertSequenceEqual(self.evaluate(gradx_2), [43.])
self.assertSequenceEqual(self.evaluate(grady_0), [55.])
self.assertSequenceEqual(self.evaluate(grady_1), [6.])
self.assertSequenceEqual(self.evaluate(grady_2), [61.])
def testNestedWhileAndTensorArray(self):
n = constant_op.constant(3.0)
def Body(row, ta, n):
def InnerBody(row, col, ta, n):
# Note: row and col are 1-based.
ta = ta.write(
math_ops.cast(n * (row - 1.) + col - 1., dtypes.int32),
row * col)
return row, col + 1., ta, n
# TODO(b/118457764): Remove n from loop_vars from both loops once fixed.
ta = while_loop_v2(lambda _, col, _1, n: col <= n,
InnerBody,
[row, constant_op.constant(1.), ta, n],
return_same_structure=False)[2]
return row + 1., ta, n
ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=9)
ta = while_loop_v2(lambda row, _, _1: row <= n,
Body, [constant_op.constant(1.), ta, n],
return_same_structure=False)[1]
output = array_ops.reshape(ta.stack(), [3, 3])
self.assertAllEqual(self.evaluate(output),
[[1., 2., 3.], [2., 4., 6.], [3., 6., 9.]])
def testDuplicateAccumulator(self):
x = constant_op.constant(2.)
tensor_list = list_ops.empty_tensor_list(
element_dtype=dtypes.float32, element_shape=ScalarShape())
def Cond(x, tl):
del tl # Unused for Cond.
return x < 5.
def Body(x, tl):
# There is an accumulator in the loop already so we should not add
# another.
tl = list_ops.tensor_list_push_back(tl, x)
return x**2., tl
ret = while_loop_v2(
Cond, Body, [x, tensor_list], return_same_structure=False)
for op in ops.get_default_graph().get_operations():
if op.type == "While":
while_op = op
body_graph = while_v2._get_graph(while_op, "body")
x_input_index = [i for i, inp in enumerate(while_op.inputs) if inp == x][0]
x_input_t = body_graph.inputs[x_input_index]
accumulator_count = len(
[c for c in x_input_t.consumers() if c.type == "TensorListPushBack"])
self.assertEqual(accumulator_count, 1)
grad = gradients_impl.gradients(ret[0], x)
with self.cached_session() as sess:
self.assertEqual(sess.run(ret[0]), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
def testSingleLoopVar(self):
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8., lambda v: v * v, [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(sess.run(ret), 16.)
self.assertSequenceEqual(sess.run(grad), [32.])
def testDuplicateAccumulator(self):
x = constant_op.constant(2.)
tensor_list = list_ops.empty_tensor_list(
element_dtype=dtypes.float32, element_shape=ScalarShape())
def Cond(x, tl):
del tl # Unused for Cond.
return x < 5.
def Body(x, tl):
# There is an accumulator in the loop already so we should not add
# another.
tl = list_ops.tensor_list_push_back(tl, x)
return x**2., tl
ret = while_loop_v2(Cond, Body, [x, tensor_list])
for op in ops.get_default_graph().get_operations():
if op.type == "While":
while_op = op
body_graph = while_v2._get_body_graph(while_op)
# body_graph.inputs: [counter_arg, x_arg, tl_arg, *accumulators]
x_input_t = body_graph.inputs[1]
accumulator_count = len(
[c for c in x_input_t.consumers() if c.type == "TensorListPushBack"])
self.assertEqual(accumulator_count, 1)
grad = gradients_impl.gradients(ret[0], x)
with self.cached_session() as sess:
self.assertEqual(sess.run(ret[0]), 16.)
self.assertSequenceEqual(sess.run(grad), [32.])
def testSingleLoopVar(self):
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8., lambda v: v * v, [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(sess.run(ret), 16.)
self.assertSequenceEqual(sess.run(grad), [32.])
def testMultipleLoopVars(self):
x = constant_op.constant(5.)
y = constant_op.constant(3.)
# x = 5.
# y = 3.
# while x < 45.:
# x = x * y
# y = x + y
ret = while_loop_v2(lambda v, _: v < 45., lambda v, w: (v * w, v + w),
[x, y])
# ret = [y*x**2 + x*y**2, x*y + x + y]
gradx_0 = gradients_impl.gradients(ret[0], [x]) # [2*x*y + y**2]
gradx_1 = gradients_impl.gradients(ret[1], [x]) # [y + 1]
gradx_2 = gradients_impl.gradients(ret, [x]) # [2*x*y + y**2 + 2*y + 1]
grady_0 = gradients_impl.gradients(ret[0], [y]) # [2*x*y + x**2]
grady_1 = gradients_impl.gradients(ret[1], [y]) # [x + 1]
grady_2 = gradients_impl.gradients(ret, [y]) # [2*x*y + x**2 + x + 1]
with self.cached_session() as sess:
self.assertSequenceEqual(sess.run(ret), [120., 23.])
self.assertSequenceEqual(sess.run(gradx_0), [39.])
self.assertSequenceEqual(sess.run(gradx_1), [4.])
self.assertSequenceEqual(sess.run(gradx_2), [43.])
self.assertSequenceEqual(sess.run(grady_0), [55.])
self.assertSequenceEqual(sess.run(grady_1), [6.])
self.assertSequenceEqual(sess.run(grady_2), [61.])
def testNestedWhileWithLegacyDefun(self):
n = constant_op.constant(3.)
m = constant_op.constant(5.)
sum_of_powers = constant_op.constant(0.)
def Body(i, previous_sum):
prod = constant_op.constant(1.)
def InnerBodyWrapper(c, v):
@function.Defun(dtypes.float32, dtypes.float32)
def InnerBody(c, v):
return c - 1., v * n
results = InnerBody(c, v)
results[0].set_shape([])
results[1].set_shape([])
return results
return i - 1., previous_sum + while_loop_v2(
lambda c, _: c > 0,
InnerBodyWrapper, [i, prod],
return_same_structure=False)[1]
result = while_loop_v2(lambda i, _: i >= 0,
Body, [m, sum_of_powers],
return_same_structure=False)[1]
grad = gradients_impl.gradients(result, [n])
self.assertEqual(self.evaluate(result), 364.)
self.assertSequenceEqual(self.evaluate(grad), [547.])
def Fn():
with backprop.GradientTape() as tape:
x = constant_op.constant(2.)
tape.watch(x)
def Body(i, x):
forward_graph = ops.get_default_graph()
@custom_gradient.custom_gradient
def SquaredWithZeroGrad(x):
def Grad(unused_g, variables=None): # pylint: disable=redefined-outer-name
del variables
gradient_graph = ops.get_default_graph()
shape = gen_array_ops.shape(x)
assert shape.graph is forward_graph
rank = gen_array_ops.rank(x)
assert rank.graph is forward_graph
size = gen_array_ops.size(x)
assert size.graph is forward_graph
zeros = array_ops.zeros(shape)
assert zeros.graph is gradient_graph
return zeros
return x * 2, Grad
return i + 1, SquaredWithZeroGrad(x)
_, result = while_loop_v2(lambda i, _: i < 2, Body, [0, x])
grad = tape.gradient(result, x)
return grad
def Fn():
def Body1(v):
x1.assign(x1)
return v * x1
ret1 = while_loop_v2(
lambda v: v < 4.,
Body1, [c],
return_same_structure=False,
name="while_1") # 2x
def Body2(v):
x1.assign(x1)
return v * x1 * x1
ret2 = while_loop_v2(
lambda v: v < 16.,
Body2, [c],
return_same_structure=False,
name="while_2") # 4x
def Body3(v):
x2.assign(x2)
return v * x2
ret3 = while_loop_v2(
lambda v: v < 4.,
Body3, [c],
return_same_structure=False,
name="while_3") # 3x
def Body4(v):
x2.assign(x2)
return v * x2 * x2
ret4 = while_loop_v2(
lambda v: v < 16.,
Body4, [c],
return_same_structure=False,
name="while_4") # 9x
ret5 = while_loop_v2(
lambda v: v < 16.,
lambda v: v * v, [c],
return_same_structure=False,
name="while_stateless") # x**2
return ret1, ret2, ret3, ret4, ret5
def testSingleLoopVar(self):
x = constant_op.constant(2.)
ret = while_loop_v2(
lambda v: v < 8., lambda v: v * v, [x], return_same_structure=False)
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
def testCaptureExternalTensorInBody(self):
x = constant_op.constant(2.)
y = constant_op.constant(3.)
ret = while_loop_v2(lambda v: v < 8., lambda v: v * y, [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(self.evaluate(ret), 18.)
self.assertSequenceEqual(self.evaluate(grad), [9.])
def Func():
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8.,
lambda v: v**2, [x],
return_same_structure=False) # x**4
grad = gradients_impl.gradients(ret, [x])[0] # 4x**3
grad_grad = gradients_impl.gradients(grad, [x])[0] # 12x**2
return ret, grad, grad_grad
def testCaptureExternalTensorInCond(self):
x = constant_op.constant(2.)
y = constant_op.constant(1.)
ret = while_loop_v2(lambda v: v + y < 9., lambda v: v * 3., [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(self.evaluate(ret), 18.)
self.assertSequenceEqual(self.evaluate(grad), [9.])
def testCaptureExternalTensorInCond(self):
x = constant_op.constant(2.)
y = constant_op.constant(1.)
ret = while_loop_v2(lambda v: v + y < 9., lambda v: v * 3., [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(sess.run(ret), 18.)
self.assertSequenceEqual(sess.run(grad), [9.])
def testCaptureExternalTensorInBody(self):
x = constant_op.constant(2.)
y = constant_op.constant(3.)
ret = while_loop_v2(lambda v: v < 8., lambda v: v * y, [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(sess.run(ret), 18.)
self.assertSequenceEqual(sess.run(grad), [9.])
def testSingleLoopVar(self):
x = constant_op.constant(2.)
ret = while_loop_v2(
lambda v: v < 8., lambda v: v * v, [x], return_same_structure=False)
grad = gradients_impl.gradients(ret, [x])
with self.cached_session():
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
def testDoubleDerivative(self):
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8., lambda v: v**2, [x]) # x**4
grad = gradients_impl.gradients(ret, [x]) # 4x**3
grad_grad = gradients_impl.gradients(grad, [x]) # 12x**2
with self.cached_session() as sess:
self.assertEqual(sess.run(ret), 16.)
self.assertSequenceEqual(sess.run(grad), [32.])
self.assertSequenceEqual(sess.run(grad_grad), [48.])
def testDoubleDerivative(self):
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8., lambda v: v**2, [x]) # x**4
grad = gradients_impl.gradients(ret, [x]) # 4x**3
grad_grad = gradients_impl.gradients(grad, [x]) # 12x**2
with self.cached_session() as sess:
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
self.assertSequenceEqual(self.evaluate(grad_grad), [48.])
def testAccumulatorElementShape(self, shape):
def MatchShape(actual_tensor_shape):
# Compare the shapes, treating None dimensions as equal. We do not
# directly check actual_tensor_shape and tf.TensorShape(shape) for
# equality because tf.Dimension.__eq__ returns None if either dimension is
# None.
if shape is None:
self.assertIsNone(actual_tensor_shape.dims)
else:
self.assertListEqual(actual_tensor_shape.as_list(), shape)
def GetAccumulatorForInputAtIndex(while_op, idx):
body_graph = while_v2._get_graph(while_op, "body")
y_input_t = body_graph.inputs[idx]
push_back_node = [
c for c in y_input_t.consumers()
if c.type == "TensorListPushBack"
][0]
output_idx = body_graph.outputs.index(push_back_node.outputs[0])
return while_op.outputs[output_idx]
x = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
y = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
# Forward pass.
ret = while_loop_v2(lambda v, u: v < 8.,
lambda v, u: (math_ops.pow(v, u), u), [x, y],
return_same_structure=True)
while_op = ret[0].op.inputs[0].op
# Gradient pass.
grad = gradients_impl.gradients(ret[0], x)
# Note: There is an Identity b/w grad[0] and the While op.
grad_while_op = grad[0].op.inputs[0].op
# Get the TensorList output of While op containing the accumulated values
# of y.
x_input_index = [
i for i, inp in enumerate(while_op.inputs) if x == inp
][0]
output = GetAccumulatorForInputAtIndex(while_op, x_input_index)
_, val = list_ops.tensor_list_pop_back(output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
# Take second derivative to generate intermediate grad_while_op outputs
gradients_impl.gradients(grad, x)
# Get the TensorList output of gradient While op containing the accumulated
# values of grad_x (note that grad_x is needed by the second derivative).
# grad_while_op.inputs:
grad_output_index = grad_while_op.outputs.index(grad[0].op.inputs[0])
grad_output = GetAccumulatorForInputAtIndex(grad_while_op,
grad_output_index)
_, val = list_ops.tensor_list_pop_back(grad_output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
def testSingleLoopVarBackPropFalse(self):
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8.,
lambda v: v * v, [x],
return_same_structure=False,
back_prop=False)
grad = gradients_impl.gradients(ret, [x])
self.assertEqual(grad, [None])
with self.cached_session():
self.assertEqual(self.evaluate(ret), 16.)
def testGradientTape(self):
with backprop.GradientTape() as t:
x = constant_op.constant(2.)
t.watch(x)
ret = while_loop_v2(lambda v: v < 4.,
lambda v: v * v, [x],
return_same_structure=False) # x**2
grad = t.gradient(ret, x)
with self.cached_session() as sess:
self.assertAllEqual(sess.run(grad), 4.0)
def testGradientTape(self):
with backprop.GradientTape() as t:
x = constant_op.constant(2.)
t.watch(x)
ret = while_loop_v2(
lambda v: v < 4., lambda v: v * v, [x],
return_same_structure=False) # x**2
grad = t.gradient(ret, x)
with self.cached_session() as sess:
self.assertAllEqual(sess.run(grad), 4.0)
def Fn():
def Body1(v):
x.assign(x)
return v * x
ret1 = while_loop_v2(lambda v: v < 4.,
Body1, [c],
return_same_structure=False,
name="while_1") # 2x
def Body2(v):
x.assign(x)
return v * x * x
ret2 = while_loop_v2(lambda v: v < 16.,
Body2, [c],
return_same_structure=False,
name="while_2") # 4x
return ret1, ret2
def testCaptureExternalTensorInCond(self):
x = constant_op.constant(2.)
y = constant_op.constant(1.)
ret = while_loop_v2(lambda v: v + y < 9.,
lambda v: v * 3., [x],
return_same_structure=False)
grad = gradients_impl.gradients(ret, [x])
with self.cached_session():
self.assertEqual(self.evaluate(ret), 18.)
self.assertSequenceEqual(self.evaluate(grad), [9.])
def testReturnSameStructureTrue(self):
x = constant_op.constant(2.)
ret = while_loop_v2(
lambda v: v < 8., lambda v: v * v, [x], return_same_structure=True)
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
eval_result = sess.run(ret)
self.assertIsInstance(eval_result, list)
self.assertLen(eval_result, 1)
self.assertEqual(16., eval_result[0])
self.assertSequenceEqual(sess.run(grad), [32.])
def testAccumulatorElementShape(self, shape):
def MatchShape(actual_tensor_shape):
# Compare the shapes, treating None dimensions as equal. We do not
# directly check actual_tensor_shape and tf.TensorShape(shape) for
# equality because tf.Dimension.__eq__ returns None if either dimension is
# None.
if shape is None:
self.assertIsNone(actual_tensor_shape.dims)
else:
self.assertListEqual(actual_tensor_shape.as_list(), shape)
def GetAccumulatorForInputAtIndex(while_op, idx):
body_graph = while_v2._get_graph(while_op, "body")
y_input_t = body_graph.inputs[idx]
push_back_node = [c for c in y_input_t.consumers()
if c.type == "TensorListPushBack"][0]
output_idx = body_graph.outputs.index(push_back_node.outputs[0])
return while_op.outputs[output_idx]
x = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
y = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
# Forward pass.
ret = while_loop_v2(lambda v, u: v < 8.,
lambda v, u: (math_ops.pow(v, u), u),
[x, y],
return_same_structure=True)
while_op = ret[0].op.inputs[0].op
# Gradient pass.
grad = gradients_impl.gradients(ret[0], x)
# Note: There is an Identity b/w grad[0] and the While op.
grad_while_op = grad[0].op.inputs[0].op
# Get the TensorList output of While op containing the accumulated values
# of y.
x_input_index = [i for i, inp in enumerate(while_op.inputs) if x == inp][0]
output = GetAccumulatorForInputAtIndex(while_op, x_input_index)
_, val = list_ops.tensor_list_pop_back(output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
# Take second derivative to generate intermediate grad_while_op outputs
gradients_impl.gradients(grad, x)
# Get the TensorList output of gradient While op containing the accumulated
# values of grad_x (note that grad_x is needed by the second derivative).
# grad_while_op.inputs:
grad_output_index = grad_while_op.outputs.index(grad[0].op.inputs[0])
grad_output = GetAccumulatorForInputAtIndex(grad_while_op,
grad_output_index)
_, val = list_ops.tensor_list_pop_back(grad_output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
def testReturnSameStructureTrue(self):
x = constant_op.constant(2.)
ret = while_loop_v2(
lambda v: v < 8., lambda v: v * v, [x], return_same_structure=True)
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
eval_result = sess.run(ret)
self.assertIsInstance(eval_result, list)
self.assertLen(eval_result, 1)
self.assertEqual(16., eval_result[0])
self.assertSequenceEqual(sess.run(grad), [32.])
def Body(row, ta, n):
def InnerBody(row, col, ta, n):
# Note: row and col are 1-based.
ta = ta.write(
math_ops.cast(n * (row - 1.) + col - 1., dtypes.int32), row * col)
return row, col + 1., ta, n
# TODO(b/118457764): Remove n from loop_vars from both loops once fixed.
ta = while_loop_v2(lambda _, col, _1, n: col <= n, InnerBody,
[row, constant_op.constant(1.), ta, n])[2]
return row + 1., ta, n
def testExternalColocationGrad(self):
external_t = constant_op.constant(2.)
v0 = constant_op.constant(2.)
def Body(v):
with ops.colocate_with(external_t):
return v * v
ret = while_loop_v2(lambda v: v < 8., Body, [v0])[0]
grad = gradients_impl.gradients(ret, [v0])[0]
self.assertAllEqual(ret, 16.)
self.assertAllEqual(grad, 32.)
def testIdentityNodeInBody(self):
def Body(v):
v = array_ops.identity(v)
v = array_ops.identity(v)
return v * v
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8., Body, [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
def Body(row, ta, n):
def InnerBody(row, col, ta, n):
# Note: row and col are 1-based.
ta = ta.write(
math_ops.cast(n * (row - 1.) + col - 1., dtypes.int32),
row * col)
return row, col + 1., ta, n
# TODO(b/118457764): Remove n from loop_vars from both loops once fixed.
ta = while_loop_v2(lambda _, col, _1, n: col <= n, InnerBody,
[row, constant_op.constant(1.), ta, n])[2]
return row + 1., ta, n
def Fn():
ret1 = while_loop_v2(lambda v: v < 4.,
lambda v: v * x1, [c],
return_same_structure=False,
name="while_1") # 2x
ret2 = while_loop_v2(lambda v: v < 16.,
lambda v: v * x1 * x1, [c],
return_same_structure=False,
name="while_2") # 4x
ret3 = while_loop_v2(lambda v: v < 4.,
lambda v: v * x2, [c],
return_same_structure=False,
name="while_3") # 3x
ret4 = while_loop_v2(lambda v: v < 16.,
lambda v: v * x2 * x2, [c],
return_same_structure=False,
name="while_4") # 9x
ret5 = while_loop_v2(lambda v: v < 16.,
lambda v: v * v, [c],
return_same_structure=False,
name="while_stateless") # x**2
return ret1, ret2, ret3, ret4, ret5
def testIdentityNodeInBody(self):
def Body(v):
v = array_ops.identity(v)
v = array_ops.identity(v)
return v * v
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8., Body, [x])
grad = gradients_impl.gradients(ret, [x])
with self.cached_session() as sess:
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
def testExternalControlDependencies(self):
with ops.Graph().as_default(), self.test_session():
v = variables.Variable(1.)
v.initializer.run()
op = v.assign_add(1.)
def body_fn(i): # pylint: disable=invalid-name
with ops.control_dependencies([op]):
return i + 1
loop = while_loop_v2(lambda i: i < 1, body_fn, [0])
loop[0].op.run()
self.assertAllEqual(self.evaluate(v), 2.0)
def testExternalControlDependencies(self):
with ops.Graph().as_default(), self.test_session():
v = variables.Variable(1.)
v.initializer.run()
op = v.assign_add(1.)
def body_fn(i): # pylint: disable=invalid-name
with ops.control_dependencies([op]):
return i + 1
loop = while_loop_v2(lambda i: i < 1, body_fn, [0])
loop[0].op.run()
self.assertAllEqual(self.evaluate(v), 2.0)
def testIdentityNodeInBody(self):
def Body(v):
v = array_ops.identity(v)
v = array_ops.identity(v)
return v * v
x = constant_op.constant(2.)
ret = while_loop_v2(lambda v: v < 8.,
Body, [x],
return_same_structure=False)
grad = gradients_impl.gradients(ret, [x])
self.assertEqual(self.evaluate(ret), 16.)
self.assertSequenceEqual(self.evaluate(grad), [32.])
def testRandomUniformShape(self):
shape = constant_op.constant([3])
def Body(i, u):
shape_extended = array_ops.concat([[5], shape], axis=0)
u = random_ops.random_uniform(shape_extended)
self.assertAllEqual(u.shape.as_list(), [5, 3])
return i + 1, u
_, _ = while_loop_v2(cond=lambda i, _: i < 3,
body=Body,
loop_vars=[
0,
array_ops.zeros([5, 3], dtype=dtypes.float32),
])
def testAccumulatorElementShape(self, shape):
def MatchShape(actual_tensor_shape):
# Compare the shapes, treating None dimensions as equal. We do not
# directly check actual_tensor_shape and tf.TensorShape(shape) for
# equality because tf.Dimension.__eq__ returns None if either dimension is
# None.
if shape is None:
self.assertIsNone(actual_tensor_shape.dims)
else:
self.assertListEqual(actual_tensor_shape.as_list(), shape)
def GetAccumulatorForInputAtIndex(while_op, idx):
body_graph = while_v2._get_body_graph(while_op)
y_input_t = body_graph.inputs[idx]
push_back_node = [c for c in y_input_t.consumers()
if c.type == "TensorListPushBack"][0]
output_idx = body_graph.outputs.index(push_back_node.outputs[0])
return while_op.outputs[output_idx]
x = constant_op.constant(2.)
y = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
# Forward pass.
ret = while_loop_v2(
lambda v, u: v < 8.,
lambda v, u: (v * v, u), [x, y],
return_same_structure=False)
while_op = ret[0].op.inputs[0].op
# Get the TensorList output of While op containing the accumulated values
# of y.
# while_op.inputs: [counter_arg, x_arg, y_arg, *accumulators]
output = GetAccumulatorForInputAtIndex(while_op, 2)
_, val = list_ops.tensor_list_pop_back(output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
# Gradient pass.
grad = gradients_impl.gradients(ret[1], y)
grad_while_op = grad[0].op.inputs[0].op
# Get the TensorList output of gradient While op containing the accumulated
# values of grad_y.
# grad_while_op.inputs:
# [counter_arg, total_iters_arg, grad_x_arg, grad_y_arg, *other_args]
grad_output = GetAccumulatorForInputAtIndex(grad_while_op, 3)
_, val = list_ops.tensor_list_pop_back(grad_output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
def testAccumulatorElementShape(self, shape):
def MatchShape(actual_tensor_shape):
# Compare the shapes, treating None dimensions as equal. We do not
# directly check actual_tensor_shape and tf.TensorShape(shape) for
# equality because tf.Dimension.__eq__ returns None if either dimension is
# None.
if shape is None:
self.assertIsNone(actual_tensor_shape.dims)
else:
self.assertListEqual(actual_tensor_shape.as_list(), shape)
def GetAccumulatorForInputAtIndex(while_op, idx):
body_graph = while_v2._get_body_graph(while_op)
y_input_t = body_graph.inputs[idx]
push_back_node = [
c for c in y_input_t.consumers()
if c.type == "TensorListPushBack"
][0]
output_idx = body_graph.outputs.index(push_back_node.outputs[0])
return while_op.outputs[output_idx]
x = constant_op.constant(2.)
y = array_ops.placeholder(dtype=dtypes.float32, shape=shape)
# Forward pass.
ret = while_loop_v2(lambda v, u: v < 8.,
lambda v, u: (v * v, u), [x, y],
return_same_structure=False)
while_op = ret[0].op.inputs[0].op
# Get the TensorList output of While op containing the accumulated values
# of y.
# while_op.inputs: [counter_arg, x_arg, y_arg, *accumulators]
output = GetAccumulatorForInputAtIndex(while_op, 2)
_, val = list_ops.tensor_list_pop_back(output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
# Gradient pass.
grad = gradients_impl.gradients(ret[1], y)
grad_while_op = grad[0].op.inputs[0].op
# Get the TensorList output of gradient While op containing the accumulated
# values of grad_y.
# grad_while_op.inputs:
# [counter_arg, total_iters_arg, grad_x_arg, grad_y_arg, *other_args]
grad_output = GetAccumulatorForInputAtIndex(grad_while_op, 3)
_, val = list_ops.tensor_list_pop_back(grad_output,
element_dtype=dtypes.float32)
MatchShape(val.shape)
def testRandomOpsShape(self, random_fn, expected_shape):
shape = constant_op.constant([3])
def Body(i, u):
shape_extended = array_ops.concat([[5], shape], axis=0)
u = random_fn(shape_extended)
assert u.shape.as_list() == expected_shape, str(u.shape.as_list())
return i + 1, u
_, _ = while_loop_v2(cond=lambda i, _: i < 3,
body=Body,
loop_vars=[
0,
array_ops.zeros(expected_shape,
dtype=dtypes.float32),
])
def testMultipleLoopVarsBasic(self):
x = constant_op.constant(5.)
y = constant_op.constant(3.)
# x = 5.
# y = 3.
# while x < 45.:
# x = x * y
ret = while_loop_v2(lambda v, _: v < 45., lambda v, w: (v * w, w), [x, y])
# ret = [x*y^2, y]
# Note: This is simply d_ret[0]/d_x since d_ret[1]/d_x is 0.
grad = gradients_impl.gradients(ret, [x]) # [2*x*y]
with self.cached_session() as sess:
self.assertSequenceEqual(sess.run(ret), [45., 3.])
self.assertSequenceEqual(sess.run(grad), [9.])
def testFillOpsShape(self, fill_fn):
shape = constant_op.constant([3, 4])
def Body(i, u):
shape_extended = array_ops.concat([[5], shape], axis=0)
u = fill_fn(shape_extended)
assert u.shape.as_list() == [5, 3, 4], str(u.shape.as_list())
return i + 1, u
_, _ = while_loop_v2(cond=lambda i, _: i < 3,
body=Body,
loop_vars=[
0,
array_ops.zeros([5, 3, 4],
dtype=dtypes.float32),
])
def testNestedWhile(self):
# Compute sum of geometric progression: n^0 + n^1 + ... + n^m
# We compute the pow using a while loop.
n = constant_op.constant(3.)
m = constant_op.constant(5.)
sum_of_powers = constant_op.constant(0.)
def Body(i, previous_sum):
prod = constant_op.constant(1.)
return i - 1., previous_sum + while_loop_v2(
lambda c, _: c > 0, lambda c, v: (c - 1., v * n), [i, prod])[1]
result = while_loop_v2(lambda i, _: i >= 0, Body, [m, sum_of_powers])[1]
grad = gradients_impl.gradients(result, [n])
with self.cached_session() as sess:
self.assertEqual(self.evaluate(result), 364.)
self.assertSequenceEqual(self.evaluate(grad), [547.])
def testLoopWithTensorListPushBack(self):
x = constant_op.constant(2.)
tensor_list = list_ops.empty_tensor_list(
element_dtype=dtypes.float32, element_shape=ScalarShape())
def Cond(x, tl):
del tl # Unused for Cond.
return x < 5.
def Body(x, tl):
tl = list_ops.tensor_list_push_back(tl, x)
tl = list_ops.tensor_list_push_back(tl, constant_op.constant(100.))
return x**2., tl
ret = while_loop_v2(Cond, Body, [x, tensor_list])
grad = gradients_impl.gradients(ret[0], x)
with self.cached_session() as sess:
self.assertEqual(sess.run(ret[0]), 16.)
self.assertSequenceEqual(sess.run(grad), [32.])
def Body(i, previous_sum):
prod = constant_op.constant(1.)
return i - 1., previous_sum + while_loop_v2(
lambda c, _: c > 0,
lambda c, v: (c - 1., v * n), [i, prod],
return_same_structure=False)[1]
def Body(i, previous_sum):
prod = constant_op.constant(1.)
return i - 1., previous_sum + while_loop_v2(
lambda c, _: c > 0, lambda c, v: (c - 1., v * n), [i, prod])[1]
