import tensorflow as tf


  def _forward(self, x):
    return self._transpose(x, self.perm)

  def _inverse(self, y):
    return self._transpose(y, tf.argsort(self.perm))

  def _inverse_log_det_jacobian(self, y):
    return tf.constant(0, dtype=y.dtype)

  def _forward_log_det_jacobian(self, x):
    return tf.constant(0, dtype=x.dtype)

  def _transpose(self, x, perm):
    sample_batch_ndims = tf.rank(x) - self.rightmost_transposed_ndims
    perm = tf.concat([
        tf.range(sample_batch_ndims),
        sample_batch_ndims + perm,
    ], axis=0)
    return tf.transpose(a=x, perm=perm)


