def reset(
self, x, y=None, weights=None, target_weights=None, num_targets=None):
"""Resets the Sinkhorn matrix for new inputs x, y, a, b."""
self.x = x
if x is None:
return
self.dtype = x.dtype
self._batch = tf.shape(x)[0]
self._set_input(x, weights)
self._set_target(y, num_targets, target_weights)
# We run sinkhorn on the rescaled input values x_s.
self.transport = sinkhorn.sinkhorn(
self._x_s, self.y, self.weights, self.target_weights, **self._kwargs)
def test_routine(self):
p = sinkhorn.sinkhorn(self.x, self.y, self.a, self.b)
self.assertEqual(p.shape.as_list(), [1, 8, 8])