def softranks(self, axis, direction):
"""Test ops.softranks for a given shape, axis and direction."""
shape = tf.TensorShape((3, 8, 6))
n = shape[axis]
p = int(np.prod(shape) / shape[axis])
# Build a target tensor of ranks, of rank 2.
# Those targets are zero based.
target = tf.constant([np.random.permutation(n) for _ in range(p)],
dtype=tf.float32)
# Turn it into a tensor of desired shape.
target = ops._postprocess(target, shape, axis)
# Apply a monotonic transformation to turn ranks into values
sign = 2 * float(direction == 'ASCENDING') - 1
x = sign * (1.2 * target - 0.4)
# The softranks of x along the axis should be close to the target.
eps = 1e-3
sinkhorn_threshold = 1e-3
tolerance = 0.5
for zero_based in [False, True]:
ranks = ops.softranks(x,
direction=direction,
axis=axis,
zero_based=zero_based,
epsilon=eps,
sinkhorn_threshold=sinkhorn_threshold)
targets = target + 1 if not zero_based else target
self.assertAllClose(ranks, targets, tolerance, tolerance)
def test_postprocess(self):
"""Tests that _postprocess is the inverse of _preprocess."""
shape = (4, 21, 7, 10)
for i in range(1, len(shape)):
x = tf.random.uniform(shape[:i])
for axis in range(x.shape.rank):
z = ops._postprocess(ops._preprocess(x, axis), x.shape, axis)
self.assertAllEqual(x, z)