import tensorflow as tf

    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),
