jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_weight_norm_parametrization(self, device):
        for dtype in [torch.float, torch.bfloat16]:
            input = torch.randn(3, 4, dtype=dtype, device=device)
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input)

            # add weight normalization
            m = torch.nn.utils.parametrizations.weight_norm(m)
            self.assertEqual(
                m.parametrizations.weight.original1.size(), m.weight.size()
            )
            self.assertEqual(m.parametrizations.weight.original0.size(), (5, 1))
            self.assertEqual(m(input), expected_output)

            # remove weight norm
            torch.nn.utils.parametrize.remove_parametrizations(m, "weight")
            self.assertFalse(hasattr(m, "parametrizations"))
            self.assertEqual(m(input), expected_output)

            # test with dim=1
            m = torch.nn.utils.parametrizations.weight_norm(m, dim=1)
            self.assertEqual(
                m.parametrizations.weight.original1.size(), m.weight.size()
            )
            self.assertEqual(m.parametrizations.weight.original0.size(), (1, 4))
            self.assertEqual(m(input), expected_output)

            # test with dim=None
            m = nn.Linear(4, 5, dtype=dtype, device=device)
            expected_output = m(input)
            m = torch.nn.utils.parametrizations.weight_norm(m, dim=None)
            self.assertEqual(m(input), expected_output)
//  [python]
def __init__(
        self,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
    ):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        self.hidden_channels = hidden_channels
        self.kernel_size = (kernel_size,)
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(
                gin_channels, 2 * hidden_channels * n_layers, 1
            )
            self.cond_layer = torch.nn.utils.weight_norm(
                cond_layer, name="weight"
            )

        for i in range(n_layers):
            dilation = dilation_rate**i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(
                hidden_channels,
                2 * hidden_channels,
                kernel_size,
                dilation=dilation,
                padding=padding,
            )
            in_layer = torch.nn.utils.weight_norm(
                in_layer, name="weight"
            )
            self.in_layers.append(in_layer)
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(
                res_skip_layer, name="weight"
            )
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(
        self,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
    ):
        super(WN, self).__init__()
        assert kernel_size % 2 == 1
        self.hidden_channels = hidden_channels
        self.kernel_size = (kernel_size,)
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(
                gin_channels, 2 * hidden_channels * n_layers, 1
            )
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(
                cond_layer, name="weight"
            )

        for i in range(n_layers):
            dilation = dilation_rate**i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(
                hidden_channels,
                2 * hidden_channels,
                kernel_size,
                dilation=dilation,
                padding=padding,
            )
            in_layer = torch.nn.utils.parametrizations.weight_norm(
                in_layer, name="weight"
            )
            self.in_layers.append(in_layer)
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(
                res_skip_layer, name="weight"
            )
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def get_ortho_mat(mat_type, batched_image_shape, device, method='bases'):
    edge_dim = batched_image_shape[1] * batched_image_shape[2] * batched_image_shape[3]
    if method == 'scipyv2':
        U = ortho_group.rvs(dim=batched_image_shape[2])
        U = torch.from_numpy(U).type(mat_type).to(device)
    elif method == 'bases':
        col_idx = torch.randperm(edge_dim)
        row_idx = torch.arange(edge_dim)
        U = torch.zeros(size=(edge_dim, edge_dim))
        ones = torch.ones(edge_dim)
        ones[torch.rand(edge_dim) > 0.5] *= -1
        U[row_idx, col_idx] = ones
        U = U.to(device)
    elif method == 'std-bases':
        col_idx = torch.randperm(edge_dim)
        row_idx = torch.arange(edge_dim)
        U = torch.zeros(size=(edge_dim, edge_dim))
        ones = torch.ones(edge_dim)
        U[row_idx, col_idx] = ones
        U = U.to(device)
    elif method == 'normal':
        U = torch.randn((edge_dim, edge_dim), device=device)
        U /= np.sqrt(edge_dim)
    elif method == 'hadamard':
        U = torch.from_numpy(get_hadamard_matrix(edge_dim)).type(mat_type).to(device)
        U /= np.sqrt(edge_dim)
    elif method == 'qr':
        U = torch.randn((edge_dim, edge_dim), device=device)
        U, _ = torch.qr(U)
    return U
//  [python]
def get_ortho_mat(mat_type, batched_image_shape, device, method='bases'):
    edge_dim = batched_image_shape[1] * batched_image_shape[2] * batched_image_shape[3]
    if method == 'scipyv2':
        U = ortho_group.rvs(dim=batched_image_shape[2])
        U = torch.from_numpy(U).type(mat_type).to(device)
    elif method == 'bases':
        col_idx = torch.randperm(edge_dim)
        row_idx = torch.arange(edge_dim)
        U = torch.zeros(size=(edge_dim, edge_dim))
        ones = torch.ones(edge_dim)
        ones[torch.rand(edge_dim) > 0.5] *= -1
        U[row_idx, col_idx] = ones
        U = U.to(device)
    elif method == 'std-bases':
        col_idx = torch.randperm(edge_dim)
        row_idx = torch.arange(edge_dim)
        U = torch.zeros(size=(edge_dim, edge_dim))
        ones = torch.ones(edge_dim)
        U[row_idx, col_idx] = ones
        U = U.to(device)
    elif method == 'normal':
        U = torch.randn((edge_dim, edge_dim), device=device)
        U /= np.sqrt(edge_dim)
    elif method == 'hadamard':
        U = torch.from_numpy(get_hadamard_matrix(edge_dim)).type(mat_type).to(device)
        U /= np.sqrt(edge_dim)
    elif method == 'qr':
        U = torch.randn((edge_dim, edge_dim), device=device)
        U, _ = torch.linalg.qr(U)
    return U
//  [python]
def forward(self, input):
        # input: (bs*t) x 512

        weight = self.weight + 1e-8
        Q, R = torch.qr(weight)  # get eignvector, orthogonal [n1, n2, n3, n4]

        if input is None:
            return Q
        else:
            input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
            out = torch.matmul(input_diag, Q.T)
            out = torch.sum(out, dim=1)

            return out
//  [python]
def forward(self, input):
        # input: (bs*t) x 512

        weight = self.weight + 1e-8
        Q, R = torch.linalg.qr(weight)  # get eignvector, orthogonal [n1, n2, n3, n4]

        if input is None:
            return Q
        else:
            input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
            out = torch.matmul(input_diag, Q.T)
            out = torch.sum(out, dim=1)

            return out
//  [python]
def _orthogonal(tensor, gain=1):

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    dt = flattened.dtype
    flattened = flattened.to(torch.float32)  # orthogonal init does not support fp16
    q, r = torch.qr(flattened)
    q, r = q.to(dtype=dt), r.to(dtype=dt)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def _orthogonal(tensor, gain=1):

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    dt = flattened.dtype
    flattened = flattened.to(torch.float32)  # orthogonal init does not support fp16
    q, r = torch.linalg.qr(flattened)
    q, r = q.to(dtype=dt), r.to(dtype=dt)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def __init__(self, num_channels, LU_decomposed):
        super().__init__()
        w_shape = [num_channels, num_channels]
        w_init = torch.qr(torch.randn(*w_shape))[0]

        if not LU_decomposed:
            self.weight = nn.Parameter(torch.Tensor(w_init))
        else:
            p, lower, upper = torch.lu_unpack(*torch.lu(w_init))
            s = torch.diag(upper)
            sign_s = torch.sign(s)
            log_s = torch.log(torch.abs(s))
            upper = torch.triu(upper, 1)
            l_mask = torch.tril(torch.ones(w_shape), -1)
            eye = torch.eye(*w_shape)

            self.register_buffer("p", p)
            self.register_buffer("sign_s", sign_s)
            self.lower = nn.Parameter(lower)
            self.log_s = nn.Parameter(log_s)
            self.upper = nn.Parameter(upper)
            self.l_mask = l_mask
            self.eye = eye

        self.w_shape = w_shape
        self.LU_decomposed = LU_decomposed
//  [python]
def __init__(self, num_channels, LU_decomposed):
        super().__init__()
        w_shape = [num_channels, num_channels]
        w_init = torch.linalg.qr(torch.randn(*w_shape))[0]

        if not LU_decomposed:
            self.weight = nn.Parameter(torch.Tensor(w_init))
        else:
            p, lower, upper = torch.lu_unpack(*torch.lu(w_init))
            s = torch.diag(upper)
            sign_s = torch.sign(s)
            log_s = torch.log(torch.abs(s))
            upper = torch.triu(upper, 1)
            l_mask = torch.tril(torch.ones(w_shape), -1)
            eye = torch.eye(*w_shape)

            self.register_buffer("p", p)
            self.register_buffer("sign_s", sign_s)
            self.lower = nn.Parameter(lower)
            self.log_s = nn.Parameter(log_s)
            self.upper = nn.Parameter(upper)
            self.l_mask = l_mask
            self.eye = eye

        self.w_shape = w_shape
        self.LU_decomposed = LU_decomposed
//  [python]
def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.qr(weight)
        weight = q
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.linalg.qr(weight)
        weight = q
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, c, memory_efficient=False):
        super().__init__(c, c, 1, bias=False) # init as nn.Conv1d(c, c, kernel_size=1, stride=1) 
        
        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
        
        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.weight.data = W
        
        if memory_efficient:
            self.efficient_forward = Conv1x1Func.apply
//  [python]
def __init__(self, c, memory_efficient=False):
        super().__init__(c, c, 1, bias=False) # init as nn.Conv1d(c, c, kernel_size=1, stride=1) 
        
        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]
        
        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.weight.data = W
        
        if memory_efficient:
            self.efficient_forward = Conv1x1Func.apply
//  [python]
def random_orthogonal(size):
    """
    Returns a random orthogonal matrix as a 2-dim tensor of shape [size, size].
    """

    # Use the QR decomposition of a random Gaussian matrix.
    x = torch.randn(size, size)
    q, _ = torch.qr(x)
    return q
//  [python]
def random_orthogonal(size):
    """
    Returns a random orthogonal matrix as a 2-dim tensor of shape [size, size].
    """

    # Use the QR decomposition of a random Gaussian matrix.
    x = torch.randn(size, size)
    q, _ = torch.linalg.qr(x)
    return q
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, lu=True, n_sqz=2, **kwargs):
        super().__init__()
        assert n_split % 2 == 0
        self.channels = channels
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.no_jacobian = no_jacobian
        w_init = torch.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.lu = lu
        if lu:
            np_p, np_l, np_u = scipy.linalg.lu(w_init)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(np.abs(np_s))
            np_u = np.triu(np_u, k=1)
            l_mask = np.tril(np.ones(w_init.shape, dtype=float), -1)
            eye = np.eye(*w_init.shape, dtype=float)
            self.register_buffer('p', torch.Tensor(np_p.astype(float)))
            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
            self.l = nn.Parameter(torch.Tensor(np_l.astype(float)), requires_grad=True)
            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)), requires_grad=True)
            self.u = nn.Parameter(torch.Tensor(np_u.astype(float)), requires_grad=True)
            self.register_buffer('l_mask', torch.Tensor(l_mask))
            self.register_buffer('eye', torch.Tensor(eye))
        else:
            self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, lu=True, n_sqz=2, **kwargs):
        super().__init__()
        assert n_split % 2 == 0
        self.channels = channels
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.no_jacobian = no_jacobian
        w_init = torch.linalg.qr(torch.randn(self.n_split, self.n_split))[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.lu = lu
        if lu:
            np_p, np_l, np_u = scipy.linalg.lu(w_init)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(np.abs(np_s))
            np_u = np.triu(np_u, k=1)
            l_mask = np.tril(np.ones(w_init.shape, dtype=float), -1)
            eye = np.eye(*w_init.shape, dtype=float)
            self.register_buffer('p', torch.Tensor(np_p.astype(float)))
            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
            self.l = nn.Parameter(torch.Tensor(np_l.astype(float)), requires_grad=True)
            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)), requires_grad=True)
            self.u = nn.Parameter(torch.Tensor(np_u.astype(float)), requires_grad=True)
            self.register_buffer('l_mask', torch.Tensor(l_mask))
            self.register_buffer('eye', torch.Tensor(eye))
        else:
            self.weight = nn.Parameter(w_init)
//  [python]
def random_rotation(dimension=3,allow_mirror=False):
    r=torch.qr(torch.randn(dimension,dimension))[0]
    f=torch.randint(2,(3,))
    if f.sum()%2==0 and not allow_mirror:
        f=1-f
    return r*(2*f-1).float()
//  [python]
def random_rotation(dimension=3,allow_mirror=False):
    r=torch.linalg.qr(torch.randn(dimension,dimension))[0]
    f=torch.randint(2,(3,))
    if f.sum()%2==0 and not allow_mirror:
        f=1-f
    return r*(2*f-1).float()
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, lu=True, n_sqz=2, **kwargs):
        super().__init__()
        assert (n_split % 2 == 0)
        self.channels = channels
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.no_jacobian = no_jacobian

        w_init = torch.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.lu = lu
        if lu:
            # LU decomposition can slightly speed up the inverse
            np_p, np_l, np_u = scipy.linalg.lu(w_init)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(np.abs(np_s))
            np_u = np.triu(np_u, k=1)
            l_mask = np.tril(np.ones(w_init.shape, dtype=float), -1)
            eye = np.eye(*w_init.shape, dtype=float)

            self.register_buffer('p', torch.Tensor(np_p.astype(float)))
            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
            self.l = nn.Parameter(torch.Tensor(np_l.astype(float)), requires_grad=True)
            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)), requires_grad=True)
            self.u = nn.Parameter(torch.Tensor(np_u.astype(float)), requires_grad=True)
            self.register_buffer('l_mask', torch.Tensor(l_mask))
            self.register_buffer('eye', torch.Tensor(eye))
        else:
            self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, lu=True, n_sqz=2, **kwargs):
        super().__init__()
        assert (n_split % 2 == 0)
        self.channels = channels
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.no_jacobian = no_jacobian

        w_init = torch.linalg.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.lu = lu
        if lu:
            # LU decomposition can slightly speed up the inverse
            np_p, np_l, np_u = scipy.linalg.lu(w_init)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(np.abs(np_s))
            np_u = np.triu(np_u, k=1)
            l_mask = np.tril(np.ones(w_init.shape, dtype=float), -1)
            eye = np.eye(*w_init.shape, dtype=float)

            self.register_buffer('p', torch.Tensor(np_p.astype(float)))
            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
            self.l = nn.Parameter(torch.Tensor(np_l.astype(float)), requires_grad=True)
            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)), requires_grad=True)
            self.u = nn.Parameter(torch.Tensor(np_u.astype(float)), requires_grad=True)
            self.register_buffer('l_mask', torch.Tensor(l_mask))
            self.register_buffer('eye', torch.Tensor(eye))
        else:
            self.weight = nn.Parameter(w_init)
//  [python]
def ortho_init(self, init_matrix, coordinate_matrix):
		_, w, h, cin =  coordinate_matrix.size()
		coordinate_matrix_v = coordinate_matrix.view(-1, w*h*cin).transpose(0, 1)
		print(coordinate_matrix_v.size())
		q, r = torch.qr(coordinate_matrix_v)
		r_d = torch.diag(r)
		print(r_d)
		sorted_d, indices = torch.sort(torch.abs(r_d), descending=True)
		eps = 1e-10
		rnk = min(coordinate_matrix_v.size()) - torch.ge(sorted_d, eps).sum()
		print('rank :' ,rnk)
		basis = q[:, indices[rnk :]].transpose(0, 1)
		basis = basis.view(-1, w, h, cin)
		stdv = 1. / math.sqrt(w*h*cin)
		nll_count = 0
		for k in range(init_matrix.size()[0]):
			if nll_count < basis.size()[0]:
				init_matrix[k, :, :, :] = torch.clamp(basis[k, :, :, :],-stdv, stdv)
				nll_count += 1
			else:
				init_matrix[k].uniform_(-stdv, stdv)
//  [python]
def ortho_init(self, init_matrix, coordinate_matrix):
		_, w, h, cin =  coordinate_matrix.size()
		coordinate_matrix_v = coordinate_matrix.view(-1, w*h*cin).transpose(0, 1)
		q, r = torch.linalg.qr(coordinate_matrix_v)
		r_d = torch.diag(r)
		sorted_d, indices = torch.sort(torch.abs(r_d), descending=True)
		eps = 1e-10
		rnk = min(coordinate_matrix_v.size()) - torch.ge(sorted_d, eps).sum()
		basis = q[:, indices[rnk :]].transpose(0, 1)
		basis = basis.view(-1, w, h, cin)
		stdv = 1. / math.sqrt(w*h*cin)
		nll_count = 0
		for k in range(init_matrix.size()[0]):
			if nll_count < basis.size()[0]:
				init_matrix[k, :, :, :] = torch.clamp(basis[k, :, :, :],-stdv, stdv)
				nll_count += 1
			else:
				init_matrix[k].uniform_(-stdv, stdv)
//  [python]
def rand(self):
        """
        Generate random Stiefel point using qr of random normally distributed
        matrix
        """
        if self._k == 1:
            X = torch.randn(self._n, self._p)
            q, r = torch.qr(X)
            return q

        X = torch.zeros((self._k, self._n, self._p))

        # TODO: update with batch implementation
        for i in range(self._k):
            X[i], r = torch.qr(torch.randn(self._n, self._p))
        return X
//  [python]
def rand(self):
        """
        Generate random Stiefel point using qr of random normally distributed
        matrix
        """
        if self._k == 1:
            X = torch.randn(self._n, self._p)
            q, r = torch.linalg.qr(X)
            return q

        X = torch.zeros((self._k, self._n, self._p))

        # TODO: update with batch implementation
        for i in range(self._k):
            X[i], r = torch.linalg.qr(torch.randn(self._n, self._p))
        return X
//  [python]
def retr(self, X, G):
        """
        Retract to the Stiefel using the qr decomposition of X + G.
        """
        if self._k == 1:
            # Calculate 'thin' qr decomposition of X + G
            q, r = torch.qr(X + G)
            # Unflip any flipped signs
            XNew = torch.matmul(q, torch.diag(
                torch.sign(torch.sign(torch.diag(r)) + .5)))
        else:
            XNew = X + G
            # TODO: update with batch implementation
            for i in range(self._k):
                q, r = torch.qr(XNew[i])
                XNew[i] = torch.matmul(q, torch.diag(
                    torch.sign(torch.sign(torch.diag(r)) + .5)))
        return XNew
//  [python]
def retr(self, X, G):
        """
        Retract to the Stiefel using the qr decomposition of X + G.
        """
        if self._k == 1:
            # Calculate 'thin' qr decomposition of X + G
            q, r = torch.linalg.qr(X + G)
            # Unflip any flipped signs
            XNew = torch.matmul(q, torch.diag(
                torch.sign(torch.sign(torch.diag(r)) + .5)))
        else:
            XNew = X + G
            # TODO: update with batch implementation
            for i in range(self._k):
                q, r = torch.linalg.qr(XNew[i])
                XNew[i] = torch.matmul(q, torch.diag(
                    torch.sign(torch.sign(torch.diag(r)) + .5)))
        return XNew
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
                                    bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1*W[:, 0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
        self.cache_inverse = cache_inverse
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
                                    bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1*W[:, 0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
        self.cache_inverse = cache_inverse
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        
        self.conv = torch.nn.Conv2d(c, c, kernel_size=(1,1), stride=1, padding=0,
                                    bias=False)
        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W.unsqueeze_(-1)
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        
        self.conv = torch.nn.Conv2d(c, c, kernel_size=(1,1), stride=1, padding=0,
                                    bias=False)
        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W.unsqueeze_(-1)
//  [python]
def orthogonal_random_matrix_(w):
    """Initialize the matrix w in-place to compute orthogonal random features.

    The matrix is initialized such that its columns are orthogonal to each
    other (in groups of size `rows`) and their norms is drawn from the
    chi-square distribution with `rows` degrees of freedom (namely the norm of
    a `rows`-dimensional vector distributed as N(0, I)).

    Arguments
    ---------
        w: float tensor of size (rows, columns)
    """
    rows, columns = w.shape
    start = 0
    while start < columns:
        end = min(start+rows, columns)
        block = torch.randn(rows, rows, device=w.device)
        norms = torch.sqrt(torch.einsum("ab,ab->a", block, block))
        Q, _ = torch.qr(block)
        w[:, start:end] = (
            Q[:, :end-start] * norms[None, :end-start]
        )
        start += rows
//  [python]
def orthogonal_random_matrix_(w):
    """Initialize the matrix w in-place to compute orthogonal random features.

    The matrix is initialized such that its columns are orthogonal to each
    other (in groups of size `rows`) and their norms is drawn from the
    chi-square distribution with `rows` degrees of freedom (namely the norm of
    a `rows`-dimensional vector distributed as N(0, I)).

    Arguments
    ---------
        w: float tensor of size (rows, columns)
    """
    rows, columns = w.shape
    start = 0
    while start < columns:
        end = min(start+rows, columns)
        block = torch.randn(rows, rows, device=w.device)
        norms = torch.sqrt(torch.einsum("ab,ab->a", block, block))
        Q, _ = torch.linalg.qr(block)
        w[:, start:end] = (
            Q[:, :end-start] * norms[None, :end-start]
        )
        start += rows
//  [python]
def __init__(self, input_dim, permutation=None):
        super(LinearPLUTransform, self).__init__()
        self.__delattr__('permutation')
        self._cache_size = 0
        # Sample a random orthogonal matrix
        W, _ = torch.qr(torch.randn(input_dim, input_dim))
        # Construct the partially pivoted LU-form and the pivots
        LU, pivots = W.lu()

        # Convert the pivots into the permutation matrix
        if permutation is None:
            P, _, _ = torch.lu_unpack(LU, pivots)
        else:
            if len(permutation) != input_dim:
                raise ValueError(
                    'Keyword argument "permutation" expected to have {} elements but {} found.'.format(
                        input_dim, len(permutation)))
            P = torch.eye(input_dim, input_dim)[permutation.type(dtype=torch.int64)]

        # We register the permutation matrix so that the model can be serialized
        self.register_buffer('permutation', P)

        # NOTE: For this implementation I have chosen to store the parameters densely, rather than
        # storing L, U, and s separately
        self.LU = torch.nn.Parameter(LU)
//  [python]
def __init__(self, input_dim, permutation=None):
        super(LinearPLUTransform, self).__init__()
        self.__delattr__('permutation')
        self._cache_size = 0
        # Sample a random orthogonal matrix
        W, _ = torch.linalg.qr(torch.randn(input_dim, input_dim))
        # Construct the partially pivoted LU-form and the pivots
        LU, pivots = W.lu()

        # Convert the pivots into the permutation matrix
        if permutation is None:
            P, _, _ = torch.lu_unpack(LU, pivots)
        else:
            if len(permutation) != input_dim:
                raise ValueError(
                    'Keyword argument "permutation" expected to have {} elements but {} found.'.format(
                        input_dim, len(permutation)))
            P = torch.eye(input_dim, input_dim)[permutation.type(dtype=torch.int64)]

        # We register the permutation matrix so that the model can be serialized
        self.register_buffer('permutation', P)

        # NOTE: For this implementation I have chosen to store the parameters densely, rather than
        # storing L, U, and s separately
        self.LU = torch.nn.Parameter(LU)
//  [python]
def test_ormqr(self, device):
        mat1 = torch.randn(7, 7)
        mat2 = torch.randn(7, 7)
        q, r = torch.qr(mat1)
        m, tau = torch.geqrf(mat1)
        out_holder = torch.empty_like(mat1)

        res1 = torch.mm(q, mat2)
        res2 = torch.ormqr(m, tau, mat2, left=True, transpose=False)
        torch.ormqr(m, tau, mat2, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)

        res1 = torch.mm(mat2, q)
        res2 = torch.ormqr(m, tau, mat2, left=False, transpose=False)
        torch.ormqr(m, tau, mat2, left=False, transpose=False, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)

        res1 = torch.mm(q.t(), mat2)
        res2 = torch.ormqr(m, tau, mat2, left=True, transpose=True)
        torch.ormqr(m, tau, mat2, left=True, transpose=True, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)

        res1 = torch.mm(mat2, q.t())
        res2 = torch.ormqr(m, tau, mat2, left=False, transpose=True)
        torch.ormqr(m, tau, mat2, left=False, transpose=True, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)
//  [python]
def test_ormqr(self, device):
        mat1 = torch.randn(7, 7)
        mat2 = torch.randn(7, 7)
        q, r = torch.linalg.qr(mat1)
        m, tau = torch.geqrf(mat1)
        out_holder = torch.empty_like(mat1)

        res1 = torch.mm(q, mat2)
        res2 = torch.ormqr(m, tau, mat2, left=True, transpose=False)
        torch.ormqr(m, tau, mat2, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)

        res1 = torch.mm(mat2, q)
        res2 = torch.ormqr(m, tau, mat2, left=False, transpose=False)
        torch.ormqr(m, tau, mat2, left=False, transpose=False, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)

        res1 = torch.mm(q.t(), mat2)
        res2 = torch.ormqr(m, tau, mat2, left=True, transpose=True)
        torch.ormqr(m, tau, mat2, left=True, transpose=True, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)

        res1 = torch.mm(mat2, q.t())
        res2 = torch.ormqr(m, tau, mat2, left=False, transpose=True)
        torch.ormqr(m, tau, mat2, left=False, transpose=True, out=out_holder)
        self.assertEqual(res1, res2)
        self.assertEqual(res2, out_holder)
//  [python]
def __init__(self, in_channel, fixed, use_lu):
        super().__init__()
        self.use_lu = use_lu
        self.fixed = fixed
        if not use_lu:
            weight = th.randn(in_channel, in_channel)
            q, _ = th.qr(weight)
            weight = q
            if fixed:
                self.register_buffer('weight', weight.data)
                self.register_buffer('weight_inverse', weight.data.inverse())
                self.register_buffer('fixed_log_det',
                                     th.slogdet(self.weight.double())[1].float())
            else:
                self.weight = nn.Parameter(weight)
        if use_lu:
            assert not fixed
            #weight = np.random.randn(in_channel, in_channel)
            weight = th.randn(in_channel, in_channel)
            #q, _ = la.qr(weight)
            q, _ = th.qr(weight)

            # w_p, w_l, w_u = la.lu(q.astype(np.float32))
            w_p, w_l, w_u = th.lu_unpack(*th.lu(q))

            #w_s = np.diag(w_u)
            w_s = th.diag(w_u)
            #w_u = np.triu(w_u, 1)
            w_u = th.triu(w_u, 1)
            #u_mask = np.triu(np.ones_like(w_u), 1)
            u_mask = th.triu(th.ones_like(w_u), 1)
            #l_mask = u_mask.T
            l_mask = u_mask.t()

            #w_p = th.from_numpy(w_p)
            #w_l = th.from_numpy(w_l)
            #w_s = th.from_numpy(w_s)
            #w_u = th.from_numpy(w_u)

            self.register_buffer('w_p', w_p)
            self.register_buffer('u_mask', u_mask)
            self.register_buffer('l_mask', l_mask)
            self.register_buffer('s_sign', th.sign(w_s))
            self.register_buffer('l_eye', th.eye(l_mask.shape[0]))
            self.w_l = nn.Parameter(w_l)
            self.w_s = nn.Parameter(th.log(th.abs(w_s)))
            self.w_u = nn.Parameter(w_u)
//  [python]
def __init__(self, in_channel, fixed, use_lu):
        super().__init__()
        self.use_lu = use_lu
        self.fixed = fixed
        if not use_lu:
            weight = th.randn(in_channel, in_channel)
            q, _ = th.linalg.qr(weight)
            weight = q
            if fixed:
                self.register_buffer('weight', weight.data)
                self.register_buffer('weight_inverse', weight.data.inverse())
                self.register_buffer('fixed_log_det',
                                     th.slogdet(self.weight.double())[1].float())
            else:
                self.weight = nn.Parameter(weight)
        if use_lu:
            assert not fixed
            #weight = np.random.randn(in_channel, in_channel)
            weight = th.randn(in_channel, in_channel)
            #q, _ = la.qr(weight)
            q, _ = th.linalg.qr(weight)

            # w_p, w_l, w_u = la.lu(q.astype(np.float32))
            w_p, w_l, w_u = th.lu_unpack(*th.lu(q))

            #w_s = np.diag(w_u)
            w_s = th.diag(w_u)
            #w_u = np.triu(w_u, 1)
            w_u = th.triu(w_u, 1)
            #u_mask = np.triu(np.ones_like(w_u), 1)
            u_mask = th.triu(th.ones_like(w_u), 1)
            #l_mask = u_mask.T
            l_mask = u_mask.t()

            #w_p = th.from_numpy(w_p)
            #w_l = th.from_numpy(w_l)
            #w_s = th.from_numpy(w_s)
            #w_u = th.from_numpy(w_u)

            self.register_buffer('w_p', w_p)
            self.register_buffer('u_mask', u_mask)
            self.register_buffer('l_mask', l_mask)
            self.register_buffer('s_sign', th.sign(w_s))
            self.register_buffer('l_eye', th.eye(l_mask.shape[0]))
            self.w_l = nn.Parameter(w_l)
            self.w_s = nn.Parameter(th.log(th.abs(w_s)))
            self.w_u = nn.Parameter(w_u)
//  [python]
def init_templates(self):
        template_shape = [self._n_caps, self._n_channels] + list(self._template_size)  # torch generally uses (N, C, H, W)
        n_elems = np.prod(template_shape[1:])

        # make each templates orthogonal to each other at init
        n = max(self._n_caps, n_elems)
        q, _ = torch.qr(torch.rand(n, n))  # TODO: test whether orthogonal init is even helpful

        col_idxs = list(range(q.shape[1]))
        random.shuffle(col_idxs)
        q = q[:, torch.tensor(col_idxs)]

        ts = q[:self._n_caps, :n_elems].reshape(template_shape)

        t_min = ts.min()
        t_max = ts.max()
        ts = (ts - t_min) / (t_max - t_min) * 2 - 1  # normalize values to [-1,1]
        if self._use_alpha_channel:
            alphas = torch.zeros(self._n_caps, 1, *self._template_size)
            ts = torch.cat([ts, alphas], dim=1)
        else:
            self.temperature_logit = torch.nn.Parameter(torch.tensor([0.]), requires_grad=True)

        self.templates = torch.nn.Parameter(ts, requires_grad=True)
//  [python]
def init_templates(self):
        template_shape = [self._n_caps, self._n_channels] + list(self._template_size)  # torch generally uses (N, C, H, W)
        n_elems = np.prod(template_shape[1:])

        # make each templates orthogonal to each other at init
        n = max(self._n_caps, n_elems)
        q, _ = torch.linalg.qr(torch.rand(n, n))  # TODO: test whether orthogonal init is even helpful

        col_idxs = list(range(q.shape[1]))
        random.shuffle(col_idxs)
        q = q[:, torch.tensor(col_idxs)]

        ts = q[:self._n_caps, :n_elems].reshape(template_shape)

        t_min = ts.min()
        t_max = ts.max()
        ts = (ts - t_min) / (t_max - t_min) * 2 - 1  # normalize values to [-1,1]
        if self._use_alpha_channel:
            alphas = torch.zeros(self._n_caps, 1, *self._template_size)
            ts = torch.cat([ts, alphas], dim=1)
        else:
            self.temperature_logit = torch.nn.Parameter(torch.tensor([0.]), requires_grad=True)

        self.templates = torch.nn.Parameter(ts, requires_grad=True)
//  [python]
def test_qr(self, device, dtype):
        def run_test(tensor_dims, some):
            A = torch.randn(*tensor_dims, dtype=dtype, device=device)
            Q, R = torch.qr(A, some=some)

            # Check0: Q[-2:] = (m, n_columns), R[-2:] = (n_columns, n)
            m, n = tensor_dims[-2:]
            n_columns = m if (not some) and m > n else min(m, n)
            self.assertEqual(Q.size(-2), m)
            self.assertEqual(R.size(-1), n)
            self.assertEqual(Q.size(-1), n_columns)

            A_ = A.cpu().numpy()
            Q_ = Q.cpu().numpy()
            R_ = R.cpu().numpy()

            # Check1: A = QR
            self.assertEqual(A_, np.matmul(Q_, R_))

            # Check2: A = QR (with out)
            Q_out, R_out = torch.full_like(Q, math.nan), torch.full_like(R, math.nan)
            torch.qr(A, some=some, out=(Q_out, R_out))
            Q_out_ = Q_out.cpu().numpy()
            R_out_ = R_out.cpu().numpy()
            self.assertEqual(A_, np.matmul(Q_out_, R_out_))

            # Check3: Q == Q_out, R == R_out
            self.assertEqual(Q_, Q_out_)
            self.assertEqual(R_, R_out_)

            # Check4: Q^{T}Q = I, triu(R) = R
            eye = torch.eye(n_columns, device=device, dtype=dtype).expand(Q.shape[:-2] + (n_columns, n_columns)).cpu().numpy()
            self.assertEqual(np.matmul(Q_.swapaxes(-1, -2).conj(), Q_), eye)
            self.assertEqual(R.triu(), R)

        tensor_dims_list = [(3, 5), (5, 5), (5, 3),  # Single matrix
                            (7, 3, 5), (7, 5, 5), (7, 5, 3),  # 3-dim Tensors
                            (7, 5, 3, 5), (7, 5, 5, 5), (7, 5, 5, 3)]  # 4-dim Tensors
        for tensor_dims, some in itertools.product(tensor_dims_list, [True, False]):
            run_test(tensor_dims, some)
//  [python]
def test_qr(self, device, dtype):
        def run_test(tensor_dims, some):
            A = torch.randn(*tensor_dims, dtype=dtype, device=device)
            Q, R = torch.linalg.qr(A, mode='reduced' if not some else 'complete')

            # Check0: Q[-2:] = (m, n_columns), R[-2:] = (n_columns, n)
            m, n = tensor_dims[-2:]
            n_columns = m if (not some) and m > n else min(m, n)
            self.assertEqual(Q.size(-2), m)
            self.assertEqual(R.size(-1), n)
            self.assertEqual(Q.size(-1), n_columns)

            A_ = A.cpu().numpy()
            Q_ = Q.cpu().numpy()
            R_ = R.cpu().numpy()

            # Check1: A = QR
            self.assertEqual(A_, np.matmul(Q_, R_))

            # Check2: A = QR (with out)
            Q_out, R_out = torch.full_like(Q, math.nan), torch.full_like(R, math.nan)
            torch.linalg.qr(A, mode='reduced' if not some else 'complete', out=(Q_out, R_out))
            Q_out_ = Q_out.cpu().numpy()
            R_out_ = R_out.cpu().numpy()
            self.assertEqual(A_, np.matmul(Q_out_, R_out_))

            # Check3: Q == Q_out, R == R_out
            self.assertEqual(Q_, Q_out_)
            self.assertEqual(R_, R_out_)

            # Check4: Q^{T}Q = I, triu(R) = R
            eye = torch.eye(n_columns, device=device, dtype=dtype).expand(Q.shape[:-2] + (n_columns, n_columns)).cpu().numpy()
            self.assertEqual(np.matmul(Q_.swapaxes(-1, -2).conj(), Q_), eye)
            self.assertEqual(R.triu(), R)

        tensor_dims_list = [(3, 5), (5, 5), (5, 3),  # Single matrix
                            (7, 3, 5), (7, 5, 5), (7, 5, 3),  # 3-dim Tensors
                            (7, 5, 3, 5), (7, 5, 5, 5), (7, 5, 5, 3)]  # 4-dim Tensors
        for tensor_dims, some in itertools.product(tensor_dims_list, [True, False]):
            run_test(tensor_dims, some)
//  [python]
def __init__(self, num_channels, use_lu=False):
        """
        Constructor
        :param num_channels: Number of channels of the data
        :param use_lu: Flag whether to parametrize weights through the LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q = torch.qr(torch.randn(self.num_channels, self.num_channels))[0]
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer('P', P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(torch.triu(U, diagonal=1))  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def __init__(self, num_channels, use_lu=False):
        """
        Constructor
        :param num_channels: Number of channels of the data
        :param use_lu: Flag whether to parametrize weights through the LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q = torch.linalg.qr(torch.randn(self.num_channels, self.num_channels))[0]
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer('P', P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(torch.triu(U, diagonal=1))  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def __init__(self, dim):
        super().__init__()

        weight = torch.randn(dim, dim)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(0)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, dim):
        super().__init__()

        weight = torch.randn(dim, dim)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(0)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, num_channels, use_lu=True):
        """
        Constructor
        :param num_channels: Number of channels of the data
        :param use_lu: Flag whether to parametrize weights through the
        LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q = torch.qr(torch.randn(self.num_channels, self.num_channels))[0]
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer('P', P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(torch.triu(U, diagonal=1))  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def __init__(self, num_channels, use_lu=True):
        """
        Constructor
        :param num_channels: Number of channels of the data
        :param use_lu: Flag whether to parametrize weights through the
        LU decomposition
        """
        super().__init__()
        self.num_channels = num_channels
        self.use_lu = use_lu
        Q = torch.linalg.qr(torch.randn(self.num_channels, self.num_channels))[0]
        if use_lu:
            P, L, U = torch.lu_unpack(*Q.lu())
            self.register_buffer('P', P)  # remains fixed during optimization
            self.L = nn.Parameter(L)  # lower triangular portion
            S = U.diag()  # "crop out" the diagonal to its own parameter
            self.register_buffer("sign_S", torch.sign(S))
            self.log_S = nn.Parameter(torch.log(torch.abs(S)))
            self.U = nn.Parameter(torch.triu(U, diagonal=1))  # "crop out" diagonal, stored in S
            self.register_buffer("eye", torch.diag(torch.ones(self.num_channels)))
        else:
            self.W = nn.Parameter(Q)
//  [python]
def orthogonal_(tensor, gain=1):
    # proper orthogonal init, see https://github.com/pytorch/pytorch/pull/10672
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = tensor.new_empty(rows, cols).normal_(0, 1)

    for i in range(0, rows, cols):
        # Compute the qr factorization
        q, r = torch.qr(flattened[i:i + cols].t())
        # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
        q *= torch.diag(r, 0).sign()
        q.t_()

        with torch.no_grad():
            tensor[i:i + cols].view_as(q).copy_(q)

    with torch.no_grad():
        tensor.mul_(gain)
    return tensor