jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self,
              mean,
              cov=None,
              *,
              decayscales,
              render_color=True):
    """Run MLP. mean is [D, *batch, 3] and cov is [D, *batch, 3, 3]."""
    # NOTE(jainajay): cov is ignored as there IPE for grid embeddings
    # is not implemented.
    batch_shape = mean.shape[:-1]
    num_points = np.prod(batch_shape)
    assert mean.shape[-1] == 3

    x = mean.reshape(num_points, 3)

    # Density network.
    emb_density = neural_fields.mask_grid_features(
        self.emb_density(x), decayscales["density"],
        self.emb_density_config["n_features_per_level"])
    features = self.layers_density(emb_density)
    d, features = features[Ellipsis, :1], features[Ellipsis, 1:]
    assert d.dtype == torch.float32

    # Parameterize canonical frame VolSDF distance as offset from sphere.
    if self.parameterization == "volsdf" and self.base_radius is not None:
      d = d + torch.linalg.norm(x, dim=-1, keepdim=True) - self.base_radius

    # Appearance network.
    if render_color:
      emb_color = neural_fields.mask_grid_features(
          self.emb_color(x), decayscales["color"],
          self.emb_color_config["n_features_per_level"])
      features = torch.cat([emb_color, features], dim=-1)
      raw_rgb = self.layers_color(features)
      raw_output = torch.cat([raw_rgb, d], dim=-1)
      return raw_output.view(*batch_shape, 4)

    return d.view(*batch_shape)
//  [python]
def quaternion_to_rotation_matrix_torch(q):
    """
    Convert a quaternion into a full three-dimensional rotation matrix.

    Input:
    :param q: A tensor of size (B, 4), where B is batch size and quaternion is in format (x, y, z, w).

    Output:
    :return: A tensor of size (B, 3, 3), where B is batch size.
    """
    # Ensure quaternion has four components
    assert q.shape[-1] == 4, "Input quaternion should have 4 components!"

    x, y, z, w = q.unbind(-1)

    # Compute quaternion norms
    q_norm = torch.norm(q, dim=-1, keepdim=True)
    # Normalize input quaternions
    q = q / q_norm

    # Compute the quaternion outer product
    q_outer = torch.einsum('...i,...j->...ij', q, q)

    # Compute rotation matrix
    rot_matrix = torch.empty(
        (*q.shape[:-1], 3, 3), dtype=q.dtype, device=q.device)
    rot_matrix[..., 0, 0] = 1 - 2 * (y**2 + z**2)
    rot_matrix[..., 0, 1] = 2 * (x*y - z*w)
    rot_matrix[..., 0, 2] = 2 * (x*z + y*w)
    rot_matrix[..., 1, 0] = 2 * (x*y + z*w)
    rot_matrix[..., 1, 1] = 1 - 2 * (x**2 + z**2)
    rot_matrix[..., 1, 2] = 2 * (y*z - x*w)
    rot_matrix[..., 2, 0] = 2 * (x*z - y*w)
    rot_matrix[..., 2, 1] = 2 * (y*z + x*w)
    rot_matrix[..., 2, 2] = 1 - 2 * (x**2 + y**2)

    return rot_matrix
//  [python]
def quaternion_to_rotation_matrix_torch(q):
    """
    Convert a quaternion into a full three-dimensional rotation matrix.

    Input:
    :param q: A tensor of size (B, 4), where B is batch size and quaternion is in format (x, y, z, w).

    Output:
    :return: A tensor of size (B, 3, 3), where B is batch size.
    """
    # Ensure quaternion has four components
    assert q.shape[-1] == 4, "Input quaternion should have 4 components!"

    x, y, z, w = q.unbind(-1)

    # Compute quaternion norms
    q_norm = torch.linalg.norm(q, dim=-1, keepdim=True)
    # Normalize input quaternions
    q = q / q_norm

    # Compute the quaternion outer product
    q_outer = torch.einsum('...i,...j->...ij', q, q)

    # Compute rotation matrix
    rot_matrix = torch.empty(
        (*q.shape[:-1], 3, 3), dtype=q.dtype, device=q.device)
    rot_matrix[..., 0, 0] = 1 - 2 * (y**2 + z**2)
    rot_matrix[..., 0, 1] = 2 * (x*y - z*w)
    rot_matrix[..., 0, 2] = 2 * (x*z + y*w)
    rot_matrix[..., 1, 0] = 2 * (x*y + z*w)
    rot_matrix[..., 1, 1] = 1 - 2 * (x**2 + z**2)
    rot_matrix[..., 1, 2] = 2 * (y*z - x*w)
    rot_matrix[..., 2, 0] = 2 * (x*z - y*w)
    rot_matrix[..., 2, 1] = 2 * (y*z + x*w)
    rot_matrix[..., 2, 2] = 1 - 2 * (x**2 + y**2)

    return rot_matrix
//  [python]
def forward(self, x):
        # x: [B, T, C]
        bs, seq_len = x.shape[:2]
        x = x.transpose(1, 2)

        x = self.inp_layer(x)
        for n in range(2):
            x = (x + self.enc_layers[2 * n](x))
            x = self.enc_layers[2 * n + 1](x)

        contact = self.contact_layer(x).sigmoid()
        assert contact.shape[1] == self.num_contact
        contact = contact.transpose(1, 2)

        residual_gravity = self.gravity_layer(x).tanh()
        assert residual_gravity.shape[1] == 3
        residual_gravity = residual_gravity.transpose(1, 2)

        gravity = torch.zeros_like(residual_gravity)
        gravity[:, :, 1] = self.g_const
        gravity = gravity + residual_gravity
        gravity = self.g_const * gravity / torch.norm(gravity, p=2, dim=-1, keepdim=True)

        inertia_force = self.inertia_force_layer(x)
        assert inertia_force.shape[1] == 3
        inertia_angular = self.inertia_angular_layer(x).tanh() * math.pi * 0.5
        assert inertia_angular.shape[1] == 3

        inertia_force = inertia_force.transpose(1, 2)
        inertia_force = self.inertia_const * inertia_force.reshape(bs, seq_len, 3, 1).tanh()
        inertia_angular = inertia_angular.transpose(1, 2)

        force = self.force_layer(x)
        force = force.transpose(1, 2).contiguous()
        assert force.shape[2] == self.num_contact * 3
        force = force.reshape(bs, seq_len, self.num_contact, 3).clone()
        force[:, :, :, 1] = -self.act(force[:, :, :, 1].clone()) * self.w_y

        # inertia_force[:, :, 2] = inertia_force[:, :, 2] * 5

        # force = self.force_layer(x)
        # force = force.transpose(1, 2).contiguous()
        # assert force.shape[2] == self.num_contact * 3
        # force = force.reshape(bs, seq_len, self.num_contact, 3).clone()
        # force[:, :, :, 1] = -self.act(force[:, :, :, 1].clone()) * self.w_y

        return contact, gravity[:, :, :, None], inertia_force, inertia_angular, force
//  [python]
def forward(self, x):
        # x: [B, T, C]
        bs, seq_len = x.shape[:2]
        x = x.transpose(1, 2)

        x = self.inp_layer(x)
        for n in range(2):
            x = (x + self.enc_layers[2 * n](x))
            x = self.enc_layers[2 * n + 1](x)

        contact = self.contact_layer(x).sigmoid()
        assert contact.shape[1] == self.num_contact
        contact = contact.transpose(1, 2)

        residual_gravity = self.gravity_layer(x).tanh()
        assert residual_gravity.shape[1] == 3
        residual_gravity = residual_gravity.transpose(1, 2)

        gravity = torch.zeros_like(residual_gravity)
        gravity[:, :, 1] = self.g_const
        gravity = gravity + residual_gravity
        gravity = self.g_const * gravity / torch.linalg.norm(gravity, p=2, dim=-1, keepdim=True)

        inertia_force = self.inertia_force_layer(x)
        assert inertia_force.shape[1] == 3
        inertia_angular = self.inertia_angular_layer(x).tanh() * math.pi * 0.5
        assert inertia_angular.shape[1] == 3

        inertia_force = inertia_force.transpose(1, 2)
        inertia_force = self.inertia_const * inertia_force.reshape(bs, seq_len, 3, 1).tanh()
        inertia_angular = inertia_angular.transpose(1, 2)

        force = self.force_layer(x)
        force = force.transpose(1, 2).contiguous()
        assert force.shape[2] == self.num_contact * 3
        force = force.reshape(bs, seq_len, self.num_contact, 3).clone()
        force[:, :, :, 1] = -self.act(force[:, :, :, 1].clone()) * self.w_y

        # inertia_force[:, :, 2] = inertia_force[:, :, 2] * 5

        # force = self.force_layer(x)
        # force = force.transpose(1, 2).contiguous()
        # assert force.shape[2] == self.num_contact * 3
        # force = force.reshape(bs, seq_len, self.num_contact, 3).clone()
        # force[:, :, :, 1] = -self.act(force[:, :, :, 1].clone()) * self.w_y

        return contact, gravity[:, :, :, None], inertia_force, inertia_angular, force
//  [python]
def forward(self, pred_batch, label_batch):
        diff = torch.norm(pred_batch - label_batch, p=2, dim=-1)
        if self.reduction == 'mean':
            metric = diff.mean()
        elif self.reduction == 'sum':
            metric = diff.sum()
        else:
            metric = diff
        return metric
//  [python]
def forward(self, pred_batch, label_batch):
        diff = torch.linalg.norm(pred_batch - label_batch, ord=2, dim=-1)
        if self.reduction == 'mean':
            metric = diff.mean()
        elif self.reduction == 'sum':
            metric = diff.sum()
        else:
            metric = diff
        return metric
//  [python]
def test_so3_cube_conv(self):
        g = rot3dOnR3()

        # cl = FourierELU(g, 3, g.fibergroup.bl_irreps(2), type='thomson_cube', N=1)
        # cl = FourierELU(g, 3, g.fibergroup.bl_irreps(2), type='cube', N=1)
        cl = QuotientFourierELU(g, (False, -1), 3, g.fibergroup.bl_irreps(2),
                                grid=g.fibergroup.sphere_grid(type='thomson_cube', N=1))
        # cl = FourierELU(g, 3, g.fibergroup.bl_irreps(2), type='thomson_cube', N=1)

        t = g.type(g.trivial_repr)

        conv1 = R3PointConv(t, cl.in_type, bias=False, width=2, n_rings=3)
        conv2 = R3PointConv(cl.out_type, t, bias=False, width=2, n_rings=3)

        P = 30
        atol = 1e-4
        rtol = 1e-5

        for _ in range(8):

            conv1.weights.data.normal_()
            # init.generalized_he_init(conv1.weights.data, conv1.basissampler)
            conv2.weights.data.normal_()
            # init.generalized_he_init(conv2.weights.data, conv2.basissampler)
            conv1.eval()
            conv2.eval()

            pos = torch.randn(P, 3)
            x = torch.randn(P, t.size)
            x = GeometricTensor(x, t, pos)

            distance = torch.norm(pos.unsqueeze(1) - pos, dim=2, keepdim=False)
            thr = sorted(distance.view(-1).tolist())[int(P ** 2 // 16)]
            edge_index = torch.nonzero(distance < thr).T.contiguous()

            for i, el in enumerate(g.fibergroup.grid('cube')):

                out1 = (conv2(cl(conv1(x, edge_index)), edge_index)).transform(el).tensor.detach().numpy()
                out2 = (conv2(cl(conv1(x.transform(el), edge_index)), edge_index)).tensor.detach().numpy()

                out1 = out1.mean(0)
                out2 = out2.mean(0)

                errs = np.abs(out1 - out2)

                esum = np.maximum(np.abs(out1), np.abs(out2))
                esum[esum <1e-7] = 1

                tol = rtol * esum + atol

                self.assertTrue(
                    np.all(errs < tol),
                    'The error found during equivariance check with element "{}" is too high: max = {}, mean = {} var ={}'
                    .format(el, errs.max(), errs.mean(), errs.var())
                )
//  [python]
def test_so3_cube_conv(self):
        g = rot3dOnR3()

        # cl = FourierELU(g, 3, g.fibergroup.bl_irreps(2), type='thomson_cube', N=1)
        # cl = FourierELU(g, 3, g.fibergroup.bl_irreps(2), type='cube', N=1)
        cl = QuotientFourierELU(g, (False, -1), 3, g.fibergroup.bl_irreps(2),
                                grid=g.fibergroup.sphere_grid(type='thomson_cube', N=1))
        # cl = FourierELU(g, 3, g.fibergroup.bl_irreps(2), type='thomson_cube', N=1)

        t = g.type(g.trivial_repr)

        conv1 = R3PointConv(t, cl.in_type, bias=False, width=2, n_rings=3)
        conv2 = R3PointConv(cl.out_type, t, bias=False, width=2, n_rings=3)

        P = 30
        atol = 1e-4
        rtol = 1e-5

        for _ in range(8):

            conv1.weights.data.normal_()
            # init.generalized_he_init(conv1.weights.data, conv1.basissampler)
            conv2.weights.data.normal_()
            # init.generalized_he_init(conv2.weights.data, conv2.basissampler)
            conv1.eval()
            conv2.eval()

            pos = torch.randn(P, 3)
            x = torch.randn(P, t.size)
            x = GeometricTensor(x, t, pos)

            distance = torch.linalg.norm(pos.unsqueeze(1) - pos, dim=2, keepdim=False)
            thr = sorted(distance.view(-1).tolist())[int(P ** 2 // 16)]
            edge_index = torch.nonzero(distance < thr).T.contiguous()

            for i, el in enumerate(g.fibergroup.grid('cube')):

                out1 = (conv2(cl(conv1(x, edge_index)), edge_index)).transform(el).tensor.detach().numpy()
                out2 = (conv2(cl(conv1(x.transform(el), edge_index)), edge_index)).tensor.detach().numpy()

                out1 = out1.mean(0)
                out2 = out2.mean(0)

                errs = np.abs(out1 - out2)

                esum = np.maximum(np.abs(out1), np.abs(out2))
                esum[esum <1e-7] = 1

                tol = rtol * esum + atol

                self.assertTrue(
                    np.all(errs < tol),
                    'The error found during equivariance check with element "{}" is too high: max = {}, mean = {} var ={}'
                    .format(el, errs.max(), errs.mean(), errs.var())
                )
//  [python]
def normalize_positions(pos, faces=None, method='mean', scale_method='max_rad'):
    # center and unit-scale positions

    if method == 'mean':
        # center using the average point position
        pos = (pos - torch.mean(pos, dim=-2, keepdim=True))
    elif method == 'bbox': 
        # center via the middle of the axis-aligned bounding box
        bbox_min = torch.min(pos, dim=-2).values
        bbox_max = torch.max(pos, dim=-2).values
        center = (bbox_max + bbox_min) / 2.
        pos -= center.unsqueeze(-2)
    else:
        raise ValueError("unrecognized method")

    if scale_method == 'max_rad':
        scale = torch.max(norm(pos), dim=-1, keepdim=True).values.unsqueeze(-1)
        pos = pos / scale
    elif scale_method == 'area': 
        if faces is None:
            raise ValueError("must pass faces for area normalization")
        coords = pos[faces]
        vec_A = coords[:, 1, :] - coords[:, 0, :]
        vec_B = coords[:, 2, :] - coords[:, 0, :]
        face_areas = torch.norm(torch.cross(vec_A, vec_B, dim=-1), dim=1) * 0.5
        total_area = torch.sum(face_areas)
        scale = (1. / torch.sqrt(total_area))
        pos = pos * scale
    else:
        raise ValueError("unrecognized scale method")
    return pos
//  [python]
def normalize_positions(pos, faces=None, method='mean', scale_method='max_rad'):
    # center and unit-scale positions

    if method == 'mean':
        # center using the average point position
        pos = (pos - torch.mean(pos, dim=-2, keepdim=True))
    elif method == 'bbox': 
        # center via the middle of the axis-aligned bounding box
        bbox_min = torch.min(pos, dim=-2).values
        bbox_max = torch.max(pos, dim=-2).values
        center = (bbox_max + bbox_min) / 2.
        pos -= center.unsqueeze(-2)
    else:
        raise ValueError("unrecognized method")

    if scale_method == 'max_rad':
        scale = torch.max(torch.linalg.norm(pos, dim=-1, keepdim=True), dim=-1, keepdim=True).values.unsqueeze(-1)
        pos = pos / scale
    elif scale_method == 'area': 
        if faces is None:
            raise ValueError("must pass faces for area normalization")
        coords = pos[faces]
        vec_A = coords[:, 1, :] - coords[:, 0, :]
        vec_B = coords[:, 2, :] - coords[:, 0, :]
        face_areas = torch.linalg.norm(torch.cross(vec_A, vec_B, dim=-1), dim=1) * 0.5
        total_area = torch.sum(face_areas)
        scale = (1. / torch.sqrt(total_area))
        pos = pos * scale
    else:
        raise ValueError("unrecognized scale method")
    return pos
//  [python]
def ev_zero(self, Kp_fr, att_fr):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        new_t = torch.sum(Kp_fr * pconf2, dim=1).view(1, 3).contiguous()

        kp_dis = torch.norm(new_t.view(-1))

        new_t *= 1000.0
        return new_t.detach().cpu().numpy()[0], att_fr, kp_dis.item()
//  [python]
def ev_zero(self, Kp_fr, att_fr):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        new_t = torch.sum(Kp_fr * pconf2, dim=1).view(1, 3).contiguous()

        kp_dis = torch.linalg.norm(new_t.view(-1))

        new_t *= 1000.0
        return new_t.detach().cpu().numpy()[0], att_fr, kp_dis.item()
//  [python]
def inf(self, Kp_fr, Kp_to):
        ori_Kp_to = Kp_to

        new_r, new_t = self.estimate_pose(Kp_fr, Kp_to)

        Kp_to = torch.bmm((ori_Kp_to - new_t), new_r)

        Kp_dis = torch.mean(torch.norm((Kp_fr - Kp_to), dim=2), dim=1)

        new_t *= 1000.0
        return new_r.detach().cpu().numpy()[0], new_t.detach().cpu().numpy()[0], Kp_dis.item()
//  [python]
def inf(self, Kp_fr, Kp_to):
        ori_Kp_to = Kp_to

        new_r, new_t = self.estimate_pose(Kp_fr, Kp_to)

        Kp_to = torch.bmm((ori_Kp_to - new_t), new_r)

        Kp_dis = torch.mean(torch.linalg.norm((Kp_fr - Kp_to), dim=2), dim=1)

        new_t *= 1000.0
        return new_r.detach().cpu().numpy()[0], new_t.detach().cpu().numpy()[0], Kp_dis.item()
//  [python]
def add_batch(self, inp, out):
        if self.metric_type == 'Weight':
            return

        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if len(inp.shape) == 3:
            inp = inp.reshape((-1, inp.shape[-1]))
        inp = inp.t()

        if self.metric_type == 'Wanda':
            self.scaler_row *= self.nsamples / (self.nsamples+tmp)
        elif self.metric_type == 'SparseGPT' or self.metric_type == 'SparseGPT-Git':
            self.Hessian *= self.nsamples / (self.nsamples + tmp)
        else:
            raise NotImplementedError(f"Invalid metric type: {self.metric_type}")

        self.nsamples += tmp

        if self.metric_type == 'Wanda':
            self.scaler_row += torch.norm(inp.float(), p=2, dim=1) ** 2  / self.nsamples
        elif self.metric_type == 'SparseGPT' or self.metric_type == 'SparseGPT-Git':
            inp = math.sqrt(2 / self.nsamples) * inp.float()
            self.Hessian += inp.matmul(inp.t())
        else:
            raise NotImplementedError(f"Invalid metric type: {self.metric_type}")
//  [python]
def add_batch(self, inp, out):
        if self.metric_type == 'Weight':
            return

        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if len(inp.shape) == 3:
            inp = inp.reshape((-1, inp.shape[-1]))
        inp = inp.t()

        if self.metric_type == 'Wanda':
            self.scaler_row *= self.nsamples / (self.nsamples+tmp)
        elif self.metric_type == 'SparseGPT' or self.metric_type == 'SparseGPT-Git':
            self.Hessian *= self.nsamples / (self.nsamples + tmp)
        else:
            raise NotImplementedError(f"Invalid metric type: {self.metric_type}")

        self.nsamples += tmp

        if self.metric_type == 'Wanda':
            self.scaler_row += torch.linalg.norm(inp.float(), ord=2, dim=1) ** 2  / self.nsamples
        elif self.metric_type == 'SparseGPT' or self.metric_type == 'SparseGPT-Git':
            inp = math.sqrt(2 / self.nsamples) * inp.float()
            self.Hessian += inp.matmul(inp.t())
        else:
            raise NotImplementedError(f"Invalid metric type: {self.metric_type}")
//  [python]
def ev(self, Kp_fr, Kp_to, att_to):
        ori_Kp_fr = Kp_fr
        ori_Kp_to = Kp_to

        new_r, new_t = self.estimate_pose(Kp_fr, Kp_to)

        Kp_to = torch.bmm((ori_Kp_to - new_t), new_r)

        Kp_dis = torch.mean(torch.norm((Kp_fr - Kp_to), dim=2), dim=1)

        new_t *= 1000.0
        return ori_Kp_fr, new_r.detach().cpu().numpy()[0], new_t.detach().cpu().numpy()[0], Kp_dis.item(), att_to
//  [python]
def ev(self, Kp_fr, Kp_to, att_to):
        ori_Kp_fr = Kp_fr
        ori_Kp_to = Kp_to

        new_r, new_t = self.estimate_pose(Kp_fr, Kp_to)

        Kp_to = torch.bmm((ori_Kp_to - new_t), new_r)

        Kp_dis = torch.mean(torch.linalg.norm((Kp_fr - Kp_to), dim=2), dim=1)

        new_t *= 1000.0
        return ori_Kp_fr, new_r.detach().cpu().numpy()[0], new_t.detach().cpu().numpy()[0], Kp_dis.item(), att_to
//  [python]
def inf_zero(self, Kp_fr):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        new_t = torch.sum(Kp_fr * pconf2, dim=1).view(1, 3).contiguous()

        Kp_dis = torch.norm(new_t.view(-1))

        new_t *= 1000.0
        return new_t.detach().cpu().numpy()[0], Kp_dis.item()
//  [python]
def inf_zero(self, Kp_fr):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        new_t = torch.sum(Kp_fr * pconf2, dim=1).view(1, 3).contiguous()

        Kp_dis = torch.linalg.norm(new_t.view(-1))

        new_t *= 1000.0
        return new_t.detach().cpu().numpy()[0], Kp_dis.item()
//  [python]
def backward(self, grad_output):
        grad_input = grad_output.clone()
        grad_input = grad_input / (torch.norm(grad_input, keepdim=True) + torch.finfo(grad_input.dtype).eps)
        return grad_input, None
//  [python]
def backward(self, grad_output):
        grad_input = grad_output.clone()
        grad_input = grad_input / (torch.linalg.norm(grad_input, keepdim=True) + torch.finfo(grad_input.dtype).eps)
        return grad_input, None
//  [python]
def change_to_ver(self, Kp):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        cent0 = torch.sum(Kp * pconf2, dim=1).view(-1).contiguous()
        
        num_kp = self.num_key
        ver_Kp_1 = Kp[:, :, 1].view(1, num_kp, 1).contiguous()

        kk_1 = Kp[:, :, 0].view(1, num_kp, 1).contiguous()
        kk_2 = Kp[:, :, 2].view(1, num_kp, 1).contiguous()
        rad = torch.cat((kk_1, kk_2), dim=2).contiguous()
        ver_Kp_2 = torch.norm(rad, dim=2).view(1, num_kp, 1).contiguous()

        tmp_aim_0 = torch.cat((Kp[:, 1:, :], Kp[:, 0:1, :]), dim=1).contiguous()
        aim_0_x = tmp_aim_0[:, :, 0].view(-1).contiguous()
        aim_0_y = tmp_aim_0[:, :, 2].view(-1).contiguous()

        aim_1_x = Kp[:, :, 0].view(-1).contiguous()
        aim_1_y = Kp[:, :, 2].view(-1).contiguous()

        angle = torch.atan2(aim_1_y, aim_1_x) - torch.atan2(aim_0_y, aim_0_x)
        angle[angle < 0] += 2 * math.pi
        ver_Kp_3 = angle.view(1, num_kp, 1).contiguous() * 0.01

        ver_Kp = torch.cat((ver_Kp_1, ver_Kp_2, ver_Kp_3), dim=2).contiguous()

        return ver_Kp, cent0
//  [python]
def change_to_ver(self, Kp):
        pconf2 = self.pconf.view(1, self.num_key, 1)
        cent0 = torch.sum(Kp * pconf2, dim=1).view(-1).contiguous()
        
        num_kp = self.num_key
        ver_Kp_1 = Kp[:, :, 1].view(1, num_kp, 1).contiguous()

        kk_1 = Kp[:, :, 0].view(1, num_kp, 1).contiguous()
        kk_2 = Kp[:, :, 2].view(1, num_kp, 1).contiguous()
        rad = torch.cat((kk_1, kk_2), dim=2).contiguous()
        ver_Kp_2 = torch.linalg.norm(rad, dim=2).view(1, num_kp, 1).contiguous()

        tmp_aim_0 = torch.cat((Kp[:, 1:, :], Kp[:, 0:1, :]), dim=1).contiguous()
        aim_0_x = tmp_aim_0[:, :, 0].view(-1).contiguous()
        aim_0_y = tmp_aim_0[:, :, 2].view(-1).contiguous()

        aim_1_x = Kp[:, :, 0].view(-1).contiguous()
        aim_1_y = Kp[:, :, 2].view(-1).contiguous()

        angle = torch.atan2(aim_1_y, aim_1_x) - torch.atan2(aim_0_y, aim_0_x)
        angle[angle < 0] += 2 * math.pi
        ver_Kp_3 = angle.view(1, num_kp, 1).contiguous() * 0.01

        ver_Kp = torch.cat((ver_Kp_1, ver_Kp_2, ver_Kp_3), dim=2).contiguous()

        return ver_Kp, cent0
//  [python]
def test_normalize_embeddings() -> None:
    """Tests the correct computation of util.normalize_embeddings"""
    embedding_size = 100
    a = torch.tensor(np.random.randn(50, embedding_size))
    a_norm = util.normalize_embeddings(a)

    for embedding in a_norm:
        assert len(embedding) == embedding_size
        emb_norm = torch.norm(embedding)
        assert abs(emb_norm.item() - 1) < 0.0001
//  [python]
def test_normalize_embeddings() -> None:
    """Tests the correct computation of util.normalize_embeddings"""
    embedding_size = 100
    a = torch.tensor(np.random.randn(50, embedding_size))
    a_norm = util.normalize_embeddings(a)

    for embedding in a_norm:
        assert len(embedding) == embedding_size
        emb_norm = torch.linalg.norm(embedding)
        assert abs(emb_norm.item() - 1) < 0.0001
//  [python]
def _f2overlap(f, mode = 'l2'):
    # f: n_scan*c
    # normalization
    f = f/torch.norm(f,dim=1,keepdim=True)
    if mode == 'inner':
        overlap = (f @ f.T + 1)/2.0
    else:
        overlap = (2 - torch.norm(f[:,None,:] - f[None,:,:], dim = -1))/2.0
    # remove self-loops
    for i in range(overlap.shape[0]):
        overlap[i,i] = 0.0
    return overlap
//  [python]
def _f2overlap(f, mode = 'l2'):
    # f: n_scan*c
    # normalization
    f = f/torch.linalg.norm(f,dim=1,keepdim=True)
    if mode == 'inner':
        overlap = (f @ f.T + 1)/2.0
    else:
        overlap = (2 - torch.linalg.norm(f[:,None,:] - f[None,:,:], dim = -1))/2.0
    # remove self-loops
    for i in range(overlap.shape[0]):
        overlap[i,i] = 0.0
    return overlap
//  [python]
def _construct_quadratic(
        device: torch.device,
    ) -> Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]:
        data = torch.arange(10, dtype=torch.float, device=device)
        data /= torch.norm(data)
        model = nn.Sequential(
            nn.Linear(10, 1, bias=False),
            nn.Linear(1, 1, bias=False),
        ).to(device=device)
        model[0].weight.data.fill_(1.0)
        model[1].weight.data.fill_(1.0)
        model[1].requires_grad = False
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=device)
        return model, loss, data, target
//  [python]
def _construct_quadratic(
        device: torch.device,
    ) -> Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]:
        data = torch.arange(10, dtype=torch.float, device=device)
        data /= torch.linalg.norm(data)
        model = nn.Sequential(
            nn.Linear(10, 1, bias=False),
            nn.Linear(1, 1, bias=False),
        ).to(device=device)
        model[0].weight.data.fill_(1.0)
        model[1].weight.data.fill_(1.0)
        model[1].requires_grad = False
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=device)
        return model, loss, data, target
//  [python]
def _calc_PVE_(verts_gts,verts_preds):
    batch_PVE = torch.norm(verts_gts-verts_preds, p=2, dim=-1).mean(-1) *1000
    return batch_PVE
//  [python]
def _calc_PVE_(verts_gts,verts_preds):
    batch_PVE = torch.linalg.norm(verts_gts-verts_preds, ord=2, dim=-1).mean(-1) *1000
    return batch_PVE
//  [python]
def forward(self, x):
        r = torch.norm(x-self.z0,dim=-1,keepdim=True)

        alpha = torch.exp(self.log_alpha)
        h = 1/(alpha+r)
        beta = -alpha+torch.log(1+torch.exp(self.beta))

        z = x+beta*h*(x-self.z0)
        
        log_det = (self.d-1)*torch.log(1+beta*h)+torch.log(1+beta*h-beta*r/(alpha+r)**2)

        return z, log_det[:,0]
//  [python]
def forward(self, x):
        r = torch.linalg.norm(x-self.z0,dim=-1,keepdim=True)

        alpha = torch.exp(self.log_alpha)
        h = 1/(alpha+r)
        beta = -alpha+torch.log(1+torch.exp(self.beta))

        z = x+beta*h*(x-self.z0)
        
        log_det = (self.d-1)*torch.log(1+beta*h)+torch.log(1+beta*h-beta*r/(alpha+r)**2)

        return z, log_det[:,0]
//  [python]
def early_stop(**args):
    stop = False
    if (args['i'] % 10) == 0:
        T_delta, grad = args['T_delta'], args['grad']
        grad_norm = torch.norm(grad.detach(), dim=-1)
        small_grad = grad_norm < 0.00001
        dR, dt = T_delta.magnitude()
        small_step = ((dt < 0.00001)
                        & (dR < 0.0001))
        if torch.all(small_step | small_grad):
            stop = True
    return stop
//  [python]
def early_stop(**args):
    stop = False
    if (args['i'] % 10) == 0:
        T_delta, grad = args['T_delta'], args['grad']
        grad_norm = torch.linalg.norm(grad.detach(), dim=-1)
        small_grad = grad_norm < 0.00001
        dR, dt = T_delta.magnitude()
        small_step = ((dt < 0.00001)
                        & (dR < 0.0001))
        if torch.all(small_step | small_grad):
            stop = True
    return stop
//  [python]
def regularization(model: nn.Module, mode: str):
    regu, counter = 0, 0
    for name, param in model.named_parameters():
        if "mask_scores" in name:
            if mode == "l1":
                regu += torch.norm(torch.sigmoid(param), p=1) / param.numel()
            elif mode == "l0":
                regu += torch.sigmoid(param - 2 / 3 * np.log(0.1 / 1.1)).sum() / param.numel()
            else:
                ValueError("Don't know this mode.")
            counter += 1
    return regu / counter
//  [python]
def regularization(model: nn.Module, mode: str):
    regu, counter = 0, 0
    for name, param in model.named_parameters():
        if "mask_scores" in name:
            if mode == "l1":
                regu += torch.linalg.norm(torch.sigmoid(param), ord=1) / param.numel()
            elif mode == "l0":
                regu += torch.sigmoid(param - 2 / 3 * np.log(0.1 / 1.1)).sum() / param.numel()
            else:
                ValueError("Don't know this mode.")
            counter += 1
    return regu / counter
//  [python]
def normal_crit(ret, batch):
    surf_normal_pred = ret['surf_normal'][ret['surf_mask']]
    surf_normal = batch['normal'][ret['surf_mask']]

    viewdir = batch['ray_d'][ret['surf_mask']]
    weights = torch.sum(-surf_normal_pred * viewdir, dim=1)
    weights = torch.clamp(weights, min=0, max=1)**2

    norm = torch.norm(surf_normal, dim=1)
    norm[norm < 1e-8] = 1e-8
    surf_normal = surf_normal / norm[..., None]

    surf_normal_pred[:, 1:] = surf_normal_pred[:, 1:] * -1

    normal_loss = torch.norm(surf_normal_pred - surf_normal, dim=1)
    normal_loss = (weights * normal_loss).mean()

    return normal_loss
//  [python]
def normal_crit(ret, batch):
    surf_normal_pred = ret['surf_normal'][ret['surf_mask']]
    surf_normal = batch['normal'][ret['surf_mask']]

    viewdir = batch['ray_d'][ret['surf_mask']]
    weights = torch.sum(-surf_normal_pred * viewdir, dim=1)
    weights = torch.clamp(weights, min=0, max=1)**2

    norm = torch.linalg.norm(surf_normal, dim=1)
    norm[norm < 1e-8] = 1e-8
    surf_normal = surf_normal / norm[..., None]

    surf_normal_pred[:, 1:] = surf_normal_pred[:, 1:] * -1

    normal_loss = torch.linalg.norm(surf_normal_pred - surf_normal, dim=1)
    normal_loss = (weights * normal_loss).mean()

    return normal_loss
//  [python]
def average_points_relative_error(points,relative_pose):
    points_x_none_zero_id=points[:,0].nonzero().squeeze()
    points=points[points_x_none_zero_id,:]
    points_y_none_zero_id=points[:,1].nonzero().squeeze()
    points=points[points_y_none_zero_id,:]
    points_z_none_zero_id=points[:,2].nonzero().squeeze()
    points=points[points_z_none_zero_id,:]

    relative_pose_=relative_pose.unsqueeze(0).expand(points.size(0),4)
    relative_points=quat_ops.qrotv(relative_pose_, points)

    relative_error=abs(torch.norm((relative_points-points))/(torch.norm(points)+torch.norm(relative_points)+ 1e-9))

    return torch.mean(relative_error)
//  [python]
def average_points_relative_error(points,relative_pose):
    points_x_none_zero_id=points[:,0].nonzero().squeeze()
    points=points[points_x_none_zero_id,:]
    points_y_none_zero_id=points[:,1].nonzero().squeeze()
    points=points[points_y_none_zero_id,:]
    points_z_none_zero_id=points[:,2].nonzero().squeeze()
    points=points[points_z_none_zero_id,:]

    relative_pose_=relative_pose.unsqueeze(0).expand(points.size(0),4)
    relative_points=quat_ops.qrotv(relative_pose_, points)

    relative_error=abs(torch.linalg.norm((relative_points-points), dim=1)/(torch.linalg.norm(points, dim=1)+torch.linalg.norm(relative_points, dim=1)+ 1e-9))

    return torch.mean(relative_error)
//  [python]
def forward(self,
              mean,
              deformation_codes,
              *,
              cov=None,
              decayscales,
              render_deformation=True,
              render_color=True):
    """Run MLP. mean is [D, *batch, 3] and cov is [D, *batch, 3, 3]."""
    batch_shape = mean.shape[:-1]
    num_points = np.prod(batch_shape)
    assert mean.shape[-1] == 3

    x = mean.reshape(num_points, 3)

    # Positional embedding. NOTE(jainajay): decayscales is unused.
    emb = self.emb(x).type(torch.float32)
    emb_density, emb_color = emb[Ellipsis, 3:], emb[Ellipsis, :3]

    # Parameterize canonical frame VolSDF distance as offset from sphere.
    if self.parameterization == "volsdf" and self.base_radius is not None:
      # TODO(jainajay): explicitly redistance?
      d = emb_density + torch.norm(x, dim=-1, keepdim=True) - self.base_radius
    else:
      d = F.relu(emb_density)
    assert d.dtype == torch.float32

    # Canonical appearance network.
    if render_color:
      # TODO(jainajay): Concatenate density embedding?
      raw_rgb = F.relu(emb_color)
      raw_output = torch.cat([raw_rgb, d], dim=-1)
      return raw_output.view(*batch_shape, 4)

    return d.view(*batch_shape)
//  [python]
def forward(self,
              mean,
              deformation_codes,
              *,
              cov=None,
              decayscales,
              render_deformation=True,
              render_color=True):
    """Run MLP. mean is [D, *batch, 3] and cov is [D, *batch, 3, 3]."""
    batch_shape = mean.shape[:-1]
    num_points = np.prod(batch_shape)
    assert mean.shape[-1] == 3

    x = mean.reshape(num_points, 3)

    # Positional embedding. NOTE(jainajay): decayscales is unused.
    emb = self.emb(x).type(torch.float32)
    emb_density, emb_color = emb[Ellipsis, 3:], emb[Ellipsis, :3]

    # Parameterize canonical frame VolSDF distance as offset from sphere.
    if self.parameterization == "volsdf" and self.base_radius is not None:
      # TODO(jainajay): explicitly redistance?
      d = emb_density + torch.linalg.norm(x, dim=-1, keepdim=True) - self.base_radius
    else:
      d = F.relu(emb_density)
    assert d.dtype == torch.float32

    # Canonical appearance network.
    if render_color:
      # TODO(jainajay): Concatenate density embedding?
      raw_rgb = F.relu(emb_color)
      raw_output = torch.cat([raw_rgb, d], dim=-1)
      return raw_output.view(*batch_shape, 4)

    return d.view(*batch_shape)
//  [python]
def make_atom_graph( xyz, mask, num_bonds, top_k=16, maxbonds=4 ):
    B,L,A = xyz.shape[:3]
    device = xyz.device

    D = torch.norm(
        xyz[:,None,None,:,:] - xyz[:,:,:,None,None], dim=-1
    )
    mask2d = mask[:,:,:,None,None]*mask[:,None,None,:,:]
    D[~mask2d] = 9999.
    D[D==0] = 9999.

    # select top K neighbors for each atom
    # keep indices as batch/res/atm indices
    D_neigh, E_idx = torch.topk(D.reshape(B,L,A,-1), top_k, largest=False) # shape of E_idx: (B, L, top_k)
    Eres, Eatm = torch.div(E_idx,A,rounding_mode='trunc'), E_idx%A
    bi,ri,ai = mask.nonzero(as_tuple=True)
    bi = bi[:,None].repeat(1,top_k).reshape(-1)
    ri = ri[:,None].repeat(1,top_k).reshape(-1)
    ai = ai[:,None].repeat(1,top_k).reshape(-1)
    rj,aj = Eres[mask].reshape(-1), Eatm[mask].reshape(-1)

    # on each edge, 1-hot encode the number of bonds (up to maxbonds) seperating each atom
    edge = torch.full(ri.shape, maxbonds, device=device)
    resmask = ri==rj
    edge[resmask] = num_bonds[bi[resmask],ri[resmask],ai[resmask],aj[resmask]]-1
    resmask = ri+1==rj
    edge[resmask] = num_bonds[bi[resmask],ri[resmask],ai[resmask],2]+num_bonds[bi[resmask],rj[resmask],0,aj[resmask]]
    resmask = ri-1==rj
    edge[resmask] = num_bonds[bi[resmask],ri[resmask],ai[resmask],0]+num_bonds[bi[resmask],rj[resmask],2,aj[resmask]]
    edge = edge.clamp(0,maxbonds-1)
    edge = F.one_hot(edge)[...,None]

    natm = torch.sum(mask)
    index = torch.zeros_like(mask, dtype=torch.long, device=device)
    index[mask] = torch.arange(natm, device=device)
    src=index[bi,ri,ai]
    tgt=index[bi,rj,aj]
    
    G = dgl.graph((src, tgt), num_nodes=natm).to(device)
    G.edata['rel_pos'] = (xyz[bi,ri,ai] - xyz[bi,rj,aj]).detach() # no gradient through basis function

    return G, edge
//  [python]
def make_atom_graph( xyz, mask, num_bonds, top_k=16, maxbonds=4 ):
    B,L,A = xyz.shape[:3]
    device = xyz.device

    D = torch.linalg.norm(
        xyz[:,None,None,:,:] - xyz[:,:,:,None,None], dim=-1
    )
    mask2d = mask[:,:,:,None,None]*mask[:,None,None,:,:]
    D[~mask2d] = 9999.
    D[D==0] = 9999.

    # select top K neighbors for each atom
    # keep indices as batch/res/atm indices
    D_neigh, E_idx = torch.topk(D.reshape(B,L,A,-1), top_k, largest=False) # shape of E_idx: (B, L, top_k)
    Eres, Eatm = torch.div(E_idx,A,rounding_mode='trunc'), E_idx%A
    bi,ri,ai = mask.nonzero(as_tuple=True)
    bi = bi[:,None].repeat(1,top_k).reshape(-1)
    ri = ri[:,None].repeat(1,top_k).reshape(-1)
    ai = ai[:,None].repeat(1,top_k).reshape(-1)
    rj,aj = Eres[mask].reshape(-1), Eatm[mask].reshape(-1)

    # on each edge, 1-hot encode the number of bonds (up to maxbonds) seperating each atom
    edge = torch.full(ri.shape, maxbonds, device=device)
    resmask = ri==rj
    edge[resmask] = num_bonds[bi[resmask],ri[resmask],ai[resmask],aj[resmask]]-1
    resmask = ri+1==rj
    edge[resmask] = num_bonds[bi[resmask],ri[resmask],ai[resmask],2]+num_bonds[bi[resmask],rj[resmask],0,aj[resmask]]
    resmask = ri-1==rj
    edge[resmask] = num_bonds[bi[resmask],ri[resmask],ai[resmask],0]+num_bonds[bi[resmask],rj[resmask],2,aj[resmask]]
    edge = edge.clamp(0,maxbonds-1)
    edge = F.one_hot(edge)[...,None]

    natm = torch.sum(mask)
    index = torch.zeros_like(mask, dtype=torch.long, device=device)
    index[mask] = torch.arange(natm, device=device)
    src=index[bi,ri,ai]
    tgt=index[bi,rj,aj]
    
    G = dgl.graph((src, tgt), num_nodes=natm).to(device)
    G.edata['rel_pos'] = (xyz[bi,ri,ai] - xyz[bi,rj,aj]).detach() # no gradient through basis function

    return G, edge
//  [python]
def get_edge_length(v, e):
    v0 = linear_gather(v, e[..., 0], dim=-2)
    v1 = linear_gather(v, e[..., 1], dim=-2)
    return torch.norm(v0 - v1, dim=-1)
//  [python]
def get_edge_length(v, e):
    v0 = linear_gather(v, e[..., 0], dim=-2)
    v1 = linear_gather(v, e[..., 1], dim=-2)
    return torch.linalg.norm(v0 - v1, dim=-1)
//  [python]
def axisangle2quat(vec, eps=1e-6):
    """
    Converts scaled axis-angle to quat.
    Args:
        vec (tensor): (..., 3) tensor where final dim is (ax,ay,az) axis-angle exponential coordinates
        eps (float): Stability value below which small values will be mapped to 0

    Returns:
        tensor: (..., 4) tensor where final dim is (x,y,z,w) vec4 float quaternion
    """
    # type: (Tensor, float) -> Tensor
    # store input shape and reshape
    input_shape = vec.shape[:-1]
    vec = vec.reshape(-1, 3)

    # Grab angle
    angle = torch.norm(vec, dim=-1, keepdim=True)

    # Create return array
    quat = torch.zeros(torch.prod(torch.tensor(input_shape)), 4, device=vec.device)
    quat[:, 3] = 1.0

    # Grab indexes where angle is not zero an convert the input to its quaternion form
    idx = angle.reshape(-1) > eps
    quat[idx, :] = torch.cat([
        vec[idx, :] * torch.sin(angle[idx, :] / 2.0) / angle[idx, :],
        torch.cos(angle[idx, :] / 2.0)
    ], dim=-1)

    # Reshape and return output
    quat = quat.reshape(list(input_shape) + [4, ])
    return quat
//  [python]
def axisangle2quat(vec, eps=1e-6):
    """
    Converts scaled axis-angle to quat.
    Args:
        vec (tensor): (..., 3) tensor where final dim is (ax,ay,az) axis-angle exponential coordinates
        eps (float): Stability value below which small values will be mapped to 0

    Returns:
        tensor: (..., 4) tensor where final dim is (x,y,z,w) vec4 float quaternion
    """
    # type: (Tensor, float) -> Tensor
    # store input shape and reshape
    input_shape = vec.shape[:-1]
    vec = vec.reshape(-1, 3)

    # Grab angle
    angle = torch.linalg.norm(vec, dim=-1, keepdim=True)

    # Create return array
    quat = torch.zeros(torch.prod(torch.tensor(input_shape)), 4, device=vec.device)
    quat[:, 3] = 1.0

    # Grab indexes where angle is not zero an convert the input to its quaternion form
    idx = angle.reshape(-1) > eps
    quat[idx, :] = torch.cat([
        vec[idx, :] * torch.sin(angle[idx, :] / 2.0) / angle[idx, :],
        torch.cos(angle[idx, :] / 2.0)
    ], dim=-1)

    # Reshape and return output
    quat = quat.reshape(list(input_shape) + [4, ])
    return quat
//  [python]
def compute(self, X, Y):
        """Compute adaptation loss

        Arguments:
            X {tensor} -- source matrix
            Y {tensor} -- target matrix

        Returns:
            [tensor] -- transfer loss
        """
        if self.loss_type == 'mmd_lin' or self.loss_type =='mmd':
            mmdloss = mmd.MMD_loss(kernel_type='linear')
            loss = mmdloss(X, Y)
        elif self.loss_type == 'coral':
            loss = CORAL(X, Y)
        elif self.loss_type == 'cosine' or self.loss_type == 'cos':
            loss = 1 - cosine(X, Y)
        elif self.loss_type == 'kl':
            loss = kl_js.kl_div(X, Y)
        elif self.loss_type == 'js':
            loss = kl_js.js(X, Y)
        elif self.loss_type == 'mine':
            mine_model = mutual_info.Mine_estimator(
                input_dim=self.input_dim, hidden_dim=60).cuda()
            loss = mine_model(X, Y)
        elif self.loss_type == 'adv':
            loss = adv_loss.adv(X, Y, input_dim=self.input_dim, hidden_dim=32)
        elif self.loss_type == 'mmd_rbf':
            mmdloss = mmd.MMD_loss(kernel_type='rbf')
            loss = mmdloss(X, Y)
        elif self.loss_type == 'pairwise':
            pair_mat = pairwise_dist(X, Y)
            import torch
            loss = torch.norm(pair_mat)

        return loss
//  [python]
def compute(self, X, Y):
        """Compute adaptation loss

        Arguments:
            X {tensor} -- source matrix
            Y {tensor} -- target matrix

        Returns:
            [tensor] -- transfer loss
        """
        if self.loss_type == 'mmd_lin' or self.loss_type =='mmd':
            mmdloss = mmd.MMD_loss(kernel_type='linear')
            loss = mmdloss(X, Y)
        elif self.loss_type == 'coral':
            loss = CORAL(X, Y)
        elif self.loss_type == 'cosine' or self.loss_type == 'cos':
            loss = 1 - cosine(X, Y)
        elif self.loss_type == 'kl':
            loss = kl_js.kl_div(X, Y)
        elif self.loss_type == 'js':
            loss = kl_js.js(X, Y)
        elif self.loss_type == 'mine':
            mine_model = mutual_info.Mine_estimator(
                input_dim=self.input_dim, hidden_dim=60).cuda()
            loss = mine_model(X, Y)
        elif self.loss_type == 'adv':
            loss = adv_loss.adv(X, Y, input_dim=self.input_dim, hidden_dim=32)
        elif self.loss_type == 'mmd_rbf':
            mmdloss = mmd.MMD_loss(kernel_type='rbf')
            loss = mmdloss(X, Y)
        elif self.loss_type == 'pairwise':
            pair_mat = pairwise_dist(X, Y)
            import torch
            loss = torch.linalg.norm(pair_mat)

        return loss
//  [python]
def get_face_areas(v: torch.Tensor, f: torch.Tensor) -> torch.Tensor:
    v0 = v[f[..., 0]]
    v1 = v[f[..., 1]]
    v2 = v[f[..., 2]]

    u = v2 - v0
    v = v1 - v0

    return torch.norm(torch.cross(u, v), dim=-1) / 2.0
//  [python]
def get_face_areas(v: torch.Tensor, f: torch.Tensor) -> torch.Tensor:
    v0 = v[f[..., 0]]
    v1 = v[f[..., 1]]
    v2 = v[f[..., 2]]

    u = v2 - v0
    v = v1 - v0

    return torch.linalg.norm(torch.cross(u, v), dim=-1) / 2.0
//  [python]
def l2n(self, x, eps=1e-6):
        return x / (torch.norm(x, p=2, dim=1, keepdim=True) + eps).expand_as(x)
//  [python]
def l2n(self, x, eps=1e-6):
        return x / (torch.linalg.norm(x, ord=2, dim=1, keepdim=True) + eps).expand_as(x)
//  [python]
def loss(self):
        m = self.A.size(0)
        return lambda z: (1.0 / m) * torch.norm(torch.abs(self.A @ z) - self.y, 1)
//  [python]
def loss(self):
        m = self.A.size(0)
        return lambda z: (1.0 / m) * torch.linalg.norm(torch.abs(self.A @ z) - self.y, 1)
//  [python]
def dist_range(self, x):
        return torch.norm(torch.linalg.lstsq(self.A, self.y - self.A @ x).solution)
//  [python]
def dist_range(self, x):
        return torch.linalg.norm(torch.linalg.lstsq(self.A, self.y - self.A @ x).solution)
//  [python]
def loss(self):
        def f(z):
            return (1.0 / self.A.size()[0]) * torch.norm(
                torch.max(self.A @ (z), torch.zeros(self.A.size()[0])) - self.y, 1
            )

        return f
//  [python]
def loss(self):
        def f(z):
            return (1.0 / self.A.size()[0]) * torch.linalg.norm(
                torch.max(self.A @ (z), torch.zeros(self.A.size()[0])) - self.y, 1
            )

        return f
//  [python]
def initializer(self, δ):
        Δ = torch.randn(self.W.numel() + self.X.numel(), dtype=torch.double)
        return (
            torch.cat([self.W.reshape(-1), self.X.reshape(-1)]) + δ * Δ / torch.norm(Δ)
        ).requires_grad_(True)
//  [python]
def initializer(self, δ):
        Δ = torch.randn(self.W.numel() + self.X.numel(), dtype=torch.double)
        return (
            torch.cat([self.W.reshape(-1), self.X.reshape(-1)]) + δ * Δ / torch.linalg.norm(Δ)
        ).requires_grad_(True)
//  [python]
def quaternion_to_axis_angle(quaternions):
    """
    Convert rotations given as quaternions to axis/angle.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles])
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48)
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def quaternion_to_axis_angle(quaternions):
    """
    Convert rotations given as quaternions to axis/angle.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.linalg.norm(quaternions[..., 1:], ord=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles])
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48)
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def _converged(self, s, trust_radius):

        if abs(norm(s) - trust_radius) <= self.kappa_easy * trust_radius:
            return True
        else:
            return False
//  [python]
def _converged(self, s, trust_radius):

        if abs(torch.linalg.norm(s) - trust_radius) <= self.kappa_easy * trust_radius:
            return True
        else:
            return False
//  [python]
def _nu_next(self, _lambda, trust_radius, s, w):

        norm_s = norm(s)
        norm_w = norm(w)

        phi = 1 / norm_s - 1 / trust_radius

        phi_prime = norm_w ** 2 / norm_s ** 3

        return _lambda - phi / phi_prime
//  [python]
def _nu_next(self, _lambda, trust_radius, s, w):

        norm_s = torch.linalg.norm(s)
        norm_w = torch.linalg.norm(w)

        phi = 1 / norm_s - 1 / trust_radius

        phi_prime = norm_w ** 2 / norm_s ** 3

        return _lambda - phi / phi_prime