jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, x, fun, k_easy=0.1, k_hard=0.2):

        super().__init__(x, fun)

        # When the trust-region shrinks in two consecutive
        # calculations (``tr_radius < previous_tr_radius``)
        # the lower bound ``lambda_lb`` may be reused,
        # facilitating  the convergence. To indicate no
        # previous value is known at first ``previous_tr_radius``
        # is set to -1  and ``lambda_lb`` to None.
        self.previous_tr_radius = -1
        self.lambda_lb = None

        self.niter = 0
        self.EPS = torch.finfo(x.dtype).eps

        # ``k_easy`` and ``k_hard`` are parameters used
        # to determine the stop criteria to the iterative
        # subproblem solver. Take a look at pp. 194-197
        # from reference _[1] for a more detailed description.
        self.k_easy = k_easy
        self.k_hard = k_hard

        # Get Lapack function for cholesky decomposition.
        # NOTE: cholesky_ex requires pytorch >= 1.9.0
        if "cholesky_ex" in dir(torch.linalg):
            self.torch_cholesky = True
        else:
            # if we don't have torch cholesky, use potrf from scipy
            (self.cholesky,) = get_lapack_funcs(("potrf",), (self.hess.cpu().numpy(),))
            self.torch_cholesky = False

        # Get info about Hessian
        self.dimension = len(self.hess)
        self.hess_gershgorin_lb, self.hess_gershgorin_ub = gershgorin_bounds(self.hess)
        self.hess_inf = torch.norm(self.hess, float("inf"))
        self.hess_fro = torch.norm(self.hess, "fro")

        # A constant such that for vectors smaler than that
        # backward substituition is not reliable. It was stabilished
        # based on Golub, G. H., Van Loan, C. F. (2013).
        # "Matrix computations". Forth Edition. JHU press., p.165.
        self.CLOSE_TO_ZERO = self.dimension * self.EPS * self.hess_inf
//  [python]
def __init__(self, x, fun, k_easy=0.1, k_hard=0.2):

        super().__init__(x, fun)

        # When the trust-region shrinks in two consecutive
        # calculations (``tr_radius < previous_tr_radius``)
        # the lower bound ``lambda_lb`` may be reused,
        # facilitating  the convergence. To indicate no
        # previous value is known at first ``previous_tr_radius``
        # is set to -1  and ``lambda_lb`` to None.
        self.previous_tr_radius = -1
        self.lambda_lb = None

        self.niter = 0
        self.EPS = torch.finfo(x.dtype).eps

        # ``k_easy`` and ``k_hard`` are parameters used
        # to determine the stop criteria to the iterative
        # subproblem solver. Take a look at pp. 194-197
        # from reference _[1] for a more detailed description.
        self.k_easy = k_easy
        self.k_hard = k_hard

        # Get Lapack function for cholesky decomposition.
        # NOTE: cholesky_ex requires pytorch >= 1.9.0
        if "cholesky_ex" in dir(torch.linalg):
            self.torch_cholesky = True
        else:
            # if we don't have torch cholesky, use potrf from scipy
            (self.cholesky,) = get_lapack_funcs(("potrf",), (self.hess.cpu().numpy(),))
            self.torch_cholesky = False

        # Get info about Hessian
        self.dimension = len(self.hess)
        self.hess_gershgorin_lb, self.hess_gershgorin_ub = gershgorin_bounds(self.hess)
        self.hess_inf = norm(self.hess, float("inf"))
        self.hess_fro = norm(self.hess, "fro")

        # A constant such that for vectors smaler than that
        # backward substituition is not reliable. It was stabilished
        # based on Golub, G. H., Van Loan, C. F. (2013).
        # "Matrix computations". Forth Edition. JHU press., p.165.
        self.CLOSE_TO_ZERO = self.dimension * self.EPS * self.hess_inf
//  [python]
def add_edge_features(self):
        """
        adds relative displacement & displacement norm as edge features
        """
        for i in range(len(self.graphs)):
            pos = self.graphs[i].ndata["pos"]
            row, col = self.graphs[i].edges()
            disp = torch.tensor(pos[row.long()] - pos[col.long()])
            disp_norm = torch.norm(disp, dim=-1, keepdim=True)
            self.graphs[i].edata["x"] = torch.cat((disp, disp_norm), dim=-1)

        return self.graphs
//  [python]
def add_edge_features(self):
        """
        adds relative displacement & displacement norm as edge features
        """
        for i in range(len(self.graphs)):
            pos = self.graphs[i].ndata["pos"]
            row, col = self.graphs[i].edges()
            disp = torch.tensor(pos[row.long()] - pos[col.long()])
            disp_norm = torch.linalg.norm(disp, dim=-1, keepdim=True)
            self.graphs[i].edata["x"] = torch.cat((disp, disp_norm), dim=-1)

        return self.graphs
//  [python]
def solve(self, tr_radius):
        g = self.jac
        gamma_0 = self.jac_mag
        (n,) = g.shape
        m = n if self.max_lanczos is None else min(n, self.max_lanczos)
        dtype = g.dtype
        device = g.device
        h_best = None
        error_best = float("inf")

        # Lanczos Q matrix buffer
        Q = torch.zeros(m, n, dtype=dtype, device=device)
        Q[0] = g / gamma_0

        # Lanczos T matrix buffers
        # a and b are the diagonal and off-diagonal entries of T, respectively
        a = torch.zeros(m, dtype=dtype, device=device)
        b = torch.zeros(m, dtype=dtype, device=device)

        # first lanczos iteration
        r = self.hessp(Q[0])
        torch.dot(Q[0], r, out=a[0])
        r.sub_(Q[0] * a[0])
        b[0] = r.norm()
        if b[0] < self.eps:
            raise RuntimeError("initial beta is zero.")

        # remaining iterations
        for i in range(1, m):
            Q[i] = r / b[i - 1]
            r = self.hessp(Q[i])
            r.sub_(Q[i - 1] * b[i - 1])
            torch.dot(Q[i], r, out=a[i])
            r.sub_(Q[i] * a[i])
            if self.ortho:
                # Re-orthogonalize with Gram-Schmidt
                r.addmv_(Q[: i + 1].T, Q[: i + 1].mv(r), alpha=-1)
            b[i] = r.norm()
            if b[i] < self.eps:
                # This should never occur when self.ortho=True
                raise RuntimeError("reducible T matrix encountered.")

            # GLTR sub-problem
            h, status, lambd = self.tridiag_subproblem(a[: i + 1], b[:i], tr_radius)

            if status >= 0:
                # convergence check; see Algorithm 1 of [1]_ and
                # Algorithm 5.1 of [2]_. Equivalent to the following:
                #     p = Q[:i+1].T.mv(h)
                #     error = torch.linalg.norm(self.hessp(p) + lambd * p + g)
                error = b[i] * h[-1].abs()
                if self._debug:
                    print(
                        "iter %3d - status: %d - lambd: %0.4e - error: %0.4e"
                        % (i + 1, status, lambd, error)
                    )
                if error < error_best:
                    # we've found a new best
                    hits_boundary = status != 0
                    h_best = h
                    error_best = error
                    if error_best <= self.tol:
                        break

            elif self._debug:
                print("iter %3d - status: %d - lambd: %0.4e" % (i + 1, status, lambd))

        if h_best is None:
            # TODO: what should we do here?
            raise RuntimeError("gltr solution not found")

        # project h back to R^n
        p_best = Q[: i + 1].T.mv(h_best)

        return p_best, hits_boundary
//  [python]
def solve(self, tr_radius):
        g = self.jac
        gamma_0 = self.jac_mag
        (n,) = g.shape
        m = n if self.max_lanczos is None else min(n, self.max_lanczos)
        dtype = g.dtype
        device = g.device
        h_best = None
        error_best = float("inf")

        # Lanczos Q matrix buffer
        Q = torch.zeros(m, n, dtype=dtype, device=device)
        Q[0] = g / gamma_0

        # Lanczos T matrix buffers
        # a and b are the diagonal and off-diagonal entries of T, respectively
        a = torch.zeros(m, dtype=dtype, device=device)
        b = torch.zeros(m, dtype=dtype, device=device)

        # first lanczos iteration
        r = self.hessp(Q[0])
        torch.dot(Q[0], r, out=a[0])
        r.sub_(Q[0] * a[0])
        torch.linalg.norm(r, out=b[0])
        if b[0] < self.eps:
            raise RuntimeError("initial beta is zero.")

        # remaining iterations
        for i in range(1, m):
            torch.div(r, b[i - 1], out=Q[i])
            r = self.hessp(Q[i])
            r.sub_(Q[i - 1] * b[i - 1])
            torch.dot(Q[i], r, out=a[i])
            r.sub_(Q[i] * a[i])
            if self.ortho:
                # Re-orthogonalize with Gram-Schmidt
                r.addmv_(Q[: i + 1].T, Q[: i + 1].mv(r), alpha=-1)
            torch.linalg.norm(r, out=b[i])
            if b[i] < self.eps:
                # This should never occur when self.ortho=True
                raise RuntimeError("reducible T matrix encountered.")

            # GLTR sub-problem
            h, status, lambd = self.tridiag_subproblem(a[: i + 1], b[:i], tr_radius)

            if status >= 0:
                # convergence check; see Algorithm 1 of [1]_ and
                # Algorithm 5.1 of [2]_. Equivalent to the following:
                #     p = Q[:i+1].T.mv(h)
                #     error = torch.linalg.norm(self.hessp(p) + lambd * p + g)
                error = b[i] * h[-1].abs()
                if self._debug:
                    print(
                        "iter %3d - status: %d - lambd: %0.4e - error: %0.4e"
                        % (i + 1, status, lambd, error)
                    )
                if error < error_best:
                    # we've found a new best
                    hits_boundary = status != 0
                    h_best = h
                    error_best = error
                    if error_best <= self.tol:
                        break

            elif self._debug:
                print("iter %3d - status: %d - lambd: %0.4e" % (i + 1, status, lambd))

        if h_best is None:
            # TODO: what should we do here?
            raise RuntimeError("gltr solution not found")

        # project h back to R^n
        p_best = Q[: i + 1].T.mv(h_best)

        return p_best, hits_boundary
//  [python]
def plot_diagnostics(plot_window, wPCA, Wall0, clu0, settings):
    # Temporal features (top left)
    p1 = plot_window.plot_widget.addPlot(
        row=0, col=0, labels={'bottom': 'Time (s)'}
        )
    p1.setTitle('Temporal Features')
    t = np.arange(wPCA.shape[1])/(settings['fs']/1000)
    for i in range(wPCA.shape[0]):
        color = _COLOR_CODES[i % len(_COLOR_CODES)]
        p1.plot(t, wPCA[i,:], pen=color)

    # Spatial features (top right)
    p2 = plot_window.plot_widget.addPlot(
        row=0, col=1, labels={'bottom': 'Channel Number', 'left': 'Unit Number'}
        )
    p2.setTitle('Spatial Features')
    features = torch.norm(Wall0, dim=2).cpu().numpy()
    img = pg.ImageItem(image=features.T)
    img.setLevels([0, 25])
    p2.addItem(img)

    # Comput spike counts and mean amplitudes
    n_units = int(clu0.max()) + 1
    spike_counts = np.zeros(n_units)
    for i in range(n_units):
        spike_counts[i] = (clu0[clu0 == i]).size
    mean_amp = torch.norm(Wall0, dim=(1,2)).cpu().numpy()

    # Unit amplitudes (bottom left)
    p3 = plot_window.plot_widget.addPlot(
        row=1, col=0, labels={'bottom': 'Unit Number', 'left': 'Amplitude (a.u.)'}
        )
    p3.setTitle('Unit Amplitudes')
    p3.plot(mean_amp)

    # Amplitude vs Spike Count (bottom right)
    p4 = plot_window.plot_widget.addPlot(
        row=1, col=1,
        labels={'bottom': 'Log(1 + Spike Count)', 'left': 'Amplitude (a.u.)'}
        )
    p4.setTitle('Amplitude vs Spike Count')
    scatter = pg.ScatterPlotItem(np.log(1 + spike_counts), mean_amp,
                                 symbol='o', size=3)
    p4.addItem(scatter)

    # Finished, draw plot
    plot_window.show()
//  [python]
def plot_diagnostics(plot_window, wPCA, Wall0, clu0, settings):
    # Temporal features (top left)
    p1 = plot_window.plot_widget.addPlot(
        row=0, col=0, labels={'bottom': 'Time (s)'}
        )
    p1.setTitle('Temporal Features')
    t = np.arange(wPCA.shape[1])/(settings['fs']/1000)
    for i in range(wPCA.shape[0]):
        color = _COLOR_CODES[i % len(_COLOR_CODES)]
        p1.plot(t, wPCA[i,:], pen=color)

    # Spatial features (top right)
    p2 = plot_window.plot_widget.addPlot(
        row=0, col=1, labels={'bottom': 'Channel Number', 'left': 'Unit Number'}
        )
    p2.setTitle('Spatial Features')
    features = torch.linalg.norm(Wall0, dim=2).cpu().numpy()
    img = pg.ImageItem(image=features.T)
    img.setLevels([0, 25])
    p2.addItem(img)

    # Comput spike counts and mean amplitudes
    n_units = int(clu0.max()) + 1
    spike_counts = np.zeros(n_units)
    for i in range(n_units):
        spike_counts[i] = (clu0[clu0 == i]).size
    mean_amp = torch.linalg.norm(Wall0, dim=(1,2)).cpu().numpy()

    # Unit amplitudes (bottom left)
    p3 = plot_window.plot_widget.addPlot(
        row=1, col=0, labels={'bottom': 'Unit Number', 'left': 'Amplitude (a.u.)'}
        )
    p3.setTitle('Unit Amplitudes')
    p3.plot(mean_amp)

    # Amplitude vs Spike Count (bottom right)
    p4 = plot_window.plot_widget.addPlot(
        row=1, col=1,
        labels={'bottom': 'Log(1 + Spike Count)', 'left': 'Amplitude (a.u.)'}
        )
    p4.setTitle('Amplitude vs Spike Count')
    scatter = pg.ScatterPlotItem(np.log(1 + spike_counts), mean_amp,
                                 symbol='o', size=3)
    p4.addItem(scatter)

    # Finished, draw plot
    plot_window.show()
//  [python]
def optimize_C(self,A,B,lr,iters,orthog,verbose):
        #parameterize mapping to be orthogonal
        n = A.shape[0]
        sim_net = LearnableSimilarityTransform(n,orthog=orthog).to(self.device)
        if orthog:
            parametrize.register_parametrization(sim_net, "C", Skew(n,self.device))
            parametrize.register_parametrization(sim_net, "C", CayleyMap(n,self.device))
        else:
            parametrize.register_parametrization(sim_net, "C", Matrix(n,self.device))
        
        simdist_loss = nn.MSELoss(reduction = 'sum')

        optimizer = optim.Adam(sim_net.parameters(), lr=lr)
        # scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999)

        losses = []
        A /= torch.norm(A)
        B /= torch.norm(B)
        for _ in range(iters):
            # Zero the gradients of the optimizer.
            optimizer.zero_grad()      
            # Compute the Frobenius norm between A and the product.
            loss = simdist_loss(A, sim_net(B))

            loss.backward()

            optimizer.step()
            # if _ % 99:
            #     scheduler.step()
            losses.append(loss.item())

        if verbose:
            print("Finished optimizing C")

        C_star = sim_net.C.detach()
        return losses, C_star,sim_net
//  [python]
def optimize_C(self,A,B,lr,iters,orthog,verbose):
        #parameterize mapping to be orthogonal
        n = A.shape[0]
        sim_net = LearnableSimilarityTransform(n,orthog=orthog).to(self.device)
        if orthog:
            parametrize.register_parametrization(sim_net, "C", Skew(n,self.device))
            parametrize.register_parametrization(sim_net, "C", CayleyMap(n,self.device))
        else:
            parametrize.register_parametrization(sim_net, "C", Matrix(n,self.device))
        
        simdist_loss = nn.MSELoss(reduction = 'sum')

        optimizer = optim.Adam(sim_net.parameters(), lr=lr)
        # scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999)

        losses = []
        A /= torch.linalg.norm(A)
        B /= torch.linalg.norm(B)
        for _ in range(iters):
            # Zero the gradients of the optimizer.
            optimizer.zero_grad()      
            # Compute the Frobenius norm between A and the product.
            loss = simdist_loss(A, sim_net(B))

            loss.backward()

            optimizer.step()
            # if _ % 99:
            #     scheduler.step()
            losses.append(loss.item())

        if verbose:
            print("Finished optimizing C")

        C_star = sim_net.C.detach()
        return losses, C_star,sim_net
//  [python]
def min_sphere_radius(rays, origin):
    rays_o, rays_d = rays[..., :3] - origin, rays[..., 3:6]
    rays_d = torch.nn.functional.normalize(rays_d, p=2.0, dim=-1)

    m = torch.cross(rays_o, rays_d, dim=-1)
    rays_o = torch.cross(rays_d, m, dim=-1)
    return torch.norm(rays_o, dim=-1)
//  [python]
def min_sphere_radius(rays, origin):
    rays_o, rays_d = rays[..., :3] - origin, rays[..., 3:6]
    rays_d = torch.nn.functional.normalize(rays_d, p=2.0, dim=-1)

    m = torch.cross(rays_o, rays_d, dim=-1)
    rays_o = torch.cross(rays_d, m, dim=-1)
    return torch.linalg.norm(rays_o, dim=-1)
//  [python]
def apply(cls, hparams, loss_dict, metric_dict, batch, out, compute_metric):
        gt_surface_config = hparams.supervision.gt_surface
        field = out['field']

        if gt_surface_config.value > 0.0 or gt_surface_config.normal > 0.0:

            if DS.GT_GEOMETRY not in batch.keys():
                ref_xyz, ref_normal = batch[DS.GT_DENSE_PC][0], batch[DS.GT_DENSE_NORMAL][0]
            else:
                ref_geometry = batch[DS.GT_GEOMETRY][0]
                ref_xyz, ref_normal, _ = ref_geometry.torch_attr()

            n_subsample = gt_surface_config.subsample
            if 0 < n_subsample < ref_xyz.size(0):
                ref_xyz_inds = (torch.rand((n_subsample,), device=ref_xyz.device) *
                                ref_xyz.size(0)).long()
            else:
                ref_xyz_inds = torch.arange(ref_xyz.size(0), device=ref_xyz.device)

            compute_grad = gt_surface_config.normal > 0.0
            eval_res = field.evaluate_f(ref_xyz[ref_xyz_inds], grad=compute_grad)

            if compute_grad:
                pd_grad = eval_res.gradient
                pd_grad = -pd_grad / (torch.norm(pd_grad, dim=-1, keepdim=True) + 1.0e-6)
                loss_dict.add_loss('gt-surface-normal',
                                   1.0 - torch.sum(pd_grad * ref_normal[ref_xyz_inds], dim=-1).mean(),
                                   gt_surface_config.normal)

            loss_dict.add_loss('gt-surface-value', torch.abs(eval_res.value).mean(), gt_surface_config.value)
//  [python]
def apply(cls, hparams, loss_dict, metric_dict, batch, out, compute_metric):
        gt_surface_config = hparams.supervision.gt_surface
        field = out['field']

        if gt_surface_config.value > 0.0 or gt_surface_config.normal > 0.0:

            if DS.GT_GEOMETRY not in batch.keys():
                ref_xyz, ref_normal = batch[DS.GT_DENSE_PC][0], batch[DS.GT_DENSE_NORMAL][0]
            else:
                ref_geometry = batch[DS.GT_GEOMETRY][0]
                ref_xyz, ref_normal, _ = ref_geometry.torch_attr()

            n_subsample = gt_surface_config.subsample
            if 0 < n_subsample < ref_xyz.size(0):
                ref_xyz_inds = (torch.rand((n_subsample,), device=ref_xyz.device) *
                                ref_xyz.size(0)).long()
            else:
                ref_xyz_inds = torch.arange(ref_xyz.size(0), device=ref_xyz.device)

            compute_grad = gt_surface_config.normal > 0.0
            eval_res = field.evaluate_f(ref_xyz[ref_xyz_inds], grad=compute_grad)

            if compute_grad:
                pd_grad = eval_res.gradient
                pd_grad = -pd_grad / (torch.linalg.norm(pd_grad, dim=-1, keepdim=True) + 1.0e-6)
                loss_dict.add_loss('gt-surface-normal',
                                   1.0 - torch.sum(pd_grad * ref_normal[ref_xyz_inds], dim=-1).mean(),
                                   gt_surface_config.normal)

            loss_dict.add_loss('gt-surface-value', torch.abs(eval_res.value).mean(), gt_surface_config.value)
//  [python]
def min_cylinder_radius(rays, origin):
    rays_o, rays_d = rays[..., 0:3] - origin, rays[..., 3:6]
    rays_o = torch.cat([rays_o[..., 0:1], torch.zeros_like(rays[..., 1:2]), rays_o[..., 2:3]], -1)
    rays_d = torch.cat([rays_d[..., 0:1], torch.zeros_like(rays[..., 1:2]), rays_d[..., 2:3]], -1)
    rays_d = torch.nn.functional.normalize(rays_d, p=2.0, dim=-1)

    m = torch.cross(rays_o, rays_d, dim=-1)
    rays_o = torch.cross(rays_d, m, dim=-1)
    return torch.norm(rays_o, dim=-1)
//  [python]
def min_cylinder_radius(rays, origin):
    rays_o, rays_d = rays[..., 0:3] - origin, rays[..., 3:6]
    rays_o = torch.cat([rays_o[..., 0:1], torch.zeros_like(rays[..., 1:2]), rays_o[..., 2:3]], -1)
    rays_d = torch.cat([rays_d[..., 0:1], torch.zeros_like(rays[..., 1:2]), rays_d[..., 2:3]], -1)
    rays_d = torch.nn.functional.normalize(rays_d, p=2.0, dim=-1)

    m = torch.cross(rays_o, rays_d, dim=-1)
    rays_o = torch.cross(rays_d, m, dim=-1)
    return torch.linalg.norm(rays_o, dim=-1)
//  [python]
def _cast_ray_to_line(
        self,
        line: Entity,
        ray_origin: Tensor,
        ray_direction: Tensor,
        max_range: float,
    ):
        """
        Inspired by https://stackoverflow.com/questions/563198/how-do-you-detect-where-two-line-segments-intersect/565282#565282
        Computes distance of ray originating from pos at angle to a line and sets distance to
        max_range if there is no intersection.
        """
        assert ray_origin.ndim == 2 and ray_direction.ndim == 1
        assert ray_origin.shape[0] == ray_direction.shape[0]
        assert isinstance(line.shape, Line)

        p = line.state.pos
        r = (
            torch.stack(
                [
                    torch.cos(line.state.rot.squeeze(1)),
                    torch.sin(line.state.rot.squeeze(1)),
                ],
                dim=-1,
            )
            * line.shape.length
        )

        q = ray_origin
        s = torch.stack(
            [
                torch.cos(ray_direction),
                torch.sin(ray_direction),
            ],
            dim=-1,
        )

        rxs = TorchUtils.cross(r, s)
        t = TorchUtils.cross(q - p, s / rxs)
        u = TorchUtils.cross(q - p, r / rxs)
        d = torch.norm(u * s, dim=-1)

        perpendicular = rxs == 0.0
        above_line = t > 0.5
        below_line = t < -0.5
        behind_line = u < 0.0
        d[perpendicular.squeeze(-1)] = max_range
        d[above_line.squeeze(-1)] = max_range
        d[below_line.squeeze(-1)] = max_range
        d[behind_line.squeeze(-1)] = max_range
        return d
//  [python]
def _cast_ray_to_line(
        self,
        line: Entity,
        ray_origin: Tensor,
        ray_direction: Tensor,
        max_range: float,
    ):
        """
        Inspired by https://stackoverflow.com/questions/563198/how-do-you-detect-where-two-line-segments-intersect/565282#565282
        Computes distance of ray originating from pos at angle to a line and sets distance to
        max_range if there is no intersection.
        """
        assert ray_origin.ndim == 2 and ray_direction.ndim == 1
        assert ray_origin.shape[0] == ray_direction.shape[0]
        assert isinstance(line.shape, Line)

        p = line.state.pos
        r = (
            torch.stack(
                [
                    torch.cos(line.state.rot.squeeze(1)),
                    torch.sin(line.state.rot.squeeze(1)),
                ],
                dim=-1,
            )
            * line.shape.length
        )

        q = ray_origin
        s = torch.stack(
            [
                torch.cos(ray_direction),
                torch.sin(ray_direction),
            ],
            dim=-1,
        )

        rxs = TorchUtils.cross(r, s)
        t = TorchUtils.cross(q - p, s / rxs)
        u = TorchUtils.cross(q - p, r / rxs)
        d = torch.linalg.norm(u * s, dim=-1)

        perpendicular = rxs == 0.0
        above_line = t > 0.5
        below_line = t < -0.5
        behind_line = u < 0.0
        d[perpendicular.squeeze(-1)] = max_range
        d[above_line.squeeze(-1)] = max_range
        d[below_line.squeeze(-1)] = max_range
        d[behind_line.squeeze(-1)] = max_range
        return d
//  [python]
def find_symm_subs(xyz,Rs,metasymm):
    com = xyz[:,:,1].mean(dim=-2)
    rcoms = torch.einsum('sij,bj->si', Rs, com)

    subsymms, nneighs = metasymm

    subs = []
    for i in range(len(subsymms)):
        drcoms = torch.norm(rcoms[0,:] - rcoms[subsymms[i],:], dim=-1)
        _,subs_i = torch.topk(drcoms,nneighs[i],largest=False)
        subs_i,_ = torch.sort( subsymms[i][subs_i] )
        subs.append(subs_i)

    subs=torch.cat(subs)
    xyz_new = torch.einsum('sij,braj->bsrai', Rs[subs], xyz).reshape(
        xyz.shape[0],-1,xyz.shape[2],3)
    return xyz_new, subs
//  [python]
def find_symm_subs(xyz,Rs,metasymm):
    com = xyz[:,:,1].mean(dim=-2)
    rcoms = torch.einsum('sij,bj->si', Rs, com)

    subsymms, nneighs = metasymm

    subs = []
    for i in range(len(subsymms)):
        drcoms = torch.linalg.norm(rcoms[0,:] - rcoms[subsymms[i],:], dim=-1)
        _,subs_i = torch.topk(drcoms,nneighs[i],largest=False)
        subs_i,_ = torch.sort( subsymms[i][subs_i] )
        subs.append(subs_i)

    subs=torch.cat(subs)
    xyz_new = torch.einsum('sij,braj->bsrai', Rs[subs], xyz).reshape(
        xyz.shape[0],-1,xyz.shape[2],3)
    return xyz_new, subs
//  [python]
def _convert_gau(self):
        """Reshape and merge tensors for Gated Attention Unit (GAU).

        This function pre-processes the gamma and beta tensors of the GAU and
        handles the position encoding if available. It also redefines the GAU's
        forward method to incorporate these pre-processed tensors, optimizing
        the computation process.
        """
        # Reshape gamma and beta tensors in advance
        gamma_q = self.gau.gamma[0].view(1, 1, 1, self.gau.gamma.size(-1))
        gamma_k = self.gau.gamma[1].view(1, 1, 1, self.gau.gamma.size(-1))
        beta_q = self.gau.beta[0].view(1, 1, 1, self.gau.beta.size(-1))
        beta_k = self.gau.beta[1].view(1, 1, 1, self.gau.beta.size(-1))

        # Adjust beta tensors with position encoding if available
        if hasattr(self, 'pos_enc'):
            pos_enc = self.pos_enc.reshape(1, 1, *self.pos_enc.shape)
            beta_q = beta_q + pos_enc
            beta_k = beta_k + pos_enc

        gamma_q = gamma_q.detach().cpu()
        gamma_k = gamma_k.detach().cpu()
        beta_q = beta_q.detach().cpu()
        beta_k = beta_k.detach().cpu()

        @torch.no_grad()
        def _forward(self, x, *args, **kwargs):
            norm = torch.norm(x, dim=-1, keepdim=True) * self.ln.scale
            x = x / norm.clamp(min=self.ln.eps) * self.ln.g

            uv = self.uv(x)
            uv = self.act_fn(uv)

            u, v, base = torch.split(uv, [self.e, self.e, self.s], dim=-1)
            if not torch.onnx.is_in_onnx_export():
                q = base * gamma_q.to(base) + beta_q.to(base)
                k = base * gamma_k.to(base) + beta_k.to(base)
            else:
                q = base * gamma_q + beta_q
                k = base * gamma_k + beta_k
            qk = torch.matmul(q, k.transpose(-1, -2))

            kernel = torch.square(torch.nn.functional.relu(qk / self.sqrt_s))
            x = u * torch.matmul(kernel, v)
            x = self.o(x)
            return x

        self.gau._forward = types.MethodType(_forward, self.gau)
//  [python]
def _convert_gau(self):
        """Reshape and merge tensors for Gated Attention Unit (GAU).

        This function pre-processes the gamma and beta tensors of the GAU and
        handles the position encoding if available. It also redefines the GAU's
        forward method to incorporate these pre-processed tensors, optimizing
        the computation process.
        """
        # Reshape gamma and beta tensors in advance
        gamma_q = self.gau.gamma[0].view(1, 1, 1, self.gau.gamma.size(-1))
        gamma_k = self.gau.gamma[1].view(1, 1, 1, self.gau.gamma.size(-1))
        beta_q = self.gau.beta[0].view(1, 1, 1, self.gau.beta.size(-1))
        beta_k = self.gau.beta[1].view(1, 1, 1, self.gau.beta.size(-1))

        # Adjust beta tensors with position encoding if available
        if hasattr(self, 'pos_enc'):
            pos_enc = self.pos_enc.reshape(1, 1, *self.pos_enc.shape)
            beta_q = beta_q + pos_enc
            beta_k = beta_k + pos_enc

        gamma_q = gamma_q.detach().cpu()
        gamma_k = gamma_k.detach().cpu()
        beta_q = beta_q.detach().cpu()
        beta_k = beta_k.detach().cpu()

        @torch.no_grad()
        def _forward(self, x, *args, **kwargs):
            norm = torch.linalg.norm(x, dim=-1, keepdim=True) * self.ln.scale
            x = x / norm.clamp(min=self.ln.eps) * self.ln.g

            uv = self.uv(x)
            uv = self.act_fn(uv)

            u, v, base = torch.split(uv, [self.e, self.e, self.s], dim=-1)
            if not torch.onnx.is_in_onnx_export():
                q = base * gamma_q.to(base) + beta_q.to(base)
                k = base * gamma_k.to(base) + beta_k.to(base)
            else:
                q = base * gamma_q + beta_q
                k = base * gamma_k + beta_k
            qk = torch.matmul(q, k.transpose(-1, -2))

            kernel = torch.square(torch.nn.functional.relu(qk / self.sqrt_s))
            x = u * torch.matmul(kernel, v)
            x = self.o(x)
            return x

        self.gau._forward = types.MethodType(_forward, self.gau)
//  [python]
def norm(x: _Numeric, ord: Optional[Any] = None, axis: Optional[int] = None):  # type: ignore
    """
    Matrix or vector norm.
    """
    return t.norm(x, ord=ord, dim=axis)
//  [python]
def norm(x: _Numeric, ord: Optional[Any] = None, axis: Optional[int] = None):  # type: ignore
    """
    Matrix or vector norm.
    """
    return torch.linalg.norm(x, ord=ord, dim=axis)
//  [python]
def test_rbf_kernel(self) -> None:
        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
        lengthscale = 2
        # test rbf
        # pyre-fixme[6]: For 3rd param expected `Tensor` but got `int`.
        res = rbf_kernel(a, b, lengthscale=lengthscale)
        actual = (
            torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float)
            .pow_(2.0)
            .mul_(-0.5 / (lengthscale**2))
            .exp()
        )
        self.assertLess(torch.norm(res - actual), 1e-3)
//  [python]
def test_rbf_kernel(self) -> None:
        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
        lengthscale = 2
        # test rbf
        # pyre-fixme[6]: For 3rd param expected `Tensor` but got `int`.
        res = rbf_kernel(a, b, lengthscale=lengthscale)
        actual = (
            torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float)
            .pow_(2.0)
            .mul_(-0.5 / (lengthscale**2))
            .exp()
        )
        self.assertLess(torch.linalg.norm(res - actual), 1e-3)
//  [python]
def test_quat_to_yaw(quats_wxyz: Tensor) -> None:
    """Test converting a quaternion to yaw.

    Args:
        quats_wxyz: (N,4) Scalar first quaternions.

    NOTE: Yaw is rotation about the vertical axis in our
        coordinate system.
    """
    quats_wxyz /= torch.norm(quats_wxyz, dim=-1, keepdim=True)
    yaw = quat_to_yaw(quats_wxyz).rad2deg()
    yaw_ = torch.as_tensor(
        R.from_quat(quats_wxyz[..., [1, 2, 3, 0]].numpy()).as_euler(
            "xyz", degrees=True
        )
    )[..., -1]
    torch.testing.assert_allclose(yaw, yaw_)
//  [python]
def test_quat_to_yaw(quats_wxyz: Tensor) -> None:
    """Test converting a quaternion to yaw.

    Args:
        quats_wxyz: (N,4) Scalar first quaternions.

    NOTE: Yaw is rotation about the vertical axis in our
        coordinate system.
    """
    quats_wxyz /= torch.linalg.norm(quats_wxyz, dim=-1, keepdim=True)
    yaw = quat_to_yaw(quats_wxyz).rad2deg()
    yaw_ = torch.as_tensor(
        R.from_quat(quats_wxyz[..., [1, 2, 3, 0]].numpy()).as_euler(
            "xyz", degrees=True
        )
    )[..., -1]
    torch.testing.assert_allclose(yaw, yaw_)
//  [python]
def test_quat_to_xyz(quats_wxyz: Tensor) -> None:
    """Test converting a quaternion to Tait-Bryan angles (roll, pitch, yaw)."""
    quats_wxyz /= torch.norm(quats_wxyz, dim=-1, keepdim=True)
    xyz = quat_to_xyz(quats_wxyz).rad2deg()
    xyz_ = torch.as_tensor(
        R.from_quat(quats_wxyz[..., [1, 2, 3, 0]].numpy()).as_euler(
            "xyz", degrees=True
        )
    )
    torch.testing.assert_allclose(xyz, xyz_)
//  [python]
def test_quat_to_xyz(quats_wxyz: Tensor) -> None:
    """Test converting a quaternion to Tait-Bryan angles (roll, pitch, yaw)."""
    quats_wxyz /= torch.linalg.norm(quats_wxyz, dim=-1, keepdim=True)
    xyz = quat_to_xyz(quats_wxyz).rad2deg()
    xyz_ = torch.as_tensor(
        R.from_quat(quats_wxyz[..., [1, 2, 3, 0]].numpy()).as_euler(
            "xyz", degrees=True
        )
    )
    torch.testing.assert_allclose(xyz, xyz_)
//  [python]
def _test_eigen_root(
        self,
        A: torch.Tensor,
        root: int,
        make_positive_semidefinite: bool,
        inverse: bool,
        epsilon: float,
        tolerance: float,
        eig_sols: Tensor,
    ) -> None:
        X, L, Q = _matrix_root_eigen(
            A=A,
            root=root,
            epsilon=epsilon,
            make_positive_semidefinite=make_positive_semidefinite,
            inverse=inverse,
        )
        if inverse:
            root = -root
        abs_error = torch.dist(torch.linalg.matrix_power(X, root), A, p=torch.inf)
        A_norm = torch.norm(A, ord=torch.inf)
        rel_error = abs_error / torch.maximum(torch.tensor(1.0), A_norm)
        torch.testing.assert_close(L, eig_sols)
        self.assertTrue(rel_error <= tolerance)
//  [python]
def _test_eigen_root(
        self,
        A: torch.Tensor,
        root: int,
        make_positive_semidefinite: bool,
        inverse: bool,
        epsilon: float,
        tolerance: float,
        eig_sols: Tensor,
    ) -> None:
        X, L, Q = _matrix_root_eigen(
            A=A,
            root=root,
            epsilon=epsilon,
            make_positive_semidefinite=make_positive_semidefinite,
            inverse=inverse,
        )
        if inverse:
            root = -root
        abs_error = torch.dist(torch.linalg.matrix_power(X, root), A, p=torch.inf)
        A_norm = torch.linalg.norm(A, ord=torch.inf)
        rel_error = abs_error / torch.maximum(torch.tensor(1.0), A_norm)
        torch.testing.assert_close(L, eig_sols)
        self.assertTrue(rel_error <= tolerance)
//  [python]
def _test_newton_root_inverse(
        self,
        A: Tensor,
        root: int,
        epsilon: float,
        max_iterations: int,
        A_tol: float,
        M_tol: float,
    ) -> None:
        X, M, flag, iteration, M_error = _matrix_inverse_root_newton(
            A, root, epsilon, max_iterations, M_tol
        )
        abs_A_error = torch.dist(torch.linalg.matrix_power(X, -root), A, p=torch.inf)
        A_norm = torch.norm(A, ord=float('inf'))
        rel_A_error = abs_A_error / torch.maximum(torch.tensor(1.0), A_norm)
        self.assertTrue(M_error <= M_tol)
        self.assertTrue(rel_A_error <= A_tol)
//  [python]
def _test_newton_root_inverse(
        self,
        A: Tensor,
        root: int,
        epsilon: float,
        max_iterations: int,
        A_tol: float,
        M_tol: float,
    ) -> None:
        X, M, flag, iteration, M_error = _matrix_inverse_root_newton(
            A, root, epsilon, max_iterations, M_tol
        )
        abs_A_error = torch.dist(torch.linalg.matrix_power(X, -root), A, p=torch.inf)
        A_norm = torch.linalg.norm(A, ord=torch.inf)
        rel_A_error = abs_A_error / torch.maximum(torch.tensor(1.0), A_norm)
        self.assertTrue(M_error <= M_tol)
        self.assertTrue(rel_A_error <= A_tol)
//  [python]
def rgbd2pcd(im, depth, w2c, k, show_depth=False, project_to_cam_w_scale=None):
    d_near = 1.5
    d_far = 6
    invk = torch.inverse(torch.tensor(k).cuda().float())
    c2w = torch.inverse(torch.tensor(w2c).cuda().float())
    radial_depth = depth[0].reshape(-1)
    def_rays = (invk @ def_pix.T).T
    def_radial_rays = def_rays / torch.norm(def_rays, 2, dim=-1)[:, None]
    pts_cam = def_radial_rays * radial_depth[:, None]
    z_depth = pts_cam[:, 2]
    if project_to_cam_w_scale is not None:
        pts_cam = project_to_cam_w_scale * pts_cam / z_depth[:, None]
    pts4 = torch.concat((pts_cam, pix_ones), 1)
    pts = (c2w @ pts4.T).T[:, :3]
    if show_depth:
        cols = ((z_depth - d_near) / (d_far - d_near))[:, None].repeat(1, 3)
    else:
        cols = torch.permute(im, (1, 2, 0)).reshape(-1, 3)
    pts = o3d.utility.Vector3dVector(pts.contiguous().double().cpu().numpy())
    cols = o3d.utility.Vector3dVector(cols.contiguous().double().cpu().numpy())
    return pts, cols
//  [python]
def rgbd2pcd(im, depth, w2c, k, show_depth=False, project_to_cam_w_scale=None):
    d_near = 1.5
    d_far = 6
    invk = torch.inverse(torch.tensor(k).cuda().float())
    c2w = torch.inverse(torch.tensor(w2c).cuda().float())
    radial_depth = depth[0].reshape(-1)
    def_rays = (invk @ def_pix.T).T
    def_radial_rays = def_rays / torch.linalg.norm(def_rays, ord=2, dim=-1)[:, None]
    pts_cam = def_radial_rays * radial_depth[:, None]
    z_depth = pts_cam[:, 2]
    if project_to_cam_w_scale is not None:
        pts_cam = project_to_cam_w_scale * pts_cam / z_depth[:, None]
    pts4 = torch.concat((pts_cam, pix_ones), 1)
    pts = (c2w @ pts4.T).T[:, :3]
    if show_depth:
        cols = ((z_depth - d_near) / (d_far - d_near))[:, None].repeat(1, 3)
    else:
        cols = torch.permute(im, (1, 2, 0)).reshape(-1, 3)
    pts = o3d.utility.Vector3dVector(pts.contiguous().double().cpu().numpy())
    cols = o3d.utility.Vector3dVector(cols.contiguous().double().cpu().numpy())
    return pts, cols
//  [python]
def test_manual_fill():
    class ProblemWithManualFill(et.Problem):
        def __init__(self):
            super().__init__(objective_sense="min", solution_length=5, dtype="int64")

        def _evaluate_batch(self, batch: et.SolutionBatch):
            batch.set_evals(torch.norm(batch.values, dim=-1))

        def _fill(self, x: torch.Tensor):
            x.zero_()

    problem = ProblemWithManualFill()

    new_batch = problem.generate_batch(10)
    assert len(new_batch) == 10
    assert torch.all(new_batch.values == 0)

    new_values = problem.generate_values(10)
    assert new_values.shape == (10, 5)
    assert torch.all(new_values == 0)
//  [python]
def test_manual_fill():
    class ProblemWithManualFill(et.Problem):
        def __init__(self):
            super().__init__(objective_sense="min", solution_length=5, dtype="int64")

        def _evaluate_batch(self, batch: et.SolutionBatch):
            batch.set_evals(torch.linalg.norm(batch.values, dim=-1))

        def _fill(self, x: torch.Tensor):
            x.zero_()

    problem = ProblemWithManualFill()

    new_batch = problem.generate_batch(10)
    assert len(new_batch) == 10
    assert torch.all(new_batch.values == 0)

    new_values = problem.generate_values(10)
    assert new_values.shape == (10, 5)
    assert torch.all(new_values == 0)
//  [python]
def test_matern_kernel(self) -> None:
        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
        lengthscale = 2
        # test matern 1/2
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `int`.
        res = matern_kernel(a, b, nu=0.5, lengthscale=lengthscale)
        actual = (
            torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float)
            .div_(-lengthscale)
            .exp()
        )
        self.assertLess(torch.norm(res - actual), 1e-3)
        # matern test 3/2
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `int`.
        res = matern_kernel(a, b, nu=1.5, lengthscale=lengthscale)
        dist = torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float).mul_(
            sqrt(3) / lengthscale
        )
        actual = (dist + 1).mul(torch.exp(-dist))
        self.assertLess(torch.norm(res - actual), 1e-3)
        # matern test 5/2
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `int`.
        res = matern_kernel(a, b, nu=2.5, lengthscale=lengthscale)
        dist = torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float).mul_(
            sqrt(5) / lengthscale
        )
        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
        actual = (dist**2 / 3 + dist + 1).mul(torch.exp(-dist))
        self.assertLess(torch.norm(res - actual), 1e-3)

        # test k(x,x) with no gradients
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `float`.
        res = matern_kernel(b, b, nu=0.5, lengthscale=2.0)
        actual = (
            torch.tensor([[0, 2], [2, 0]], dtype=torch.float).div_(-lengthscale).exp()
        )
        self.assertLess(torch.norm(res - actual), 1e-3)

        # test unsupported nu
        with self.assertRaises(AxError):
            # pyre-fixme[6]: For 4th param expected `Tensor` but got `float`.
            matern_kernel(b, b, nu=0.0, lengthscale=2.0)
//  [python]
def test_matern_kernel(self) -> None:
        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
        lengthscale = 2
        # test matern 1/2
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `int`.
        res = matern_kernel(a, b, nu=0.5, lengthscale=lengthscale)
        actual = (
            torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float)
            .div_(-lengthscale)
            .exp()
        )
        self.assertLess(torch.linalg.norm(res - actual), 1e-3)
        # matern test 3/2
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `int`.
        res = matern_kernel(a, b, nu=1.5, lengthscale=lengthscale)
        dist = torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float).mul_(
            sqrt(3) / lengthscale
        )
        actual = (dist + 1).mul(torch.exp(-dist))
        self.assertLess(torch.linalg.norm(res - actual), 1e-3)
        # matern test 5/2
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `int`.
        res = matern_kernel(a, b, nu=2.5, lengthscale=lengthscale)
        dist = torch.tensor([[4, 2], [2, 0], [8, 6]], dtype=torch.float).mul_(
            sqrt(5) / lengthscale
        )
        # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
        actual = (dist**2 / 3 + dist + 1).mul(torch.exp(-dist))
        self.assertLess(torch.linalg.norm(res - actual), 1e-3)

        # test k(x,x) with no gradients
        # pyre-fixme[6]: For 4th param expected `Tensor` but got `float`.
        res = matern_kernel(b, b, nu=0.5, lengthscale=2.0)
        actual = (
            torch.tensor([[0, 2], [2, 0]], dtype=torch.float).div_(-lengthscale).exp()
        )
        self.assertLess(torch.linalg.norm(res - actual), 1e-3)

        # test unsupported nu
        with self.assertRaises(AxError):
            # pyre-fixme[6]: For 4th param expected `Tensor` but got `float`.
            matern_kernel(b, b, nu=0.0, lengthscale=2.0)
//  [python]
def get_token_attributions(
    model: LudwigModel,
    feature_name: str,
    input_ids: torch.Tensor,
    token_attributions: torch.Tensor,
) -> List[List[Tuple[str, float]]]:
    """Convert token-level attributions to an array of token-attribution pairs of shape.

    [batch_size, sequence_length, 2].

    Args:
        model: The LudwigModel used to generate the attributions.
        feature_name: The name of the feature for which the attributions were generated.
        input_ids: The input ids of shape [batch_size, sequence_length].
        token_attributions: The token-level attributions of shape [batch_size, sequence_length].

    Returns:
        An array of token-attribution pairs of shape [batch_size, sequence_length, 2].
    """
    assert (
        input_ids.dtype == torch.int8
        or input_ids.dtype == torch.int16
        or input_ids.dtype == torch.int32
        or input_ids.dtype == torch.int64
    )

    # Normalize token-level attributions to visualize the relative importance of each token.
    norm = torch.norm(token_attributions, dim=1)
    # Safe divide by zero by setting the norm to 1 if the norm is 0.
    norm = torch.where(norm == 0, torch.ones_like(norm), norm)
    token_attributions = token_attributions / norm.unsqueeze(-1)

    # map input ids to input tokens via the vocabulary
    feature = model.training_set_metadata[feature_name]
    vocab = feature.get("idx2str", feature.get("word_idx2str"))
    idx2str = np.vectorize(lambda idx: vocab[idx])
    input_tokens = idx2str(input_ids)

    # add attribution to the input tokens
    tok_attrs = [
        list(zip(t, a)) for t, a in zip(input_tokens, token_attributions.tolist())
    ]  # [batch_size, sequence_length, 2]

    return tok_attrs
//  [python]
def get_token_attributions(
    model: LudwigModel,
    feature_name: str,
    input_ids: torch.Tensor,
    token_attributions: torch.Tensor,
) -> List[List[Tuple[str, float]]]:
    """Convert token-level attributions to an array of token-attribution pairs of shape.

    [batch_size, sequence_length, 2].

    Args:
        model: The LudwigModel used to generate the attributions.
        feature_name: The name of the feature for which the attributions were generated.
        input_ids: The input ids of shape [batch_size, sequence_length].
        token_attributions: The token-level attributions of shape [batch_size, sequence_length].

    Returns:
        An array of token-attribution pairs of shape [batch_size, sequence_length, 2].
    """
    assert (
        input_ids.dtype == torch.int8
        or input_ids.dtype == torch.int16
        or input_ids.dtype == torch.int32
        or input_ids.dtype == torch.int64
    )

    # Normalize token-level attributions to visualize the relative importance of each token.
    norm = torch.linalg.norm(token_attributions, dim=1)
    # Safe divide by zero by setting the norm to 1 if the norm is 0.
    norm = torch.where(norm == 0, torch.ones_like(norm), norm)
    token_attributions = token_attributions / norm.unsqueeze(-1)

    # map input ids to input tokens via the vocabulary
    feature = model.training_set_metadata[feature_name]
    vocab = feature.get("idx2str", feature.get("word_idx2str"))
    idx2str = np.vectorize(lambda idx: vocab[idx])
    input_tokens = idx2str(input_ids)

    # add attribution to the input tokens
    tok_attrs = [
        list(zip(t, a)) for t, a in zip(input_tokens, token_attributions.tolist())
    ]  # [batch_size, sequence_length, 2]

    return tok_attrs
//  [python]
def _compute_dq(
    step_size,
    coords_x: torch.Tensor,
    grad_q: torch.Tensor,
    hess_q: torch.Tensor,
    ic: tico.ic.IC,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the delta in internal coordinates to take."""
    grad_q_len = len(grad_q)

    if isinstance(ic, tico.ic.DLC) and ic.constr is not None:
        grad_q, hess_q = ic.augment_hess_q(coords_x, grad_q, hess_q)

    hess_q = hess_q + step_size * torch.eye(len(hess_q))

    for i in range(grad_q_len, len(grad_q)):
        # don't shift the constrained degrees of freedom
        hess_q[i, i] = 0.0

    hess_q_inv = tico.utils.pinv(hess_q, 1e-12)

    dq_with_const = -hess_q_inv @ grad_q
    dq = dq_with_const[:grad_q_len]

    d_prime = -hess_q_inv @ dq_with_const
    dq_prime = dq @ d_prime[:grad_q_len] / torch.norm(dq)

    return dq, dq_prime
//  [python]
def _compute_dq(
    step_size,
    coords_x: torch.Tensor,
    grad_q: torch.Tensor,
    hess_q: torch.Tensor,
    ic: tico.ic.IC,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute the delta in internal coordinates to take."""
    grad_q_len = len(grad_q)

    if isinstance(ic, tico.ic.DLC) and ic.constr is not None:
        grad_q, hess_q = ic.augment_hess_q(coords_x, grad_q, hess_q)

    hess_q = hess_q + step_size * torch.eye(len(hess_q))

    for i in range(grad_q_len, len(grad_q)):
        # don't shift the constrained degrees of freedom
        hess_q[i, i] = 0.0

    hess_q_inv = tico.utils.pinv(hess_q, 1e-12)

    dq_with_const = -hess_q_inv @ grad_q
    dq = dq_with_const[:grad_q_len]

    d_prime = -hess_q_inv @ dq_with_const
    dq_prime = dq @ d_prime[:grad_q_len] / torch.linalg.norm(dq)

    return dq, dq_prime