code: def restore(self, x):
      
      with tf.name_scope("pad_reduce/restore"):
        x = tf.scatter_nd(
            indices=self.nonpad_ids,
            updates=x,
            shape=tf.concat([self.dim_origin, tf.shape(x)[1:]], axis=0),
        )
      return x
docs: Add padding back to the given tensor.
  
      Args:
        x (tf.Tensor): of shape [dim_compressed,...]
  
      Returns:
        a tensor of shape [dim_origin,...] with dim_compressed >= dim_origin. The
        dim is restored from the original reference tensor