Fix base sampling budget in adaptive

apd10 · apd10 · commit 9cc29968ac42 · 2025-07-21T13:39:06.000-05:00
diff --git a/sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/adaptive_sampling.py b/sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/adaptive_sampling.py
@@ -300,13 +300,13 @@ def add_mask(
             num_base_samples,
             previous_mask.dtype,
         )
-
         # Compute denominators and budget
         sampled_denominator = apply_inv_mask_sum(expwts, base_sampling_mask)
         estimated_denominator = static_denominator + sampled_denominator
         budget = self._compute_adaptive_budget(
             std_estimate, estimated_denominator, sampling_range
         )
+        budget = torch.clamp(budget, min=num_base_samples, max=sampling_range)
 
         # Create adaptive sampling mask
         sampling_probabilities = (budget / sampling_range).to(previous_mask.dtype)

Original file line number	Diff line number	Diff line change
`@@ -300,13 +300,13 @@ def add_mask(`
`300`	`300`	`num_base_samples,`
`301`	`301`	`previous_mask.dtype,`
`302`	`302`	`)`
`303`		`-`
`304`	`303`	`# Compute denominators and budget`
`305`	`304`	`sampled_denominator = apply_inv_mask_sum(expwts, base_sampling_mask)`
`306`	`305`	`estimated_denominator = static_denominator + sampled_denominator`
`307`	`306`	`budget = self._compute_adaptive_budget(`
`308`	`307`	`std_estimate, estimated_denominator, sampling_range`
`309`	`308`	`)`
	`309`	`+ budget = torch.clamp(budget, min=num_base_samples, max=sampling_range)`
`310`	`310`
`311`	`311`	`# Create adaptive sampling mask`
`312`	`312`	`sampling_probabilities = (budget / sampling_range).to(previous_mask.dtype)`