jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def spherical_to_camera(theta, azimuth, z):
    # Calculate the camera's position in Cartesian coordinates
    X = z * torch.cos(theta) * torch.cos(azimuth)
    Y = z * torch.cos(theta) * torch.sin(azimuth)
    Z = z * torch.sin(theta)

    camera_T_c2w = torch.cat([X, Y, Z], dim=-1)
    # Calculate the camera's look-at direction vector
    camera_z = - camera_T_c2w
    camera_z /= torch.norm(camera_z, dim=-1, keepdim=True)
    # Define the camera's up direction vector
    up_direction = torch.zeros((X.size(0), 3)).to(z.device)
    up_direction[:, 2] = - 1.0
    camera_x = torch.cross(up_direction, camera_z, dim=-1)
    camera_x[camera_x.sum(dim=-1) == 0] = torch.tensor([0, 1, 0]).type_as(camera_z)
    camera_x /= torch.norm(camera_x, dim=-1, keepdim=True)
    camera_y = torch.cross(camera_z, camera_x, dim=-1)
    camera_y =  camera_y / torch.norm( camera_y, dim=-1, keepdim=True)
    # Construct the translation vector as the negative of the camera position
    C2W = torch.eye(4).unsqueeze(0).repeat(X.size(0),1,1).to(z.device)
    C2W[:, :3, 0] = camera_x
    C2W[:, :3, 1] = camera_y
    C2W[:, :3, 2] = camera_z

    C2W[:, :3, 3] = camera_T_c2w

    # world to cam
    W2C = torch.inverse(C2W)

    return W2C
//  [python]
def spherical_to_camera(theta, azimuth, z):
    # Calculate the camera's position in Cartesian coordinates
    X = z * torch.cos(theta) * torch.cos(azimuth)
    Y = z * torch.cos(theta) * torch.sin(azimuth)
    Z = z * torch.sin(theta)

    camera_T_c2w = torch.cat([X, Y, Z], dim=-1)
    # Calculate the camera's look-at direction vector
    camera_z = - camera_T_c2w
    camera_z /= torch.linalg.norm(camera_z, dim=-1, keepdim=True)
    # Define the camera's up direction vector
    up_direction = torch.zeros((X.size(0), 3)).to(z.device)
    up_direction[:, 2] = - 1.0
    camera_x = torch.cross(up_direction, camera_z, dim=-1)
    camera_x[camera_x.sum(dim=-1) == 0] = torch.tensor([0, 1, 0]).type_as(camera_z)
    camera_x /= torch.linalg.norm(camera_x, dim=-1, keepdim=True)
    camera_y = torch.cross(camera_z, camera_x, dim=-1)
    camera_y =  camera_y / torch.linalg.norm( camera_y, dim=-1, keepdim=True)
    # Construct the translation vector as the negative of the camera position
    C2W = torch.eye(4).unsqueeze(0).repeat(X.size(0),1,1).to(z.device)
    C2W[:, :3, 0] = camera_x
    C2W[:, :3, 1] = camera_y
    C2W[:, :3, 2] = camera_z

    C2W[:, :3, 3] = camera_T_c2w

    # world to cam
    W2C = torch.linalg.inv(C2W)

    return W2C
//  [python]
def backward(ctx, y_grad):
        if not ctx.needs_input_grad[0]:
            return None, None, None, None, None

        x, y = ctx.saved_tensors
        B, C = x.shape[0], x.shape[1]

        x_minus_y = y.view(B, C, 1) - x.view(B, C, -1)
        z = torch.norm(x_minus_y, dim=1, keepdim=True) + 1.0e-9

        k1, k2 = ctx.penalty.kappa(z, ctx.alpha)
        if torch.all(k2 == 0.0):
            return (k1 * (y_grad / k1.sum(dim=2)).view(B, C, 1)).reshape(x.shape), None, None, None, None

        H = k1.sum(dim=2).view(B, 1, 1) * torch.eye(C, dtype=x.dtype, device=x.device).view(1, C, C) + \
            torch.einsum("bik,bjk->bij", x_minus_y, k2 * x_minus_y)
        try:
            L = torch.cholesky(H + ctx.hess_reg * torch.eye(C, dtype=x.dtype, device=x.device).view(1, C, C))
            v = torch.cholesky_solve(y_grad.view(B, C, -1), L).view(B, C)
        except:
            warnings.warn("backward pass encountered a singular matrix for penalty function {}".format(ctx.penalty.__name__))
            v = torch.empty_like(y_grad)
            for b in range(B):
                try:
                    L = torch.cholesky(H[b, :, :])
                    v[b, :] = torch.cholesky_solve(y_grad[b, :].view(C, 1), L).view(1, C)
                except:
                    v[b, :] = torch.lstsq(y_grad[b, :].view(C, 1), H[b, :, :])[0].view(1, C)


        w = torch.einsum("bi,bik->bk", v, k2 * x_minus_y)
        x_grad = (k1 * v.view(B, C, 1) + torch.einsum("bk,bik->bik", w, x_minus_y)).reshape(x.shape)

        return x_grad, None, None, None, None
//  [python]
def backward(ctx, y_grad):
        if not ctx.needs_input_grad[0]:
            return None, None, None, None, None

        x, y = ctx.saved_tensors
        B, C = x.shape[0], x.shape[1]

        x_minus_y = y.view(B, C, 1) - x.view(B, C, -1)
        z = torch.linalg.norm(x_minus_y, dim=1, keepdim=True) + 1.0e-9

        k1, k2 = ctx.penalty.kappa(z, ctx.alpha)
        if torch.all(k2 == 0.0):
            return (k1 * (y_grad / k1.sum(dim=2)).view(B, C, 1)).reshape(x.shape), None, None, None, None

        H = k1.sum(dim=2).view(B, 1, 1) * torch.eye(C, dtype=x.dtype, device=x.device).view(1, C, C) + \
            torch.einsum("bik,bjk->bij", x_minus_y, k2 * x_minus_y)
        try:
            L = torch.cholesky(H + ctx.hess_reg * torch.eye(C, dtype=x.dtype, device=x.device).view(1, C, C))
            v = torch.cholesky_solve(y_grad.view(B, C, -1), L).view(B, C)
        except:
            warnings.warn("backward pass encountered a singular matrix for penalty function {}".format(ctx.penalty.__name__))
            v = torch.empty_like(y_grad)
            for b in range(B):
                try:
                    L = torch.cholesky(H[b, :, :])
                    v[b, :] = torch.cholesky_solve(y_grad[b, :].view(C, 1), L).view(1, C)
                except:
                    v[b, :] = torch.lstsq(y_grad[b, :].view(C, 1), H[b, :, :])[0].view(1, C)


        w = torch.einsum("bi,bik->bk", v, k2 * x_minus_y)
        x_grad = (k1 * v.view(B, C, 1) + torch.einsum("bk,bik->bik", w, x_minus_y)).reshape(x.shape)

        return x_grad, None, None, None, None
//  [python]
def test_norms():

    v1 = variable(torch.zeros(100, requires_grad=True))
    v2 = variable(torch.zeros(100, requires_grad=True))
    assert torch.equal(torch.norm(v1 - v2)(), torch.tensor(0.0))
//  [python]
def test_norms():

    v1 = variable(torch.zeros(100, requires_grad=True))
    v2 = variable(torch.zeros(100, requires_grad=True))
    assert torch.equal(torch.linalg.norm(v1 - v2)(), torch.tensor(0.0))
//  [python]
def get_weights(A, B, weight):
    # get weights
    if weight == "L2":
        weights_A = torch.norm(A, dim=2, p=2) 
        weights_B = torch.norm(B, dim=2, p=2) 
    elif weight == "L1":
        weights_A = torch.norm(A, dim=2, p=1) 
        weights_B = torch.norm(B, dim=2, p=1) 
    elif weight == "no":
        weights_A = torch.ones(A.size(0), A.size(1)).to(A.device)
        weights_B = torch.ones(B.size(0), B.size(1)).to(B.device)
    else:
        raise NotImplementedError
    return weights_A, weights_B
//  [python]
def get_weights(A, B, weight):
    # get weights
    if weight == "L2":
        weights_A = torch.linalg.norm(A, dim=2) 
        weights_B = torch.linalg.norm(B, dim=2) 
    elif weight == "L1":
        weights_A = torch.linalg.norm(A, dim=2, ord=1) 
        weights_B = torch.linalg.norm(B, dim=2, ord=1)
    elif weight == "no":
        weights_A = torch.ones(A.size(0), A.size(1)).to(A.device)
        weights_B = torch.ones(B.size(0), B.size(1)).to(B.device)
    else:
        raise NotImplementedError
    return weights_A, weights_B
//  [python]
for w in ws:
    _, dim = w.shape
    w = w.float().to(device)
    A, B, r, _ = svd_approx(w, r)
    As.append(A)
    Bs.append(B)


Wp = torch.zeros_like(ws[0])

for i in range(experts):
    Wd = As[i] @ Bs[i]
    Wp += ws[i] - Wd
Wp = Wp / experts

for idx, W, A, B in enumerate(zip(ws, As, Bs)):
    Wapx = Wp + A @ B
    print(idx, torch.norm(W - Wapx))
//  [python]
for w in ws:
    _, dim = w.shape
    w = w.float().to(device)
    A, B, r, _ = svd_approx(w, r)
    As.append(A)
    Bs.append(B)


Wp = torch.zeros_like(ws[0])

for i in range(experts):
    Wd = As[i] @ Bs[i]
    Wp += ws[i] - Wd
Wp = Wp / experts

for idx, W, A, B in enumerate(zip(ws, As, Bs)):
    Wapx = Wp + A @ B
    print(idx, torch.linalg.norm(W - Wapx))
//  [python]
def grad_and_value(operator, x, x_hat, measurement):
    difference = measurement - operator(x_hat)
    norm = torch.norm(difference)
    norm_grad = torch.autograd.grad(outputs=norm, inputs=x)[0]
    return norm_grad,  norm
//  [python]
def grad_and_value(operator, x, x_hat, measurement):
    difference = measurement - operator(x_hat)
    norm = torch.linalg.norm(difference)
    norm_grad = torch.autograd.grad(outputs=norm, inputs=x)[0]
    return norm_grad,  norm
//  [python]
def solve_inverse_kinematics(
        self,
        position: torch.Tensor,
        orientation: torch.Tensor,
        q0: torch.Tensor,
        tol: float = 1e-3,
    ) -> Tuple[torch.Tensor, bool]:
        """Compute inverse kinematics given desired EE pose"""
        # Call IK
        joint_pos_output = self.robot_model.inverse_kinematics(
            position, orientation, rest_pose=q0
        )

        # Check result
        pos_output, quat_output = self.robot_model.forward_kinematics(joint_pos_output)
        pose_desired = T.from_rot_xyz(R.from_quat(orientation), position)
        pose_output = T.from_rot_xyz(R.from_quat(quat_output), pos_output)
        err = torch.norm((pose_desired * pose_output.inv()).as_twist())
        ik_sol_found = err < tol

        return joint_pos_output, ik_sol_found
//  [python]
def solve_inverse_kinematics(
        self,
        position: torch.Tensor,
        orientation: torch.Tensor,
        q0: torch.Tensor,
        tol: float = 1e-3,
    ) -> Tuple[torch.Tensor, bool]:
        """Compute inverse kinematics given desired EE pose"""
        # Call IK
        joint_pos_output = self.robot_model.inverse_kinematics(
            position, orientation, rest_pose=q0
        )

        # Check result
        pos_output, quat_output = self.robot_model.forward_kinematics(joint_pos_output)
        pose_desired = T.from_rot_xyz(R.from_quat(orientation), position)
        pose_output = T.from_rot_xyz(R.from_quat(quat_output), pos_output)
        err = torch.linalg.norm((pose_desired * pose_output.inv()).as_twist())
        ik_sol_found = err < tol

        return joint_pos_output, ik_sol_found
//  [python]
def forward(self, input):
        eps = self.eps

        norm = torch.norm(input, dim=2, keepdim=True)
        x = input / (norm + eps)
        x_U = self.conv1d_U(x)
        x_V = self.conv1d_V(x)
        output = self.relu(x_U) * self.sigmoid(x_V)

        return output
//  [python]
def forward(self, input):
        eps = self.eps

        norm = torch.linalg.norm(input, dim=2, keepdim=True)
        x = input / (norm + eps)
        x_U = self.conv1d_U(x)
        x_V = self.conv1d_V(x)
        output = self.relu(x_U) * self.sigmoid(x_V)

        return output
//  [python]
def omp(X, D, s):
    """
    X: data matrix of size (n_samples, n_features)
    D: dictionary of size (n_atoms, n_features)
    s: desired sparsity level
    Returns the sparse code matrix S of size (n_samples, n_atoms)
    """
    n_samples, n_features = X.shape
    n_atoms, _ = D.shape

    # Normalize rows
    norms = torch.norm(D, dim=1) + 1e-6
    D2 = D / norms[:, None]

    residuals = X.clone()
    ids = torch.zeros((n_samples, s), dtype=int, device=X.device)

    for step in range(s):
        corrs = torch.abs(residuals @ D2.T)
        corrs.scatter_(1, ids[:, :step], -1)  # Mask ids we've already picked
        ids[:, step] = torch.argmax(corrs, axis=1)
        subDs = D2[ids[:, : step + 1]].mT
        m = torch.linalg.lstsq(subDs, X, rcond=None).solution
        residuals = X - (subDs @ m.unsqueeze(2)).squeeze(2)

    m /= norms[ids]
    return ids, m
//  [python]
def omp(X, D, s):
    """
    X: data matrix of size (n_samples, n_features)
    D: dictionary of size (n_atoms, n_features)
    s: desired sparsity level
    Returns the sparse code matrix S of size (n_samples, n_atoms)
    """
    n_samples, n_features = X.shape
    n_atoms, _ = D.shape

    # Normalize rows
    norms = torch.linalg.norm(D, axis=1) + 1e-6
    D2 = D / norms[:, None]

    residuals = X.clone()
    ids = torch.zeros((n_samples, s), dtype=int, device=X.device)

    for step in range(s):
        corrs = torch.abs(residuals @ D2.T)
        corrs.scatter_(1, ids[:, :step], -1)  # Mask ids we've already picked
        ids[:, step] = torch.argmax(corrs, axis=1)
        subDs = D2[ids[:, : step + 1]].mT
        m = torch.linalg.lstsq(subDs, X, rcond=None).solution
        residuals = X - (subDs @ m.unsqueeze(2)).squeeze(2)

    m /= norms[ids]
    return ids, m
//  [python]
def test_trajopt_single_js(trajopt_solver):
    q_start = trajopt_solver.retract_config
    q_goal = q_start.clone() + 0.2
    goal_state = JointState.from_position(q_goal)
    current_state = JointState.from_position(q_start)
    # do single planning:
    js_goal = Goal(goal_state=goal_state, current_state=current_state)
    result = trajopt_solver.solve_single(js_goal)

    traj = result.solution.position[..., -1, :].view(q_goal.shape)
    assert torch.norm((goal_state.position - traj)).item() < 5e-3
    assert result.success.item()
//  [python]
def test_trajopt_single_js(trajopt_solver):
    q_start = trajopt_solver.retract_config
    q_goal = q_start.clone() + 0.2
    goal_state = JointState.from_position(q_goal)
    current_state = JointState.from_position(q_start)
    # do single planning:
    js_goal = Goal(goal_state=goal_state, current_state=current_state)
    result = trajopt_solver.solve_single(js_goal)

    traj = result.solution.position[..., -1, :].view(q_goal.shape)
    assert torch.linalg.norm((goal_state.position - traj)).item() < 5e-3
    assert result.success.item()
//  [python]
def test_trajopt_batch_js(trajopt_solver):
    # run goalset planning:
    q_start = trajopt_solver.retract_config.repeat(2, 1)
    q_goal = q_start.clone()
    q_goal[0] += 0.1
    q_goal[1] -= 0.1
    kin_state = trajopt_solver.fk(q_goal)
    # goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
    goal_state = JointState.from_position(q_goal)
    current_state = JointState.from_position(q_start)

    js_goal = Goal(goal_state=goal_state, current_state=current_state)
    result = trajopt_solver.solve_batch(js_goal)
    traj = result.solution.position
    interpolated_traj = result.interpolated_solution.position
    assert torch.count_nonzero(result.success) > 0
    assert torch.norm((goal_state.position - traj[:, -1, :])).item() < 0.05
    assert torch.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.05
//  [python]
def test_trajopt_batch_js(trajopt_solver):
    # run goalset planning:
    q_start = trajopt_solver.retract_config.repeat(2, 1)
    q_goal = q_start.clone()
    q_goal[0] += 0.1
    q_goal[1] -= 0.1
    kin_state = trajopt_solver.fk(q_goal)
    # goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
    goal_state = JointState.from_position(q_goal)
    current_state = JointState.from_position(q_start)

    js_goal = Goal(goal_state=goal_state, current_state=current_state)
    result = trajopt_solver.solve_batch(js_goal)
    traj = result.solution.position
    interpolated_traj = result.interpolated_solution.position
    assert torch.count_nonzero(result.success) > 0
    assert torch.linalg.norm((goal_state.position - traj[:, -1, :])).item() < 0.05
    assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.05
//  [python]
def grad_penalty(self, obs: torch.Tensor, action: torch.Tensor, ebm: nn.Module) -> torch.Tensor:
        """
        Overview:
            Calculate gradient penalty.
        Arguments:
            - obs (:obj:`torch.Tensor`): Observations.
            - action (:obj:`torch.Tensor`): Actions.
            - ebm (:obj:`torch.nn.Module`): Energy based model.
        Returns:
            - loss (:obj:`torch.Tensor`): Gradient penalty.
        Shapes:
            - obs (:obj:`torch.Tensor`): :math:`(B, N+1, O)`.
            - action (:obj:`torch.Tensor`): :math:`(B, N+1, A)`.
            - ebm (:obj:`torch.nn.Module`): :math:`(B, N+1, O)`.
            - loss (:obj:`torch.Tensor`): :math:`(B, )`.
        """
        if not self.add_grad_penalty:
            return 0.
        # (B, N+1, A), this gradient is differentiable w.r.t model parameters
        de_dact = MCMC._gradient_wrt_act(obs, action, ebm, create_graph=True)

        def compute_grad_norm(grad_norm_type, de_dact) -> torch.Tensor:
            # de_deact: B, N+1, A
            # return:   B, N+1
            grad_norm_type_to_ord = {
                '1': 1,
                '2': 2,
                'inf': float('inf'),
            }
            ord = grad_norm_type_to_ord[grad_norm_type]
            return torch.norm(de_dact, ord, dim=-1)

        # (B, N+1)
        grad_norms = compute_grad_norm(self.grad_norm_type, de_dact)
        grad_norms = grad_norms - self.grad_margin
        grad_norms = grad_norms.clamp(min=0., max=1e10)
        grad_norms = grad_norms.pow(2)

        grad_loss = grad_norms.mean()
        return grad_loss * self.grad_loss_weight
//  [python]
def grad_penalty(self, obs: torch.Tensor, action: torch.Tensor, ebm: nn.Module) -> torch.Tensor:
        """
        Overview:
            Calculate gradient penalty.
        Arguments:
            - obs (:obj:`torch.Tensor`): Observations.
            - action (:obj:`torch.Tensor`): Actions.
            - ebm (:obj:`torch.nn.Module`): Energy based model.
        Returns:
            - loss (:obj:`torch.Tensor`): Gradient penalty.
        Shapes:
            - obs (:obj:`torch.Tensor`): :math:`(B, N+1, O)`.
            - action (:obj:`torch.Tensor`): :math:`(B, N+1, A)`.
            - ebm (:obj:`torch.nn.Module`): :math:`(B, N+1, O)`.
            - loss (:obj:`torch.Tensor`): :math:`(B, )`.
        """
        if not self.add_grad_penalty:
            return 0.
        # (B, N+1, A), this gradient is differentiable w.r.t model parameters
        de_dact = MCMC._gradient_wrt_act(obs, action, ebm, create_graph=True)

        def compute_grad_norm(grad_norm_type, de_dact) -> torch.Tensor:
            # de_deact: B, N+1, A
            # return:   B, N+1
            grad_norm_type_to_ord = {
                '1': 1,
                '2': 2,
                'inf': float('inf'),
            }
            ord = grad_norm_type_to_ord[grad_norm_type]
            return torch.linalg.norm(de_dact, ord, dim=-1)

        # (B, N+1)
        grad_norms = compute_grad_norm(self.grad_norm_type, de_dact)
        grad_norms = grad_norms - self.grad_margin
        grad_norms = grad_norms.clamp(min=0., max=1e10)
        grad_norms = grad_norms.pow(2)

        grad_loss = grad_norms.mean()
        return grad_loss * self.grad_loss_weight
//  [python]
def test_computes_downsampling_function(self):
        a = torch.tensor([0.1, 0.2]).view(2, 1)
        b = torch.tensor([0.2, 0.4]).view(2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel()
        kernel.initialize(power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[0.72, 0.54], [0.64, 0.48]])
        actual = offset + diff.pow(1 + power)
        res = kernel(a, b).to_dense()

        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_computes_downsampling_function(self):
        a = torch.tensor([0.1, 0.2]).view(2, 1)
        b = torch.tensor([0.2, 0.4]).view(2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel()
        kernel.initialize(power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[0.72, 0.54], [0.64, 0.48]])
        actual = offset + diff.pow(1 + power)
        res = kernel(a, b).to_dense()

        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def test_trajopt_batch_env_js(trajopt_solver_batch_env):
    # run goalset planning:
    q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
    q_goal = q_start.clone()
    q_goal += 0.1
    q_goal[2] += 0.1
    q_goal[1] -= 0.2
    # q_goal[2, -1] += 0.1
    kin_state = trajopt_solver_batch_env.fk(q_goal)
    goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
    goal_state = JointState.from_position(q_goal)
    current_state = JointState.from_position(q_start)

    js_goal = Goal(goal_state=goal_state, current_state=current_state)
    result = trajopt_solver_batch_env.solve_batch_env(js_goal)

    traj = result.solution.position
    interpolated_traj = result.interpolated_solution.position
    assert torch.count_nonzero(result.success) == 3
    assert torch.norm((goal_state.position - traj[:, -1, :])).item() < 0.005
    assert torch.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.005
    assert len(result) == 3
//  [python]
def test_trajopt_batch_env_js(trajopt_solver_batch_env):
    # run goalset planning:
    q_start = trajopt_solver_batch_env.retract_config.repeat(3, 1)
    q_goal = q_start.clone()
    q_goal += 0.1
    q_goal[2] += 0.1
    q_goal[1] -= 0.2
    # q_goal[2, -1] += 0.1
    kin_state = trajopt_solver_batch_env.fk(q_goal)
    goal_pose = Pose(kin_state.ee_position, kin_state.ee_quaternion)
    goal_state = JointState.from_position(q_goal)
    current_state = JointState.from_position(q_start)

    js_goal = Goal(goal_state=goal_state, current_state=current_state)
    result = trajopt_solver_batch_env.solve_batch_env(js_goal)

    traj = result.solution.position
    interpolated_traj = result.interpolated_solution.position
    assert torch.count_nonzero(result.success) == 3
    assert torch.linalg.norm((goal_state.position - traj[:, -1, :])).item() < 0.005
    assert torch.linalg.norm((goal_state.position - interpolated_traj[:, -1, :])).item() < 0.005
    assert len(result) == 3
//  [python]
def test_subset_active_compute_downsampling_function(self):
        a = torch.tensor([0.1, 0.2]).view(2, 1)
        a_p = torch.tensor([0.3, 0.4]).view(2, 1)
        a = torch.cat((a, a_p), 1)
        b = torch.tensor([0.2, 0.4]).view(2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel(active_dims=[0])
        kernel.initialize(power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[0.72, 0.54], [0.64, 0.48]])
        actual = offset + diff.pow(1 + power)
        res = kernel(a, b).to_dense()

        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_subset_active_compute_downsampling_function(self):
        a = torch.tensor([0.1, 0.2]).view(2, 1)
        a_p = torch.tensor([0.3, 0.4]).view(2, 1)
        a = torch.cat((a, a_p), 1)
        b = torch.tensor([0.2, 0.4]).view(2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel(active_dims=[0])
        kernel.initialize(power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[0.72, 0.54], [0.64, 0.48]])
        actual = offset + diff.pow(1 + power)
        res = kernel(a, b).to_dense()

        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def test_initialize_offset(self):
        kernel = DownsamplingKernel()
        kernel.initialize(offset=1)
        actual_value = torch.tensor(1.0).view_as(kernel.offset)
        self.assertLess(torch.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_offset(self):
        kernel = DownsamplingKernel()
        kernel.initialize(offset=1)
        actual_value = torch.tensor(1.0).view_as(kernel.offset)
        self.assertLess(torch.linalg.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_offset_batch(self):
        kernel = DownsamplingKernel(batch_shape=torch.Size([2]))
        off_init = torch.tensor([1.0, 2.0])
        kernel.initialize(offset=off_init)
        actual_value = off_init.view_as(kernel.offset)
        self.assertLess(torch.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_offset_batch(self):
        kernel = DownsamplingKernel(batch_shape=torch.Size([2]))
        off_init = torch.tensor([1.0, 2.0])
        kernel.initialize(offset=off_init)
        actual_value = off_init.view_as(kernel.offset)
        self.assertLess(torch.linalg.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def get_rays_torch(focal, c2w, H=64,W=64):
    """Computes rays using a General Pinhole Camera Model
    Assumes self.h, self.w, self.focal, and self.cam_to_world exist
    """
    x, y = torch.meshgrid(
        torch.arange(W),  # X-Axis (columns)
        torch.arange(H),  # Y-Axis (rows)
        indexing='xy')
    camera_directions = torch.stack(
        [(x - W * 0.5 + 0.5) / focal,
            -(y - H * 0.5 + 0.5) / focal,
            -torch.ones_like(x)],
        dim=-1).to(c2w)

    # Rotate ray directions from camera frame to the world frame
    directions = ((camera_directions[ None,..., None, :] * c2w[None,None, None, :3, :3]).sum(axis=-1))  # Translate camera frame's origin to the world frame
    origins = torch.broadcast_to(c2w[ None,None, None, :3, -1], directions.shape)
    viewdirs = directions / torch.norm(directions, axis=-1, keepdims=True)

    return torch.cat((origins,viewdirs),dim=-1)
//  [python]
def get_rays_torch(focal, c2w, H=64,W=64):
    """Computes rays using a General Pinhole Camera Model
    Assumes self.h, self.w, self.focal, and self.cam_to_world exist
    """
    x, y = torch.meshgrid(
        torch.arange(W),  # X-Axis (columns)
        torch.arange(H),  # Y-Axis (rows)
        indexing='xy')
    camera_directions = torch.stack(
        [(x - W * 0.5 + 0.5) / focal,
            -(y - H * 0.5 + 0.5) / focal,
            -torch.ones_like(x)],
        dim=-1).to(c2w)

    # Rotate ray directions from camera frame to the world frame
    directions = ((camera_directions[ None,..., None, :] * c2w[None,None, None, :3, :3]).sum(axis=-1))  # Translate camera frame's origin to the world frame
    origins = torch.broadcast_to(c2w[ None,None, None, :3, -1], directions.shape)
    viewdirs = directions / torch.linalg.norm(directions, axis=-1, keepdims=True)

    return torch.cat((origins,viewdirs),dim=-1)
//  [python]
def test_initialize_power(self):
        kernel = DownsamplingKernel()
        kernel.initialize(power=1)
        actual_value = torch.tensor(1.0).view_as(kernel.power)
        self.assertLess(torch.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power(self):
        kernel = DownsamplingKernel()
        kernel.initialize(power=1)
        actual_value = torch.tensor(1.0).view_as(kernel.power)
        self.assertLess(torch.linalg.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power_batch(self):
        kernel = DownsamplingKernel(batch_shape=torch.Size([2]))
        power_init = torch.tensor([1.0, 2.0])
        kernel.initialize(power=power_init)
        actual_value = power_init.view_as(kernel.power)
        self.assertLess(torch.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power_batch(self):
        kernel = DownsamplingKernel(batch_shape=torch.Size([2]))
        power_init = torch.tensor([1.0, 2.0])
        kernel.initialize(power=power_init)
        actual_value = power_init.view_as(kernel.power)
        self.assertLess(torch.linalg.norm(kernel.power - actual_value), 1e-5)
//  [python]
def cal_align_uniform(model,dataloader,args):
    results = []
    labels = []
    embs = []
    embas, embbs = [], []
    for batch in dataloader:
        with torch.no_grad():
            batch = [ele.to(model.device) for ele in batch]
            
            # print(batch[0].size(),batch[1].size(),batch[2].size())
            # print(batch[3].size(),batch[4].size(),batch[5].size())
            if len(batch) == 5:
                outputs = model(input_ids=batch[0],attention_mask=batch[1],output_hidden_states=True, return_dict=True)
                emba = pooling(outputs,batch[1],args)
                # emba = outputs.pooler_output
                outputs = model(input_ids=batch[2],attention_mask=batch[3],output_hidden_states=True, return_dict=True)
                embb = pooling(outputs,batch[3],args)
                label = batch[4].detach().cpu().numpy().tolist()
            else:
                outputs = model(input_ids=batch[0],token_type_ids=batch[1],attention_mask=batch[2],output_hidden_states=True, return_dict=True)
                emba = pooling(outputs,batch[2],args)
                outputs = model(input_ids=batch[3],token_type_ids=batch[4],attention_mask=batch[5],output_hidden_states=True, return_dict=True)
                embb = pooling(outputs,batch[5],args)
                label = batch[6].detach().cpu().numpy().tolist()
            emba = F.normalize(emba,dim=-1)
            embb = F.normalize(embb,dim=-1)
            scores = torch.norm(emba-embb,dim=-1).pow(2)
            embs.append(emba)
            embs.append(embb)
            # embas.append(emba)
            # embbs.append(embb)
            labels += label
        scores = scores.detach().cpu().numpy().tolist()
        results += scores
        
        labels += label
    align = []
    uniform = []
    
    embs = torch.cat(embs,dim=0)
    '''
    kernel, bias = compute_kernel_bias(embs.numpy())
    kernel = torch.Tensor(kernel)
    bias = torch.Tensor(bias).squeeze()
    embs = (embs + bias).matmul(kernel)
    embs = F.normalize(embs,dim=-1)

    embas = torch.cat(embas,dim=0)
    embbs = torch.cat(embbs,dim=0)
    embas = (embas + bias).matmul(kernel)
    embbs = (embbs + bias).matmul(kernel)
    embas = F.normalize(embas,dim=-1)
    embbs = F.normalize(embbs,dim=-1)

    scores = torch.norm(embas-embbs,dim=-1).pow(2)
    '''
    uniform = F.pdist(embs).pow(2).neg().exp().mean().log().item()
    for score,label in zip(results,labels):
        if label > 4.0:
            align.append(score)
        # uniform.append(math.exp(-2 * score))
    align = sum(align) / len(align)
    metrics = {'align':align,'uniform':uniform}
    return metrics
//  [python]
def cal_align_uniform(model,dataloader,args):
    results = []
    labels = []
    embs = []
    embas, embbs = [], []
    for batch in dataloader:
        with torch.no_grad():
            batch = [ele.to(model.device) for ele in batch]
            
            # print(batch[0].size(),batch[1].size(),batch[2].size())
            # print(batch[3].size(),batch[4].size(),batch[5].size())
            if len(batch) == 5:
                outputs = model(input_ids=batch[0],attention_mask=batch[1],output_hidden_states=True, return_dict=True)
                emba = pooling(outputs,batch[1],args)
                # emba = outputs.pooler_output
                outputs = model(input_ids=batch[2],attention_mask=batch[3],output_hidden_states=True, return_dict=True)
                embb = pooling(outputs,batch[3],args)
                label = batch[4].detach().cpu().numpy().tolist()
            else:
                outputs = model(input_ids=batch[0],token_type_ids=batch[1],attention_mask=batch[2],output_hidden_states=True, return_dict=True)
                emba = pooling(outputs,batch[2],args)
                outputs = model(input_ids=batch[3],token_type_ids=batch[4],attention_mask=batch[5],output_hidden_states=True, return_dict=True)
                embb = pooling(outputs,batch[5],args)
                label = batch[6].detach().cpu().numpy().tolist()
            emba = F.normalize(emba,dim=-1)
            embb = F.normalize(embb,dim=-1)
            scores = torch.linalg.norm(emba-embb,dim=-1).pow(2)
            embs.append(emba)
            embs.append(embb)
            # embas.append(emba)
            # embbs.append(embb)
            labels += label
        scores = scores.detach().cpu().numpy().tolist()
        results += scores
        
        labels += label
    align = []
    uniform = []
    
    embs = torch.cat(embs,dim=0)
    '''
    kernel, bias = compute_kernel_bias(embs.numpy())
    kernel = torch.Tensor(kernel)
    bias = torch.Tensor(bias).squeeze()
    embs = (embs + bias).matmul(kernel)
    embs = F.normalize(embs,dim=-1)

    embas = torch.cat(embas,dim=0)
    embbs = torch.cat(embbs,dim=0)
    embas = (embas + bias).matmul(kernel)
    embbs = (embbs + bias).matmul(kernel)
    embas = F.normalize(embas,dim=-1)
    embbs = F.normalize(embbs,dim=-1)

    scores = torch.linalg.norm(embas-embbs,dim=-1).pow(2)
    '''
    uniform = F.pdist(embs).pow(2).neg().exp().mean().log().item()
    for score,label in zip(results,labels):
        if label > 4.0:
            align.append(score)
        # uniform.append(math.exp(-2 * score))
    align = sum(align) / len(align)
    metrics = {'align':align,'uniform':uniform}
    return metrics
//  [python]
def test_computes_downsampling_function_batch(self):
        a = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.5]]).view(3, 2, 1)
        b = torch.tensor([[0.5, 0.6], [0.7, 0.8], [0.6, 0.6]]).view(3, 2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel(batch_shape=torch.Size([3]))
        kernel.initialize(power=power, offset=offset)
        kernel.eval()
        res = kernel(a, b).to_dense()

        actual = torch.zeros(3, 2, 2)

        diff = torch.tensor([[0.45, 0.36], [0.4, 0.32]])
        actual[0, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.21, 0.14], [0.18, 0.12]])
        actual[1, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.2, 0.2], [0.2, 0.2]])
        actual[2, :, :] = offset + diff.pow(1 + power)
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_computes_downsampling_function_batch(self):
        a = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.5]]).view(3, 2, 1)
        b = torch.tensor([[0.5, 0.6], [0.7, 0.8], [0.6, 0.6]]).view(3, 2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel(batch_shape=torch.Size([3]))
        kernel.initialize(power=power, offset=offset)
        kernel.eval()
        res = kernel(a, b).to_dense()

        actual = torch.zeros(3, 2, 2)

        diff = torch.tensor([[0.45, 0.36], [0.4, 0.32]])
        actual[0, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.21, 0.14], [0.18, 0.12]])
        actual[1, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.2, 0.2], [0.2, 0.2]])
        actual[2, :, :] = offset + diff.pow(1 + power)
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def convert2mip(pts):
    # pts
    pts_norm = torch.norm(pts, p=2, dim=-1)
    outside_mask = (pts_norm >= 1.0)
    mip_pts = pts.clone()
    mip_pts[outside_mask,:] = (2 - 1.0 / pts_norm[outside_mask, None]) * (pts[outside_mask,:] / pts_norm[outside_mask, None])
    return mip_pts
//  [python]
def convert2mip(pts):
    # pts
    pts_norm = torch.linalg.norm(pts, ord=2, dim=-1)
    outside_mask = (pts_norm >= 1.0)
    mip_pts = pts.clone()
    mip_pts[outside_mask,:] = (2 - 1.0 / pts_norm[outside_mask, None]) * (pts[outside_mask,:] / pts_norm[outside_mask, None])
    return mip_pts
//  [python]
def test_diag_calculation(self):
        a = torch.tensor([0.1, 0.2]).view(2, 1)
        b = torch.tensor([0.2, 0.4]).view(2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel()
        kernel.initialize(power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[0.72, 0.54], [0.64, 0.48]])
        actual = offset + diff.pow(1 + power)
        res = kernel(a, b, diag=True)

        self.assertLess(torch.norm(res - torch.diag(actual)), 1e-5)
//  [python]
def test_diag_calculation(self):
        a = torch.tensor([0.1, 0.2]).view(2, 1)
        b = torch.tensor([0.2, 0.4]).view(2, 1)
        power = 1
        offset = 1

        kernel = DownsamplingKernel()
        kernel.initialize(power=power, offset=offset)
        kernel.eval()

        diff = torch.tensor([[0.72, 0.54], [0.64, 0.48]])
        actual = offset + diff.pow(1 + power)
        res = kernel(a, b, diag=True)

        self.assertLess(torch.linalg.norm(res - torch.diag(actual)), 1e-5)
//  [python]
def test_last_dim_is_batch(self):
        a = (
            torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.5]])
            .view(3, 2)
            .transpose(-1, -2)
        )
        b = (
            torch.tensor([[0.5, 0.6], [0.7, 0.8], [0.6, 0.6]])
            .view(3, 2)
            .transpose(-1, -2)
        )
        power = 1
        offset = 1

        kernel = DownsamplingKernel()
        kernel.initialize(power=power, offset=offset)
        kernel.eval()
        res = kernel(a, b, last_dim_is_batch=True).to_dense()

        actual = torch.zeros(3, 2, 2)

        diff = torch.tensor([[0.45, 0.36], [0.4, 0.32]])
        actual[0, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.21, 0.14], [0.18, 0.12]])
        actual[1, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.2, 0.2], [0.2, 0.2]])
        actual[2, :, :] = offset + diff.pow(1 + power)
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_last_dim_is_batch(self):
        a = (
            torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.5]])
            .view(3, 2)
            .transpose(-1, -2)
        )
        b = (
            torch.tensor([[0.5, 0.6], [0.7, 0.8], [0.6, 0.6]])
            .view(3, 2)
            .transpose(-1, -2)
        )
        power = 1
        offset = 1

        kernel = DownsamplingKernel()
        kernel.initialize(power=power, offset=offset)
        kernel.eval()
        res = kernel(a, b, last_dim_is_batch=True).to_dense()

        actual = torch.zeros(3, 2, 2)

        diff = torch.tensor([[0.45, 0.36], [0.4, 0.32]])
        actual[0, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.21, 0.14], [0.18, 0.12]])
        actual[1, :, :] = offset + diff.pow(1 + power)

        diff = torch.tensor([[0.2, 0.2], [0.2, 0.2]])
        actual[2, :, :] = offset + diff.pow(1 + power)
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def _estimate_posterior_stable(self, x_test):
        """Estimate Posterior (mu, cov) of GP conditioned on data

        x_test : test points
        """
        cov11 = self._kernel(self.x_data, self.x_data)
        cov12 = self._kernel(self.x_data, x_test)
        cov22 = self._kernel(x_test, x_test)
        # cholesky decomposition of cov11
        L = cholesky(cov11)
        A = torch.triangular_solve(cov12, L, upper=False)[0]
        V = torch.triangular_solve(self.y_data.view(-1, 1), L, upper=False)[0]
        # sufficient statistics
        mu = A.t() @ V
        cov = cov22 - A.t() @ A
        return mu.view(-1), cov
//  [python]
def _estimate_posterior_stable(self, x_test):
        """Estimate Posterior (mu, cov) of GP conditioned on data

        x_test : test points
        """
        cov11 = self._kernel(self.x_data, self.x_data)
        cov12 = self._kernel(self.x_data, x_test)
        cov22 = self._kernel(x_test, x_test)
        # cholesky decomposition of cov11
        L = cholesky(cov11)
        A = torch.linalg.solve_triangular(cov12, L, upper=False)
        V = torch.linalg.solve_triangular(self.y_data.view(-1, 1), L, upper=False)
        # sufficient statistics
        mu = A.t() @ V
        cov = cov22 - A.t() @ A
        return mu.view(-1), cov
//  [python]
def precision_matrix(self):
        # We use "Woodbury matrix identity" to take advantage of low rank form::
        #     inv(W @ W.T + D) = inv(D) - inv(D) @ W @ inv(C) @ W.T @ inv(D)
        # where :math:`C` is the capacitance matrix.
        Wt_Dinv = (self._unbroadcasted_cov_factor.transpose(-1, -2)
                   / self._unbroadcasted_cov_diag.unsqueeze(-2))
        A = torch.triangular_solve(Wt_Dinv, self._capacitance_tril, upper=False)[0]
        precision_matrix = (torch.diag_embed(self._unbroadcasted_cov_diag.reciprocal())
                            - torch.matmul(A.transpose(-1, -2), A))
        return precision_matrix.expand(self._batch_shape + self._event_shape +
                                       self._event_shape)
//  [python]
def precision_matrix(self):
        # We use "Woodbury matrix identity" to take advantage of low rank form::
        #     inv(W @ W.T + D) = inv(D) - inv(D) @ W @ inv(C) @ W.T @ inv(D)
        # where :math:`C` is the capacitance matrix.
        Wt_Dinv = (self._unbroadcasted_cov_factor.transpose(-1, -2)
                   / self._unbroadcasted_cov_diag.unsqueeze(-2))
        A = torch.linalg.solve_triangular(Wt_Dinv, self._capacitance_tril, upper=False)
        precision_matrix = (torch.diag_embed(self._unbroadcasted_cov_diag.reciprocal())
                            - torch.matmul(A.transpose(-1, -2), A))
        return precision_matrix.expand(self._batch_shape + self._event_shape +
                                       self._event_shape)
//  [python]
def maha(self, mean: ch.Tensor, mean_other: ch.Tensor, std: ch.Tensor):
        diff = (mean - mean_other)[..., None]
        return ch.triangular_solve(diff, std, upper=False)[0].pow(2).sum([-2, -1])
//  [python]
def maha(self, mean: ch.Tensor, mean_other: ch.Tensor, std: ch.Tensor):
        diff = (mean - mean_other)[..., None]
        return ch.linalg.solve_triangular(diff, std, upper=False)[0].pow(2).sum([-2, -1])
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs, _ = torch.triangular_solve(outputs.t(), lower, upper=False, unitriangular=True)
        outputs, _ = torch.triangular_solve(outputs, upper, upper=True, unitriangular=False)
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs, _ = torch.linalg.solve_triangular(outputs.t(), lower, upper=False, unitriangular=True)
        outputs, _ = torch.linalg.solve_triangular(outputs, upper, upper=True, unitriangular=False)
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def _kl_multivariatenormal_multivariatenormal(p, q):
    # From https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Multivariate Normals with\
                          different event shapes cannot be computed")

    half_term1 = (q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
                  p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    combined_batch_shape = torch._C._infer_size(q._unbroadcasted_scale_tril.shape[:-2],
                                                p._unbroadcasted_scale_tril.shape[:-2])
    n = p.event_shape[0]
    q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    p_scale_tril = p._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    term2 = _batch_trace_XXT(torch.triangular_solve(p_scale_tril, q_scale_tril, upper=False)[0])
    term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
    return half_term1 + 0.5 * (term2 + term3 - n)
//  [python]
def _kl_multivariatenormal_multivariatenormal(p, q):
    # From https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Multivariate Normals with\
                          different event shapes cannot be computed")

    half_term1 = (q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
                  p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    combined_batch_shape = torch._C._infer_size(q._unbroadcasted_scale_tril.shape[:-2],
                                                p._unbroadcasted_scale_tril.shape[:-2])
    n = p.event_shape[0]
    q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    p_scale_tril = p._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    term2 = _batch_trace_XXT(torch.linalg.solve_triangular(p_scale_tril, q_scale_tril, upper=False))
    term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
    return half_term1 + 0.5 * (term2 + term3 - n)
//  [python]
def _kl_multivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             2 * p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ p_tril @ p_tril.T
    #                  = [inv(qD) - A.T @ A] @ p_tril @ p_tril.T
    qWt_qDinv = (q._unbroadcasted_cov_factor.transpose(-1, -2) /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.triangular_solve(qWt_qDinv, q._capacitance_tril, upper=False)[0]
    term21 = _batch_trace_XXT(p._unbroadcasted_scale_tril *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term22 = _batch_trace_XXT(A.matmul(p._unbroadcasted_scale_tril))
    term2 = term21 - term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_multivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             2 * p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ p_tril @ p_tril.T
    #                  = [inv(qD) - A.T @ A] @ p_tril @ p_tril.T
    qWt_qDinv = (q._unbroadcasted_cov_factor.transpose(-1, -2) /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.linalg.solve_triangular(q._capacitance_tril, qWt_qDinv, upper=False)
    term21 = _batch_trace_XXT(p._unbroadcasted_scale_tril *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term22 = _batch_trace_XXT(A.matmul(p._unbroadcasted_scale_tril))
    term2 = term21 - term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def forward_no_cache(self, x, logpx=None, reverse=False, context=None):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        if not reverse:
            y = F.linear(x, upper)
            y = F.linear(y, lower, self.bias)
            delta_logp = self.logabsdet() * x.new_ones(y.shape[0])
        else:
            y = x - self.bias
            y, _ = torch.triangular_solve(y.t(), lower, upper=False, unitriangular=True)
            y, _ = torch.triangular_solve(y, upper, upper=True, unitriangular=False)
            y = y.t()

            delta_logp = -self.logabsdet()
            delta_logp = delta_logp * x.new_ones(y.shape[0])

        if logpx is None:
            return y
        else:
            return y, logpx + delta_logp[:,None]