def BackwardLoopBody(*args):
"""Backward loop body function."""
t, dev_t = args[0], args[1]
(theta, orig_state0, inputs, acc_state, acc_extras, d_theta, d_state1,
d_inputs, d_acc_state) = _Pack(args[2:], bakloop_sig)
# The input recurrent state for time step t is previous time step's
# output, or the original state0 when on time step 0.
state_from_acc = _Index(acc_state, math_ops.maximum(0, t - 1))
state0 = functional_ops.If(
math_ops.equal(t, array_ops.constant(0, dtypes.int32)),
_Flatten([state_from_acc, orig_state0]), ReturnOrigState0,
ReturnAccState)
state0 = nest.pack_sequence_as(orig_state0, state0)
# The external inputs for time step t.
inputs_t = _Index(inputs, t)
# The extras for time step t.
extras_t = _Index(acc_extras, t)
d_state1 = _Add(_Index(d_acc_state, t), d_state1)
(d_theta_t, d_state0, d_inputs_t) = _Pack(
Bak(*_Flatten([theta, state0, inputs_t, extras_t, d_state1])),
[self._theta, self._state, self._inputs])
d_theta = _Add(d_theta, d_theta_t)
d_inputs = _Update(d_inputs, d_inputs_t, dev_t)
return [math_ops.subtract(dev_t, 1)] + _Flatten([
theta, orig_state0, inputs, acc_state, acc_extras, d_theta, d_state0,
d_inputs, d_acc_state
])
def ReorderIndicesByPhi(anchor, bboxes):
"""Sort bboxes based their angles relative to the anchor point.
Args:
anchor: A vector of (x0, y0).
bboxes: A matrix of shape [N, 4].
Returns:
A permutation of tf.range(n) which can be used to reshuffle bboxes to the
sorted order. (e.g., tf.gather(bboxes, indices)).
"""
@tf.Defun(anchor.dtype, bboxes.dtype)
def _True(anchor, bboxes):
"""True branch when num of bboxes is non-zero."""
n = tf.shape(bboxes)[0]
centroid = BBoxesCentroid(bboxes)
# Computed dot products between centroid and the anchor point.
dot = tf.squeeze(tf.matmul(centroid, tf.expand_dims(anchor, 1)),
axis=1)
# Normalize dot to get the cosine of the angles.
norm = tf.norm(anchor) * tf.norm(centroid, axis=1)
cosine = tf.where(tf.greater(norm, 0), dot / norm,
tf.zeros([n], norm.dtype))
# Disambiguates the angle anchor--O--point is positive or negative by the
# sign of cross products between angle and points. tf.linalg.cross takes
# 3-vector (x, y, z), so we set z to 0. tf.linalg.cross does not support
# broadcasting, so we tile anchor to shape [n, 3].
cross = tf.linalg.cross(
tf.tile(tf.pad(tf.expand_dims(anchor, 0), [[0, 0], [0, 1]]),
[n, 1]), tf.pad(centroid, [[0, 0], [0, 1]]))
# If the sign is positive, the points lie on the clockwise side of
# O-->anchor. Hence, -1 - cosine moves the cosine values to [-2, 0]. If the
# sign is negative, the points lie on the counter-clockwise side of
# O-->anchor. 1 + cosine moves the cosine values to [0, 2].
#
# The car dataset shows that the points are scanned in the counter-clockwise
# fashion. Therefore, top-k orders the points in the same order in which
# bboxes appears in the spin.
score = tf.where(tf.greater(cross, 0)[:, 2], -1 - cosine, 1 + cosine)
_, indices = tf.nn.top_k(score, n, sorted=True)
return indices
@tf.Defun(anchor.dtype, bboxes.dtype)
def _False(anchor, bboxes):
del anchor, bboxes
return tf.zeros([0], dtype=tf.int32)
n = tf.shape(bboxes)[0]
return functional_ops.If(tf.greater(n, 0), [anchor, bboxes], _True,
_False)[0]
def testIfWithDefun(self):
@function.Defun(dtypes.float32)
def Then(x):
return x + 1
@function.Defun(dtypes.float32)
def Else(x):
return x - 1
with self.cached_session():
inputs = [10.]
result = self.evaluate(functional_ops.If(False, inputs, Then,
Else))
self.assertEqual([9.0], result)
def testIfWithDefun(self):
# Defun should only be used in graph mode
with ops.Graph().as_default():
@function.Defun(dtypes.float32)
def Then(x):
return x + 1
@function.Defun(dtypes.float32)
def Else(x):
return x - 1
inputs = [10.]
result = self.evaluate(functional_ops.If(False, inputs, Then, Else))
self.assertEqual([9.0], result)
def testIfWithFunction(self):
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def Then(x):
return x + 1
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def Else(x):
return x - 1
with self.cached_session():
inputs = [10.]
result = self.evaluate(
functional_ops.If(False, inputs, Then.get_concrete_function(),
Else.get_concrete_function()))
self.assertEqual([9.0], result)
def testIfWithFunctionComposite(self):
signature = [tensor_spec.TensorSpec([], dtypes.float32)]
@def_function.function(input_signature=signature)
def Then(x):
return sparse_tensor.SparseTensor([[0]], [x + 1], [1])
@def_function.function(input_signature=signature)
def Else(x):
return sparse_tensor.SparseTensor([[0]], [x - 1], [1])
inputs = [10.]
then_cf = Then.get_concrete_function()
else_cf = Else.get_concrete_function()
result = functional_ops.If(False, inputs, then_cf, else_cf)
self.assertIsInstance(result, sparse_tensor.SparseTensor)
self.assertAllEqual([9.0], result.values)
def testIfWithFunction(self):
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def Then(x):
return x + 1
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def Else(x):
return x - 1
inputs = [10.]
then_cf = Then.get_concrete_function()
else_cf = Else.get_concrete_function()
result = self.evaluate(functional_ops.If(False, inputs, then_cf, else_cf))
self.assertEqual([9.0], result)
def testIf(self):
@function.Defun(dtypes.float32)
def Twice(x):
return x * 2
@function.Defun(dtypes.float32)
def Thrice(x):
return x * 3 + 1
with self.test_session(use_gpu=False) as sess:
x = array_ops.placeholder(dtypes.float32)
ret = functional_ops.If(math_ops.greater(x, 0), [x], Twice, Thrice)[0]
self.assertAllEqual(sess.run(ret, feed_dict={x: 9.}), 18.)
self.assertAllEqual(sess.run(ret, feed_dict={x: -8.}), -23.)
self.assertAllEqual(sess.run(ret, feed_dict={x: 0.}), 1.)
def BackwardLoopBody(t, limit, *args):
"""Backward loop body function."""
(
theta,
orig_state0,
inputs,
acc_state,
acc_extras,
# End of forward params
d_theta,
d_state1,
d_inputs,
d_acc_state,
d_captured) = _Pack(args, bakloop_sig)
# The input recurrent state for time step t is previous time step's
# output, or the original state0 when on time step 0.
state_from_acc = _Index(acc_state,
tf.maximum(tf.constant(0, t.dtype), t - 1))
state0 = functional_ops.If(tf.equal(t, tf.constant(0, t.dtype)),
_Flatten([state_from_acc, orig_state0]),
ReturnOrigState0, ReturnAccState)
state0 = orig_state0.Pack(state0)
# The external inputs for time step t.
inputs_t = _Index(inputs, t)
# The extras for time step t.
extras_t = _Index(acc_extras, t)
d_state1 = _Add(_Index(d_acc_state, t), d_state1)
(d_theta_t, d_state0, d_inputs_t, d_captured_t) = _Pack(
Bak(*_Flatten([theta, state0, inputs_t, extras_t, d_state1])),
[
self._theta, self._state, self._inputs,
self._implicit_captures
])
if self._unused_acc_state:
# XLA IF op requires the same shape for if and else branches.
d_state0 = d_state0.Transform(tf.reduce_sum)
d_theta = _Add(d_theta, d_theta_t)
d_inputs = _Update(d_inputs, d_inputs_t, t)
d_captured = _Add(d_captured, d_captured_t)
# Make sure this function didn't capture anything different than the
# cell_fn when reflected on at the beginning. Must come after the call
# to Bak() which adds to the captured list.
_AssertSameTensors(function.get_extra_inputs(),
self._implicit_captures.Flatten())
return [tf.subtract(t, 1), limit] + _Flatten([
theta,
orig_state0,
inputs,
acc_state,
acc_extras,
# End of forward params
d_theta,
d_state0,
d_inputs,
d_acc_state,
d_captured,
])
def Run(x):
return sess.run(
functional_ops.If(math_ops.greater(x, 0), [x], Twice,
Thrice))[0]
