Automated Code Change

examples/visualization_demo.py

@@ -52,14 +52,15 @@ def convert_tree_to_graph(
   chex.assert_rank(tree.node_values, 2)
   batch_size = tree.node_values.shape[0]
   if action_labels is None:
-    action_labels = range(tree.num_actions)
+    action_labels = range(tree.num_actions)  # pyrefly: ignore[bad-assignment]
   elif len(action_labels) != tree.num_actions:
     raise ValueError(
         f"action_labels {action_labels} has the wrong number of actions "
         f"({len(action_labels)}). "
         f"Expecting {tree.num_actions}.")
 
   def node_to_str(node_i, reward=0, discount=1):
+    # pyrefly: ignore[bad-index]
     return (f"{node_i}\n"
             f"Reward: {reward:.2f}\n"
             f"Discount: {discount:.2f}\n"
@@ -68,7 +69,8 @@ def node_to_str(node_i, reward=0, discount=1):
 
   def edge_to_str(node_i, a_i):
     node_index = jnp.full([batch_size], node_i)
-    probs = jax.nn.softmax(tree.children_prior_logits[batch_index, node_i])
+    probs = jax.nn.softmax(tree.children_prior_logits[batch_index, node_i])  # pyrefly: ignore[bad-index]
+    # pyrefly: ignore[unsupported-operation]
     return (f"{action_labels[a_i]}\n"
             f"Q: {tree.qvalues(node_index)[batch_index, a_i]:.2f}\n"  # pytype: disable=unsupported-operands  # always-use-return-annotations
             f"p: {probs[a_i]:.2f}\n")
@@ -81,14 +83,14 @@ def edge_to_str(node_i, a_i):
   for node_i in range(tree.num_simulations):
     for a_i in range(tree.num_actions):
       # Index of children, or -1 if not expanded
-      children_i = tree.children_index[batch_index, node_i, a_i]
+      children_i = tree.children_index[batch_index, node_i, a_i]  # pyrefly: ignore[bad-index]
       if children_i >= 0:
         graph.add_node(
             children_i,
             label=node_to_str(
                 node_i=children_i,
-                reward=tree.children_rewards[batch_index, node_i, a_i],
-                discount=tree.children_discounts[batch_index, node_i, a_i]),
+                reward=tree.children_rewards[batch_index, node_i, a_i],  # pyrefly: ignore[bad-index]
+                discount=tree.children_discounts[batch_index, node_i, a_i]),  # pyrefly: ignore[bad-index]
             color="red")
         graph.add_edge(node_i, children_i, label=edge_to_str(node_i, a_i))
 
@@ -166,8 +168,8 @@ def _make_batched_env_model(
   root_state = 0
   root = mctx.RootFnOutput(
       prior_logits=jnp.full([batch_size, num_actions],
-                            prior_logits[root_state]),
-      value=jnp.full([batch_size], values[root_state]),
+                            prior_logits[root_state]),  # pyrefly: ignore[bad-index]
+      value=jnp.full([batch_size], values[root_state]),  # pyrefly: ignore[bad-index]
       # The embedding will hold the state index.
       embedding=jnp.zeros([batch_size], dtype=jnp.int32),
   )
@@ -177,11 +179,11 @@ def recurrent_fn(params, rng_key, action, embedding):
     chex.assert_shape(action, [batch_size])
     chex.assert_shape(embedding, [batch_size])
     recurrent_fn_output = mctx.RecurrentFnOutput(
-        reward=rewards[embedding, action],
-        discount=discounts[embedding, action],
-        prior_logits=prior_logits[embedding],
-        value=values[embedding])
-    next_embedding = transition_matrix[embedding, action]
+        reward=rewards[embedding, action],  # pyrefly: ignore[bad-index]
+        discount=discounts[embedding, action],  # pyrefly: ignore[bad-index]
+        prior_logits=prior_logits[embedding],  # pyrefly: ignore[bad-index]
+        value=values[embedding])  # pyrefly: ignore[bad-index]
+    next_embedding = transition_matrix[embedding, action]  # pyrefly: ignore[bad-index]
     return recurrent_fn_output, next_embedding
 
   return root, recurrent_fn

mctx/_src/action_selection.py

@@ -71,14 +71,16 @@ def muzero_action_selection(
   Returns:
     action: the action selected from the given node.
   """
-  visit_counts = tree.children_visits[node_index]
-  node_visit = tree.node_visits[node_index]
+  visit_counts = tree.children_visits[node_index]  # pyrefly: ignore[bad-index]
+  node_visit = tree.node_visits[node_index]  # pyrefly: ignore[bad-index]
   pb_c = pb_c_init + jnp.log((node_visit + pb_c_base + 1.) / pb_c_base)
+  # pyrefly: ignore[bad-index]
   prior_logits = tree.children_prior_logits[node_index]
   prior_probs = jax.nn.softmax(prior_logits)
   policy_score = jnp.sqrt(node_visit) * pb_c * prior_probs / (visit_counts + 1)
   chex.assert_shape([node_index, node_visit], ())
   chex.assert_equal_shape([prior_probs, visit_counts, policy_score])
+  # pyrefly: ignore[bad-argument-type]
   value_score = qtransform(tree, node_index)
 
   # Add tiny bit of randomness for tie break
@@ -129,9 +131,12 @@ def gumbel_muzero_root_action_selection(
   """
   del rng_key
   chex.assert_shape([node_index], ())
+  # pyrefly: ignore[bad-index]
   visit_counts = tree.children_visits[node_index]
+  # pyrefly: ignore[bad-index]
   prior_logits = tree.children_prior_logits[node_index]
   chex.assert_equal_shape([visit_counts, prior_logits])
+  # pyrefly: ignore[bad-argument-type]
   completed_qvalues = qtransform(tree, node_index)
 
   table = jnp.array(seq_halving.get_table_of_considered_visits(
@@ -180,9 +185,11 @@ def gumbel_muzero_interior_action_selection(
   """
   del rng_key, depth
   chex.assert_shape([node_index], ())
-  visit_counts = tree.children_visits[node_index]
+  visit_counts = tree.children_visits[node_index]  # pyrefly: ignore[bad-index]
+  # pyrefly: ignore[bad-index]
   prior_logits = tree.children_prior_logits[node_index]
   chex.assert_equal_shape([visit_counts, prior_logits])
+  # pyrefly: ignore[bad-argument-type]
   completed_qvalues = qtransform(tree, node_index)
 
   # The `prior_logits + completed_qvalues` provide an improved policy,

mctx/_src/policies.py

@@ -86,7 +86,7 @@ def muzero_policy(
           jax.nn.softmax(root.prior_logits),
           dirichlet_fraction=dirichlet_fraction,
           dirichlet_alpha=dirichlet_alpha))
-  root = root.replace(
+  root = root.replace(  # pyrefly: ignore[missing-attribute]
       prior_logits=_mask_invalid_actions(noisy_logits, invalid_actions))
 
   # Running the search.
@@ -176,7 +176,7 @@ def gumbel_muzero_policy(
     search tree.
   """
   # Masking invalid actions.
-  root = root.replace(
+  root = root.replace(  # pyrefly: ignore[missing-attribute]
       prior_logits=_mask_invalid_actions(root.prior_logits, invalid_actions))
 
   # Generating Gumbel.
@@ -301,7 +301,7 @@ def stochastic_muzero_policy(
           dirichlet_fraction=dirichlet_fraction,
           dirichlet_alpha=dirichlet_alpha))
 
-  root = root.replace(
+  root = root.replace(  # pyrefly: ignore[missing-attribute]
       prior_logits=_mask_invalid_actions(noisy_logits, invalid_actions))
 
   # construct a dummy afterstate embedding
@@ -509,7 +509,7 @@ def _take_slice(x):
     else:
       raise ValueError(f'Unknown mode: {mode}.')
 
-  return tree.replace(
+  return tree.replace(  # pyrefly: ignore[missing-attribute]
       children_index=_take_slice(tree.children_index),
       children_prior_logits=_take_slice(tree.children_prior_logits),
       children_visits=_take_slice(tree.children_visits),
@@ -531,7 +531,8 @@ def _chance_node_selection_fn(
       tree: search.Tree,
       node_index: chex.Array,
   ) -> chex.Array:
-    num_chance = tree.children_visits[node_index]
+    num_chance = tree.children_visits[node_index]  # pyrefly: ignore[bad-index]
+    # pyrefly: ignore[bad-index]
     chance_logits = tree.children_prior_logits[node_index]
     prob_chance = jax.nn.softmax(chance_logits)
     argmax_chance = jnp.argmax(prob_chance / (num_chance + 1), axis=-1).astype(

mctx/_src/qtransforms.py

@@ -44,8 +44,9 @@ def qtransform_by_min_max(
   """
   chex.assert_shape(node_index, ())
   qvalues = tree.qvalues(node_index)
-  visit_counts = tree.children_visits[node_index]
+  visit_counts = tree.children_visits[node_index]  # pyrefly: ignore[bad-index]
   value_score = jnp.where(visit_counts > 0, qvalues, min_value)
+  # pyrefly: ignore[unsupported-operation]
   value_score = (value_score - min_value) / ((max_value - min_value))
   return value_score
 
@@ -69,9 +70,9 @@ def qtransform_by_parent_and_siblings(
   """
   chex.assert_shape(node_index, ())
   qvalues = tree.qvalues(node_index)
-  visit_counts = tree.children_visits[node_index]
+  visit_counts = tree.children_visits[node_index]  # pyrefly: ignore[bad-index]
   chex.assert_rank([qvalues, visit_counts, node_index], [1, 1, 0])
-  node_value = tree.node_values[node_index]
+  node_value = tree.node_values[node_index]  # pyrefly: ignore[bad-index]
   safe_qvalues = jnp.where(visit_counts > 0, qvalues, node_value)
   chex.assert_equal_shape([safe_qvalues, qvalues])
   min_value = jnp.minimum(node_value, jnp.min(safe_qvalues, axis=-1))
@@ -119,12 +120,12 @@ def qtransform_completed_by_mix_value(
   """
   chex.assert_shape(node_index, ())
   qvalues = tree.qvalues(node_index)
-  visit_counts = tree.children_visits[node_index]
+  visit_counts = tree.children_visits[node_index]  # pyrefly: ignore[bad-index]
 
   # Computing the mixed value and producing completed_qvalues.
-  raw_value = tree.raw_values[node_index]
+  raw_value = tree.raw_values[node_index]  # pyrefly: ignore[bad-index]
   prior_probs = jax.nn.softmax(
-      tree.children_prior_logits[node_index])
+      tree.children_prior_logits[node_index])  # pyrefly: ignore[bad-index]
   if use_mixed_value:
     value = _compute_mixed_value(
         raw_value,

mctx/_src/search.py

@@ -98,6 +98,7 @@ def body_fun(sim, loop_state):
     next_node_index = jnp.where(next_node_index == Tree.UNVISITED,
                                 sim + 1, next_node_index)
     tree = expand(
+        # pyrefly: ignore[bad-argument-type]
         params, expand_key, tree, recurrent_fn, parent_index,
         action, next_node_index)
     tree = backward(tree, next_node_index)
@@ -155,6 +156,7 @@ def body_fun(state):
     rng_key, action_selection_key = jax.random.split(state.rng_key)
     action = action_selection_fn(action_selection_key, tree, node_index,
                                  state.depth)
+    # pyrefly: ignore[bad-index]
     next_node_index = tree.children_index[node_index, action]
     # The returned action will be visited.
     depth = state.depth + 1
@@ -164,10 +166,10 @@ def body_fun(state):
     return _SimulationState(  # pytype: disable=wrong-arg-types  # jax-types
         rng_key=rng_key,
         node_index=node_index,
-        action=action,
-        next_node_index=next_node_index,
+        action=action,  # pyrefly: ignore[bad-argument-type]
+        next_node_index=next_node_index,  # pyrefly: ignore[bad-argument-type]
         depth=depth,
-        is_continuing=is_continuing)
+        is_continuing=is_continuing)  # pyrefly: ignore[bad-argument-type]
 
   node_index = jnp.array(Tree.ROOT_INDEX, dtype=jnp.int32)
   depth = jnp.zeros((), dtype=tree.children_prior_logits.dtype)
@@ -232,7 +234,7 @@ def expand(
       tree, next_node_index, step.prior_logits, step.value, embedding)
 
   # Return updated tree topology.
-  return tree.replace(
+  return tree.replace(  # pyrefly: ignore[missing-attribute]
       children_index=batch_update(
           tree.children_index, next_node_index, parent_index, action),
       children_rewards=batch_update(
@@ -286,6 +288,7 @@ def body_fun(loop_state):
     return tree, leaf_value, parent
 
   leaf_index = jnp.asarray(leaf_index, dtype=jnp.int32)
+  # pyrefly: ignore[bad-index]
   loop_state = (tree, tree.node_values[leaf_index], leaf_index)
   tree, _, _ = jax.lax.while_loop(cond_fun, body_fun, loop_state)
 
@@ -325,6 +328,7 @@ def update_tree_node(
   chex.assert_shape(prior_logits, (batch_size, tree.num_actions))
 
   # When using max_depth, a leaf can be expanded multiple times.
+  # pyrefly: ignore[bad-index]
   new_visit = tree.node_visits[batch_range, node_index] + 1
   updates = dict(  # pylint: disable=use-dict-literal
       children_prior_logits=batch_update(
@@ -339,7 +343,7 @@ def update_tree_node(
           lambda t, s: batch_update(t, s, node_index),
           tree.embeddings, embedding))
 
-  return tree.replace(**updates)
+  return tree.replace(**updates)  # pyrefly: ignore[missing-attribute]
 
 
 def instantiate_tree_from_root(

mctx/_src/tests/policies_test.py

@@ -250,7 +250,7 @@ def test_gumbel_muzero_policy(self):
         jax.tree.map(lambda x: x[0], policy_output.search_tree),
         policy_output.search_tree.ROOT_INDEX)
     self.assertEqual(max_depth, max_found_depth)
-    self.assertEqual(6, policy_output.search_tree.node_visits[0, leaf])
+    self.assertEqual(6, policy_output.search_tree.node_visits[0, leaf])  # pyrefly: ignore[bad-index]
 
   def test_gumbel_muzero_policy_without_invalid_actions(self):
     root_value = jnp.array([-5.0])

mctx/_src/tests/tree_test.py

@@ -106,9 +106,9 @@ def tree_to_pytree(tree: mctx.Tree, batch_i: int = 0):
     for a_i in range(tree.num_actions):
       prior = children_prior_probs[batch_i, node_i, a_i]
       # Index of children, or -1 if not expanded
-      child_i = int(tree.children_index[batch_i, node_i, a_i])
+      child_i = int(tree.children_index[batch_i, node_i, a_i])  # pyrefly: ignore[bad-index]
       if child_i >= 0:
-        reward = tree.children_rewards[batch_i, node_i, a_i]
+        reward = tree.children_rewards[batch_i, node_i, a_i]  # pyrefly: ignore[bad-index]
         child = _create_pynode(
             tree, batch_i, child_i, prior=prior, action=a_i, reward=reward)
         nodes[child_i] = child

mctx/_src/tree.py

@@ -98,11 +98,12 @@ def summary(self) -> SearchSummary:
     """Extract summary statistics for the root node."""
     # Get state and action values for the root nodes.
     chex.assert_rank(self.node_values, 2)
-    value = self.node_values[:, Tree.ROOT_INDEX]
+    value = self.node_values[:, Tree.ROOT_INDEX]  # pyrefly: ignore[bad-index]
     batch_size, = value.shape
     root_indices = jnp.full((batch_size,), Tree.ROOT_INDEX)
     qvalues = self.qvalues(root_indices)
     # Extract visit counts and induced probabilities for the root nodes.
+    # pyrefly: ignore[bad-index]
     visit_counts = self.children_visits[:, Tree.ROOT_INDEX].astype(value.dtype)
     total_counts = jnp.sum(visit_counts, axis=-1, keepdims=True)
     visit_probs = visit_counts / jnp.maximum(total_counts, 1)
@@ -137,6 +138,8 @@ class SearchSummary:
 def _unbatched_qvalues(tree: Tree, index: int) -> int:
   chex.assert_rank(tree.children_discounts, 2)
   return (  # pytype: disable=bad-return-type  # numpy-scalars
+      # pyrefly: ignore[bad-index, bad-return]
       tree.children_rewards[index]
+      # pyrefly: ignore[bad-index]
       + tree.children_discounts[index] * tree.children_values[index]
   )