Fix pytype failures related to teaching pytype about NumPy scalar types.

PiperOrigin-RevId: 519128722

mctx/_src/policies.py

@@ -214,7 +214,7 @@ def gumbel_muzero_policy(
   # a smaller number of valid actions.
   considered_visit = jnp.max(summary.visit_counts, axis=-1, keepdims=True)
   # The completed_qvalues include imputed values for unvisited actions.
-  completed_qvalues = jax.vmap(qtransform, in_axes=[0, None])(
+  completed_qvalues = jax.vmap(qtransform, in_axes=[0, None])(  # pytype: disable=wrong-arg-types  # numpy-scalars  # pylint: disable=line-too-long
       search_tree, search_tree.ROOT_INDEX)
   to_argmax = seq_halving.score_considered(
       considered_visit, gumbel, root.prior_logits, completed_qvalues,

mctx/_src/tree.py

@@ -106,7 +106,7 @@ def summary(self) -> SearchSummary:
     visit_probs = visit_counts / jnp.maximum(total_counts, 1)
     visit_probs = jnp.where(total_counts > 0, visit_probs, 1 / self.num_actions)
     # Return relevant stats.
-    return SearchSummary(
+    return SearchSummary(  # pytype: disable=wrong-arg-types  # numpy-scalars
         visit_counts=visit_counts,
         visit_probs=visit_probs,
         value=value,
@@ -134,6 +134,7 @@ class SearchSummary:
 
 def _unbatched_qvalues(tree: Tree, index: int) -> int:
   chex.assert_rank(tree.children_discounts, 2)
-  return (tree.children_rewards[index] +
-          tree.children_discounts[index] * tree.children_values[index])
-
+  return (  # pytype: disable=bad-return-type  # numpy-scalars
+      tree.children_rewards[index]
+      + tree.children_discounts[index] * tree.children_values[index]
+  )