Add Pyrefly type ignore in FrozenLake data.

examples/frozenlake/data.py

@@ -95,12 +95,12 @@ def create_dataset(
     **kwargs
 ) -> grain.MapDataset:
   """Lively generates the dataset, saves it to a local directory, and returns a MapDataset.
-  
+
   This function acts as the entry point for the Tunix GRPO pipeline.
   """
   os.makedirs(data_dir, exist_ok=True)
   filepath = os.path.join(data_dir, f"{split}.parquet")
-  
+
   if not os.path.exists(filepath):
     size = train_size if split == "train" else test_size
     print(f"Dynamically generating {size} instances for '{split}' split...")
@@ -112,15 +112,15 @@ def create_dataset(
     save_dataset(data, filepath)
   else:
     print(f"Loading existing dataset from {filepath}")
-    
+
   df = pd.read_parquet(filepath)
   hf_ds = datasets_lib.Dataset.from_pandas(df)
-  
+
   def process_item(item):
     item["prompts"] = ""
     return item
-    
-  return grain.MapDataset.source(hf_ds).map(process_item)
+
+  return grain.MapDataset.source(hf_ds).map(process_item)  # pyrefly: ignore[bad-argument-type]
 
 
 def main():

tunix/cli/config.py

@@ -199,7 +199,7 @@ def __init__(self, argv: list[str], **kwargs):
 
     # Handle relative paths used in example scripts as a special case.
     # TODO(noghabi): Remove this once the example scripts are updated.
-    if argv[1] == "base_config.yaml":
+    if argv[1] in ("base_config.yaml", "base_agentic_config.yaml"):
       base_config_file = pathlib.Path(__file__).parent / argv[1]
     else:
       base_config_file = argv[1]

tunix/generate/sampler.py

@@ -84,7 +84,7 @@ class _SamplingState:
   sampling_mode: str = flax.struct.field(pytree_node=False)
 
   # Number of input tokens with padding.
-  num_input_tokens: jnp.int32 = flax.struct.field(pytree_node=False)
+  num_input_tokens: int = flax.struct.field(pytree_node=False)
 
   # Tempurature for top_p sampling.
   temperature: float = flax.struct.field(pytree_node=False)
@@ -446,7 +446,7 @@ def init_sample_state(
 
     return _SamplingState(
         decoding_step=num_input_tokens - 1,
-        num_input_tokens=jnp.array(num_input_tokens, dtype=jnp.int32),
+        num_input_tokens=int(num_input_tokens),
         token_buffer=token_buffer,
         positions=positions,
         logits_buffer=logits_buffer,

tunix/models/gemma4/attention.py

@@ -780,19 +780,7 @@ def init_cache(
       cache_len = min(max_seq_len, sliding_window_size)
 
     cache_shape = (batch_size, cache_len, self.num_kv_heads, self.head_dim)
-    k = shard(
-        np.zeros(cache_shape, dtype),
-        self.config.shd_config.act_btnh,
-        eager=True,
-    )
-    v = shard(
-        np.zeros(cache_shape, dtype),
-        self.config.shd_config.act_btnh,
-        eager=True,
-    )
-    end_index = shard(
-        np.zeros((batch_size,), np.int32),
-        self.config.shd_config.act_btnh[:1],
-        eager=True,
-    )
+    k = jnp.zeros(cache_shape, dtype=dtype)
+    v = jnp.zeros(cache_shape, dtype=dtype)
+    end_index = jnp.zeros((batch_size,), dtype=jnp.int32)
     return {'k': k, 'v': v, 'end_index': end_index}