address greptile comment

tdophung · tdophung · commit d313046ffaa1 · 2026-01-06T10:27:56.000-08:00
Signed-off-by: tdophung &lt;tdophung@nvidia.com&gt;
diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py
@@ -98,7 +98,7 @@ def reference_make_row_id_map(
     # Compute total tokens per expert and expert offsets
     tokens_per_expert = jnp.sum(routing_map, axis=0)
     expert_offsets = jnp.concatenate(
-        [jnp.array([0], dtype=jnp.int32), jnp.cumsum(tokens_per_expert)[:-1]]
+        [jnp.array([0], dtype=jnp.int32), jnp.cumsum(tokens_per_expert)[:-1].astype(jnp.int32)]
     )
 
     # Compute destination rows for all (token, expert) pairs

Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ def reference_make_row_id_map(`
`98`	`98`	`# Compute total tokens per expert and expert offsets`
`99`	`99`	`tokens_per_expert = jnp.sum(routing_map, axis=0)`
`100`	`100`	`expert_offsets = jnp.concatenate(`
`101`		`- [jnp.array([0], dtype=jnp.int32), jnp.cumsum(tokens_per_expert)[:-1]]`
	`101`	`+ [jnp.array([0], dtype=jnp.int32), jnp.cumsum(tokens_per_expert)[:-1].astype(jnp.int32)]`
`102`	`102`	`)`
`103`	`103`
`104`	`104`	`# Compute destination rows for all (token, expert) pairs`