def testLogPartitionFunctionInfeasible(self):
"""Tests the log partition function on infeasible scores."""
with self.test_session():
for forest in [False, True]:
# The scores form cycles of various sizes. Note that one can compute
# the partition function for infeasible scores---it's the gradient that
# may be impacted by numerical error.
pad = 12345.6
scores = tf.constant([[[ 0, 1, pad, pad],
[ 1, 0, pad, pad],
[pad, pad, pad, pad],
[pad, pad, pad, pad]],
[[ 0, 1, 0, pad],
[ 0, 0, 1, pad],
[ 1, 0, 0, pad],
[pad, pad, pad, pad]],
[[ 0, 1, 0, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 1],
[ 1, 0, 0, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([2, 3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(tf.exp(log_partition_functions[0]).eval(), 0.0)
self.assertAlmostEqual(tf.exp(log_partition_functions[1]).eval(), 0.0)
self.assertAlmostEqual(tf.exp(log_partition_functions[2]).eval(), 0.0)
def testLogPartitionFunctionTwoTreesScaled(self):
"""Tests the log partition function with two feasible trees."""
with self.test_session():
for forest in [False, True]:
# Each score matrix supports exactly two trees with varying score, and
# the rest with score=0. Thus the log partition function will equal
# the sum of scores of those two trees in each case.
pad = 12345.6
scores = tf.constant([[[ 2, 0, 0, pad],
[ 3, 0, 0, pad],
[ 5, 7, 0, pad],
[pad, pad, pad, pad]],
[[ 0, 11, 0, 13],
[ 0, 17, 0, 0],
[ 0, 19, 0, 0],
[ 0, 23, 0, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(
tf.exp(log_partition_functions[0]).eval(),
2.0 * 3.0 * 5.0 + 2.0 * 3.0 * 7.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[1]).eval(),
11.0 * 17.0 * 19.0 * 23.0 + 13.0 * 17.0 * 19.0 * 23.0)
def testLogPartitionFunctionGradientErrorOkIfInfeasibleWithClipping(self):
"""Tests that the log partition function gradient is OK after clipping."""
with self.test_session():
for forest in [False, True]:
# The scores form cycles of various sizes.
pad = 12345.6
scores_raw = [[[0, 1, pad, pad], [1, 0, pad, pad],
[pad, pad, pad, pad], [pad, pad, pad, pad]],
[[0, 1, 0, pad], [0, 0, 1, pad], [1, 0, 0, pad],
[pad, pad, pad, pad]],
[[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1],
[1, 0, 0, 0]]] # pyformat: disable
scores = tf.log(scores_raw)
init_scores = np.log(np.array(scores_raw))
num_nodes = tf.constant([2, 3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest, max_dynamic_range=10)
gradient_error = tf.test.compute_gradient_error(
scores, [3, 4, 4], log_partition_functions, [3],
init_scores)
tf.logging.info('forest=%s gradient_error=%s', forest,
gradient_error)
# There's still a lot of error.
self.assertLessEqual(gradient_error, 1e-3)
def testLogPartitionFunctionOneTree(self):
"""Tests the log partition function with one feasible tree with score 1."""
with self.test_session():
for forest in [False, True]:
# Each score matrix supports exactly one tree with score=1*1*1, and
# the rest with score=0. Thus the log partition function will be 1.0
# in each case.
pad = 12345.6
scores = tf.constant([[[ 1, pad, pad],
[pad, pad, pad],
[pad, pad, pad]],
[[ 1, 0, pad],
[ 1, 0, pad],
[pad, pad, pad]],
[[ 1, 0, 0],
[ 1, 0, 0],
[ 0, 1, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([1, 2, 3], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(tf.exp(log_partition_functions[0]).eval(), 1.0)
self.assertAlmostEqual(tf.exp(log_partition_functions[1]).eval(), 1.0)
self.assertAlmostEqual(tf.exp(log_partition_functions[2]).eval(), 1.0)
def testLogPartitionFunctionGradientError(self):
"""Validates the log partition function gradient."""
with self.test_session():
for forest in [False, True]:
# To avoid numerical issues, provide score matrices that are weighted
# towards feasible trees or forests.
scores_raw = [[[0, 0, 0, 0], [1, 0, 0, 0], [1, 2, 0, 0],
[1, 2, 3, 4]],
[[4, 3, 2, 9], [0, 0, 2, 9], [0, 0, 0, 9],
[9, 9, 9, 9]]] # pyformat: disable
scores = tf.constant(scores_raw, tf.float64)
init_scores = np.array(scores_raw)
num_nodes = tf.constant([4, 3], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
gradient_error = tf.test.compute_gradient_error(
scores, [2, 4, 4], log_partition_functions, [2],
init_scores)
tf.logging.info('forest=%s gradient_error=%s', forest,
gradient_error)
self.assertLessEqual(gradient_error, 1e-7)
def testLogPartitionFunctionGradientErrorFailsIfInfeasible(self):
"""Tests that the partition function gradient fails on infeasible scores."""
with self.test_session():
for forest in [False, True]:
# The scores form cycles of various sizes.
pad = 12345.6
scores_raw = [[[0, 1, pad, pad], [1, 0, pad, pad],
[pad, pad, pad, pad], [pad, pad, pad, pad]],
[[0, 1, 0, pad], [0, 0, 1, pad], [1, 0, 0, pad],
[pad, pad, pad, pad]],
[[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1],
[1, 0, 0, 0]]] # pyformat: disable
scores = tf.log(scores_raw)
init_scores = np.log(np.array(scores_raw))
num_nodes = tf.constant([2, 3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
with self.assertRaises(Exception):
tf.test.compute_gradient_error(scores, [3, 4, 4],
log_partition_functions,
[3], init_scores)
def testLogPartitionFunctionTwoTreesScaled(self):
"""Tests the log partition function with two feasible trees."""
with self.test_session():
for forest in [False, True]:
# Each score matrix supports exactly two trees with varying score, and
# the rest with score=0. Thus the log partition function will equal
# the sum of scores of those two trees in each case.
pad = 12345.6
scores = tf.constant([[[2, 0, 0, pad], [3, 0, 0, pad],
[5, 7, 0, pad], [pad, pad, pad, pad]],
[[0, 11, 0, 13], [0, 17, 0, 0],
[0, 19, 0, 0], [0, 23, 0, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(
tf.exp(log_partition_functions[0]).eval(),
2.0 * 3.0 * 5.0 + 2.0 * 3.0 * 7.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[1]).eval(),
11.0 * 17.0 * 19.0 * 23.0 + 13.0 * 17.0 * 19.0 * 23.0)
def testLogPartitionFunctionInfeasible(self):
"""Tests the log partition function on infeasible scores."""
with self.test_session():
for forest in [False, True]:
# The scores form cycles of various sizes. Note that one can compute
# the partition function for infeasible scores---it's the gradient that
# may be impacted by numerical error.
pad = 12345.6
scores = tf.constant(
[[[0, 1, pad, pad], [1, 0, pad, pad], [pad, pad, pad, pad],
[pad, pad, pad, pad]],
[[0, 1, 0, pad], [0, 0, 1, pad], [1, 0, 0, pad],
[pad, pad, pad, pad]],
[[0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([2, 3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(
tf.exp(log_partition_functions[0]).eval(), 0.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[1]).eval(), 0.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[2]).eval(), 0.0)
def testLogPartitionFunctionOneTree(self):
"""Tests the log partition function with one feasible tree with score 1."""
with self.test_session():
for forest in [False, True]:
# Each score matrix supports exactly one tree with score=1*1*1, and
# the rest with score=0. Thus the log partition function will be 1.0
# in each case.
pad = 12345.6
scores = tf.constant(
[[[1, pad, pad], [pad, pad, pad], [pad, pad, pad]],
[[1, 0, pad], [1, 0, pad], [pad, pad, pad]],
[[1, 0, 0], [1, 0, 0], [0, 1, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([1, 2, 3], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(
tf.exp(log_partition_functions[0]).eval(), 1.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[1]).eval(), 1.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[2]).eval(), 1.0)
def testLogPartitionFunctionOneTreeScaled(self):
"""Tests the log partition function with one feasible tree."""
with self.test_session():
for forest in [False, True]:
# Each score matrix supports exactly one tree with varying score, and
# the rest with score=0. Thus the log partition function will equal
# the score of that single tree in each case.
pad = 12345.6
scores = tf.constant(
[[[2, pad, pad], [pad, pad, pad], [pad, pad, pad]],
[[3, 0, pad], [5, 0, pad], [pad, pad, pad]],
[[7, 0, 0], [11, 0, 0], [0, 13, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([1, 2, 3], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(
tf.exp(log_partition_functions[0]).eval(), 2.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[1]).eval(), 3.0 * 5.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[2]).eval(),
7.0 * 11.0 * 13.0)
def testLogPartitionFunctionGradientErrorOkIfInfeasibleWithClipping(self):
"""Tests that the log partition function gradient is OK after clipping."""
with self.test_session():
for forest in [False, True]:
# The scores form cycles of various sizes.
pad = 12345.6
scores_raw = [[[ 0, 1, pad, pad],
[ 1, 0, pad, pad],
[pad, pad, pad, pad],
[pad, pad, pad, pad]],
[[ 0, 1, 0, pad],
[ 0, 0, 1, pad],
[ 1, 0, 0, pad],
[pad, pad, pad, pad]],
[[ 0, 1, 0, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 1],
[ 1, 0, 0, 0]]] # pyformat: disable
scores = tf.log(scores_raw)
init_scores = np.log(np.array(scores_raw))
num_nodes = tf.constant([2, 3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest, max_dynamic_range=10)
gradient_error = tf.test.compute_gradient_error(
scores, [3, 4, 4], log_partition_functions, [3], init_scores)
tf.logging.info('forest=%s gradient_error=%s', forest, gradient_error)
# There's still a lot of error.
self.assertLessEqual(gradient_error, 1e-3)
def testLogPartitionFunctionOneTreeScaled(self):
"""Tests the log partition function with one feasible tree."""
with self.test_session():
for forest in [False, True]:
# Each score matrix supports exactly one tree with varying score, and
# the rest with score=0. Thus the log partition function will equal
# the score of that single tree in each case.
pad = 12345.6
scores = tf.constant([[[ 2, pad, pad],
[pad, pad, pad],
[pad, pad, pad]],
[[ 3, 0, pad],
[ 5, 0, pad],
[pad, pad, pad]],
[[ 7, 0, 0],
[ 11, 0, 0],
[ 0, 13, 0]]],
tf.float64) # pyformat: disable
scores = tf.log(scores)
num_nodes = tf.constant([1, 2, 3], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
self.assertAlmostEqual(tf.exp(log_partition_functions[0]).eval(), 2.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[1]).eval(), 3.0 * 5.0)
self.assertAlmostEqual(
tf.exp(log_partition_functions[2]).eval(), 7.0 * 11.0 * 13.0)
def testLogPartitionFunctionGradientErrorFailsIfInfeasible(self):
"""Tests that the partition function gradient fails on infeasible scores."""
with self.test_session():
for forest in [False, True]:
# The scores form cycles of various sizes.
pad = 12345.6
scores_raw = [[[ 0, 1, pad, pad],
[ 1, 0, pad, pad],
[pad, pad, pad, pad],
[pad, pad, pad, pad]],
[[ 0, 1, 0, pad],
[ 0, 0, 1, pad],
[ 1, 0, 0, pad],
[pad, pad, pad, pad]],
[[ 0, 1, 0, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 1],
[ 1, 0, 0, 0]]] # pyformat: disable
scores = tf.log(scores_raw)
init_scores = np.log(np.array(scores_raw))
num_nodes = tf.constant([2, 3, 4], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
with self.assertRaises(Exception):
tf.test.compute_gradient_error(
scores, [3, 4, 4], log_partition_functions, [3], init_scores)
def testLogPartitionFunctionGradientError(self):
"""Validates the log partition function gradient."""
with self.test_session():
for forest in [False, True]:
# To avoid numerical issues, provide score matrices that are weighted
# towards feasible trees or forests.
scores_raw = [[[0, 0, 0, 0],
[1, 0, 0, 0],
[1, 2, 0, 0],
[1, 2, 3, 4]],
[[4, 3, 2, 9],
[0, 0, 2, 9],
[0, 0, 0, 9],
[9, 9, 9, 9]]] # pyformat: disable
scores = tf.constant(scores_raw, tf.float64)
init_scores = np.array(scores_raw)
num_nodes = tf.constant([4, 3], tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
gradient_error = tf.test.compute_gradient_error(
scores, [2, 4, 4], log_partition_functions, [2], init_scores)
tf.logging.info('forest=%s gradient_error=%s', forest, gradient_error)
self.assertLessEqual(gradient_error, 1e-7)
def _compute_crf_loss(self, lengths, scores, gold):
"""Computes the negative CRF log-probability for a batch."""
# The |scores| are assumed to be in the log domain.
log_gold_scores_b = tf.reduce_sum(scores * gold, axis=[1, 2])
log_partition_functions_b = mst_ops.log_partition_function(
num_nodes=lengths,
scores=scores,
forest=self._attrs['forest'],
max_dynamic_range=self._attrs['crf_max_dynamic_range'])
return log_partition_functions_b - log_gold_scores_b # negative log-prob
def _compute_crf_loss(self, lengths, scores, gold):
"""Computes the negative CRF log-probability for a batch."""
# The |scores| are assumed to be in the log domain.
log_gold_scores_b = tf.reduce_sum(scores * gold, axis=[1, 2])
log_partition_functions_b = mst_ops.log_partition_function(
num_nodes=lengths,
scores=scores,
forest=self._attrs['forest'],
max_dynamic_range=self._attrs['crf_max_dynamic_range'])
return log_partition_functions_b - log_gold_scores_b # negative log-prob
def testLogPartitionFunctionWithVeryLowValues(self):
"""Tests the underflow protection in the log partition function."""
with self.test_session():
for forest in [False, True]:
# Set the scores to very low values to test underflow protection.
scores = -1000 * tf.ones([10, 10, 10], tf.float64)
num_nodes = tf.range(1, 11, dtype=tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
base_offset = 1 if forest else 0 # n+1 for forest, n for tree
for size in range(1, 11):
self.assertAlmostEqual(
log_partition_functions[size - 1].eval(),
(size - 1) * math.log(size + base_offset) - size * 1000)
def testLogPartitionFunctionWithVeryLowValues(self):
"""Tests the underflow protection in the log partition function."""
with self.test_session():
for forest in [False, True]:
# Set the scores to very low values to test underflow protection.
scores = -1000 * tf.ones([10, 10, 10], tf.float64)
num_nodes = tf.range(1, 11, dtype=tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
base_offset = 1 if forest else 0 # n+1 for forest, n for tree
for size in range(1, 11):
self.assertAlmostEqual(
log_partition_functions[size - 1].eval(),
(size - 1) * math.log(size + base_offset) -
size * 1000)
def testLogPartitionFunctionAllTrees(self):
"""Tests the log partition function with all trees feasible."""
with self.test_session():
for forest in [False, True]:
# The scores allow all trees. Using Cayley's formula, the
# number of directed spanning trees and forests in a complete
# digraph of n nodes is n^{n-1} and (n+1)^{n-1}, respectively.
# https://en.wikipedia.org/wiki/Cayley%27s_formula
scores = tf.zeros([10, 10, 10], tf.float64) # = 1 in log domain
num_nodes = tf.range(1, 11, dtype=tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
base_offset = 1 if forest else 0 # n+1 for forest, n for tree
for size in range(1, 11):
self.assertAlmostEqual(log_partition_functions[size - 1].eval(),
(size - 1) * math.log(size + base_offset))
def testLogPartitionFunctionAllTrees(self):
"""Tests the log partition function with all trees feasible."""
with self.test_session():
for forest in [False, True]:
# The scores allow all trees. Using Cayley's formula, the
# number of directed spanning trees and forests in a complete
# digraph of n nodes is n^{n-1} and (n+1)^{n-1}, respectively.
# https://en.wikipedia.org/wiki/Cayley%27s_formula
scores = tf.zeros([10, 10, 10],
tf.float64) # = 1 in log domain
num_nodes = tf.range(1, 11, dtype=tf.int32)
log_partition_functions = mst_ops.log_partition_function(
num_nodes, scores, forest=forest)
base_offset = 1 if forest else 0 # n+1 for forest, n for tree
for size in range(1, 11):
self.assertAlmostEqual(
log_partition_functions[size - 1].eval(),
(size - 1) * math.log(size + base_offset))