jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _dq_to_x_cached(
    coords_x: torch.Tensor,
    dq: torch.Tensor,
    ic: typing.Union["RIC", "DLC"],
) -> tuple[torch.Tensor, torch.Tensor, bool]:
    """A cached version of the ``_dq_to_x`` function.

    While a bit of a hack, this can significantly speed up ops like the brent
    line search.
    """
    if ic._dq_cache is not None:
        (
            cached_coords_x,
            cached_dq,
            cached_coords_x_new,
            cached_coords_q_new,
            cached_q_converged,
        ) = ic._dq_cache

        if (
            torch.linalg.norm(cached_coords_x - coords_x) < 1e-10
            and torch.linalg.norm(cached_dq - dq) < 1e-10
        ):
            return cached_coords_x_new, cached_coords_q_new, cached_q_converged

    coords_x_new, coords_q_new, q_converged = _dq_to_x(coords_x, dq, ic)
    ic._dq_cache = (coords_x, dq, coords_x_new, coords_q_new, q_converged)

    return coords_x_new, coords_q_new, q_converged
//  [python]
def forward(
        cls,
        embeds: TensorType["bs":..., "num_samples", "num_classes"],
        weights: TensorType["bs":..., "num_samples", 1],
    ) -> TensorType["bs":..., "num_classes"]:
        """Calculate semantics along the ray."""
        output = torch.sum(weights * embeds, dim=-2)
        output = output / torch.norm(output, dim=-1, keepdim=True)
        return output
//  [python]
def forward(
        cls,
        embeds: TensorType["bs":..., "num_samples", "num_classes"],
        weights: TensorType["bs":..., "num_samples", 1],
    ) -> TensorType["bs":..., "num_classes"]:
        """Calculate semantics along the ray."""
        output = torch.sum(weights * embeds, dim=-2)
        output = output / torch.linalg.norm(output, dim=-1, keepdim=True)
        return output
//  [python]
def main(self) -> None:
        """Main function."""
        torch.set_float32_matmul_precision(self.torch_precision)
        assert str(self.output_path)[-4:] == ".ply"
        self.output_path.parent.mkdir(parents=True, exist_ok=True)

        _, pipeline, _ = eval_setup(self.load_config)

        CONSOLE.print("Extract mesh with marching cubes and may take a while")

        if self.create_visibility_mask:
            assert self.resolution % 512 == 0

            coarse_mask = pipeline.get_visibility_mask(
                self.visibility_grid_resolution, self.valid_points_thres, self.sub_sample_factor
            )

            def inv_contract(x):
                mag = torch.norm(x, p=pipeline.model.scene_contraction.order, dim=-1)
                mask = mag >= 1
                x_new = x.clone()
                x_new[mask] = (1 / (2 - mag[mask][..., None])) * (x[mask] / mag[mask][..., None])
                return x_new

            if self.save_visibility_grid:
                offset = torch.linspace(-2.0, 2.0, 512)
                x, y, z = torch.meshgrid(offset, offset, offset, indexing="ij")
                offset_cube = torch.stack([x, y, z], dim=-1).reshape(-1, 3).to(coarse_mask.device)
                points = offset_cube[coarse_mask.reshape(-1) > 0]
                points = inv_contract(points)
                save_points("mask.ply", points.cpu().numpy())
                torch.save(coarse_mask, "coarse_mask.pt")

            get_surface_sliding_with_contraction(
                sdf=lambda x: (
                    pipeline.model.field.forward_geonetwork(x)[:, 0] - self.marching_cube_threshold
                ).contiguous(),
                resolution=self.resolution,
                bounding_box_min=self.bounding_box_min,
                bounding_box_max=self.bounding_box_max,
                coarse_mask=coarse_mask,
                output_path=self.output_path,
                simplify_mesh=self.simplify_mesh,
                inv_contraction=inv_contract,
            )
            return

        if self.is_occupancy:
            # for unisurf
            get_surface_occupancy(
                occupancy_fn=lambda x: torch.sigmoid(
                    10 * pipeline.model.field.forward_geonetwork(x)[:, 0].contiguous()
                ),
                resolution=self.resolution,
                bounding_box_min=self.bounding_box_min,
                bounding_box_max=self.bounding_box_max,
                level=0.5,
                device=pipeline.model.device,
                output_path=self.output_path,
            )
        else:
            assert self.resolution % 512 == 0
            # for sdf we can multi-scale extraction.
            get_surface_sliding(
                sdf=lambda x: pipeline.model.field.forward_geonetwork(x)[:, 0].contiguous(),
                resolution=self.resolution,
                bounding_box_min=self.bounding_box_min,
                bounding_box_max=self.bounding_box_max,
                coarse_mask=pipeline.model.scene_box.coarse_binary_gird,
                output_path=self.output_path,
                simplify_mesh=self.simplify_mesh,
            )
//  [python]
def main(self) -> None:
        """Main function."""
        torch.set_float32_matmul_precision(self.torch_precision)
        assert str(self.output_path)[-4:] == ".ply"
        self.output_path.parent.mkdir(parents=True, exist_ok=True)

        _, pipeline, _ = eval_setup(self.load_config)

        CONSOLE.print("Extract mesh with marching cubes and may take a while")

        if self.create_visibility_mask:
            assert self.resolution % 512 == 0

            coarse_mask = pipeline.get_visibility_mask(
                self.visibility_grid_resolution, self.valid_points_thres, self.sub_sample_factor
            )

            def inv_contract(x):
                mag = torch.linalg.norm(x, ord=pipeline.model.scene_contraction.order, dim=-1)
                mask = mag >= 1
                x_new = x.clone()
                x_new[mask] = (1 / (2 - mag[mask][..., None])) * (x[mask] / mag[mask][..., None])
                return x_new

            if self.save_visibility_grid:
                offset = torch.linspace(-2.0, 2.0, 512)
                x, y, z = torch.meshgrid(offset, offset, offset, indexing="ij")
                offset_cube = torch.stack([x, y, z], dim=-1).reshape(-1, 3).to(coarse_mask.device)
                points = offset_cube[coarse_mask.reshape(-1) > 0]
                points = inv_contract(points)
                save_points("mask.ply", points.cpu().numpy())
                torch.save(coarse_mask, "coarse_mask.pt")

            get_surface_sliding_with_contraction(
                sdf=lambda x: (
                    pipeline.model.field.forward_geonetwork(x)[:, 0] - self.marching_cube_threshold
                ).contiguous(),
                resolution=self.resolution,
                bounding_box_min=self.bounding_box_min,
                bounding_box_max=self.bounding_box_max,
                coarse_mask=coarse_mask,
                output_path=self.output_path,
                simplify_mesh=self.simplify_mesh,
                inv_contraction=inv_contract,
            )
            return

        if self.is_occupancy:
            # for unisurf
            get_surface_occupancy(
                occupancy_fn=lambda x: torch.sigmoid(
                    10 * pipeline.model.field.forward_geonetwork(x)[:, 0].contiguous()
                ),
                resolution=self.resolution,
                bounding_box_min=self.bounding_box_min,
                bounding_box_max=self.bounding_box_max,
                level=0.5,
                device=pipeline.model.device,
                output_path=self.output_path,
            )
        else:
            assert self.resolution % 512 == 0
            # for sdf we can multi-scale extraction.
            get_surface_sliding(
                sdf=lambda x: pipeline.model.field.forward_geonetwork(x)[:, 0].contiguous(),
                resolution=self.resolution,
                bounding_box_min=self.bounding_box_min,
                bounding_box_max=self.bounding_box_max,
                coarse_mask=pipeline.model.scene_box.coarse_binary_gird,
                output_path=self.output_path,
                simplify_mesh=self.simplify_mesh,
            )
//  [python]
def _project_constr(
        cls, v: torch.Tensor, constr: ConstraintDict, ric: RIC
    ) -> torch.Tensor:
        """Attempts to project the constrained primitive internal coordinates onto the
        non-redundant basis.

        Notes:
            This uses the 'old' method used in geomeTRIC, which is a bit more
            battle tested than the more 'correct new' method.
        """

        n_v = v.shape[1]

        constr_to_ic_idx = _match_constr(constr, ric.idxs)

        v_constr = v[constr_to_ic_idx, :]
        v_constr = v_constr / t.norm(v_constr, dim=-1, keepdim=True)

        v_proj = v @ v_constr.T
        v_proj = v_proj / t.norm(v_proj, dim=0, keepdim=True)

        v = torch.hstack((v_proj, v))

        threshold = 1e-6

        while True:
            u = []
            for i in range(v.shape[1]):
                v_i = v[:, i].flatten()
                u.append(v_i.clone())

                for u_i in u[:-1]:
                    u[-1] = u[-1] - u_i * torch.dot(u_i, v_i)

                if t.norm(u[-1]) < threshold:
                    u = u[:-1]
                    continue

                u[-1] = u[-1] / t.norm(u[-1])

            if len(u) > n_v:
                threshold *= 10.0
            elif len(u) == n_v:
                break
            elif len(u) < n_v:
                raise RuntimeError("Gram-Schmidt orthogonalization failed")

        return torch.vstack(u).T
//  [python]
def _project_constr(
        cls, v: torch.Tensor, constr: ConstraintDict, ric: RIC
    ) -> torch.Tensor:
        """Attempts to project the constrained primitive internal coordinates onto the
        non-redundant basis.

        Notes:
            This uses the 'old' method used in geomeTRIC, which is a bit more
            battle tested than the more 'correct new' method.
        """

        n_v = v.shape[1]

        constr_to_ic_idx = _match_constr(constr, ric.idxs)

        v_constr = v[constr_to_ic_idx, :]
        v_constr = v_constr / torch.linalg.norm(v_constr, dim=-1, keepdim=True)

        v_proj = v @ v_constr.T
        v_proj = v_proj / torch.linalg.norm(v_proj, dim=0, keepdim=True)

        v = torch.hstack((v_proj, v))

        threshold = 1e-6

        while True:
            u = []
            for i in range(v.shape[1]):
                v_i = v[:, i].flatten()
                u.append(v_i.clone())

                for u_i in u[:-1]:
                    u[-1] = u[-1] - u_i * torch.dot(u_i, v_i)

                if torch.linalg.norm(u[-1]) < threshold:
                    u = u[:-1]
                    continue

                u[-1] = u[-1] / torch.linalg.norm(u[-1])

            if len(u) > n_v:
                threshold *= 10.0
            elif len(u) == n_v:
                break
            elif len(u) < n_v:
                raise RuntimeError("Gram-Schmidt orthogonalization failed")

        return torch.vstack(u).T
//  [python]
def _dq_to_x(
    coords_x: torch.Tensor, dq: torch.Tensor, ic: typing.Union["RIC", "DLC"]
) -> tuple[torch.Tensor, torch.Tensor, bool]:
    """Perturb a set of cartesian coordinates by a set of internal coordinate
    displacements.

    Args:
        coords_x: The coordinates with ``shape=(n_atoms, 3)``.
        dq: The internal coordinate displacements with ``shape=(n_ic,)``.

    Returns:
        The perturbed cartesian (with ``shape=(n_atoms, 3)``) and internal
        coordinates (with ``shape=(n_ic,)``), and a boolean indicating whether the
        perturbation was successful.
    """

    iteration, failures = 0, 0

    damp = 1.0

    coords_iter_1, coords_final = None, None

    dq_delta_prev = None

    while True:
        iteration += 1

        b_matrix = ic.compute_b(coords_x)
        g_matrix_inv = tico.utils.pinv(b_matrix @ b_matrix.T)

        dx = damp * torch.linalg.multi_dot([b_matrix.T, g_matrix_inv, dq])

        coords_x_new = coords_x + dx.reshape(-1, 3)
        coords_q_new = ic.compute_q(coords_x_new)

        if iteration == 1:
            coords_final = coords_iter_1 = coords_x_new, coords_q_new

        dq_actual = ic.compute_dq(coords_x, coords_x_new)
        dq_delta = torch.norm(dq - dq_actual)

        rmsd = torch.sqrt(torch.mean((coords_x_new - coords_x) ** 2))

        if dq_delta_prev is not None:
            if dq_delta > dq_delta_prev:
                damp *= 0.5
                failures += 1
            else:
                damp = min(damp * 1.2, 1.0)
                failures = 0

                coords_final = coords_x_new, coords_q_new
                dq_delta_prev = dq_delta
        else:
            dq_delta_prev = dq_delta

        if (rmsd < 1e-6 or dq_delta < 1e-6) or failures >= 5 or iteration >= 50:
            converged = dq_delta <= 1e-1

            coords_x, coords_q = coords_final if converged else coords_iter_1
            return coords_x, coords_q, converged

        dq = dq - dq_actual
        coords_x, coords_q = coords_x_new, coords_q_new
//  [python]
def _dq_to_x(
    coords_x: torch.Tensor, dq: torch.Tensor, ic: typing.Union["RIC", "DLC"]
) -> tuple[torch.Tensor, torch.Tensor, bool]:
    """Perturb a set of cartesian coordinates by a set of internal coordinate
    displacements.

    Args:
        coords_x: The coordinates with ``shape=(n_atoms, 3)``.
        dq: The internal coordinate displacements with ``shape=(n_ic,)``.

    Returns:
        The perturbed cartesian (with ``shape=(n_atoms, 3)``) and internal
        coordinates (with ``shape=(n_ic,)``), and a boolean indicating whether the
        perturbation was successful.
    """

    iteration, failures = 0, 0

    damp = 1.0

    coords_iter_1, coords_final = None, None

    dq_delta_prev = None

    while True:
        iteration += 1

        b_matrix = ic.compute_b(coords_x)
        g_matrix_inv = tico.utils.pinv(b_matrix @ b_matrix.T)

        dx = damp * torch.linalg.multi_dot([b_matrix.T, g_matrix_inv, dq])

        coords_x_new = coords_x + dx.reshape(-1, 3)
        coords_q_new = ic.compute_q(coords_x_new)

        if iteration == 1:
            coords_final = coords_iter_1 = coords_x_new, coords_q_new

        dq_actual = ic.compute_dq(coords_x, coords_x_new)
        dq_delta = torch.linalg.norm(dq - dq_actual)

        rmsd = torch.sqrt(torch.mean((coords_x_new - coords_x) ** 2))

        if dq_delta_prev is not None:
            if dq_delta > dq_delta_prev:
                damp *= 0.5
                failures += 1
            else:
                damp = min(damp * 1.2, 1.0)
                failures = 0

                coords_final = coords_x_new, coords_q_new
                dq_delta_prev = dq_delta
        else:
            dq_delta_prev = dq_delta

        if (rmsd < 1e-6 or dq_delta < 1e-6) or failures >= 5 or iteration >= 50:
            converged = dq_delta <= 1e-1

            coords_x, coords_q = coords_final if converged else coords_iter_1
            return coords_x, coords_q, converged

        dq = dq - dq_actual
        coords_x, coords_q = coords_x_new, coords_q_new
//  [python]
def forward(self, X: Tensor) -> Tensor:
        r"""
        Args:
            X: A "batch_shape x q x dim" representing the points to be evaluated.

        Returns:
            A tensor of size "batch_shape" representing the acqfn for each q-batch.
        """
        regularization_term = (
            torch.norm((X - self.init_point), p=2, dim=-1).max(dim=-1).values
            ** 2
        )
        return regularization_term
//  [python]
def forward(self, X: Tensor) -> Tensor:
        r"""
        Args:
            X: A "batch_shape x q x dim" representing the points to be evaluated.

        Returns:
            A tensor of size "batch_shape" representing the acqfn for each q-batch.
        """
        regularization_term = (
            torch.linalg.norm((X - self.init_point), ord=2, dim=-1).max(dim=-1).values
            ** 2
        )
        return regularization_term
//  [python]
def forward(self, X: Tensor) -> Tensor:
        r"""
        Args:
            X: A "batch_shape x q x dim" representing the points to be evaluated.

        Returns:
            A tensor of size "batch_shape" representing the acqfn for each q-batch.
        """
        sq_diff = torch.norm((X - self.init_point), p=2, dim=-1) ** 2
        pdf = torch.exp(sq_diff / 2 / self.sigma**2)
        regularization_term = pdf.max(dim=-1).values
        return regularization_term
//  [python]
def forward(self, X: Tensor) -> Tensor:
        r"""
        Args:
            X: A "batch_shape x q x dim" representing the points to be evaluated.

        Returns:
            A tensor of size "batch_shape" representing the acqfn for each q-batch.
        """
        sq_diff = torch.linalg.norm((X - self.init_point), ord=2, dim=-1) ** 2
        pdf = torch.exp(sq_diff / 2 / self.sigma**2)
        regularization_term = pdf.max(dim=-1).values
        return regularization_term
//  [python]
def group_lasso_regularizer(X: Tensor, groups: List[List[int]]) -> Tensor:
    r"""Computes the group lasso regularization function for the given point.

    Args:
        X: A bxd tensor representing the points to evaluate the regularization at.
        groups: List of indices of different groups.

    Returns:
        Computed group lasso norm of at the given points.
    """
    return torch.sum(
        torch.stack(
            [
                math.sqrt(len(g)) * torch.norm(X[..., g], p=2, dim=-1)
                for g in groups
            ],
            dim=-1,
        ),
        dim=-1,
    )
//  [python]
def group_lasso_regularizer(X: Tensor, groups: List[List[int]]) -> Tensor:
    r"""Computes the group lasso regularization function for the given point.

    Args:
        X: A bxd tensor representing the points to evaluate the regularization at.
        groups: List of indices of different groups.

    Returns:
        Computed group lasso norm of at the given points.
    """
    return torch.sum(
        torch.stack(
            [
                math.sqrt(len(g)) * torch.linalg.norm(X[..., g], ord=2, dim=-1)
                for g in groups
            ],
            dim=-1,
        ),
        dim=-1,
    )
//  [python]
def _apply_constraints(
        self, coords_x: torch.Tensor, tol: float = 1e-6
    ) -> tuple[torch.Tensor, torch.Tensor, bool]:
        """Attempt to take steps in ICs to enforce any internal coordinate constraints.

        Args:
            coords_x: The coordinates with shape=(n_atoms, 3).
            tol: The tolerance for the constraint satisfaction.

        Returns:
            The modified coordinates and internal coordinates, and a boolean indicating
            whether the conversion of cartesian to internal coordinates was successful.
        """
        iteration = 0

        constr_to_ic_idx = _match_constr(self.constr, self.idxs)

        dq_norm = 0.0
        dq_norm_best, soln_best = None, None

        while True:
            dq = smee.utils.zeros_like(self.v.shape[-1], self.v)

            # because of how we construct v, the first n_v columns should be the
            # projected constrained DoF
            constr_delta = -self.compute_constr_delta(coords_x)

            dq[: len(constr_delta)] = constr_delta
            dq[: len(constr_delta)] /= self.v[
                constr_to_ic_idx, list(range(len(constr_delta)))
            ]

            dq_norm_prev = dq_norm
            dq_norm = torch.norm(dq)

            soln = self.dq_to_x(coords_x, dq, enforce_constr=False)
            coords_x = soln[0]

            if dq_norm_best is None or dq_norm < dq_norm_best:
                dq_norm_best, soln_best = dq_norm, soln

            if dq_norm < tol:
                return soln

            if iteration > 0 and dq_norm > dq_norm_prev:
                _LOGGER.warning("constraint satisfaction failed to converge")
                return soln_best

            iteration += 1
//  [python]
def _apply_constraints(
        self, coords_x: torch.Tensor, tol: float = 1e-6
    ) -> tuple[torch.Tensor, torch.Tensor, bool]:
        """Attempt to take steps in ICs to enforce any internal coordinate constraints.

        Args:
            coords_x: The coordinates with shape=(n_atoms, 3).
            tol: The tolerance for the constraint satisfaction.

        Returns:
            The modified coordinates and internal coordinates, and a boolean indicating
            whether the conversion of cartesian to internal coordinates was successful.
        """
        iteration = 0

        constr_to_ic_idx = _match_constr(self.constr, self.idxs)

        dq_norm = 0.0
        dq_norm_best, soln_best = None, None

        while True:
            dq = smee.utils.zeros_like(self.v.shape[-1], self.v)

            # because of how we construct v, the first n_v columns should be the
            # projected constrained DoF
            constr_delta = -self.compute_constr_delta(coords_x)

            dq[: len(constr_delta)] = constr_delta
            dq[: len(constr_delta)] /= self.v[
                constr_to_ic_idx, list(range(len(constr_delta)))
            ]

            dq_norm_prev = dq_norm
            dq_norm = torch.linalg.norm(dq)

            soln = self.dq_to_x(coords_x, dq, enforce_constr=False)
            coords_x = soln[0]

            if dq_norm_best is None or dq_norm < dq_norm_best:
                dq_norm_best, soln_best = dq_norm, soln

            if dq_norm < tol:
                return soln

            if iteration > 0 and dq_norm > dq_norm_prev:
                _LOGGER.warning("constraint satisfaction failed to converge")
                return soln_best

            iteration += 1
//  [python]
def forward(self, X: Tensor) -> Tensor:
        r"""
        Args:
            X: A "batch_shape x q x dim" representing the points to be evaluated.

        Returns:
            A "1 x batch_shape x q" tensor representing the penalty for each point.
            The first dimension corresponds to the dimension of MC samples.
        """
        return torch.norm((X - self.init_point), p=1, dim=-1).unsqueeze(dim=0)
//  [python]
def forward(self, X: Tensor) -> Tensor:
        r"""
        Args:
            X: A "batch_shape x q x dim" representing the points to be evaluated.

        Returns:
            A "1 x batch_shape x q" tensor representing the penalty for each point.
            The first dimension corresponds to the dimension of MC samples.
        """
        return torch.linalg.norm((X - self.init_point), ord=1, dim=-1).unsqueeze(dim=0)
//  [python]
def find_symm_subs(xyz,Rs,metasymm):
    com = xyz[:,:,1].sum(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].sum(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 evaluate_true(self, X: Tensor) -> Tensor:
        norm = torch.norm(X, dim=-1)
        part1 = 1.0 + torch.cos(12.0 * norm)
        part2 = 0.5 * norm.pow(2) + 2.0
        return -part1 / part2
//  [python]
def evaluate_true(self, X: Tensor) -> Tensor:
        norm = torch.linalg.norm(X, dim=-1)
        part1 = 1.0 + torch.cos(12.0 * norm)
        part2 = 0.5 * norm.pow(2) + 2.0
        return -part1 / part2
//  [python]
def evaluate_true(self, X: Tensor) -> Tensor:
        a, b, c = self.a, self.b, self.c
        part1 = -a * torch.exp(-b / math.sqrt(self.dim) * torch.norm(X, dim=-1))
        part2 = -(torch.exp(torch.mean(torch.cos(c * X), dim=-1)))
        return part1 + part2 + a + math.e
//  [python]
def evaluate_true(self, X: Tensor) -> Tensor:
        a, b, c = self.a, self.b, self.c
        part1 = -a * torch.exp(-b / math.sqrt(self.dim) * torch.linalg.norm(X, dim=-1))
        part2 = -(torch.exp(torch.mean(torch.cos(c * X), dim=-1)))
        return part1 + part2 + a + math.e
//  [python]
def l2_square(bn_mean: Tensor,
              input_mean: Tensor,
              bn_std: Tensor,
              input_std: Tensor) -> Tensor:
    """
    Compute the L2 Square loss for batch normalization alignment.

    Args:
        bn_mean (Tensor): The mean of the batch normalization layer from the original statistics.
        input_mean (Tensor): The mean of the batch normalization layer from the current batch statistics.
        bn_std (Tensor): The standard deviation of the batch normalization layer from the original statistics.
        input_std (Tensor): The standard deviation of the batch normalization layer from the current batch statistics.

    Returns:
        Tensor: The L2 Square loss value for batch normalization alignment.
    """
    return torch.norm(input_mean - bn_mean) ** 2 / bn_mean.size(0) + \
           torch.norm(input_std - bn_std) ** 2 / bn_std.size(0)
//  [python]
def l2_square(bn_mean: Tensor,
              input_mean: Tensor,
              bn_std: Tensor,
              input_std: Tensor) -> Tensor:
    """
    Compute the L2 Square loss for batch normalization alignment.

    Args:
        bn_mean (Tensor): The mean of the batch normalization layer from the original statistics.
        input_mean (Tensor): The mean of the batch normalization layer from the current batch statistics.
        bn_std (Tensor): The standard deviation of the batch normalization layer from the original statistics.
        input_std (Tensor): The standard deviation of the batch normalization layer from the current batch statistics.

    Returns:
        Tensor: The L2 Square loss value for batch normalization alignment.
    """
    return torch.linalg.norm(input_mean - bn_mean) ** 2 / bn_mean.size(0) + \
           torch.linalg.norm(input_std - bn_std) ** 2 / bn_std.size(0)
//  [python]
def evaluate_true(self, X: Tensor) -> Tensor:
        term = torch.abs(1 - torch.norm(X, dim=-1) / math.pi)
        return -(
            torch.abs(torch.sin(X[..., 0]) * torch.cos(X[..., 1]) * torch.exp(term))
        )
//  [python]
def evaluate_true(self, X: Tensor) -> Tensor:
        term = torch.abs(1 - torch.linalg.norm(X, dim=-1) / math.pi)
        return -(
            torch.abs(torch.sin(X[..., 0]) * torch.cos(X[..., 1]) * torch.exp(term))
        )
//  [python]
def normalize(x):
    return x / torch.norm(x, dim=-2, keepdim=True)
//  [python]
def normalize(x):
    return x / torch.linalg.norm(x, dim=-2, keepdim=True)
//  [python]
def smallest_eigenvector_power_method(V, n_iter=10):

    assert V.shape[-2] == V.shape[-1]

    # compute inverse of input matrix
    V_inv = inv_loaded(V, load=1e-5)

    # initial point x0 = [1, 0, ..., 0]
    x = V_inv[..., :, 0]
    x = normalize(x)

    for epoch in range(n_iter - 1):
        x = torch.einsum("...fcd,...fd->...fc", V_inv, x)
        ev = torch.norm(x, dim=-1)
        x = x / ev[..., None]

    # fix sign of first element to positive
    s = torch.conj(torch.sgn(x[..., 0]))
    x = x * s[..., None]

    return ev, x
//  [python]
def smallest_eigenvector_power_method(V, n_iter=10):

    assert V.shape[-2] == V.shape[-1]

    # compute inverse of input matrix
    V_inv = inv_loaded(V, load=1e-5)

    # initial point x0 = [1, 0, ..., 0]
    x = V_inv[..., :, 0]
    x = normalize(x)

    for epoch in range(n_iter - 1):
        x = torch.einsum("...fcd,...fd->...fc", V_inv, x)
        ev = torch.linalg.norm(x, dim=-1)
        x = x / ev[..., None]

    # fix sign of first element to positive
    s = torch.conj(torch.sgn(x[..., 0]))
    x = x * s[..., None]

    return ev, x
//  [python]
def test_initialize_power(self):
        kernel = LinearTruncatedFidelityKernel(fidelity_dims=[1, 2], dimension=3)
        kernel.initialize(power=1)
        actual_value = torch.tensor(1, dtype=torch.float).view_as(kernel.power)
        self.assertLess(torch.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power(self):
        kernel = LinearTruncatedFidelityKernel(fidelity_dims=[1, 2], dimension=3)
        kernel.initialize(power=1)
        actual_value = torch.tensor(1, dtype=torch.float).view_as(kernel.power)
        self.assertLess(torch.linalg.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power_batch(self):
        kernel = LinearTruncatedFidelityKernel(
            fidelity_dims=[1, 2], dimension=3, batch_shape=torch.Size([2])
        )
        power_init = torch.tensor([1, 2], dtype=torch.float)
        kernel.initialize(power=power_init)
        actual_value = power_init.view_as(kernel.power)
        self.assertLess(torch.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power_batch(self):
        kernel = LinearTruncatedFidelityKernel(
            fidelity_dims=[1, 2], dimension=3, batch_shape=torch.Size([2])
        )
        power_init = torch.tensor([1, 2], dtype=torch.float)
        kernel.initialize(power=power_init)
        actual_value = power_init.view_as(kernel.power)
        self.assertLess(torch.linalg.norm(kernel.power - actual_value), 1e-5)
//  [python]
def gen_rays_from_single_image(H, W, image, intrinsic, c2w, depth=None, mask=None):
    """
    generate rays in world space, for image image
    :param H:
    :param W:
    :param intrinsics: [3,3]
    :param c2ws: [4,4]
    :return:
    """
    device = image.device
    ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                            torch.linspace(0, W - 1, W))  # pytorch's meshgrid has indexing='ij'
    p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

    # normalized ndc uv coordinates, (-1, 1)
    ndc_u = 2 * xs / (W - 1) - 1
    ndc_v = 2 * ys / (H - 1) - 1
    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)

    intrinsic_inv = torch.inverse(intrinsic)

    p = p.view(-1, 3).float().to(device)  # N_rays, 3
    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
    rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)  # N_rays, 3
    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3

    image = image.permute(1, 2, 0)
    color = image.view(-1, 3)
    depth = depth.view(-1, 1) if depth is not None else None
    mask = mask.view(-1, 1) if mask is not None else torch.ones([H * W, 1]).to(device)
    sample = {
        'rays_o': rays_o,
        'rays_v': rays_v,
        'rays_ndc_uv': rays_ndc_uv,
        'rays_color': color,
        # 'rays_depth': depth,
        'rays_mask': mask,
        'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
    }
    if depth is not None:
        sample['rays_depth'] = depth

    return sample
//  [python]
def gen_rays_from_single_image(H, W, image, intrinsic, c2w, depth=None, mask=None):
    """
    generate rays in world space, for image image
    :param H:
    :param W:
    :param intrinsics: [3,3]
    :param c2ws: [4,4]
    :return:
    """
    device = image.device
    ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
                            torch.linspace(0, W - 1, W))  # pytorch's meshgrid has indexing='ij'
    p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3

    # normalized ndc uv coordinates, (-1, 1)
    ndc_u = 2 * xs / (W - 1) - 1
    ndc_v = 2 * ys / (H - 1) - 1
    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)

    intrinsic_inv = torch.inverse(intrinsic)

    p = p.view(-1, 3).float().to(device)  # N_rays, 3
    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays, 3
    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3

    image = image.permute(1, 2, 0)
    color = image.view(-1, 3)
    depth = depth.view(-1, 1) if depth is not None else None
    mask = mask.view(-1, 1) if mask is not None else torch.ones([H * W, 1]).to(device)
    sample = {
        'rays_o': rays_o,
        'rays_v': rays_v,
        'rays_ndc_uv': rays_ndc_uv,
        'rays_color': color,
        # 'rays_depth': depth,
        'rays_mask': mask,
        'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
    }
    if depth is not None:
        sample['rays_depth'] = depth

    return sample
//  [python]
def test_project_gaussians_forward():
    num_points = 100

    means3d = torch.randn((num_points, 3), device=device, requires_grad=True)
    scales = torch.rand((num_points, 3), device=device) + 0.2
    glob_scale = 0.1
    quats = torch.randn((num_points, 4), device=device)
    quats /= torch.norm(quats, dim=-1, keepdim=True)

    H, W = 512, 512
    cx, cy = W / 2, H / 2
    # 90 degree FOV
    fx, fy = W / 2, W / 2
    clip_thresh = 0.01
    viewmat = torch.tensor(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 8.0],
            [0.0, 0.0, 0.0, 1.0],
        ],
        device=device,
    )
    viewmat[:3, :3] = _torch_impl.quat_to_rotmat(torch.randn(4))
    BLOCK_SIZE = 16

    (
        xys,
        depths,
        radii,
        conics,
        compensation,
        num_tiles_hit,
        cov3d,
    ) = project_gaussians(
        means3d,
        scales,
        glob_scale,
        quats,
        viewmat,
        fx,
        fy,
        cx,
        cy,
        H,
        W,
        BLOCK_SIZE,
        clip_thresh,
    )
    masks = num_tiles_hit > 0

    with torch.no_grad():
        (
            _cov3d,
            _,
            _xys,
            _depths,
            _radii,
            _conics,
            _compensation,
            _num_tiles_hit,
            _masks,
        ) = _torch_impl.project_gaussians_forward(
            means3d,
            scales,
            glob_scale,
            quats,
            viewmat,
            (fx, fy, cx, cy),
            (W, H),
            BLOCK_SIZE,
            clip_thresh,
        )

    check_close(masks, _masks, atol=1e-5, rtol=1e-5)
    check_close(cov3d, _cov3d)
    check_close(xys, _xys)
    check_close(depths, _depths)
    check_close(radii, _radii)
    check_close(conics, _conics)
    check_close(compensation, _compensation)
    check_close(num_tiles_hit, _num_tiles_hit)
    print("passed project_gaussians_forward test")
//  [python]
def test_project_gaussians_forward():
    num_points = 100

    means3d = torch.randn((num_points, 3), device=device, requires_grad=True)
    scales = torch.rand((num_points, 3), device=device) + 0.2
    glob_scale = 0.1
    quats = torch.randn((num_points, 4), device=device)
    quats /= torch.linalg.norm(quats, dim=-1, keepdim=True)

    H, W = 512, 512
    cx, cy = W / 2, H / 2
    # 90 degree FOV
    fx, fy = W / 2, W / 2
    clip_thresh = 0.01
    viewmat = torch.tensor(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 8.0],
            [0.0, 0.0, 0.0, 1.0],
        ],
        device=device,
    )
    viewmat[:3, :3] = _torch_impl.quat_to_rotmat(torch.randn(4))
    BLOCK_SIZE = 16

    (
        xys,
        depths,
        radii,
        conics,
        compensation,
        num_tiles_hit,
        cov3d,
    ) = project_gaussians(
        means3d,
        scales,
        glob_scale,
        quats,
        viewmat,
        fx,
        fy,
        cx,
        cy,
        H,
        W,
        BLOCK_SIZE,
        clip_thresh,
    )
    masks = num_tiles_hit > 0

    with torch.no_grad():
        (
            _cov3d,
            _,
            _xys,
            _depths,
            _radii,
            _conics,
            _compensation,
            _num_tiles_hit,
            _masks,
        ) = _torch_impl.project_gaussians_forward(
            means3d,
            scales,
            glob_scale,
            quats,
            viewmat,
            (fx, fy, cx, cy),
            (W, H),
            BLOCK_SIZE,
            clip_thresh,
        )

    check_close(masks, _masks, atol=1e-5, rtol=1e-5)
    check_close(cov3d, _cov3d)
    check_close(xys, _xys)
    check_close(depths, _depths)
    check_close(radii, _radii)
    check_close(conics, _conics)
    check_close(compensation, _compensation)
    check_close(num_tiles_hit, _num_tiles_hit)
    print("passed project_gaussians_forward test")
//  [python]
def norm(array, axis=None, keepdims=False, ord=None):
    return torch.norm(array, dim=axis, keepdim=keepdims, p=ord)
//  [python]
def norm(array, axis=None, keepdims=False, ord=None):
    return torch.linalg.norm(array, dim=axis, keepdim=keepdims, ord=ord)
//  [python]
def test_active_dims_list(self):
        kernel = LinearTruncatedFidelityKernel(
            fidelity_dims=[1, 2], dimension=10, active_dims=[0, 2, 4, 6]
        )
        x = self.create_data_no_batch()
        covar_mat = kernel(x).evaluate_kernel().to_dense()
        kernel_basic = LinearTruncatedFidelityKernel(fidelity_dims=[1, 2], dimension=4)
        covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().to_dense()
        self.assertLess(
            torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(),
            1e-4,
        )
//  [python]
def test_active_dims_list(self):
        kernel = LinearTruncatedFidelityKernel(
            fidelity_dims=[1, 2], dimension=10, active_dims=[0, 2, 4, 6]
        )
        x = self.create_data_no_batch()
        covar_mat = kernel(x).evaluate_kernel().to_dense()
        kernel_basic = LinearTruncatedFidelityKernel(fidelity_dims=[1, 2], dimension=4)
        covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().to_dense()
        self.assertLess(
            torch.linalg.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(),
            1e-4,
        )
//  [python]
def test_compute_linear_truncated_kernel_no_batch(self):
        x1 = torch.tensor([[1, 0.1, 0.2], [2, 0.3, 0.4]])
        x2 = torch.tensor([[3, 0.5, 0.6], [4, 0.7, 0.8]])
        t_1 = torch.tensor([[0.3584, 0.1856], [0.2976, 0.1584]])
        for nu, fidelity_dims in itertools.product({0.5, 1.5, 2.5}, ([2], [1, 2])):
            kernel = LinearTruncatedFidelityKernel(
                fidelity_dims=fidelity_dims, dimension=3, nu=nu
            )
            kernel.power = 1
            n_fid = len(fidelity_dims)
            if n_fid > 1:
                active_dimsM = [0]
                t_2 = torch.tensor([[0.4725, 0.2889], [0.4025, 0.2541]])
                t_3 = torch.tensor([[0.1685, 0.0531], [0.1168, 0.0386]])
                t = 1 + t_1 + t_2 + t_3
            else:
                active_dimsM = [0, 1]
                t = 1 + t_1

            matern_ker = MaternKernel(nu=nu, active_dims=active_dimsM)
            matern_term = matern_ker(x1, x2).to_dense()
            actual = t * matern_term
            res = kernel(x1, x2).to_dense()
            self.assertLess(torch.norm(res - actual), 1e-4)
            # test diagonal mode
            res_diag = kernel(x1, x2, diag=True)
            self.assertLess(torch.norm(res_diag - actual.diag()), 1e-4)
        # make sure that we error out if last_dim_is_batch=True
        with self.assertRaises(NotImplementedError):
            kernel(x1, x2, diag=True, last_dim_is_batch=True)
//  [python]
def test_compute_linear_truncated_kernel_no_batch(self):
        x1 = torch.tensor([[1, 0.1, 0.2], [2, 0.3, 0.4]])
        x2 = torch.tensor([[3, 0.5, 0.6], [4, 0.7, 0.8]])
        t_1 = torch.tensor([[0.3584, 0.1856], [0.2976, 0.1584]])
        for nu, fidelity_dims in itertools.product({0.5, 1.5, 2.5}, ([2], [1, 2])):
            kernel = LinearTruncatedFidelityKernel(
                fidelity_dims=fidelity_dims, dimension=3, nu=nu
            )
            kernel.power = 1
            n_fid = len(fidelity_dims)
            if n_fid > 1:
                active_dimsM = [0]
                t_2 = torch.tensor([[0.4725, 0.2889], [0.4025, 0.2541]])
                t_3 = torch.tensor([[0.1685, 0.0531], [0.1168, 0.0386]])
                t = 1 + t_1 + t_2 + t_3
            else:
                active_dimsM = [0, 1]
                t = 1 + t_1

            matern_ker = MaternKernel(nu=nu, active_dims=active_dimsM)
            matern_term = matern_ker(x1, x2).to_dense()
            actual = t * matern_term
            res = kernel(x1, x2).to_dense()
            self.assertLess(torch.linalg.norm(res - actual), 1e-4)
            # test diagonal mode
            res_diag = kernel(x1, x2, diag=True)
            self.assertLess(torch.linalg.norm(res_diag - actual.diag()), 1e-4)
        # make sure that we error out if last_dim_is_batch=True
        with self.assertRaises(NotImplementedError):
            kernel(x1, x2, diag=True, last_dim_is_batch=True)
//  [python]
def test_active_dims_range(self):
        active_dims = list(range(3, 9))
        kernel = LinearTruncatedFidelityKernel(
            fidelity_dims=[1, 2], dimension=10, active_dims=active_dims
        )
        x = self.create_data_no_batch()
        covar_mat = kernel(x).evaluate_kernel().to_dense()
        kernel_basic = LinearTruncatedFidelityKernel(fidelity_dims=[1, 2], dimension=6)
        covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().to_dense()
        self.assertLess(
            torch.norm(covar_mat - covar_mat_actual) / torch.norm(covar_mat_actual),
            1e-4,
        )
//  [python]
def test_active_dims_range(self):
        active_dims = list(range(3, 9))
        kernel = LinearTruncatedFidelityKernel(
            fidelity_dims=[1, 2], dimension=10, active_dims=active_dims
        )
        x = self.create_data_no_batch()
        covar_mat = kernel(x).evaluate_kernel().to_dense()
        kernel_basic = LinearTruncatedFidelityKernel(fidelity_dims=[1, 2], dimension=6)
        covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().to_dense()
        self.assertLess(
            torch.linalg.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(),
            1e-4,
        )
//  [python]
def normalize(v):
    if isinstance(v, torch.Tensor):
        norm = torch.norm(v)
    else:
        norm = np.linalg.norm(v)
    if norm == 0.:
        return v
    return v / norm
//  [python]
def normalize(v):
    if isinstance(v, torch.Tensor):
        norm = torch.linalg.norm(v)
    else:
        norm = np.linalg.norm(v)
    if norm == 0.:
        return v
    return v / norm
//  [python]
def test_model_normalization(self):
        names = self.open_clip_test_model

        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-6

        for name in names:
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

            for sentence in sentences:
                output = model.encode(sentence, normalize=True)
                output = _convert_vectorized_output(output)
                max_output_norm = max(torch.norm(FloatTensor(output), dim=1))
                min_output_norm = min(torch.norm(FloatTensor(output), dim=1))

                assert abs(max_output_norm - 1) < eps, f"{name}, {sentence}"
                assert abs(min_output_norm - 1) < eps, f"{name}, {sentence}"

            clear_loaded_models()
//  [python]
def test_model_normalization(self):
        names = self.open_clip_test_model

        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-6

        for name in names:
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

            for sentence in sentences:
                output = model.encode(sentence, normalize=True)
                output = _convert_vectorized_output(output)
                max_output_norm = max(torch.linalg.norm(FloatTensor(output), dim=1))
                min_output_norm = min(torch.linalg.norm(FloatTensor(output), dim=1))

                assert abs(max_output_norm - 1) < eps, f"{name}, {sentence}"
                assert abs(min_output_norm - 1) < eps, f"{name}, {sentence}"

            clear_loaded_models()
//  [python]
def test_model_normalization(self):
        names = ["fp16/ViT-B/32", "onnx16/open_clip/ViT-B-32/laion400m_e32", 'onnx32/open_clip/ViT-B-32-quickgelu/laion400m_e32',
                 'RN50', "ViT-B/16", "all-MiniLM-L6-v1",
                 "all_datasets_v4_MiniLM-L6", "hf/all-MiniLM-L6-v1", "hf/all_datasets_v4_MiniLM-L6",
                 "onnx/all-MiniLM-L6-v1", "onnx/all_datasets_v4_MiniLM-L6"]

        names_e5 = ["hf/e5-small", "hf/e5-base", "hf/e5-small-unsupervised", "hf/e5-base-unsupervised", "hf/e5-base-v2",
                    "hf/multilingual-e5-small"]
        names += names_e5
                 
        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-6

        for name in names:
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

            for sentence in sentences:
                output = model.encode(sentence, normalize=True)
                output = _convert_vectorized_output(output)
                max_output_norm = max(torch.norm(FloatTensor(output), dim=1))
                min_output_norm = min(torch.norm(FloatTensor(output), dim=1))

                assert abs(max_output_norm - 1) < eps, f"{name}, {sentence}"
                assert abs(min_output_norm - 1) < eps, f"{name}, {sentence}"

            clear_loaded_models()
//  [python]
def test_model_normalization(self):
        names = ["fp16/ViT-B/32", "onnx16/open_clip/ViT-B-32/laion400m_e32", 'onnx32/open_clip/ViT-B-32-quickgelu/laion400m_e32',
                 'RN50', "ViT-B/16", "all-MiniLM-L6-v1",
                 "all_datasets_v4_MiniLM-L6", "hf/all-MiniLM-L6-v1", "hf/all_datasets_v4_MiniLM-L6",
                 "onnx/all-MiniLM-L6-v1", "onnx/all_datasets_v4_MiniLM-L6"]

        names_e5 = ["hf/e5-small", "hf/e5-base", "hf/e5-small-unsupervised", "hf/e5-base-unsupervised", "hf/e5-base-v2",
                    "hf/multilingual-e5-small"]
        names += names_e5
                 
        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-6

        for name in names:
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

            for sentence in sentences:
                output = model.encode(sentence, normalize=True)
                output = _convert_vectorized_output(output)
                max_output_norm = max(torch.linalg.norm(FloatTensor(output), dim=1))
                min_output_norm = min(torch.linalg.norm(FloatTensor(output), dim=1))

                assert abs(max_output_norm - 1) < eps, f"{name}, {sentence}"
                assert abs(min_output_norm - 1) < eps, f"{name}, {sentence}"

            clear_loaded_models()