From 11a8e72eeff86a6f31805e9c6a942c2ea69ffa74 Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Sat, 25 Jul 2026 10:01:29 +0530 Subject: [PATCH] Fix wrong parameter name in diffusion attention-mask error message `_make_global_attention_mask` raises when its `num_valid_tokens` argument is None, but the ValueError text referred to `num_valid_samples`, a name that does not exist in the function signature. The sibling guard in `_make_block_local_attention_mask` already uses the correct name. Align the message with the actual parameter so the error is actionable. --- gemma/diffusion/_sampler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gemma/diffusion/_sampler.py b/gemma/diffusion/_sampler.py index 5ead7050..8e586628 100644 --- a/gemma/diffusion/_sampler.py +++ b/gemma/diffusion/_sampler.py @@ -685,7 +685,7 @@ def _make_global_attention_mask( if num_valid_tokens is None: raise ValueError( - 'num_valid_samples must be provided if cache_length is set.' + 'num_valid_tokens must be provided if cache_length is set.' ) total_valid = jnp.minimum(num_valid_tokens + canvas_length, cache_length)