jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self):
        N0 = self.matA.shape[0]
        coeff = self.theta / torch.norm(self.theta)
        if self.use_complex:
            coeff = coeff[:N0] + 1j*coeff[N0:]
        matH = (coeff @ self.matA.reshape(self.matA.shape[0],-1)).reshape(self.matA.shape[1], self.matA.shape[2])
        self.matH = matH
        loss = torch.linalg.svd(matH, full_matrices=False)[1][1]
        # loss = torch.linalg.eigvalsh(matH @ matH.T)[-2]
        return loss
//  [python]
def forward(self):
        N0 = self.matA.shape[0]
        coeff = self.theta / torch.linalg.norm(self.theta)
        if self.use_complex:
            coeff = coeff[:N0] + 1j*coeff[N0:]
        matH = (coeff @ self.matA.reshape(self.matA.shape[0],-1)).reshape(self.matA.shape[1], self.matA.shape[2])
        self.matH = matH
        loss = torch.linalg.svd(matH, full_matrices=False)[1][1]
        # loss = torch.linalg.eigvalsh(matH @ matH.T)[-2]
        return loss
//  [python]
def forward(self):
        vecY = self.theta0 / torch.norm(self.theta0)
        tmp1 = self.matB @ vecY
        loss = torch.linalg.eigvalsh(tmp1.T @ tmp1)[0] #if not full rank, then exist vecX
        return loss
//  [python]
def forward(self):
        vecY = self.theta0 / torch.linalg.norm(self.theta0)
        tmp1 = self.matB @ vecY
        loss = torch.linalg.eigvalsh(tmp1.T @ tmp1)[0] #if not full rank, then exist vecX
        return loss
//  [python]
def _exp_impl_helper(tangent_vector: torch.Tensor):
    theta = torch.norm(tangent_vector, dim=-1, keepdim=True).unsqueeze(-1)
    theta2 = theta**2
    # Compute the approximations when theta ~ 0
    near_zero = theta < LIE_PARAMS.get_eps("so3", "near_zero", tangent_vector.dtype)
    theta_nz = torch.where(near_zero, constants._NON_ZERO, theta)
    theta2_nz = torch.where(near_zero, constants._NON_ZERO, theta2)

    cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
    sine = theta.sin()
    sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
    one_minus_cosine_by_theta2 = torch.where(
        near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz
    )

    size = tangent_vector.shape[:-1]
    ret = (
        one_minus_cosine_by_theta2
        * tangent_vector.view(*size, 3, 1)
        @ tangent_vector.view(*size, 1, 3)
    )
    ret[..., 0, 0] += cosine.view(size)
    ret[..., 1, 1] += cosine.view(size)
    ret[..., 2, 2] += cosine.view(size)
    sine_axis = sine_by_theta.view(*size, 1) * tangent_vector
    ret[..., 0, 1] -= sine_axis[..., 2]
    ret[..., 1, 0] += sine_axis[..., 2]
    ret[..., 0, 2] += sine_axis[..., 1]
    ret[..., 2, 0] -= sine_axis[..., 1]
    ret[..., 1, 2] -= sine_axis[..., 0]
    ret[..., 2, 1] += sine_axis[..., 0]

    return ret, (
        theta,
        theta2,
        theta_nz,
        theta2_nz,
        sine,
        cosine,
        sine_by_theta,
        one_minus_cosine_by_theta2,
    )
//  [python]
def _exp_impl_helper(tangent_vector: torch.Tensor):
    theta = torch.linalg.norm(tangent_vector, dim=-1, keepdim=True).unsqueeze(-1)
    theta2 = theta**2
    # Compute the approximations when theta ~ 0
    near_zero = theta < LIE_PARAMS.get_eps("so3", "near_zero", tangent_vector.dtype)
    theta_nz = torch.where(near_zero, constants._NON_ZERO, theta)
    theta2_nz = torch.where(near_zero, constants._NON_ZERO, theta2)

    cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
    sine = theta.sin()
    sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
    one_minus_cosine_by_theta2 = torch.where(
        near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz
    )

    size = tangent_vector.shape[:-1]
    ret = (
        one_minus_cosine_by_theta2
        * tangent_vector.view(*size, 3, 1)
        @ tangent_vector.view(*size, 1, 3)
    )
    ret[..., 0, 0] += cosine.view(size)
    ret[..., 1, 1] += cosine.view(size)
    ret[..., 2, 2] += cosine.view(size)
    sine_axis = sine_by_theta.view(*size, 1) * tangent_vector
    ret[..., 0, 1] -= sine_axis[..., 2]
    ret[..., 1, 0] += sine_axis[..., 2]
    ret[..., 0, 2] += sine_axis[..., 1]
    ret[..., 2, 0] -= sine_axis[..., 1]
    ret[..., 1, 2] -= sine_axis[..., 0]
    ret[..., 2, 1] += sine_axis[..., 0]

    return ret, (
        theta,
        theta2,
        theta_nz,
        theta2_nz,
        sine,
        cosine,
        sine_by_theta,
        one_minus_cosine_by_theta2,
    )
//  [python]
def get_vecX_vecY(self):
        N0 = self.matA.shape[0]
        tmp0 = self.theta.detach().numpy()
        coeff = tmp0 / np.linalg.norm(tmp0)
        if self.use_complex:
            coeff = coeff[:N0] + 1j*coeff[N0:]
        matH = (coeff @ self.matA.cpu().numpy().reshape(self.matA.shape[0],-1)).reshape(self.matA.shape[1], self.matA.shape[2])
        u,s,v = np.linalg.svd(matH)
        npX = u[:,0]
        npY = v[0]
        return coeff, npX, npY
//  [python]
def get_vecX_vecY(self):
        N0 = self.matA.shape[0]
        tmp0 = self.theta.detach().numpy()
        coeff = tmp0 / torch.linalg.norm(tmp0)
        if self.use_complex:
            coeff = coeff[:N0] + 1j*coeff[N0:]
        matH = (coeff @ self.matA.cpu().numpy().reshape(self.matA.shape[0],-1)).reshape(self.matA.shape[1], self.matA.shape[2])
        u,s,v = np.linalg.svd(matH)
        npX = u[:,0]
        npY = v[0]
        return coeff, npX, npY
//  [python]
def so3exp(self, phi):
        angle = torch.norm(phi)
        # Near phi==0, use first order Taylor expansion
        if torch.abs(angle) < 1e-8:
            skew_phi = self.skew(phi)
            return self.Id3 + skew_phi
        axis = phi / angle
        skew_axis = self.skew(axis)
        s = torch.sin(angle)
        c = torch.cos(angle)
        return c * self.Id3 + (1 - c) * torch.outer(axis, axis) + s * skew_axis
//  [python]
def so3exp(self, phi):
        angle = torch.linalg.norm(phi)
        # Near phi==0, use first order Taylor expansion
        if torch.abs(angle) < 1e-8:
            skew_phi = self.skew(phi)
            return self.Id3 + skew_phi
        axis = phi / angle
        skew_axis = self.skew(axis)
        s = torch.sin(angle)
        c = torch.cos(angle)
        return c * self.Id3 + (1 - c) * torch.outer(axis, axis) + s * skew_axis
//  [python]
def sen3exp(self, xi):
        phi = xi[:3]
        angle = torch.norm(phi)
        # Near |phi|==0, use first order Taylor expansion
        if torch.abs(angle) < 1e-8:
            skew_phi = self.skew(phi)
            J = self.Id3 + 0.5 * skew_phi
            Rot = self.Id3 + skew_phi
        else:
            axis = phi / angle
            skew_axis = self.skew(axis)
            s = torch.sin(angle)
            c = torch.cos(angle)
            J = (s / angle) * self.Id3 + (1 - s / angle) * torch.outer(axis, axis) + ((1 - c) / angle) * skew_axis
            Rot = c * self.Id3 + (1 - c) * torch.outer(axis, axis) + s * skew_axis
        x = J.mm(xi[3:].view(-1, 3).t())
        return Rot, x
//  [python]
def sen3exp(self, xi):
        phi = xi[:3]
        angle = torch.linalg.norm(phi)
        # Near |phi|==0, use first order Taylor expansion
        if torch.abs(angle) < 1e-8:
            skew_phi = self.skew(phi)
            J = self.Id3 + 0.5 * skew_phi
            Rot = self.Id3 + skew_phi
        else:
            axis = phi / angle
            skew_axis = self.skew(axis)
            s = torch.sin(angle)
            c = torch.cos(angle)
            J = (s / angle) * self.Id3 + (1 - s / angle) * torch.outer(axis, axis) + ((1 - c) / angle) * skew_axis
            Rot = c * self.Id3 + (1 - c) * torch.outer(axis, axis) + s * skew_axis
        x = J.mm(xi[3:].view(-1, 3).t())
        return Rot, x
//  [python]
def forward(self):
        space_char = self.space_char
        if space_char in {'R_T', 'C_H'}:
            tmp0 = [
                self.EVL_free,
                None if (self.EVL_positive is None) else self.EVL_positive(),
                None if (self.EVL_negative is None) else (-self.EVL_negative()),
            ]
            tmp1 = torch.cat([x for x in tmp0 if x is not None])
            EVL = tmp1 / torch.norm(tmp1)
            unitary = self.manifold_U()
            matH = (unitary.conj()*EVL) @ unitary.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char=='R_A':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary = self.manifold_U()
            tmp0 = unitary[:,::2]
            tmp1 = unitary[:,1::2]
            matH = (tmp0 * EVL) @ tmp1.T - (tmp1 * EVL) @ tmp0.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char=='C_T':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary = self.manifold_U()
            matH = (unitary*EVL) @ unitary.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char in {'R','C'}:
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary0 = self.manifold_U0()
            unitary1 = self.manifold_U1()
            matH = (unitary0.conj()*EVL) @ unitary1.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char=='R_cT':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary = self.manifold_U()
            matH = (unitary*EVL) @ unitary.T
            tmp0 = numqi.utils.hf_complex_to_real(matH)
            loss = hf_torch_norm_square(self.basis_orth_conj @ tmp0.reshape(-1))
        elif space_char=='R_c':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary0 = self.manifold_U0()
            unitary1 = self.manifold_U1()
            matH = (unitary0.conj()*EVL) @ unitary1.T
            tmp0 = numqi.utils.hf_complex_to_real(matH)
            loss = hf_torch_norm_square(self.basis_orth_conj @ tmp0.reshape(-1))
        self.matH = matH.detach()
        return loss
//  [python]
def forward(self):
        space_char = self.space_char
        if space_char in {'R_T', 'C_H'}:
            tmp0 = [
                self.EVL_free,
                None if (self.EVL_positive is None) else self.EVL_positive(),
                None if (self.EVL_negative is None) else (-self.EVL_negative()),
            ]
            tmp1 = torch.cat([x for x in tmp0 if x is not None])
            EVL = tmp1 / torch.linalg.norm(tmp1)
            unitary = self.manifold_U()
            matH = (unitary.conj()*EVL) @ unitary.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char=='R_A':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary = self.manifold_U()
            tmp0 = unitary[:,::2]
            tmp1 = unitary[:,1::2]
            matH = (tmp0 * EVL) @ tmp1.T - (tmp1 * EVL) @ tmp0.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char=='C_T':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary = self.manifold_U()
            matH = (unitary*EVL) @ unitary.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char in {'R','C'}:
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary0 = self.manifold_U0()
            unitary1 = self.manifold_U1()
            matH = (unitary0.conj()*EVL) @ unitary1.T
            loss = hf_torch_norm_square(self.basis_orth_conj @ matH.reshape(-1))
        elif space_char=='R_cT':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary = self.manifold_U()
            matH = (unitary*EVL) @ unitary.T
            tmp0 = numqi.utils.hf_complex_to_real(matH)
            loss = hf_torch_norm_square(self.basis_orth_conj @ tmp0.reshape(-1))
        elif space_char=='R_c':
            EVL = 1 if (self.manifold_EVL is None) else self.manifold_EVL()
            unitary0 = self.manifold_U0()
            unitary1 = self.manifold_U1()
            matH = (unitary0.conj()*EVL) @ unitary1.T
            tmp0 = numqi.utils.hf_complex_to_real(matH)
            loss = hf_torch_norm_square(self.basis_orth_conj @ tmp0.reshape(-1))
        self.matH = matH.detach()
        return loss
//  [python]
def compute(self, simulation_output: SimulationOutputCLE) -> torch.Tensor:
        """Compute the metric on all frames of the scene.

        :param simulation_output: the output from the closed-loop simulation
        :returns: driven miles per each frame
        """
        observed_ego_states_centroid = simulation_output.recorded_ego_states[:, :2]  # [Timesteps, 2]
        observed_ego_states_centroid = observed_ego_states_centroid.to(torch.float64)

        drive_meters = \
            torch.norm(observed_ego_states_centroid[1:]
                              - observed_ego_states_centroid[0:-1], dim=1)
        pad = torch.as_tensor([0.], device=observed_ego_states_centroid.device)
        pad_drive_meters = torch.cat((pad, drive_meters))
        driven_miles = pad_drive_meters * self.METER_TO_MILES
        return driven_miles
//  [python]
def compute(self, simulation_output: SimulationOutputCLE) -> torch.Tensor:
        """Compute the metric on all frames of the scene.

        :param simulation_output: the output from the closed-loop simulation
        :returns: driven miles per each frame
        """
        observed_ego_states_centroid = simulation_output.recorded_ego_states[:, :2]  # [Timesteps, 2]
        observed_ego_states_centroid = observed_ego_states_centroid.to(torch.float64)

        drive_meters = \
            torch.linalg.norm(observed_ego_states_centroid[1:]
                              - observed_ego_states_centroid[0:-1], dim=1)
        pad = torch.as_tensor([0.], device=observed_ego_states_centroid.device)
        pad_drive_meters = torch.cat((pad, drive_meters))
        driven_miles = pad_drive_meters * self.METER_TO_MILES
        return driven_miles
//  [python]
def compute(self, simulation_output: SimulationOutputCLE) -> torch.Tensor:
        """Compute the metric on all frames of the scene.

        :param simulation_output: the output from the closed-loop simulation
        :returns: driven miles per each frame
        """
        simulated_scene_ego_state = simulation_output.simulated_ego_states
        simulated_centroid = simulated_scene_ego_state[:, :2]  # [Timesteps, 2]
        simulated_centroid = simulated_centroid.to(torch.float64)

        drive_meters = torch.norm(simulated_centroid[1:] - simulated_centroid[0:-1], dim=1)
        pad = torch.as_tensor([0.], device=simulated_centroid.device)
        pad_drive_meters = torch.cat((pad, drive_meters))
        driven_miles = pad_drive_meters * self.METER_TO_MILES
        return driven_miles
//  [python]
def compute(self, simulation_output: SimulationOutputCLE) -> torch.Tensor:
        """Compute the metric on all frames of the scene.

        :param simulation_output: the output from the closed-loop simulation
        :returns: driven miles per each frame
        """
        simulated_scene_ego_state = simulation_output.simulated_ego_states
        simulated_centroid = simulated_scene_ego_state[:, :2]  # [Timesteps, 2]
        simulated_centroid = simulated_centroid.to(torch.float64)

        drive_meters = torch.linalg.norm(simulated_centroid[1:] - simulated_centroid[0:-1], dim=1)
        pad = torch.as_tensor([0.], device=simulated_centroid.device)
        pad_drive_meters = torch.cat((pad, drive_meters))
        driven_miles = pad_drive_meters * self.METER_TO_MILES
        return driven_miles
//  [python]
def compute(self, simulation_output: SimulationOutputCLE) -> torch.Tensor:
        # calculate the speed error between simulated and recorded ego over the course of the simulation.
        simulated_centroid = simulation_output.simulated_ego_states[:, :2]
        simulated_velocity = simulated_centroid[1:] - simulated_centroid[:-1]
        simulated_speed = torch.norm(simulated_velocity, dim=1)

        recorded_centroid = simulation_output.recorded_ego_states[:, :2]
        recorded_velocity = recorded_centroid[1:] - recorded_centroid[:-1]
        recorded_speed = torch.norm(recorded_velocity, dim=1)

        # TODO: Replace 10 by (1 / step_time)
        return (simulated_speed - recorded_speed) * 10
//  [python]
def compute(self, simulation_output: SimulationOutputCLE) -> torch.Tensor:
        # calculate the speed error between simulated and recorded ego over the course of the simulation.
        simulated_centroid = simulation_output.simulated_ego_states[:, :2]
        simulated_velocity = simulated_centroid[1:] - simulated_centroid[:-1]
        simulated_speed = torch.linalg.norm(simulated_velocity, dim=1)

        recorded_centroid = simulation_output.recorded_ego_states[:, :2]
        recorded_velocity = recorded_centroid[1:] - recorded_centroid[:-1]
        recorded_speed = torch.linalg.norm(recorded_velocity, dim=1)

        # TODO: Replace 10 by (1 / step_time)
        return (simulated_speed - recorded_speed) * 10
//  [python]
def calculate_diversity(activation, diversity_times):
    assert len(activation.shape) == 2
    assert activation.shape[0] > diversity_times
    num_samples = activation.shape[0]

    first_indices = np.random.choice(num_samples,
                                     diversity_times,
                                     replace=False)
    second_indices = np.random.choice(num_samples,
                                      diversity_times,
                                      replace=False)
    dist = np.linalg.norm(activation[first_indices] - activation[second_indices],
                       axis=1)
    return dist.mean()
//  [python]
def calculate_diversity(activation, diversity_times):
    assert len(activation.shape) == 2
    assert activation.shape[0] > diversity_times
    num_samples = activation.shape[0]

    first_indices = np.random.choice(num_samples,
                                     diversity_times,
                                     replace=False)
    second_indices = np.random.choice(num_samples,
                                      diversity_times,
                                      replace=False)
    dist = linalg.norm(activation[first_indices] - activation[second_indices],
                       axis=1)
    return dist.mean()
//  [python]
def calc_accel(preds, target):
    """
    Mean joint acceleration error
    often referred to as "Protocol #1" in many papers.
    """
    assert preds.shape == target.shape, print(preds.shape,
                                              target.shape)  # BxJx3
    assert preds.dim() == 3
    # Expects BxJx3
    # valid_mask = torch.BoolTensor(target[:, :, 0].shape)
    accel_gt = target[:-2] - 2 * target[1:-1] + target[2:]
    accel_pred = preds[:-2] - 2 * preds[1:-1] + preds[2:]
    normed = torch.norm(accel_pred - accel_gt, dim=-1)
    accel_seq = normed.mean(1)
    return accel_seq
//  [python]
def calc_accel(preds, target):
    """
    Mean joint acceleration error
    often referred to as "Protocol #1" in many papers.
    """
    assert preds.shape == target.shape, print(preds.shape,
                                              target.shape)  # BxJx3
    assert preds.dim() == 3
    # Expects BxJx3
    # valid_mask = torch.BoolTensor(target[:, :, 0].shape)
    accel_gt = target[:-2] - 2 * target[1:-1] + target[2:]
    accel_pred = preds[:-2] - 2 * preds[1:-1] + preds[2:]
    normed = torch.linalg.norm(accel_pred - accel_gt, dim=-1)
    accel_seq = normed.mean(1)
    return accel_seq
//  [python]
def compute_Li_for_linear_regression(self):
        # ==============================================================================================================
        # Compute L, Li for linear regression
        # ==============================================================================================================
        A = self.data
        self.L = ((2/A.shape[0]) * torch.norm(A, 2)**2).item()
        self.Li_all_clients = []

        for c in range(self.num_clients):
            self.set_client(c)
            subdata = self.data[int(self.client_id) * self.n_client_samples:
                                (int(self.client_id) + 1) * self.n_client_samples, ...]

            Li = ((2/subdata.shape[0]) * torch.norm(subdata, 2) ** 2).item()
            self.Li_all_clients.append(Li)

        assert max(self.Li_all_clients) + 1.0e+3 >= self.L
        assert max(self.Li_all_clients) - 1.0e+3 <= self.L * self.num_clients
//  [python]
def compute_Li_for_linear_regression(self):
        # ==============================================================================================================
        # Compute L, Li for linear regression
        # ==============================================================================================================
        A = self.data
        self.L = ((2/A.shape[0]) * torch.linalg.norm(A, 2)**2).item()
        self.Li_all_clients = []

        for c in range(self.num_clients):
            self.set_client(c)
            subdata = self.data[int(self.client_id) * self.n_client_samples:
                                (int(self.client_id) + 1) * self.n_client_samples, ...]

            Li = ((2/subdata.shape[0]) * torch.linalg.norm(subdata, 2) ** 2).item()
            self.Li_all_clients.append(Li)

        assert max(self.Li_all_clients) + 1.0e+3 >= self.L
        assert max(self.Li_all_clients) - 1.0e+3 <= self.L * self.num_clients
//  [python]
def relative_lp_error(pred, y, p=2):
    """
    Calculate relative L2 error norm
    Parameters:
    -----------
    pred: torch.Tensor
        Prediction
    y: torch.Tensor
        Ground truth
    Returns:
    --------
    error: float
        Calculated relative L2 error norm (percentage) on cpu
    """

    error = (
        torch.mean(torch.norm(pred - y, p=p) / torch.norm(y, p=p))
        .cpu()
        .numpy()
    )
    return error * 100
//  [python]
def relative_lp_error(pred, y, p=2):
    """
    Calculate relative L2 error norm
    Parameters:
    -----------
    pred: torch.Tensor
        Prediction
    y: torch.Tensor
        Ground truth
    Returns:
    --------
    error: float
        Calculated relative L2 error norm (percentage) on cpu
    """

    error = (
        torch.mean(torch.linalg.norm(pred - y, ord=p) / torch.linalg.norm(y, ord=p))
        .cpu()
        .numpy()
    )
    return error * 100
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: The tensor after applying scale norm.
        """

        if torch.onnx.is_in_onnx_export() and \
                digit_version(TORCH_VERSION) >= digit_version('1.12'):

            norm = torch.norm(x, dim=-1, keepdim=True)

        else:
            norm = torch.norm(x, dim=-1, keepdim=True)
        norm = norm * self.scale
        return x / norm.clamp(min=self.eps) * self.g
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: The tensor after applying scale norm.
        """

        if torch.onnx.is_in_onnx_export() and \
                digit_version(TORCH_VERSION) >= digit_version('1.12'):

            norm = torch.linalg.norm(x, dim=-1, keepdim=True)

        else:
            norm = torch.norm(x, dim=-1, keepdim=True)
        norm = norm * self.scale
        return x / norm.clamp(min=self.eps) * self.g
//  [python]
def SE3_logmap(T, eps=1e-6):
    w = SO3_logmap(T[:,:3,:3])
    theta = torch.norm(w,dim=(1))
    theta = torch.clamp(theta, min=eps)
    w_norm = w/theta
    tan = torch.tan(0.5*theta)

    t = T[:,:3,3]
    wnorm_x_t = torch.cross(w_norm,t)
    V_inv_t = t - (0.5 * t) * wnorm_x_t + (1.0 - theta/(2.0*tan)) * torch.cross(w_norm, wnorm_x_t)
    xi = torch.cat((w, V_inv_t),dim=-1)

    return xi
//  [python]
def SE3_logmap(T, eps=1e-6):
    w = SO3_logmap(T[:,:3,:3])
    theta = torch.linalg.norm(w,dim=(1))
    theta = torch.clamp(theta, min=eps)
    w_norm = w/theta
    tan = torch.tan(0.5*theta)

    t = T[:,:3,3]
    wnorm_x_t = torch.cross(w_norm,t)
    V_inv_t = t - (0.5 * t) * wnorm_x_t + (1.0 - theta/(2.0*tan)) * torch.cross(w_norm, wnorm_x_t)
    xi = torch.cat((w, V_inv_t),dim=-1)

    return xi
//  [python]
def training_step(self, batch, batch_idx):
        """
        Training step, runs once per batch
        """
        loss_terms = self._get_loss_terms(batch)
        avg_loss = torch.mean(loss_terms)

        # L1 loss implementation
        if self.l1_lambda > 0:
            l1_penalty = sum(torch.norm(p, 1) for p in self.parameters())
            avg_loss += self.l1_lambda * l1_penalty

        pseudo_ft_names = (
            (self.ft_names + ["pairwise_dist_loss"])
            if self.use_pairwise_dist_loss
            else self.ft_names
        )
        assert len(loss_terms) == len(pseudo_ft_names)
        loss_dict = {
            f"train_loss_{val_name}": val
            for val_name, val in zip(pseudo_ft_names, loss_terms)
        }
        loss_dict["train_loss"] = avg_loss
        self.log_dict(loss_dict)  # Don't seem to need rank zero or sync dist

        return avg_loss
//  [python]
def training_step(self, batch, batch_idx):
        """
        Training step, runs once per batch
        """
        loss_terms = self._get_loss_terms(batch)
        avg_loss = torch.mean(loss_terms)

        # L1 loss implementation
        if self.l1_lambda > 0:
            l1_penalty = sum(torch.linalg.norm(p, 1) for p in self.parameters())
            avg_loss += self.l1_lambda * l1_penalty

        pseudo_ft_names = (
            (self.ft_names + ["pairwise_dist_loss"])
            if self.use_pairwise_dist_loss
            else self.ft_names
        )
        assert len(loss_terms) == len(pseudo_ft_names)
        loss_dict = {
            f"train_loss_{val_name}": val
            for val_name, val in zip(pseudo_ft_names, loss_terms)
        }
        loss_dict["train_loss"] = avg_loss
        self.log_dict(loss_dict)  # Don't seem to need rank zero or sync dist

        return avg_loss
//  [python]
def __init__(self,
                 rot_mats: Optional[torch.Tensor] = None,
                 quats: Optional[torch.Tensor] = None,
                 normalize_quats: bool = True,
                 ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if ((rot_mats is None and quats is None) or
                (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

        if ((rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or
                (quats is not None and quats.shape[-1] != 4)):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if (quats is not None):
            quats = quats.to(dtype=torch.float32)
        if (rot_mats is not None):
            rot_mats = rot_mats.to(dtype=torch.float32)

        if (quats is not None and normalize_quats):
            quats = quats / torch.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def __init__(self,
                 rot_mats: Optional[torch.Tensor] = None,
                 quats: Optional[torch.Tensor] = None,
                 normalize_quats: bool = True,
                 ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if ((rot_mats is None and quats is None) or
                (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

        if ((rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or
                (quats is not None and quats.shape[-1] != 4)):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if (quats is not None):
            quats = quats.to(dtype=torch.float32)
        if (rot_mats is not None):
            rot_mats = rot_mats.to(dtype=torch.float32)

        if (quats is not None and normalize_quats):
            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def est_directional_similarity(xs: torch.Tensor, n_est: int = 1000) -> torch.Tensor:
    """xs: (batch, nx). Returns (n_est, ) between 0 and 1."""
    # xs: (batch, nx)
    batch, nx = xs.shape

    # Center first.
    xs = xs - torch.mean(xs, dim=0, keepdim=True)

    rand_idxs1 = torch.randint(batch, [n_est], dtype=torch.long)
    rand_idxs2 = torch.randint(batch, [n_est], dtype=torch.long)

    # (n_est, nx)
    xs1 = xs[rand_idxs1]
    # (n_est, nx)
    xs2 = xs[rand_idxs2]

    # Normalize to unit vector.
    xs1 /= torch.norm(xs1, dim=1, keepdim=True)
    xs2 /= torch.norm(xs2, dim=1, keepdim=True)

    # (n_est, )
    cos_angle = torch.sum(xs1 * xs2, dim=1).clip(-1.0, 1.0)
    assert cos_angle.shape == (n_est,)

    # Should be in [0, pi).
    angle = torch.acos(cos_angle)
    assert (0 <= angle).all()
    assert (angle <= torch.pi).all()

    D_ij = 1.0 - angle / torch.pi
    assert D_ij.shape == (n_est,)

    return D_ij
//  [python]
def est_directional_similarity(xs: torch.Tensor, n_est: int = 1000) -> torch.Tensor:
    """xs: (batch, nx). Returns (n_est, ) between 0 and 1."""
    # xs: (batch, nx)
    batch, nx = xs.shape

    # Center first.
    xs = xs - torch.mean(xs, dim=0, keepdim=True)

    rand_idxs1 = torch.randint(batch, [n_est], dtype=torch.long)
    rand_idxs2 = torch.randint(batch, [n_est], dtype=torch.long)

    # (n_est, nx)
    xs1 = xs[rand_idxs1]
    # (n_est, nx)
    xs2 = xs[rand_idxs2]

    # Normalize to unit vector.
    xs1 /= torch.linalg.norm(xs1, dim=1, keepdim=True)
    xs2 /= torch.linalg.norm(xs2, dim=1, keepdim=True)

    # (n_est, )
    cos_angle = torch.sum(xs1 * xs2, dim=1).clip(-1.0, 1.0)
    assert cos_angle.shape == (n_est,)

    # Should be in [0, pi).
    angle = torch.acos(cos_angle)
    assert (0 <= angle).all()
    assert (angle <= torch.pi).all()

    D_ij = 1.0 - angle / torch.pi
    assert D_ij.shape == (n_est,)

    return D_ij
//  [python]
def build_xis(T, D) -> torch.Tensor:
    # Generate random unit vectors.
    rng = np.random.default_rng(seed=4078213)
    xis = rng.standard_normal([T, D])

    # Construct a xis that has some degree of "continuous" over time, as a brownian motion.
    xi = xis[0]
    bm_xis = [xi]
    std = 0.4
    dt = 1.0 / T
    for t in range(1, T):
        xi = xi - (2.0 * xi) * dt + std * math.sqrt(dt) * xis[t]
        bm_xis.append(xi)
    assert len(bm_xis) == xis.shape[0]

    xis = torch.Tensor(np.stack(bm_xis))
    xis /= torch.norm(xis, dim=-1, keepdim=True)

    # Just safeguard if the self.xis becomes different.
    print("USING BM XI! xis.sum(): {}".format(torch.sum(xis)))
    assert xis.shape == (T, D)
    return xis
//  [python]
def build_xis(T, D) -> torch.Tensor:
    # Generate random unit vectors.
    rng = np.random.default_rng(seed=4078213)
    xis = rng.standard_normal([T, D])

    # Construct a xis that has some degree of "continuous" over time, as a brownian motion.
    xi = xis[0]
    bm_xis = [xi]
    std = 0.4
    dt = 1.0 / T
    for t in range(1, T):
        xi = xi - (2.0 * xi) * dt + std * math.sqrt(dt) * xis[t]
        bm_xis.append(xi)
    assert len(bm_xis) == xis.shape[0]

    xis = torch.Tensor(np.stack(bm_xis))
    xis /= torch.linalg.norm(xis, dim=-1, keepdim=True)

    # Just safeguard if the self.xis becomes different.
    print("USING BM XI! xis.sum(): {}".format(torch.sum(xis)))
    assert xis.shape == (T, D)
    return xis
//  [python]
def forward(self, text):
        tokens = clip.tokenize(text).to(self.device)
        z = self.model.encode_text(tokens)
        if self.normalize:
            z = z / torch.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def forward(self, text):
        tokens = clip.tokenize(text).to(self.device)
        z = self.model.encode_text(tokens)
        if self.normalize:
            z = z / torch.linalg.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def cholesky(self, quat):
        # Return the covariance matrix of the noise.
        with torch.enable_grad():
            tmp = quat.detach().clone()
            N, _ = tmp.shape
            tmp = tmp[:, None, :].expand((N, 3, 4))
            tmp.requires_grad_(True)
            eul = quat2eul(tmp.reshape(-1, 4)).reshape(N, -1, 3)
            mask = torch.eye(3, device=eul.device).repeat(N, 1, 1)
            quat_jacobian = torch.autograd.grad(eul, tmp, mask, create_graph=False)[0].detach().float()
            quat_jacobian = quat_jacobian.transpose(1, 2)
            ret_chol = torch.matmul(quat_jacobian, torch.diag(torch.stack([self.r_std, self.p_std, self.y_std])).float())
            ret_chol = torch.norm(ret_chol, dim=2)
            ret_chol = torch.diag_embed(ret_chol)
        return ret_chol
//  [python]
def cholesky(self, quat):
        # Return the covariance matrix of the noise.
        with torch.enable_grad():
            tmp = quat.detach().clone()
            N, _ = tmp.shape
            tmp = tmp[:, None, :].expand((N, 3, 4))
            tmp.requires_grad_(True)
            eul = quat2eul(tmp.reshape(-1, 4)).reshape(N, -1, 3)
            mask = torch.eye(3, device=eul.device).repeat(N, 1, 1)
            quat_jacobian = torch.autograd.grad(eul, tmp, mask, create_graph=False)[0].detach().float()
            quat_jacobian = quat_jacobian.transpose(1, 2)
            ret_chol = torch.matmul(quat_jacobian, torch.diag(torch.stack([self.r_std, self.p_std, self.y_std])).float())
            ret_chol = torch.linalg.norm(ret_chol, dim=2)
            ret_chol = torch.diag_embed(ret_chol)
        return ret_chol
//  [python]
def compute_mean_drift_term(mf_x: torch.Tensor, xi: torch.Tensor) -> torch.Tensor:
    """Decompose the polarize dynamic Eq (18) in paper into 2 parts for faster computation:
          f_polarize(x,p,ξ)
        = E_{y~p}[a(x,y,ξ) * bar_y],                  where a(x,y,ξ) = sign(<x,ξ>)*sign(<y,ξ>)
                                                      and      bar_y = y / |y|^{0.5}
        = sign(<x,ξ>) * E_{y~p}[sign(<y,ξ>) * bar_y], since sign(<x,ξ>) is independent of y
        = A(x,ξ)      * B(p,ξ)
    Hence, bar_f_polarize = bar_A(x,ξ) * bar_B(p,ξ)
    This function computes only bar_B(p,ξ).
    """
    # mf_x: (B, *, D), xi: (*, D)
    # output: (*, D)

    B, Ts, D = mf_x.shape[0], mf_x.shape[1:-1], mf_x.shape[-1]
    assert xi.shape == (*Ts, D)

    mf_x_norm = torch.norm(mf_x, dim=-1, keepdim=True)
    assert torch.all(mf_x_norm > 0.0)

    normalized_mf_x = mf_x / torch.sqrt(mf_x_norm)
    assert normalized_mf_x.shape == (B, *Ts, D)

    # Compute the mean drift term:   1/J sum_j a(y_j) y_j / sqrt(| y_j |).
    mf_agree_j = opinion_thresh(torch.sum(mf_x * xi, dim=-1, keepdim=True))
    assert mf_agree_j.shape == (B, *Ts, 1)

    mean_drift_term = torch.mean(mf_agree_j * normalized_mf_x, dim=0)
    assert mean_drift_term.shape == (*Ts, D)

    mean_drift_term_norm = torch.norm(mean_drift_term, dim=-1, keepdim=True)
    mean_drift_term = mean_drift_term / torch.sqrt(mean_drift_term_norm)
    assert mean_drift_term.shape == (*Ts, D)

    return mean_drift_term
//  [python]
def compute_mean_drift_term(mf_x: torch.Tensor, xi: torch.Tensor) -> torch.Tensor:
    """Decompose the polarize dynamic Eq (18) in paper into 2 parts for faster computation:
          f_polarize(x,p,ξ)
        = E_{y~p}[a(x,y,ξ) * bar_y],                  where a(x,y,ξ) = sign(<x,ξ>)*sign(<y,ξ>)
                                                      and      bar_y = y / |y|^{0.5}
        = sign(<x,ξ>) * E_{y~p}[sign(<y,ξ>) * bar_y], since sign(<x,ξ>) is independent of y
        = A(x,ξ)      * B(p,ξ)
    Hence, bar_f_polarize = bar_A(x,ξ) * bar_B(p,ξ)
    This function computes only bar_B(p,ξ).
    """
    # mf_x: (B, *, D), xi: (*, D)
    # output: (*, D)

    B, Ts, D = mf_x.shape[0], mf_x.shape[1:-1], mf_x.shape[-1]
    assert xi.shape == (*Ts, D)

    mf_x_norm = torch.linalg.norm(mf_x, dim=-1, keepdim=True)
    assert torch.all(mf_x_norm > 0.0)

    normalized_mf_x = mf_x / torch.sqrt(mf_x_norm)
    assert normalized_mf_x.shape == (B, *Ts, D)

    # Compute the mean drift term:   1/J sum_j a(y_j) y_j / sqrt(| y_j |).
    mf_agree_j = opinion_thresh(torch.sum(mf_x * xi, dim=-1, keepdim=True))
    assert mf_agree_j.shape == (B, *Ts, 1)

    mean_drift_term = torch.mean(mf_agree_j * normalized_mf_x, dim=0)
    assert mean_drift_term.shape == (*Ts, D)

    mean_drift_term_norm = torch.linalg.norm(mean_drift_term, dim=-1, keepdim=True)
    mean_drift_term = mean_drift_term / torch.sqrt(mean_drift_term_norm)
    assert mean_drift_term.shape == (*Ts, D)

    return mean_drift_term
//  [python]
def forward(self, initial_states, controls):
        # The dynamics model is a simple random walk model.
        # The controls are null.
        N, state_dim = initial_states.shape
        self.state_dim = state_dim

        accel = controls[:, :3] - initial_states[:,10:13]
        accel_mag = torch.norm(accel, dim=1)
#         print("Accel mag: ", accel_mag)
        accel_body = torch.zeros(N, 3)
#         accel_body[:, 1] = accel_mag
        accel_enu = tf.quaternion_apply(tf.quaternion_invert(initial_states[:,6:10].detach().clone()), accel_body)
#         print("Accel enu: ", accel_enu)
        omega = controls[:, 3:6] - initial_states[:,13:]
#         omega = torch.zeros(N, 3)

        pos_vel, pos_vel_cov = self.pos_vel_model(initial_states[:,:6], accel_enu)
        quat, quat_cov = self.quat_model(initial_states[:,6:10], omega)
        imu_bias, imu_bias_cov = self.imu_bias_model(initial_states[:,10:], None)

        # Combine the covariance matrices
        cholseky_matrix = torch.zeros((N, 16, 16))
        cholseky_matrix[:, 0:6, 0:6] = pos_vel_cov
        cholseky_matrix[:, 6:10, 6:10] = quat_cov
        cholseky_matrix[:, 10:, 10:] = imu_bias_cov

        return torch.cat((pos_vel, quat, imu_bias), axis=1), cholseky_matrix
//  [python]
def forward(self, initial_states, controls):
        # The dynamics model is a simple random walk model.
        # The controls are null.
        N, state_dim = initial_states.shape
        self.state_dim = state_dim

        accel = controls[:, :3] - initial_states[:,10:13]
        accel_mag = torch.linalg.norm(accel, dim=1)
#         print("Accel mag: ", accel_mag)
        accel_body = torch.zeros(N, 3)
#         accel_body[:, 1] = accel_mag
        accel_enu = tf.quaternion_apply(tf.quaternion_invert(initial_states[:,6:10].detach().clone()), accel_body)
#         print("Accel enu: ", accel_enu)
        omega = controls[:, 3:6] - initial_states[:,13:]
#         omega = torch.zeros(N, 3)

        pos_vel, pos_vel_cov = self.pos_vel_model(initial_states[:,:6], accel_enu)
        quat, quat_cov = self.quat_model(initial_states[:,6:10], omega)
        imu_bias, imu_bias_cov = self.imu_bias_model(initial_states[:,10:], None)

        # Combine the covariance matrices
        cholseky_matrix = torch.zeros((N, 16, 16))
        cholseky_matrix[:, 0:6, 0:6] = pos_vel_cov
        cholseky_matrix[:, 6:10, 6:10] = quat_cov
        cholseky_matrix[:, 10:, 10:] = imu_bias_cov

        return torch.cat((pos_vel, quat, imu_bias), axis=1), cholseky_matrix
//  [python]
def get_k_voxel_grid(k_cutoff, delta_k, num_k_rbf):

    # Get indices for a cube of k-lattice sites containing the cutoff sphere
    num_k = k_cutoff / delta_k
    k_index_product_set, _ = get_k_index_product_set(num_k, num_k, num_k)

    # Orthogonal k-space basis, norm delta_k
    k_cell = torch.tensor(
        [[delta_k, 0, 0], [0, delta_k, 0], [0, 0, delta_k]], dtype=torch.float
    )

    # Translate lattice indices into k-vectors
    k_grid = torch.matmul(k_index_product_set, k_cell)

    # Prune all k-vectors outside the cutoff sphere
    k_grid = k_grid[torch.sum(k_grid**2, dim=-1) <= k_cutoff**2]

    # Probably quite arbitrary, for backwards compatibility with scaling
    # yaml files produced with old Ewald Message Passing code
    k_offset = 0.1 if num_k_rbf <= 48 else 0.25

    # Evaluate a basis of Gaussian RBF on the k-vectors
    k_rbf_values = RadialBasis(
        num_radial=num_k_rbf,
        # Avoids zero or extremely small RBF values (there are k-points until
        # right at the cutoff, where all RBF would otherwise drop to 0)
        cutoff=k_cutoff + k_offset,
        rbf={"name": "gaussian"},
        envelope={"name": "polynomial", "exponent": 5},
    )(
        torch.norm(k_grid, dim=-1)
    )  # Tensor of shape (N_k, N_RBF)

    num_k_degrees_of_freedom = k_rbf_values.shape[-1]

    return k_grid, k_rbf_values, num_k_degrees_of_freedom
//  [python]
def get_k_voxel_grid(k_cutoff, delta_k, num_k_rbf):

    # Get indices for a cube of k-lattice sites containing the cutoff sphere
    num_k = k_cutoff / delta_k
    k_index_product_set, _ = get_k_index_product_set(num_k, num_k, num_k)

    # Orthogonal k-space basis, norm delta_k
    k_cell = torch.tensor(
        [[delta_k, 0, 0], [0, delta_k, 0], [0, 0, delta_k]], dtype=torch.float
    )

    # Translate lattice indices into k-vectors
    k_grid = torch.matmul(k_index_product_set, k_cell)

    # Prune all k-vectors outside the cutoff sphere
    k_grid = k_grid[torch.sum(k_grid**2, dim=-1) <= k_cutoff**2]

    # Probably quite arbitrary, for backwards compatibility with scaling
    # yaml files produced with old Ewald Message Passing code
    k_offset = 0.1 if num_k_rbf <= 48 else 0.25

    # Evaluate a basis of Gaussian RBF on the k-vectors
    k_rbf_values = RadialBasis(
        num_radial=num_k_rbf,
        # Avoids zero or extremely small RBF values (there are k-points until
        # right at the cutoff, where all RBF would otherwise drop to 0)
        cutoff=k_cutoff + k_offset,
        rbf={"name": "gaussian"},
        envelope={"name": "polynomial", "exponent": 5},
    )(
        torch.linalg.norm(k_grid, dim=-1)
    )  # Tensor of shape (N_k, N_RBF)

    num_k_degrees_of_freedom = k_rbf_values.shape[-1]

    return k_grid, k_rbf_values, num_k_degrees_of_freedom
//  [python]
def test_norm_vector(self, device, dtype):
        def run_test_case(input, p, dim, keepdim):
            result = torch.norm(input, p, dim, keepdim)
            input_numpy = input.cpu().numpy()
            result_numpy = np.linalg.norm(input_numpy, p, dim, keepdim)

            msg = f'input.size()={input.size()}, ord={p}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            self.assertEqual(result, result_numpy, msg=msg)

            result_out = torch.empty_like(result)
            torch.norm(input, p, dim, keepdim, out=result_out)
            self.assertEqual(result, result_out, msg=msg)

        ord_vector = [0, 1, -1, 2, -2, 3, -3, 4.5, -4.5, inf, -inf, None]
        S = 10
        test_cases = [
            # input size, p settings, dim
            ((S, ), ord_vector, None),
            ((S, ), ord_vector, 0),
            ((S, S, S), ord_vector, 0),
            ((S, S, S), ord_vector, 1),
            ((S, S, S), ord_vector, 2),
            ((S, S, S), ord_vector, -1),
            ((S, S, S), ord_vector, -2),
        ]
        L = 1_000_000
        if dtype == torch.double:
            test_cases.append(((L, ), ord_vector, None))
        for keepdim in [True, False]:
            for input_size, ord_settings, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_settings:
                    run_test_case(input, ord, dim, keepdim)
//  [python]
def test_norm_vector(self, device, dtype):
        def run_test_case(input, p, dim, keepdim):
            result = torch.linalg.norm(input, ord, dim, keepdim)
            input_numpy = input.cpu().numpy()
            result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)

            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            self.assertEqual(result, result_numpy, msg=msg)

            result_out = torch.empty_like(result)
            torch.linalg.norm(input, ord, dim, keepdim, out=result_out)
            self.assertEqual(result, result_out, msg=msg)

        ord_vector = [0, 1, -1, 2, -2, 3, -3, 4.5, -4.5, inf, -inf, None]
        S = 10
        test_cases = [
            # input size, p settings, dim
            ((S, ), ord_vector, None),
            ((S, ), ord_vector, 0),
            ((S, S, S), ord_vector, 0),
            ((S, S, S), ord_vector, 1),
            ((S, S, S), ord_vector, 2),
            ((S, S, S), ord_vector, -1),
            ((S, S, S), ord_vector, -2),
        ]
        L = 1_000_000
        if dtype == torch.double:
            test_cases.append(((L, ), ord_vector, None))
        for keepdim in [True, False]:
            for input_size, ord_settings, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_settings:
                    run_test_case(input, ord, dim, keepdim)
//  [python]
def test_norm_matrix(self, device, dtype):
        def run_test_case(input, p, dim, keepdim):
            result = t.norm(input, ord, dim, keepdim)
            input_numpy = input.cpu().numpy()
            result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)

            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            self.assertEqual(result, result_numpy, msg=msg)

            result_out = torch.empty_like(result)
            t.norm(input, ord, dim, keepdim, out=result_out)
            self.assertEqual(result, result_out, msg=msg)

        ord_matrix = [1, -1, 2, -2, inf, -inf, 'nuc', 'fro', None]
        S = 10
        test_cases = [
            # input size, p settings, dim
            ((S, S), ord_matrix, None),
            ((S, S), ord_matrix, (0, 1)),
            ((S, S), ord_matrix, (1, 0)),
            ((S, S, S, S), ord_matrix, (2, 0)),
            ((S, S, S, S), ord_matrix, (-1, -2)),
            ((S, S, S, S), ord_matrix, (-1, -3)),
            ((S, S, S, S), ord_matrix, (-3, 2)),
        ]
        L = 1_000
        if dtype == torch.double:
            test_cases.append(((L, L), ord_matrix, None))
        for keepdim in [True, False]:
            for input_size, ord_settings, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_settings:
                    run_test_case(input, ord, dim, keepdim)