jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def randomwalk(idx, val):
    dim = idx.shape[0]
    rows = torch.arange(idx.shape[0]).repeat_interleave(idx.shape[1])
    cols = idx.reshape(1, -1).squeeze()
    val = val.reshape(1, -1).squeeze()

    split = cols > rows
    rows, cols = torch.cat([rows[split], cols[~split]], dim=0), torch.cat(
        [cols[split], rows[~split]], dim=0)
    idx = torch.stack([rows, cols], dim=0)
    val = torch.cat([val[split], val[~split]])
    idx, val = coalesce(idx, val, reduce='mean')
    rows = idx[0, :]
    cols = idx[1, :]

    idx = torch.cat((torch.arange(dim).repeat(2, 1), idx, torch.stack((cols,rows),dim=0)), dim=1)
    val = torch.cat([torch.zeros(dim), val.repeat(2)])
    
    val = val.div(-4*epsilon.square()).exp().squeeze()

    Lsparse = torch.sparse_coo_tensor(idx, val, (dim, dim))
    L = Lsparse.to_dense()# + Lsparse.to_dense().t()
    D = L.sum(dim=1).pow(-1).diag()
    L = torch.mm(D, torch.mm(L, D))

    D = L.sum(dim=1)  # .pow(-1).diag()
    L = torch.solve(L, D.diag())[0]  # torch.mm(D.pow(-1).diag(), L)

    # L = torch.mm(D, L)
    # L *= D.unsqueeze(-1)

    # return L, D
    return (torch.eye(dim)-L)/epsilon.square(), D
//  [python]
def randomwalk(idx, val):
    dim = idx.shape[0]
    rows = torch.arange(idx.shape[0]).repeat_interleave(idx.shape[1])
    cols = idx.reshape(1, -1).squeeze()
    val = val.reshape(1, -1).squeeze()

    split = cols > rows
    rows, cols = torch.cat([rows[split], cols[~split]], dim=0), torch.cat(
        [cols[split], rows[~split]], dim=0)
    idx = torch.stack([rows, cols], dim=0)
    val = torch.cat([val[split], val[~split]])
    idx, val = coalesce(idx, val, reduce='mean')
    rows = idx[0, :]
    cols = idx[1, :]

    idx = torch.cat((torch.arange(dim).repeat(2, 1), idx, torch.stack((cols,rows),dim=0)), dim=1)
    val = torch.cat([torch.zeros(dim), val.repeat(2)])
    
    val = val.div(-4*epsilon.square()).exp().squeeze()

    Lsparse = torch.sparse_coo_tensor(idx, val, (dim, dim))
    L = Lsparse.to_dense()# + Lsparse.to_dense().t()
    D = L.sum(dim=1).pow(-1).diag()
    L = torch.mm(D, torch.mm(L, D))

    D = L.sum(dim=1)  # .pow(-1).diag()
    L = torch.linalg.solve(D.diag(), L)  # torch.mm(D.pow(-1).diag(), L)

    # L = torch.mm(D, L)
    # L *= D.unsqueeze(-1)

    # return L, D
    return (torch.eye(dim)-L)/epsilon.square(), D
//  [python]
def fit(self, X: Tensor, y: Tensor, freq: Tensor):
        assert X.shape[:2] == freq.shape[:2]
        assert len(X.shape) == 3
        assert len(y.shape) == 2
        self.X_ = X
        self.y_ = y
        self._freq = freq

        k = self.kernel(X, X, freq, freq)
        # klmbda = k + torch.eye(k.shape[0]) * self.lmbda
        # Below is the same as above but avoids the creation of a new tensor on a different device
        klmbda = (k - k.diag().diag()) + (self.lmbda * self.X_.shape[0] + k.diag()).diag()
        self.k_ = k
        self.alpha_ = torch.solve(y, klmbda)[0]
        return self
//  [python]
def fit(self, X: Tensor, y: Tensor, freq: Tensor):
        assert X.shape[:2] == freq.shape[:2]
        assert len(X.shape) == 3
        assert len(y.shape) == 2
        self.X_ = X
        self.y_ = y
        self._freq = freq

        k = self.kernel(X, X, freq, freq)
        # klmbda = k + torch.eye(k.shape[0]) * self.lmbda
        # Below is the same as above but avoids the creation of a new tensor on a different device
        klmbda = (k - k.diag().diag()) + (self.lmbda * self.X_.shape[0] + k.diag()).diag()
        self.k_ = k
        self.alpha_ = LA.solve(klmbda, y)
        return self
//  [python]
def solve(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:  # pragma: no cover
    """
    Like torch.linalg.solve, tries to return X
    such that AX=B, with A square.
    """
    if hasattr(torch.linalg, "solve"):
        # PyTorch version >= 1.8.0
        return torch.linalg.solve(A, B)

    return torch.solve(B, A).solution
//  [python]
def fit(self, X: Tensor, ZX: Tensor, y: Tensor):
        assert len(X.shape) == 3
        assert len(ZX.shape) == 3
        assert len(y.shape) == 2
        self.X_ = X
        self.ZX_ = ZX
        self.y_ = y

        k = self.kernel(X, ZX, X, ZX)
        # klmbda = k + torch.eye(k.shape[0]) * self.lmbda
        # Below is the same as above but avoids the creation of a new tensor on a different device
        klmbda = (k - k.diag().diag()) + (self.lmbda * self.X_.shape[0] + k.diag()).diag()
        self.k_ = k
        self.alpha_ = torch.solve(y, klmbda)[0]
        return self
//  [python]
def fit(self, X: Tensor, ZX: Tensor, y: Tensor):
        assert len(X.shape) == 3
        assert len(ZX.shape) == 3
        assert len(y.shape) == 2
        self.X_ = X
        self.ZX_ = ZX
        self.y_ = y

        k = self.kernel(X, ZX, X, ZX)
        # klmbda = k + torch.eye(k.shape[0]) * self.lmbda
        # Below is the same as above but avoids the creation of a new tensor on a different device
        klmbda = (k - k.diag().diag()) + (self.lmbda * self.X_.shape[0] + k.diag()).diag()
        self.k_ = k
        self.alpha_ = LA.solve(klmbda, y)
        return self
//  [python]
def _get_projection_matrix(W_: Tensor, sigma_: Tensor, device: Any):
    """
    This is currently just implemented for PPCA due to nicer formula. Will
    modify for broader application later.

    Computes the parameters for p(S|X)

    Description in future to be released paper

    Parameters
    ----------
    W_ : Tensor(n_components,p)
        The loadings

    sigma_ : Tensor
        Isotropic noise

    Returns
    -------
    Proj_X : Tensor(n_components,p)
        Beta such that np.dot(Proj_X, X) = E[S|X]

    Cov : Tensor(n_components,n_components)
        Var(S|X)
    """
    n_components = int(W_.shape[0])
    eye = torch.tensor(np.eye(n_components).astype(np.float32), device=device)
    M_init = torch.matmul(W_, torch.transpose(W_, 0, 1))
    M_end = sigma_ * eye
    M = M_init + M_end
    Proj_X = torch.solve(W_,M)[1]
    Cov = torch.inverse(M) * sigma_
    return Proj_X, Cov
//  [python]
def _get_projection_matrix(W_: Tensor, sigma_: Tensor, device: Any):
    """
    This is currently just implemented for PPCA due to nicer formula. Will
    modify for broader application later.

    Computes the parameters for p(S|X)

    Description in future to be released paper

    Parameters
    ----------
    W_ : Tensor(n_components,p)
        The loadings

    sigma_ : Tensor
        Isotropic noise

    Returns
    -------
    Proj_X : Tensor(n_components,p)
        Beta such that np.dot(Proj_X, X) = E[S|X]

    Cov : Tensor(n_components,n_components)
        Var(S|X)
    """
    n_components = int(W_.shape[0])
    eye = torch.tensor(np.eye(n_components).astype(np.float32), device=device)
    M_init = torch.matmul(W_, torch.transpose(W_, 0, 1))
    M_end = sigma_ * eye
    M = M_init + M_end
    Proj_X = torch.linalg.solve(M, W_)
    Cov = torch.linalg.inv(M) * sigma_
    return Proj_X, Cov
//  [python]
def line_intersection(line1, line2):
	"""
	line in format np.array[[x1,y1],[x2,y2]]
	returns intersection point
	"""
	try:
		t, s = torch.solve(line2[0] - line1[0], torch.vstack([line1[1] - line1[0], line2[0] - line2[1]]).T)
		return (1 - t) * line1[0] + t * line1[1]
	except:  # np.linalg.LinAlgError:
		print("no intersection")
		return None
//  [python]
def line_intersection(line1, line2):
	"""
	line in format np.array[[x1,y1],[x2,y2]]
	returns intersection point
	"""
	try:
		t, s = torch.linalg.solve(torch.vstack([line1[1] - line1[0], line2[0] - line2[1]]).T, line2[0] - line1[0])
		return (1 - t) * line1[0] + t * line1[1]
	except:  # np.linalg.LinAlgError:
		print("no intersection")
		return None
//  [python]
def _stable_solve(b, a, eps=1e-06):
    try:
        return torch.solve(b, a)[0]
    except RuntimeError:
        a = condition_scm(a, eps)
        return torch.solve(b, a)[0]
//  [python]
def _stable_solve(b, a, eps=1e-06):
    try:
        return torch.linalg.solve(a, b)
    except RuntimeError:
        a = condition_scm(a, eps)
        return torch.linalg.solve(a, b)
//  [python]
def check_line_intersects(line1, line2):
	"""
	line in format np.array[[x1,y1],[x2,y2]]
	:return if there is an intersection
	"""
	try:
		A = torch.vstack([line1[1] - line1[0], line2[0] - line2[1]]).T
		b = line2[0] - line1[0]
		t, s = torch.solve(b,A)[0]
		return (0 <= t <= 1) and (0 <= s <= 1)
	except:  # torch.linalg.LinAlgError:
		# print("no intersection while checking")
		return False
//  [python]
def check_line_intersects(line1, line2):
	"""
	line in format np.array[[x1,y1],[x2,y2]]
	:return if there is an intersection
	"""
	try:
		t, s = torch.linalg.solve(torch.vstack([line1[1] - line1[0], line2[0] - line2[1]]).T, line2[0] - line1[0])
		return (0 <= t <= 1) and (0 <= s <= 1)
	except:  # torch.linalg.LinAlgError:
		# print("no intersection while checking")
		return False
//  [python]
def coeff(nodes, order=1):
    '''
    This function is used to calculate the differential coefficient of given nodes
    '''
    m = len(nodes)
    device=nodes.device

    factor, pownodes = torch.ones(m, device=device), torch.ones(m, m, device=device)
    b = torch.zeros(m, 1, device=device)
    b[order] = 1

    for i in range(1, m):
        factor[i] = i * factor[i-1]
    for i in range(1,m):
        pownodes[:,i] = nodes * pownodes[:,i-1]
        
    A = pownodes / factor
    A = A.T
    x = torch.solve(b, A)[0]
    return x
//  [python]
def coeff(nodes, order=1):
    '''
    This function is used to calculate the differential coefficient of given nodes
    '''
    m = len(nodes)
    device=nodes.device

    factor, pownodes = torch.ones(m, device=device), torch.ones(m, m, device=device)
    b = torch.zeros(m, 1, device=device)
    b[order] = 1

    for i in range(1, m):
        factor[i] = i * factor[i-1]
    for i in range(1,m):
        pownodes[:,i] = nodes * pownodes[:,i-1]
        
    A = pownodes / factor
    A = A.T
    x = torch.linalg.solve(A, b)
    return x
//  [python]
def get_wout(s, sd,sdd, y0,y0dot,m1,m2,k1,k2, t):


    Lmat = torch.tensor([[m1, 0.], [0., m2]])
    Rmat = torch.tensor([[k1 + k2, -k2], [-k2, k1 + k2]])
    Amatrix = torch.linalg.inv(Lmat)@Rmat
    hddothat = torch.block_diag(sdd,sdd)
    hdothat = torch.block_diag(sd,sd)
    hhat = torch.block_diag(s,s)
    Amatrixhat = torch.zeros((hdothat.shape[0],hdothat.shape[0]))
    for i in range(Amatrix.shape[0]):
        for j in range(Amatrix.shape[1]):
            Amatrixhat[i*s.shape[0]:(i+1)*s.shape[0],j*s.shape[0]:(j+1)*s.shape[0]]=torch.eye(s.shape[0],s.shape[0])*Amatrix[i,j]

    DH = hddothat + Amatrixhat@hhat


    h0= torch.block_diag(s[0,:].reshape(1,-1),s[0,:].reshape(1,-1))

    h0dot = torch.block_diag(sd[0, :].reshape(1, -1), sd[0, :].reshape(1, -1))

    W0 = torch.solve(h0.t()@y0.t() + h0dot.t()@y0dot.t(), DH.t()@DH + h0.t()@h0 + h0dot.t()@h0dot)[0]
    return W0
//  [python]
def get_wout(s, sd,sdd, y0,y0dot,m1,m2,k1,k2, t):


    Lmat = torch.tensor([[m1, 0.], [0., m2]])
    Rmat = torch.tensor([[k1 + k2, -k2], [-k2, k1 + k2]])
    Amatrix = torch.linalg.inv(Lmat)@Rmat
    hddothat = torch.block_diag(sdd,sdd)
    hdothat = torch.block_diag(sd,sd)
    hhat = torch.block_diag(s,s)
    Amatrixhat = torch.zeros((hdothat.shape[0],hdothat.shape[0]))
    for i in range(Amatrix.shape[0]):
        for j in range(Amatrix.shape[1]):
            Amatrixhat[i*s.shape[0]:(i+1)*s.shape[0],j*s.shape[0]:(j+1)*s.shape[0]]=torch.eye(s.shape[0],s.shape[0])*Amatrix[i,j]

    DH = hddothat + Amatrixhat@hhat


    h0= torch.block_diag(s[0,:].reshape(1,-1),s[0,:].reshape(1,-1))

    h0dot = torch.block_diag(sd[0, :].reshape(1, -1), sd[0, :].reshape(1, -1))

    W0 = torch.linalg.solve(DH.t()@DH + h0.t()@h0 + h0dot.t()@h0dot, h0.t()@y0.t() + h0dot.t()@y0dot.t() )
    return W0
//  [python]
def triangle_ray_intersects(A: Point, B: Point, C: Point, O: Point, D: Point) -> bool:
	'''
	A: shape (3,), one vertex of the triangle
	B: shape (3,), second vertex of the triangle
	C: shape (3,), third vertex of the triangle
	O: shape (3,), origin point
	D: shape (3,), direction point

	Return True if the ray and the triangle intersect.
	'''
	s, u, v = t.solve(
		t.stack([-D, B - A, C - A], dim=1), 
		O - A
	)
	return ((u >= 0) & (v >= 0) & (u + v <= 1)).item()
//  [python]
def triangle_ray_intersects(A: Point, B: Point, C: Point, O: Point, D: Point) -> bool:
	'''
	A: shape (3,), one vertex of the triangle
	B: shape (3,), second vertex of the triangle
	C: shape (3,), third vertex of the triangle
	O: shape (3,), origin point
	D: shape (3,), direction point

	Return True if the ray and the triangle intersect.
	'''
	s, u, v = t.linalg.solve(
		t.stack([-D, B - A, C - A], dim=1), 
		O - A
	)
	return ((u >= 0) & (v >= 0) & (u + v <= 1)).item()
//  [python]
def recon_fastssss(y, S, T, order, sd=0.5):
    """A primary implementation of reconstruction method associated with
    "FastSSS" sampling algorithm.

    Parameters
    ----------
    y:  Tensor
        The measurements on sampling set :obj:S:. If the localization operator :obj:`T`
        has a density greater than the threshold :obj:`sd`, :obj:`y` has a shape of
        either :obj:`(M,)`，:obj:`(M,1)`，or :obj:`(M,C)`; otherwise :obj:`y` could only
        be  either :obj:`(M,)` or :obj:`(M,1)`.
    S:  List
        A list consisting of all indices of sampled nodes.
    T:  Tensor, SparseTensor
        The localization operator
    order: int
    sd: float
        The threshold of :obj:`T`'s density that controls when we use a dense or sparse
        linear solver.

    Returns
    -------
    Tensor
        The signal recovered from the measurements :obj:`y` .
    """
    T_k = matrix_power(T, order)
    dt, dv, density, on_gpu = get_ddd(T_k)
    if density > sd:
        T_k = T_k.to_dense()
        tmp = torch.solve(T_k[np.ix_(S, S)], torch.as_tensor(y))[0]
        x_hat = T_k[:, S] @ tmp

    else:
        T_k = to_xcipy(T_k)
        xp, xcipy, xsplin = get_array_module(on_gpu)
        tmp = xsplin.spsolve(T_k[xp.ix_(S, S)], xp.asarray(y))
        x_hat = T_k[:, S] * tmp
    return torch.as_tensor(x_hat, device=dv)
//  [python]
def recon_fastssss(y, S, T, order, sd=0.5):
    """A primary implementation of reconstruction method associated with
    "FastSSS" sampling algorithm.

    Parameters
    ----------
    y:  Tensor
        The measurements on sampling set :obj:S:. If the localization operator :obj:`T`
        has a density greater than the threshold :obj:`sd`, :obj:`y` has a shape of
        either :obj:`(M,)`，:obj:`(M,1)`，or :obj:`(M,C)`; otherwise :obj:`y` could only
        be  either :obj:`(M,)` or :obj:`(M,1)`.
    S:  List
        A list consisting of all indices of sampled nodes.
    T:  Tensor, SparseTensor
        The localization operator
    order: int
    sd: float
        The threshold of :obj:`T`'s density that controls when we use a dense or sparse
        linear solver.

    Returns
    -------
    Tensor
        The signal recovered from the measurements :obj:`y` .
    """
    T_k = matrix_power(T, order)
    dt, dv, density, on_gpu = get_ddd(T_k)
    if density > sd:
        T_k = T_k.to_dense()
        tmp = torch.linalg.solve(T_k[np.ix_(S, S)], torch.as_tensor(y))
        x_hat = T_k[:, S] @ tmp

    else:
        T_k = to_xcipy(T_k)
        xp, xcipy, xsplin = get_array_module(on_gpu)
        tmp = xsplin.spsolve(T_k[xp.ix_(S, S)], xp.asarray(y))
        x_hat = T_k[:, S] * tmp
    return torch.as_tensor(x_hat, device=dv)
//  [python]
def update_hypernet(self, param_shifts: Dict[str, torch.FloatTensor]):
        
        self.opt.zero_grad()
        for module_idx, module_name in enumerate(self.config.model.edit_modules):
            shape = get_shape(get_module(self.model, module_name))
            net = self.net[str(shape)]
            layer_idx = torch.LongTensor([self.name2idx[module_name]]).to(self.config.editor_device)
            keys = torch.cat([
                torch.load(f"{self.config.editor.cache_dir}/{module_idx}_{idx}_keys.pth")
                for idx in range(math.ceil(self.config.data.n_edits / self.config.data.batch_size))
            ])
            values_grad = torch.cat([
                torch.load(f"{self.config.editor.cache_dir}/{module_idx}_{idx}_values_grad.pth")
                for idx in range(math.ceil(self.config.data.n_edits / self.config.data.batch_size))
            ])
            module = get_module(self.model, module_name)
            module_grad = module.weight.grad.to(torch.float32).to(self.config.editor_device)
            param_shift = param_shifts[module_name].to(self.config.editor_device)
            if isinstance(module, nn.Linear):
                module_grad = module_grad.T
            with torch.no_grad():
                mat = torch.solve(module_grad, keys.T @ keys + net.lamda(layer_idx).exp() * torch.eye(net.key_size, device = self.config.editor_device))[1]
                lamda_grad = - net.lamda(layer_idx).exp() * (mat * param_shift).sum()
            value_diffs_grad = keys @ mat
            (lamda_grad * net.lamda(layer_idx)).backward()
            for start_idx in range(0, keys.shape[0], self.config.editor.batch_size):
                end_idx = start_idx + self.config.editor.batch_size
                pesudo_keys, pesudo_values_grad = net(
                    keys[start_idx:end_idx],
                    values_grad[start_idx:end_idx],
                    layer_idx
                )
                coeffs = - net.lr(layer_idx) * (keys[start_idx:end_idx] * pesudo_keys).sum(-1).unsqueeze(-1)
                value_diff = coeffs * pesudo_values_grad
                (value_diffs_grad[start_idx:end_idx] * value_diff).sum().backward()
            
        clip_grad_norm_(
            self.net.parameters(),
            self.config.editor.max_grad_norm
        )
        self.opt.step()
//  [python]
def update_hypernet(self, param_shifts: Dict[str, torch.FloatTensor]):
        
        self.opt.zero_grad()
        for module_idx, module_name in enumerate(self.config.model.edit_modules):
            shape = get_shape(get_module(self.model, module_name))
            net = self.net[str(shape)]
            layer_idx = torch.LongTensor([self.name2idx[module_name]]).to(self.config.editor_device)
            keys = torch.cat([
                torch.load(f"{self.config.editor.cache_dir}/{module_idx}_{idx}_keys.pth")
                for idx in range(math.ceil(self.config.data.n_edits / self.config.data.batch_size))
            ])
            values_grad = torch.cat([
                torch.load(f"{self.config.editor.cache_dir}/{module_idx}_{idx}_values_grad.pth")
                for idx in range(math.ceil(self.config.data.n_edits / self.config.data.batch_size))
            ])
            module = get_module(self.model, module_name)
            module_grad = module.weight.grad.to(torch.float32).to(self.config.editor_device)
            param_shift = param_shifts[module_name].to(self.config.editor_device)
            if isinstance(module, nn.Linear):
                module_grad = module_grad.T
            with torch.no_grad():
                mat = torch.linalg.solve(keys.T @ keys + net.lamda(layer_idx).exp() * torch.eye(net.key_size, device = self.config.editor_device), module_grad)
                lamda_grad = - net.lamda(layer_idx).exp() * (mat * param_shift).sum()
            value_diffs_grad = keys @ mat
            (lamda_grad * net.lamda(layer_idx)).backward()
            for start_idx in range(0, keys.shape[0], self.config.editor.batch_size):
                end_idx = start_idx + self.config.editor.batch_size
                pesudo_keys, pesudo_values_grad = net(
                    keys[start_idx:end_idx],
                    values_grad[start_idx:end_idx],
                    layer_idx
                )
                coeffs = - net.lr(layer_idx) * (keys[start_idx:end_idx] * pesudo_keys).sum(-1).unsqueeze(-1)
                value_diff = coeffs * pesudo_values_grad
                (value_diffs_grad[start_idx:end_idx] * value_diff).sum().backward()
            
        clip_grad_norm_(
            self.net.parameters(),
            self.config.editor.max_grad_norm
        )
        self.opt.step()
//  [python]
def intersect_ray_1d(ray: t.Tensor, segment: t.Tensor) -> bool:
	'''
	ray: shape (n_points=2, n_dim=3)  # O, D points
	segment: shape (n_points=2, n_dim=3)  # L_1, L_2 points

	Return True if the ray intersects the segment.
	'''
	# Get the x and y coordinates (ignore z)
	ray = ray[..., :2]
	segment = segment[..., :2]

	# Ray is [[Ox, Oy], [Dx, Dy]]
	O, D = ray
	# Segment is [[L1x, L1y], [L2x, L2y]]
	L_1, L_2 = segment

	# Create matrix and vector, and solve equation
	mat = t.stack([D, L_1 - L_2], dim=-1)
	vec = L_1 - O

	# Solve equation (return False if no solution)
	try:
		sol = t.solve(vec, mat)[0]
	except:
		return False

	# If there is a solution, check the soln is in the correct range for there to be an intersection
	u = sol[0].item()
	v = sol[1].item()
	return (u >= 0.0) and (v >= 0.0) and (v <= 1.0)
//  [python]
def intersect_ray_1d(ray: t.Tensor, segment: t.Tensor) -> bool:
	'''
	ray: shape (n_points=2, n_dim=3)  # O, D points
	segment: shape (n_points=2, n_dim=3)  # L_1, L_2 points

	Return True if the ray intersects the segment.
	'''
	# Get the x and y coordinates (ignore z)
	ray = ray[..., :2]
	segment = segment[..., :2]

	# Ray is [[Ox, Oy], [Dx, Dy]]
	O, D = ray
	# Segment is [[L1x, L1y], [L2x, L2y]]
	L_1, L_2 = segment

	# Create matrix and vector, and solve equation
	mat = t.stack([D, L_1 - L_2], dim=-1)
	vec = L_1 - O

	# Solve equation (return False if no solution)
	try:
		sol = t.linalg.solve(mat, vec)
	except:
		return False

	# If there is a solution, check the soln is in the correct range for there to be an intersection
	u = sol[0].item()
	v = sol[1].item()
	return (u >= 0.0) and (v >= 0.0) and (v <= 1.0)
//  [python]
def fit_poly(self, kernel_func=None, target='hypo', fkine=None): #epsilon=None, 
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.rbf_kernel = kernel_func
        kmat = self.rbf_kernel(X, X)

        self.rbf_nodes = torch.solve(y[:, None], kmat)[0].reshape(-1)
        # print(kmat@self.rbf_nodes) # DEBUG
        if self._cuda:
            self.cuda()
//  [python]
def fit_poly(self, kernel_func=None, target='hypo', fkine=None): #epsilon=None, 
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.rbf_kernel = kernel_func
        kmat = self.rbf_kernel(X, X)

        self.rbf_nodes = torch.linalg.solve(kmat, y[:, None]).reshape(-1)
        # print(kmat@self.rbf_nodes) # DEBUG
        if self._cuda:
            self.cuda()
//  [python]
def raytrace_triangle(
	rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
	triangle: Float[Tensor, "trianglePoints=3 dims=3"]
) -> Bool[Tensor, "nrays"]:
	'''
	For each ray, return True if the triangle intersects that ray.
	'''
	NR = rays.size(0)

	# Triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
	A, B, C = einops.repeat(triangle, "pts dims -> pts NR dims", NR=NR)
	assert A.shape == (NR, 3)

	# Each element of `rays` is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
	O, D = rays.unbind(dim=1)
	assert O.shape == (NR, 3)

	# Define matrix on left hand side of equation
	mat: Float[Tensor, "NR 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
	
	# Get boolean of where matrix is singular, and replace it with the identity in these positions
	# Note - this works because mat[is_singular] has shape (NR_where_singular, 3, 3), so we
	# can broadcast the identity matrix to that shape.
	dets: Float[Tensor, "NR"] = t.linalg.det(mat)
	is_singular = dets.abs() < 1e-8
	mat[is_singular] = t.eye(3)

	# Define vector on the right hand side of equation
	vec = O - A

	# Solve eqns
	sol: Float[Tensor, "NR 3"] = t.solve(mat, vec)
	s, u, v = sol.unbind(dim=-1)

	# Return boolean of (matrix is nonsingular, and solution is in correct range implying intersection)
	return ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
//  [python]
def raytrace_triangle(
	rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
	triangle: Float[Tensor, "trianglePoints=3 dims=3"]
) -> Bool[Tensor, "nrays"]:
	'''
	For each ray, return True if the triangle intersects that ray.
	'''
	NR = rays.size(0)

	# Triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
	A, B, C = einops.repeat(triangle, "pts dims -> pts NR dims", NR=NR)
	assert A.shape == (NR, 3)

	# Each element of `rays` is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
	O, D = rays.unbind(dim=1)
	assert O.shape == (NR, 3)

	# Define matrix on left hand side of equation
	mat: Float[Tensor, "NR 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
	
	# Get boolean of where matrix is singular, and replace it with the identity in these positions
	# Note - this works because mat[is_singular] has shape (NR_where_singular, 3, 3), so we
	# can broadcast the identity matrix to that shape.
	dets: Float[Tensor, "NR"] = t.linalg.det(mat)
	is_singular = dets.abs() < 1e-8
	mat[is_singular] = t.eye(3)

	# Define vector on the right hand side of equation
	vec = O - A

	# Solve eqns
	sol: Float[Tensor, "NR 3"] = t.linalg.solve(mat, vec)
	s, u, v = sol.unbind(dim=-1)

	# Return boolean of (matrix is nonsingular, and solution is in correct range implying intersection)
	return ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
//  [python]
def _maxvol(M):
    """
    Maxvol

    Args:
        M (torch.tensor): input matrix.

    Returns:
        torch.tensor: indices of tha maxvol submatrix.
    """
     
    if M.shape[1] >= M.shape[0]:
        # more cols than row -> return all the row indices 
        idx = tn.tensor(range(M.shape[0]),dtype = tn.int64)
        return idx
    else:
        L, U, P = _LU(M)
        idx = P[:M.shape[1]]
    
    Msub = M[idx,:]
   
    Mat = tn.solve(Msub.T,M.T).t()
   
    for i in range(100): 
        val_max, idx_max = _max_matrix(tn.abs(Mat)) 
        idx_max = idx_max[0]
        if val_max<=1+5e-2:
            idx = tn.sort(idx)[0]
            return idx
        Mat += tn.outer(Mat[:,idx_max[1]],Mat[idx[idx_max[1]]]-Mat[idx_max[0],:])/Mat[idx_max[0],idx_max[1]]
        idx[idx_max[1]]=idx_max[0]
    return idx
//  [python]
def _maxvol(M):
    """
    Maxvol

    Args:
        M (torch.tensor): input matrix.

    Returns:
        torch.tensor: indices of tha maxvol submatrix.
    """
     
    if M.shape[1] >= M.shape[0]:
        # more cols than row -> return all the row indices 
        idx = tn.tensor(range(M.shape[0]),dtype = tn.int64)
        return idx
    else:
        L, U, P = _LU(M)
        idx = P[:M.shape[1]]
    
    Msub = M[idx,:]
   
    Mat = tn.linalg.solve(Msub.T,M.T).t()
   
    for i in range(100): 
        val_max, idx_max = _max_matrix(tn.abs(Mat)) 
        idx_max = idx_max[0]
        if val_max<=1+5e-2:
            idx = tn.sort(idx)[0]
            return idx
        Mat += tn.outer(Mat[:,idx_max[1]],Mat[idx[idx_max[1]]]-Mat[idx_max[0],:])/Mat[idx_max[0],idx_max[1]]
        idx[idx_max[1]]=idx_max[0]
    return idx
//  [python]
def raytrace_mesh(
	rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
	triangles: Float[Tensor, "ntriangles trianglePoints=3 dims=3"]
) -> Float[Tensor, "nrays"]:
	'''
	For each ray, return the distance to the closest intersecting triangle, or infinity.
	'''
	NR = rays.size(0)
	NT = triangles.size(0)

	# Each triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
	triangles = einops.repeat(triangles, "NT pts dims -> pts NR NT dims", NR=NR)
	A, B, C = triangles
	assert A.shape == (NR, NT, 3)

	# Each ray is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
	rays = einops.repeat(rays, "NR pts dims -> pts NR NT dims", NT=NT)
	O, D = rays
	assert O.shape == (NR, NT, 3)

	# Define matrix on left hand side of equation
	mat: Float[Tensor, "NR NT 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
	# Get boolean of where matrix is singular, and replace it with the identity in these positions
	dets: Float[Tensor, "NR NT"] = t.linalg.det(mat)
	is_singular = dets.abs() < 1e-8
	mat[is_singular] = t.eye(3)

	# Define vector on the right hand side of equation
	vec: Float[Tensor, "NR NT 3"] = O - A

	# Solve eqns (note, s is the distance along ray)
	sol: Float[Tensor, "NR NT 3"] = t.solve(vec, mat)
	s, u, v = sol.unbind(-1)

	# Get boolean of intersects, and use it to set distance to infinity wherever there is no intersection
	intersects = ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
	s[~intersects] = float("inf") # t.inf

	# Get the minimum distance (over all triangles) for each ray
	return s.min(dim=-1).values
//  [python]
def raytrace_mesh(
	rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
	triangles: Float[Tensor, "ntriangles trianglePoints=3 dims=3"]
) -> Float[Tensor, "nrays"]:
	'''
	For each ray, return the distance to the closest intersecting triangle, or infinity.
	'''
	NR = rays.size(0)
	NT = triangles.size(0)

	# Each triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
	triangles = einops.repeat(triangles, "NT pts dims -> pts NR NT dims", NR=NR)
	A, B, C = triangles
	assert A.shape == (NR, NT, 3)

	# Each ray is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
	rays = einops.repeat(rays, "NR pts dims -> pts NR NT dims", NT=NT)
	O, D = rays
	assert O.shape == (NR, NT, 3)

	# Define matrix on left hand side of equation
	mat: Float[Tensor, "NR NT 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
	# Get boolean of where matrix is singular, and replace it with the identity in these positions
	dets: Float[Tensor, "NR NT"] = t.linalg.det(mat)
	is_singular = dets.abs() < 1e-8
	mat[is_singular] = t.eye(3)

	# Define vector on the right hand side of equation
	vec: Float[Tensor, "NR NT 3"] = O - A

	# Solve eqns (note, s is the distance along ray)
	sol: Float[Tensor, "NR NT 3"] = t.linalg.solve(mat, vec)
	s, u, v = sol.unbind(-1)

	# Get boolean of intersects, and use it to set distance to infinity wherever there is no intersection
	intersects = ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
	s[~intersects] = float("inf") # t.inf

	# Get the minimum distance (over all triangles) for each ray
	return s.min(dim=-1).values
//  [python]
def rmsd_pbc(batched_pos1, batched_pos2, cell, lig_mask, num_moveable):
    delta_pos = (batched_pos2.unsqueeze(2) - batched_pos1.unsqueeze(1)).float() # B x N1 x N2 x T x 3
    cell = cell.unsqueeze(1).unsqueeze(2).float() # B x 1 x 1 x 3 x 3
    delta_pos_solve = torch.solve(delta_pos.transpose(-1, -2), cell.transpose(-1, -2))[0].transpose(-1, -2)
    delta_pos_solve[:, :, :, :, 0] %= 1.0
    delta_pos_solve[:, :, :, :, 0] %= 1.0
    delta_pos_solve[:, :, :, :, 1] %= 1.0
    delta_pos_solve[:, :, :, :, 1] %= 1.0
    delta_pos_solve[delta_pos_solve > 0.5] -= 1.0
    min_delta_pos = torch.matmul(delta_pos_solve, cell)
    rmsds = torch.sqrt(torch.sum(torch.sum(min_delta_pos ** 2, dim=-1) * lig_mask, dim=-1) / num_moveable.unsqueeze(1).unsqueeze(2))
    return rmsds
//  [python]
def rmsd_pbc(batched_pos1, batched_pos2, cell, lig_mask, num_moveable):
    delta_pos = (batched_pos2.unsqueeze(2) - batched_pos1.unsqueeze(1)).float() # B x N1 x N2 x T x 3
    cell = cell.unsqueeze(1).unsqueeze(2).float() # B x 1 x 1 x 3 x 3
    delta_pos_solve = torch.linalg.solve(cell.transpose(-1, -2), delta_pos.transpose(-1, -2)).transpose(-1, -2)
    delta_pos_solve[:, :, :, :, 0] %= 1.0
    delta_pos_solve[:, :, :, :, 0] %= 1.0
    delta_pos_solve[:, :, :, :, 1] %= 1.0
    delta_pos_solve[:, :, :, :, 1] %= 1.0
    delta_pos_solve[delta_pos_solve > 0.5] -= 1.0
    min_delta_pos = torch.matmul(delta_pos_solve, cell)
    rmsds = torch.sqrt(torch.sum(torch.sum(min_delta_pos ** 2, dim=-1) * lig_mask, dim=-1) / num_moveable.unsqueeze(1).unsqueeze(2))
    return rmsds
//  [python]
def log_likelihood(self, x: torch.Tensor, u: torch.Tensor, 
                       extra_info: Dict[str, torch.Tensor], 
                       theta=None, 
                       mu_r=0, lam_r=0):
        intermediate_poses = extra_info['intermediate_poses']
        if theta is None:
            theta = torch.cat([torch.tensor(self.gt_theta), torch.tensor([-np.inf])])
        
        batch_size = x.shape[0]

        ll_val = 0.0
        for batch in range(batch_size):
            g_hat, g_tilde = AF.jacobian(partial(self.gt_reward_fn, 
                                                 theta=theta[:-1], 
                                                 extra_info={'intermediate_poses': intermediate_poses[[batch]]}), 
                                         (x[batch], u[batch]),
                                         create_graph=True)

            g = g_tilde + self.J @ g_hat
            H = self.hessian(x[batch], u[batch], theta, {'intermediate_poses': intermediate_poses[[batch]]})

            ll_val += 0.5 * g @ torch.solve(g, H)[0] + 0.5 * torch.logdet(-H)
        
        if np.isinf(mu_r):
            # We've switched off the constrained variable (auglagcost.m:17)
            return ll_val/batch_size
        else:
            return ll_val/batch_size - 0.5*mu_r*(torch.exp(theta[-1])**2) + lam_r*torch.exp(theta[-1])
//  [python]
def log_likelihood(self, x: torch.Tensor, u: torch.Tensor, 
                       extra_info: Dict[str, torch.Tensor], 
                       theta=None, 
                       mu_r=0, lam_r=0):
        intermediate_poses = extra_info['intermediate_poses']
        if theta is None:
            theta = torch.cat([torch.tensor(self.gt_theta), torch.tensor([-np.inf])])
        
        batch_size = x.shape[0]

        ll_val = 0.0
        for batch in range(batch_size):
            g_hat, g_tilde = AF.jacobian(partial(self.gt_reward_fn, 
                                                 theta=theta[:-1], 
                                                 extra_info={'intermediate_poses': intermediate_poses[[batch]]}), 
                                         (x[batch], u[batch]),
                                         create_graph=True)

            g = g_tilde + self.J @ g_hat
            H = self.hessian(x[batch], u[batch], theta, {'intermediate_poses': intermediate_poses[[batch]]})

            ll_val += 0.5 * g @ torch.linalg.solve(H, g) + 0.5 * torch.logdet(-H)
        
        if np.isinf(mu_r):
            # We've switched off the constrained variable (auglagcost.m:17)
            return ll_val/batch_size
        else:
            return ll_val/batch_size - 0.5*mu_r*(torch.exp(theta[-1])**2) + lam_r*torch.exp(theta[-1])
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")

        self.RTOL = 1e-3

        self.A_shape = (4, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A = self.A + self.A.t() + 10.0 * torch.eye(4, device=self.device)
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (4, 2)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.solve(self.B, self.A)[0]
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")

        self.RTOL = 1e-3

        self.A_shape = (4, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A = self.A + self.A.t() + 10.0 * torch.eye(4, device=self.device)
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (4, 2)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.linalg.solve(self.A, self.B)
//  [python]
def test_on_quadratic(seed, dim, device):
    """This function sets up and runs the `HessianFree` optimizer on a quadratic
    function. In this particular case, it has to converge in a single Newton
    step.

    NOTE: This test case is resticted to lower dimensions because for
    `dim > 10`, Martens' convergence criterion is triggered in `cg` which
    prevents `cg` from running until actual convergence.
    """

    msg = f"seed={seed}, dim={dim}, device={device}"
    print("\n===== TEST `HessianFree` on a quadratic =====\n" + msg)

    # Create test problem
    torch.manual_seed(seed)
    init_params = (torch.rand((dim, 1)) - 0.5).to(device)

    A, b, _ = get_linear_system(dim, seed=seed, device=device)
    assert torch.all(torch.linalg.eigvalsh(A) > 0), "Matrix A is not pos. def."
    b = b.reshape(dim, 1)
    c = (torch.rand(1) - 0.5).to(device)

    def quadratic(x):
        return 0.5 * x.T @ A @ x + x.T @ b + c

    model = TargetFuncModel(quadratic, init_params)
    opt_params = torch.solve(b, A)[0]
    print("\nopt_params = ", opt_params.T)

    def forward():
        return model.eval_loss(), None

    # Set up optimizer
    opt = HessianFree(
        [model.params],
        curvature_opt="hessian",  # use the Hessian
        lr=1.0,
        use_linesearch=False,  # fixed lerning rate
        damping=0.0,  # no damping
        adapt_damping=False,
        use_cg_backtracking=False,  # no cg-backtracking
        verbose=True,
    )

    # Training
    def eval_dist_to_opt():
        return torch.linalg.norm(model.params - opt_params).detach().item()

    init_dist = eval_dist_to_opt()
    print(f"\nInitial distance to optimum = {init_dist}")

    opt.step(forward=forward)

    final_dist = eval_dist_to_opt()
    print(f"\nFinal distance to optimum = {final_dist}")
    assert torch.allclose(model.params, opt_params, atol=1e-3)
//  [python]
def test_on_quadratic(seed, dim, device):
    """This function sets up and runs the `HessianFree` optimizer on a quadratic
    function. In this particular case, it has to converge in a single Newton
    step.

    NOTE: This test case is resticted to lower dimensions because for
    `dim > 10`, Martens' convergence criterion is triggered in `cg` which
    prevents `cg` from running until actual convergence.
    """

    msg = f"seed={seed}, dim={dim}, device={device}"
    print("\n===== TEST `HessianFree` on a quadratic =====\n" + msg)

    # Create test problem
    torch.manual_seed(seed)
    init_params = (torch.rand((dim, 1)) - 0.5).to(device)

    A, b, _ = get_linear_system(dim, seed=seed, device=device)
    assert torch.all(torch.linalg.eigvalsh(A) > 0), "Matrix A is not pos. def."
    b = b.reshape(dim, 1)
    c = (torch.rand(1) - 0.5).to(device)

    def quadratic(x):
        return 0.5 * x.T @ A @ x + x.T @ b + c

    model = TargetFuncModel(quadratic, init_params)
    opt_params = torch.linalg.solve(A, -b)
    print("\nopt_params = ", opt_params.T)

    def forward():
        return model.eval_loss(), None

    # Set up optimizer
    opt = HessianFree(
        [model.params],
        curvature_opt="hessian",  # use the Hessian
        lr=1.0,
        use_linesearch=False,  # fixed lerning rate
        damping=0.0,  # no damping
        adapt_damping=False,
        use_cg_backtracking=False,  # no cg-backtracking
        verbose=True,
    )

    # Training
    def eval_dist_to_opt():
        return torch.linalg.norm(model.params - opt_params).detach().item()

    init_dist = eval_dist_to_opt()
    print(f"\nInitial distance to optimum = {init_dist}")

    opt.step(forward=forward)

    final_dist = eval_dist_to_opt()
    print(f"\nFinal distance to optimum = {final_dist}")
    assert torch.allclose(model.params, opt_params, atol=1e-3)
//  [python]
def test_generic_solver_gradient_default(self):
        # Sparse solver:
        As1 = self.A_csr.detach().clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = sparse_generic_solve(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2)[0]
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_generic_solver_gradient_default(self):
        # Sparse solver:
        As1 = self.A_csr.detach().clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = sparse_generic_solve(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve(Ad2, Bd2)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_generic_solver_gradient_cg(self):
        # Sparse solver:
        As1 = self.A_csr.detach().clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = sparse_generic_solve(As1, Bd1, solve=linear_cg, transpose_solve=linear_cg)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2)[0]
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_generic_solver_gradient_cg(self):
        # Sparse solver:
        As1 = self.A_csr.detach().clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = sparse_generic_solve(As1, Bd1, solve=linear_cg, transpose_solve=linear_cg)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve(Ad2, Bd2)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def training_step(self, batch):
        ratings = (batch[self.frating] > 0).float()
        if batch[self.fuid].dim() == 1:  # user model, updating user embedding
            item_embed = self.item_encoder(self._get_item_feat(batch))  # B x N x D
            QuQ = torch.bmm(item_embed.transpose(1, 2),
                            item_embed) * self.config['train']['alpha'] + (self.QtQ + self.eye)  # BxDxD
            r = torch.bmm(item_embed.transpose(1, 2), ratings.unsqueeze(-1)).squeeze(-1)  # BxD
            output = torch.solve(r * (self.config['train']['alpha'] + 1), QuQ)[0]
            if self.config['model']['item_bias']:
                output[:, -1] = 1.0
            self.query_encoder.weight[batch[self.fuid]] = output  # B x D
            user_embed = self.query_encoder(self._get_query_feat(batch))  # BxD
            pred = self.score_func(user_embed, item_embed)  # BxD   BxNxD -> BxN
            reg1 = torch.multiply(user_embed @ self.QtQ,  user_embed).sum(-1)
        else:
            user_embed = self.query_encoder(self._get_query_feat(batch))
            PiP = torch.bmm(user_embed.transpose(1, 2), user_embed) * self.config['train']['alpha'] + (self.PtP + self.eye)
            r = torch.bmm(user_embed.transpose(1, 2), ratings.unsqueeze(-1)).squeeze(-1)
            output = torch.solve(r * (self.config['train']['alpha'] + 1), PiP)[0]
            self.item_encoder.weight[batch[self.fiid]] = output
            item_embed = self.item_encoder(self._get_item_feat(batch))
            pred = self.score_func(item_embed, user_embed)  # BxD BxNxD  -> BxN
            reg1 = torch.multiply(item_embed @ self.PtP, item_embed).sum(-1)
        loss = torch.sum((ratings - pred) ** 2, dim=-1) * (self.config['train']['alpha'] + 1)
        loss -= (pred**2).sum(-1)
        loss += reg1

        return {'loss': loss}
//  [python]
def training_step(self, batch):
        ratings = (batch[self.frating] > 0).float()
        if batch[self.fuid].dim() == 1:  # user model, updating user embedding
            item_embed = self.item_encoder(self._get_item_feat(batch))  # B x N x D
            QuQ = torch.bmm(item_embed.transpose(1, 2),
                            item_embed) * self.config['train']['alpha'] + (self.QtQ + self.eye)  # BxDxD
            r = torch.bmm(item_embed.transpose(1, 2), ratings.unsqueeze(-1)).squeeze(-1)  # BxD
            output = torch.linalg.solve(QuQ, r * (self.config['train']['alpha'] + 1))
            if self.config['model']['item_bias']:
                output[:, -1] = 1.0
            self.query_encoder.weight[batch[self.fuid]] = output  # B x D
            user_embed = self.query_encoder(self._get_query_feat(batch))  # BxD
            pred = self.score_func(user_embed, item_embed)  # BxD   BxNxD -> BxN
            reg1 = torch.multiply(user_embed @ self.QtQ,  user_embed).sum(-1)
        else:
            user_embed = self.query_encoder(self._get_query_feat(batch))
            PiP = torch.bmm(user_embed.transpose(1, 2), user_embed) * self.config['train']['alpha'] + (self.PtP + self.eye)
            r = torch.bmm(user_embed.transpose(1, 2), ratings.unsqueeze(-1)).squeeze(-1)
            output = torch.linalg.solve(PiP, r * (self.config['train']['alpha'] + 1))
            self.item_encoder.weight[batch[self.fiid]] = output
            item_embed = self.item_encoder(self._get_item_feat(batch))
            pred = self.score_func(item_embed, user_embed)  # BxD BxNxD  -> BxN
            reg1 = torch.multiply(item_embed @ self.PtP, item_embed).sum(-1)
        loss = torch.sum((ratings - pred) ** 2, dim=-1) * (self.config['train']['alpha'] + 1)
        loss -= (pred**2).sum(-1)
        loss += reg1

        return {'loss': loss}