From ef1bb6f40105ed7523f377a0da53535971fe31f2 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Thu, 6 Aug 2026 21:14:04 +0800
Subject: [PATCH] refactor: unify greedy check with _is_greedy helper

- Replace batch-scattered temperature==0 checks with (temperature == 0).all()

- Reuse _is_greedy in standalone sample() function
---
 astrai/inference/sample.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/astrai/inference/sample.py b/astrai/inference/sample.py
index 36186ce..df6a240 100644
--- a/astrai/inference/sample.py
+++ b/astrai/inference/sample.py
@@ -266,7 +266,7 @@ class SamplingPipeline(BaseSamplingStrategy):
     @staticmethod
     def _is_greedy(temperature: Union[float, Tensor]) -> bool:
         if isinstance(temperature, Tensor):
-            return temperature.numel() == 1 and temperature.item() == 0
+            return bool((temperature == 0).all())
         return temperature == 0
 
     @torch.inference_mode()
@@ -364,11 +364,7 @@ def sample(
         ``chosen_logprobs`` has shape ``[batch]``.
     """
     greedy = (
-        (
-            isinstance(temperature, Tensor)
-            and temperature.numel() == 1
-            and temperature.item() == 0
-        )
+        bool((temperature == 0).all())
         if isinstance(temperature, Tensor)
         else temperature == 0
     )