Replace # pytype: disable suppressions with # pyrefly: ignore

vmoe/evaluate/ensemble.py

@@ -70,7 +70,7 @@ def ensemble_softmax_xent_train(repeated_logits: Array, labels: Array,
     ensemble softmax cross entropy used at training time.
   """
   repeated_labels = jnp.repeat(labels, ensemble_size, axis=0)
-  return optax.softmax_cross_entropy(repeated_logits, repeated_labels)  # pytype: disable=bad-return-type  # numpy-scalars
+  return optax.softmax_cross_entropy(repeated_logits, repeated_labels)
 
 
 def ensemble_sigmoid_xent_train(repeated_logits: Array, labels: Array,

vmoe/evaluate/evaluator.py

@@ -138,7 +138,7 @@ def _make_callback_fn(cls, *, apply_fn, loss_fn, label_pred_fn, datasets,
     # Note: We create the eval_step_pjit here to avoid multiple compilation
     # steps. If the shapes of inputs/outputs for all datasets is the same, this
     # will be only compiled once.
-    eval_state_dtype_struct = EvalState(  # pytype: disable=wrong-arg-types  # dataclass_transform
+    eval_state_dtype_struct = EvalState(
         num=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),  # pyrefly: ignore[bad-argument-type]
         sum_correct=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),  # pyrefly: ignore[bad-argument-type]
         sum_loss=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),  # pyrefly: ignore[bad-argument-type]
@@ -160,7 +160,7 @@ def _make_callback_fn(cls, *, apply_fn, loss_fn, label_pred_fn, datasets,
         in_shardings=(), out_shardings=None, static_argnums=(0,))
     def make_eval_state_pjit(seed):
       rngs = utils.make_rngs(rng_keys, seed)
-      return EvalState(  # pytype: disable=wrong-arg-types  # jnp-type
+      return EvalState(
           num=jnp.zeros((), dtype=jnp.float32),  # pyrefly: ignore[bad-argument-type]
           sum_correct=jnp.zeros((), dtype=jnp.float32),  # pyrefly: ignore[bad-argument-type]
           sum_loss=jnp.zeros((), dtype=jnp.float32),  # pyrefly: ignore[bad-argument-type]
@@ -172,7 +172,7 @@ def compile_for_dataset(name, params, train_step):
       # Note: This is not the initial EvalState, this only serves to compile the
       # eval step for a given dataset.
       t0 = time.time()
-      eval_step_pjit_ds = eval_step_pjit.lower(  # pytype: disable=attribute-error  # always-use-return-annotations
+      eval_step_pjit_ds = eval_step_pjit.lower(  # pyrefly: ignore[missing-attribute]
           eval_state_dtype_struct,
           params,
           datasets_element_shape_dtype[name]['image'],

vmoe/evaluate/fewshot.py

@@ -334,7 +334,7 @@ def _make_fewshot_step_pjit(
     rng_keys: Sequence[str],
 ):
   """Wraps _fewshot_step with pjit."""
-  state_axis_resources = FewShotState(  # pytype: disable=wrong-arg-types
+  state_axis_resources = FewShotState(
       rngs={key: PartitionSpec() for key in rng_keys})  # pyrefly: ignore[bad-argument-type]
   fewshot_step_pjit = jax.experimental.pjit.pjit(
       functools.partial(_fewshot_step, apply_fn=apply_fn),

vmoe/initialization/initialization.py

@@ -53,7 +53,7 @@ def structure(self, directory: epath.PathLike) -> Optional[Any]:
     """See superclass documentation."""
     directory = epath.Path(directory)
     try:
-      return self._handler.structure(directory)  # pytype: disable=attribute-error
+      return self._handler.structure(directory)  # pyrefly: ignore[missing-attribute]
     except NotImplementedError:
       return
 

vmoe/nn/ensemble_routing.py

@@ -87,7 +87,7 @@ def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
     dispatcher = self._create_dispatcher(gates_softmax)
     return dispatcher, metrics
 
-  def _compute_gates_softmax_and_metrics(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+  def _compute_gates_softmax_and_metrics(  # pyrefly: ignore[bad-override]
       self, inputs: Array) -> Tuple[Array, Metrics]:
     if inputs.ndim != 3:
       raise ValueError(f'inputs.ndim must be 3, but it is {inputs.ndim}')
@@ -122,7 +122,7 @@ def _compute_gates_softmax_and_metrics(  # pytype: disable=signature-mismatch  #
 
   # Wrap the super's _compute_gates_softmax_and_metrics with vmap over both
   # inputs and parameters.
-  @functools.partial(  # pyrefly: ignore[bad-specialization]
+  @functools.partial(
       nn.vmap,
       variable_axes={'params': 1},
       split_rngs={'params': True, 'gating': True},

vmoe/nn/external.py

@@ -29,7 +29,7 @@ def __post_init__(self):
     super().__post_init__()
 
   @nn.compact
-  def __call__(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-name-checks
+  def __call__(self, inputs):
     return super().__call__(inputs, train=not self.deterministic), {}
 
 
@@ -45,5 +45,5 @@ def __post_init__(self):
     super().__post_init__()
 
   @nn.compact
-  def __call__(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-name-checks
+  def __call__(self, inputs):
     return super().__call__(inputs, train=not self.deterministic), {}

vmoe/nn/routing.py

@@ -95,7 +95,7 @@ def _compute_gates_softmax_and_metrics(
           key=self.make_rng("gating"), shape=gates_logits.shape)
       gates_logits_noisy = gates_logits + logits_noise
       gates_softmax_noisy = jax.nn.softmax(gates_logits_noisy)
-      load_loss = jax.vmap(  # pytype: disable=wrong-arg-types
+      load_loss = jax.vmap(
           functools.partial(
               self._load_auxiliary_loss,
               num_selected_experts=self.num_selected_experts,
@@ -207,7 +207,7 @@ def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
     gates_softmax = self._compute_gates_softmax(inputs, self.num_experts)
     dispatcher, metrics = self._create_dispatcher_and_metrics(gates_softmax)
     metrics["auxiliary_loss"] = 0.
-    return dispatcher, metrics  # pytype: disable=bad-return-type
+    return dispatcher, metrics
 
   @nn.nowrap
   def _compute_gates_softmax(self, inputs: Array, num_experts: int) -> Array:

vmoe/nn/vit_moe.py

@@ -41,7 +41,7 @@ class MlpBlock(models_vit.MlpBlock):
   deterministic: bool = False
 
   @nn.compact
-  def __call__(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-name-checks
+  def __call__(self, inputs):
     return super().__call__(inputs, deterministic=self.deterministic)
 
 

vmoe/nn/vit_moe_test.py

@@ -181,7 +181,7 @@ def init():
         }
     }
     if not position_emb_kwargs:
-      expected_shapes['params']['Encoder']['posembed_input'] = {  # pyrefly: ignore[bad-assignment]
+      expected_shapes['params']['Encoder']['posembed_input'] = {
           'pos_embedding': (1, 4, 8),  # pyrefly: ignore[bad-assignment]
       }
     self.assertDictEqual(shapes, expected_shapes)

vmoe/projects/adversarial_attacks/attacks.py

@@ -20,7 +20,7 @@
 import jax.numpy as jnp
 
 
-class AttackState(flax.struct.PyTreeNode):  # pytype: disable=invalid-function-definition  # dataclass_transform
+class AttackState(flax.struct.PyTreeNode):
   """State used to aggregate attack statistics over several batches of images.
 
   Attributes:

vmoe/projects/adversarial_attacks/lib.py

@@ -102,7 +102,7 @@ def run_pgd_attack(
   del mesh
   # Setup dataset and get the global shape of the image array.
   dataset = get_dataset(config.dataset)
-  element_spec: ArraySpecDict = dataset.element_spec  # pytype: disable=annotation-type-mismatch
+  element_spec: ArraySpecDict = dataset.element_spec  # pyrefly: ignore[bad-assignment]
   image_shape = tuple(element_spec['image'].shape)
   image_shape = (config.dataset.batch_size,) + image_shape[1:]
   num_examples = input_pipeline.get_data_num_examples(config.dataset)

vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py

@@ -99,7 +99,7 @@ def _compute_gates_softmax_and_metrics(  # pyrefly: ignore[bad-override]
               (self.importance_loss_weight, importance_loss)),
           "importance_loss": importance_loss,
       }
-      return gates_plan, gates_softmax, metrics  # pytype: disable=bad-return-type  # jax-ndarray
+      return gates_plan, gates_softmax, metrics  # pyrefly: ignore[bad-return]
     else:
       noise_std = (1.0 / num_experts) * self.noise_std
       logits_noise = noise_std * jax.random.normal(
@@ -111,7 +111,7 @@ def _compute_gates_softmax_and_metrics(  # pyrefly: ignore[bad-override]
         # no need to compute the loss in this case
         load_loss = 0.0
       else:
-        load_loss = jax.vmap(  # pytype: disable=wrong-arg-types
+        load_loss = jax.vmap(
             functools.partial(
                 self._load_auxiliary_loss,
                 num_selected_experts=self.num_selected_experts,
@@ -123,7 +123,7 @@ def _compute_gates_softmax_and_metrics(  # pyrefly: ignore[bad-override]
           "importance_loss": importance_loss,
           "load_loss": load_loss,
       }
-      return gates_plan_noisy, gates_softmax_noisy, metrics  # pytype: disable=bad-return-type  # jax-ndarray
+      return gates_plan_noisy, gates_softmax_noisy, metrics  # pyrefly: ignore[bad-return]
 
   @nn.nowrap
   def _get_kl_projection_algorithm(self, group_size):
@@ -197,7 +197,7 @@ def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
       dispatcher, metrics = self._create_dispatcher_and_metrics(
           gates_dispatch=gates_ot)
     metrics["auxiliary_loss"] = 0.
-    return dispatcher, metrics  # pytype: disable=bad-return-type
+    return dispatcher, metrics
 
   @nn.nowrap
   def _compute_gates(self, inputs: Array, num_experts: int) -> Array:
@@ -212,15 +212,15 @@ def _compute_gates(self, inputs: Array, num_experts: int) -> Array:
     if self.deterministic or self.noise_std == 0.0:
       gates_ot = ot_algorithm(sim_matrix=gates_logits)
       gates_softmax = jax.nn.softmax(gates_logits)
-      return gates_ot, gates_softmax  # pytype: disable=bad-return-type  # jax-ndarray
+      return gates_ot, gates_softmax  # pyrefly: ignore[bad-return]
     else:
       noise_std = (1.0 / num_experts) * self.noise_std
       logits_noise = noise_std * jax.random.normal(
           key=self.make_rng("gating"), shape=gates_logits.shape)
       gates_logits_noisy = gates_logits + logits_noise
       gates_ot_noisy = ot_algorithm(sim_matrix=gates_logits_noisy)
       gates_softmax_noisy = jax.nn.softmax(gates_logits_noisy)
-      return gates_ot_noisy, gates_softmax_noisy  # pytype: disable=bad-return-type  # jax-ndarray
+      return gates_ot_noisy, gates_softmax_noisy  # pyrefly: ignore[bad-return]
 
   @nn.nowrap
   def _get_kl_projection_algorithm(self, group_size):

vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py

@@ -86,7 +86,7 @@ def _compute_gates_softmax_and_metrics(  # pyrefly: ignore[bad-override]
     metrics["num_experts_per_item_min"] = jnp.min(num_experts_per_item, axis=1)  # pyrefly: ignore[bad-assignment]
     metrics["num_experts_per_item_max"] = jnp.max(num_experts_per_item, axis=1)  # pyrefly: ignore[bad-assignment]
     metrics["num_experts_per_item_avg"] = jnp.mean(num_experts_per_item, axis=1)  # pyrefly: ignore[bad-assignment]
-    return gates_ot, gates_softmax, metrics  # pytype: disable=bad-return-type  # jax-ndarray
+    return gates_ot, gates_softmax, metrics  # pyrefly: ignore[bad-return]
 
   @nn.nowrap
   def _get_k_sparse_ot_func(self, group_size):
@@ -133,7 +133,7 @@ def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
     metrics["num_items_per_expert_max"] = jnp.max(num_items_per_expert, axis=1)
     metrics["num_items_per_expert_avg"] = jnp.mean(num_items_per_expert, axis=1)
     metrics["auxiliary_loss"] = 0.
-    return dispatcher, metrics  # pytype: disable=bad-return-type
+    return dispatcher, metrics
 
   @nn.nowrap
   def _compute_gates(

vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py

@@ -62,7 +62,7 @@ def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
       dispatcher, metrics = self._create_dispatcher_and_metrics(
           gates_dispatch=gates_ot)
     metrics["auxiliary_loss"] = 0.0
-    return dispatcher, metrics  # pytype: disable=bad-return-type
+    return dispatcher, metrics
 
   @nn.nowrap
   def _compute_gates(