def test_one_dimensional(self):
x = tf.constant([0.3, 0.1, 0.2, -0.5, 42.0])
top_1 = self.evaluate(metrics_utils._filter_top_k(x=x, k=1))
top_2 = self.evaluate(metrics_utils._filter_top_k(x=x, k=2))
top_3 = self.evaluate(metrics_utils._filter_top_k(x=x, k=3))
self.assertAllClose(
top_1,
[
metrics_utils.NEG_INF,
metrics_utils.NEG_INF,
metrics_utils.NEG_INF,
metrics_utils.NEG_INF,
42.0,
],
)
self.assertAllClose(
top_2,
[
0.3,
metrics_utils.NEG_INF,
metrics_utils.NEG_INF,
metrics_utils.NEG_INF,
42.0,
],
)
self.assertAllClose(
top_3,
[0.3, metrics_utils.NEG_INF, 0.2, metrics_utils.NEG_INF, 42.0],
)
def test_three_dimensional(self):
x = tf.constant([[[.3, .1, .2], [-.3, -.2, -.1]],
[[5., .2, 42.], [-.3, -.6, -.99]]])
top_2 = self.evaluate(metrics_utils._filter_top_k(x=x, k=2))
self.assertAllClose(top_2, [[[.3, metrics_utils.NEG_INF, .2],
[metrics_utils.NEG_INF, -.2, -.1]],
[[5., metrics_utils.NEG_INF, 42.],
[-.3, -.6, metrics_utils.NEG_INF]]])
def test_three_dimensional(self):
x = tf.constant([
[[0.3, 0.1, 0.2], [-0.3, -0.2, -0.1]],
[[5.0, 0.2, 42.0], [-0.3, -0.6, -0.99]],
])
top_2 = self.evaluate(metrics_utils._filter_top_k(x=x, k=2))
self.assertAllClose(
top_2,
[
[
[0.3, metrics_utils.NEG_INF, 0.2],
[metrics_utils.NEG_INF, -0.2, -0.1],
],
[
[5.0, metrics_utils.NEG_INF, 42.0],
[-0.3, -0.6, metrics_utils.NEG_INF],
],
],
)
def _filter_top_k(x):
# This loses the static shape.
x = tf.numpy_function(_identity, (x, ), tf.float32)
return metrics_utils._filter_top_k(x=x, k=2)