def test_softranks(self):
layer = layers.SoftRanksLayer(axis=self._axis, epsilon=1e-4)
outputs = layer(self._inputs)
self.assertAllEqual(outputs.shape, self._inputs.shape)
ranks = tf.argsort(tf.argsort(self._inputs, axis=self._axis),
axis=self._axis)
self.assertAllClose(ranks, outputs, atol=0.5)
def test_rankslayer_in_model(self):
inputs = tf.random.uniform((32, 10))
outputs = self.take_model_output(layers.SoftRanksLayer(), inputs)
self.assertAllEqual([inputs.shape[0], 1], outputs.shape)