integrations.kd.utils

integrations.kd.utils

Helper KD utils

Functions

Name	Description
normalize_logprobs	Re-normalizes top-k raw logprobs as probabilities, and converts back to logprobs.
strided_chunk_views	Split a tensor into chunks along a dimension with striding, prioritizing views over copies.

integrations.kd.utils.normalize_logprobs(logprobs, topk)

Re-normalizes top-k raw logprobs as probabilities, and converts back to logprobs.

integrations.kd.utils.strided_chunk_views(
    tensor,
    chunks,
    dim=0,
    stride=1,
    chunk_size=None,
)

Split a tensor into chunks along a dimension with striding, prioritizing views over copies.

Name	Type	Description	Default
tensor	Union[np.ndarray, torch.Tensor]	Input tensor (numpy array or torch tensor)	required
chunks	int	Number of chunks to create	required
dim	int	Dimension along which to chunk (default: 0)	`0`
stride	int	Stride between chunk starting positions (default: 1)	`1`
chunk_size	int \| None	Size of each chunk. If None, calculated automatically (default: None)	`None`

Name	Type	Description
	List[Union[np.ndarray, torch.Tensor]]	List of tensor chunks (views when possible, copies when necessary)