def repeat_rows(x, num_reps): """Each row of tensor `x` is repeated `num_reps` times along leading dimension.""" if not utils.is_positive_int(num_reps): raise TypeError("Number of repetitions must be a positive integer.") shape = x.shape x = x.unsqueeze(1) x = x.expand(shape[0], num_reps, *shape[1:]) return merge_leading_dims(x, num_dims=2)
def tile(x, n): if not utils.is_positive_int(n): raise TypeError("Argument `n` must be a positive integer.") x_ = x.reshape(-1) x_ = x_.repeat(n) x_ = x_.reshape(n, -1) x_ = x_.transpose(1, 0) x_ = x_.reshape(-1) return x_
def merge_leading_dims(x, num_dims): """Reshapes the tensor `x` such that the first `num_dims` dimensions are merged to one.""" if not utils.is_positive_int(num_dims): raise TypeError("Number of leading dims must be a positive integer.") if num_dims > x.dim(): raise ValueError( "Number of leading dims can't be greater than total number of dims." ) new_shape = torch.Size([-1]) + x.shape[num_dims:] return torch.reshape(x, new_shape)