Skip to content

Instantly share code, notes, and snippets.

@YouJiacheng
Created May 30, 2025 09:09
Show Gist options
  • Select an option

  • Save YouJiacheng/096753a5ad12260c73d3e166a6ad7f52 to your computer and use it in GitHub Desktop.

Select an option

Save YouJiacheng/096753a5ad12260c73d3e166a6ad7f52 to your computer and use it in GitHub Desktop.
supports CUDA, no graph breaks
def cdf(t: Tensor, thresholds: list[float]):
level = torch.bucketize(t, t.new_tensor(thresholds, dtype=torch.float32), out_int32=True) # sum(x > v for v in thresholds)
count = t.new_zeros(len(thresholds) + 1, dtype=torch.int32)
count.index_put_((level,), count.new_ones(1), True)
return count.cumsum(0) / t.numel()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment