Paired gene-embedding perturbation statistics

Your task

Function signature

embedding_perturbation_comparison(
    reference_embeddings: torch.Tensor,
    edited_embeddings: torch.Tensor,
    state_embeddings: torch.Tensor,
    reference_indices: torch.Tensor,
    edited_indices: torch.Tensor,
    pair_mask: torch.Tensor,
    cell_mask: torch.Tensor,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]

Worked example

Paired reencoded genes and signed state shifts

import torch

reference_embeddings = torch.tensor([[[1.0, 0.0], [0.0, 1.0], [7.0, 9.0]], [[1.0, 0.0], [3.0, 8.0], [9.0, 2.0]]], dtype=torch.float64, requires_grad=True)
edited_embeddings = torch.tensor([[[0.0, 1.0], [1.0, 0.0], [8.0, 9.0], [2.0, 5.0]], [[7.0, 2.0], [0.0, 1.0], [4.0, 6.0], [1.0, 8.0]]], dtype=torch.float64, requires_grad=True)
state_embeddings = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64, requires_grad=True)
reference_indices = torch.tensor([[0, 1, -1], [0, -8, 99]], dtype=torch.int64)
edited_indices = torch.tensor([[0, 1, 30], [1, -7, -1]], dtype=torch.int64)
pair_mask = torch.tensor([[True, True, False], [True, False, False]], dtype=torch.bool)
cell_mask = torch.tensor([True, True], dtype=torch.bool)
eps = 1e-08

result = embedding_perturbation_comparison(reference_embeddings=reference_embeddings, edited_embeddings=edited_embeddings, state_embeddings=state_embeddings, reference_indices=reference_indices, edited_indices=edited_indices, pair_mask=pair_mask, cell_mask=cell_mask, eps=eps)

Expected return

{
    'gene_cosine': torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=torch.float64),
    'gene_squared_distance': torch.tensor([[2.0, 2.0, 0.0], [2.0, 0.0, 0.0]], dtype=torch.float64),
    'mean_gene_cosine_distance': torch.tensor([1.0, 1.0], dtype=torch.float64),
    'reference_cell': torch.tensor([[0.5, 0.5], [1.0, 0.0]], dtype=torch.float64),
    'edited_cell': torch.tensor([[0.5, 0.5], [0.0, 1.0]], dtype=torch.float64),
    'cell_cosine_distance': torch.tensor([2.220446049250313e-16, 1.0], dtype=torch.float64),
    'state_cosine_shift': torch.tensor([[0.0, 0.0], [-1.0, 1.0]], dtype=torch.float64),
    'state_shift_mean': torch.tensor([-0.5, 0.5], dtype=torch.float64),
    'state_shift_variance': torch.tensor([0.25, 0.25], dtype=torch.float64),
    'comparison_count': torch.tensor([2, 1], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Supply reference_embeddings [B,R,D], edited_embeddings [B,E,D] and state_embeddings [S,D] of one floating dtype (\mathrm{float32}/\mathrm{float64}), with B,R,E,D,S\ge 1. The two sequences have already been encoded separately with identical chosen model/layer settings.
• Supply int64 reference_indices [B,K], edited_indices [B,K], Boolean pair_mask [B,K] and cell_mask [B], K\ge 1. Effective correspondence mask M=\mathrm{pair\_mask} \land \mathrm{cell\_mask}[:,\mathrm{None}].
• A true pair aligns the SAME unperturbed gene, even if its rank position changed; the caller excludes removed/inserted/promoted genes, padding and boundary specials. Only active indices must lie in range. Masked/unreferenced embedding payloads and inactive indices are ignored, even NaNs or arbitrary integer sentinels.
• Active gathered vectors and all state prototypes must be finite. Repeated active correspondence indices are permitted, count as repeated observations and must accumulate gradients; a biological one-gene-one-pair protocol should normally supply unique alignments. Define:
• C(u,v)=\sum_{d}(u[d]\cdot v[d]) / (\sqrt{\operatorname{max}(\sum_{d} u[d]^{2},\varepsilon^{2})}\cdot \sqrt{\operatorname{max}(\sum_{d} v[d]^{2},\varepsilon^{2})}). \varepsilon defaults to 10^{-8} and must lie in [10^{-8},1]. Thus a zero vector has cosine zero, and a valid zero/zero pair has cosine distance one.
• On each active pair output \mathrm{gene\_cosine}=C(r,e) and \mathrm{gene\_squared\_distance}=\sum_{d}(r-e)^{2}; output zero at inactive pairs. \mathrm{comparison\_count}[b]=n_{b}=\sum_{k} M[b,k].
• Pool BOTH reference_cell [b,D] and edited_cell [b,D] by averaging the corresponding matched vectors over n_b, or zero when n_{b}=0. mean_gene_cosine_distance [b] is the average of 1-\mathrm{gene\_cosine} over active pairs; \mathrm{cell\_cosine\_distance}[b]=1-C(\mathrm{reference\_cell},\mathrm{edited\_cell}).
• These two distances differ in general; output zero for either when n_{b}=0.
• For nonempty cells, \mathrm{state\_cosine\_shift}[b,s]=C(\mathrm{edited\_cell}[b],\mathrm{state\_embeddings}[s])-C(\mathrm{reference\_cell}[b],\mathrm{state\_embeddings}[s]); positive means closer in cosine similarity to that supplied prototype. Empty/excluded rows are zero.
• Over the N cells with n_{b}>0, return state_shift_mean [S] and \mathrm{state\_shift\_variance}[S]=\sum_{b}(\mathrm{shift}[b,s]-\operatorname{mean}[s])^{2}/N (POPULATION variance, not N-1). When N=0 both are zero, with finite zero gradients. comparison_count is int64; every other output preserves floating dtype.
• No data-dependent resampling, p-value, encoder execution or learned classifier is involved. Active out-of-range indices, nonfinite active values/prototypes and invalid \varepsilon raise ValueError.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def embedding_perturbation_comparison(
    reference_embeddings: torch.Tensor,  # (B, R, D)
    edited_embeddings: torch.Tensor,  # (B, E, D)
    state_embeddings: torch.Tensor,  # (S, D)
    reference_indices: torch.Tensor,  # (B, K)
    edited_indices: torch.Tensor,  # (B, K)
    pair_mask: torch.Tensor,  # (B, K)
    cell_mask: torch.Tensor,  # (B,)
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    gene_cosine: (B, K).
    gene_squared_distance: (B, K).
    mean_gene_cosine_distance: (B,).
    reference_cell: (B, D).
    edited_cell: (B, D).
    cell_cosine_distance: (B,).
    state_cosine_shift: (B, S).
    state_shift_mean: (S,).
    state_shift_variance: (S,).
    comparison_count: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement embedding_perturbation_comparison")

Sources and further reading

• www.nature.com
• huggingface.co
• huggingface.co