Corpus-normalized gene rank tokenization

Your task

Function signature

rank_value_tokenization(
    counts: torch.Tensor,
    library_size: torch.Tensor,
    gene_medians: torch.Tensor,
    gene_mask: torch.Tensor,
    gene_token_ids: torch.Tensor,
    max_tokens: int = 2048,
    target_sum: float = 10000.0,
    pad_token_id: int = 0,
) -> dict[str, torch.Tensor]

Worked example

Rank expression before mapping gene IDs

import torch

counts = torch.tensor([[3.0, 1.0, 4.0, 0.0], [0.0, 2.0, 1.0, 0.0]], dtype=torch.float64, requires_grad=True)
library_size = torch.tensor([10.0, 5.0], dtype=torch.float64, requires_grad=True)
gene_medians = torch.tensor([3.0, 1.0, 2.0, float('nan')], dtype=torch.float64, requires_grad=True)
gene_mask = torch.tensor([True, True, True, False], dtype=torch.bool)
gene_token_ids = torch.tensor([41, 7, 29, -1], dtype=torch.int64)
max_tokens = 4
target_sum = 10.0
pad_token_id = 0

result = rank_value_tokenization(counts=counts, library_size=library_size, gene_medians=gene_medians, gene_mask=gene_mask, gene_token_ids=gene_token_ids, max_tokens=max_tokens, target_sum=target_sum, pad_token_id=pad_token_id)

Expected return

{
    'normalized_expression': torch.tensor([[1.0, 1.0, 2.0, 0.0], [0.0, 4.0, 1.0, 0.0]], dtype=torch.float64),
    'ranked_expression': torch.tensor([[2.0, 1.0, 1.0, 0.0], [4.0, 1.0, 0.0, 0.0]], dtype=torch.float64),
    'tokens': torch.tensor([[29, 41, 7, 0], [7, 29, 0, 0]], dtype=torch.int64),
    'gene_indices': torch.tensor([[2, 0, 1, -1], [1, 2, -1, -1]], dtype=torch.int64),
    'valid_mask': torch.tensor([[True, True, True, False], [True, True, False, False]], dtype=torch.bool),
    'lengths': torch.tensor([3, 2], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• For B\ge 1 cells and G\ge 1 supplied gene columns, compute normalized expression s[b,g] = \mathrm{counts}[b,g] / \mathrm{library\_size}[b] \cdot \mathrm{target\_sum} / \mathrm{gene\_medians}[g] when gene_mask [g] is true, otherwise zero.
• Counts are finite nonnegative \mathrm{float32}/\mathrm{float64}; all three floating inputs share dtype. The positive finite library_size [B] is supplied from the original full transcriptome before vocabulary filtering and may include genes outside these columns; do not replace it by a sum over retained genes.
• Selected gene medians must be finite and strictly positive. Excluded median/ID payloads are ignored, even if a median is NaN. The selected gene_token_ids [G] must be unique nonnegative int64 values unequal to the pad ID. gene_mask [G] is Boolean. Within each cell, keep only selected genes whose counts are strictly positive.
• Sort them by decreasing s, breaking exact ties by LOWER INPUT COLUMN INDEX (not gene token ID). Keep the first L=\mathrm{max\_tokens}, and pad to exactly L entries.
• Return normalized_expression[B,G], ranked_expression[B,L] (zero on padding), tokens[B,L] (pad ID on padding), gene_indices[B,L] (original column or -1), valid_mask[B,L] (true for real tokens), and lengths [B] (retained count after truncation).
• Tokens, indices and lengths are int64; masks are Boolean; expression outputs preserve floating dtype. Empty detected or selected sets yield only padding and zero ranked expression.
• Config: max_tokens defaults to 2048 and must be an integer between 1 and 2048, inclusive; target_sum defaults to 10000.0 and must be finite positive; pad_token_id defaults to 0 and must be a nonnegative integer. Invalid support/resources/configuration raise ValueError. No special CLS/EOS tokens are added: this is the V1 interface.
• Lower-index stable ties are an explicit deterministic teaching convention, not a claim about upstream NumPy's unspecified sort ties. Normalize, rank, select and map IDs from scratch; no sorting API exception is provided.

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 rank_value_tokenization(
    counts: torch.Tensor,  # (B, G)
    library_size: torch.Tensor,  # (B,)
    gene_medians: torch.Tensor,  # (G,)
    gene_mask: torch.Tensor,  # (G,)
    gene_token_ids: torch.Tensor,  # (G,)
    max_tokens: int = 2048,
    target_sum: float = 10000.0,
    pad_token_id: int = 0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    normalized_expression: (B, G).
    ranked_expression: (B, L).
    tokens: (B, L).
    gene_indices: (B, L).
    valid_mask: (B, L).
    lengths: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rank_value_tokenization")

Sources and further reading

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