From ebeb8040f8ef7448e830f70fc7b043766e0e782b Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Mon, 5 Oct 2026 08:05:55 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 993691087 --- vmoe/evaluate/ensemble.py | 2 +- vmoe/evaluate/evaluator.py | 6 +++--- vmoe/evaluate/fewshot.py | 2 +- vmoe/initialization/initialization.py | 2 +- vmoe/nn/ensemble_routing.py | 4 ++-- vmoe/nn/external.py | 4 ++-- vmoe/nn/routing.py | 4 ++-- vmoe/nn/vit_moe.py | 2 +- vmoe/nn/vit_moe_test.py | 2 +- vmoe/projects/adversarial_attacks/attacks.py | 2 +- vmoe/projects/adversarial_attacks/lib.py | 2 +- .../sparsity_constrained_ot/kl_projection_routing.py | 12 ++++++------ .../ksparse_projection_routing.py | 4 ++-- .../sparse_projection_routing.py | 2 +- 14 files changed, 25 insertions(+), 25 deletions(-) diff --git a/vmoe/evaluate/ensemble.py b/vmoe/evaluate/ensemble.py index e1272cc..4cd8da7 100644 --- a/vmoe/evaluate/ensemble.py +++ b/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, diff --git a/vmoe/evaluate/evaluator.py b/vmoe/evaluate/evaluator.py index 1656b46..dcf3ffb 100644 --- a/vmoe/evaluate/evaluator.py +++ b/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'], diff --git a/vmoe/evaluate/fewshot.py b/vmoe/evaluate/fewshot.py index 22cd228..a104ee3 100644 --- a/vmoe/evaluate/fewshot.py +++ b/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), diff --git a/vmoe/initialization/initialization.py b/vmoe/initialization/initialization.py index bc970c4..5869f0c 100644 --- a/vmoe/initialization/initialization.py +++ b/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 diff --git a/vmoe/nn/ensemble_routing.py b/vmoe/nn/ensemble_routing.py index b2c24e9..7733dc9 100644 --- a/vmoe/nn/ensemble_routing.py +++ b/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}, diff --git a/vmoe/nn/external.py b/vmoe/nn/external.py index 9c334d0..f31a2a9 100644 --- a/vmoe/nn/external.py +++ b/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), {} diff --git a/vmoe/nn/routing.py b/vmoe/nn/routing.py index d36307f..c20af51 100644 --- a/vmoe/nn/routing.py +++ b/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: diff --git a/vmoe/nn/vit_moe.py b/vmoe/nn/vit_moe.py index 7bb97e1..44920d0 100644 --- a/vmoe/nn/vit_moe.py +++ b/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) diff --git a/vmoe/nn/vit_moe_test.py b/vmoe/nn/vit_moe_test.py index 584ad07..6183113 100644 --- a/vmoe/nn/vit_moe_test.py +++ b/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) diff --git a/vmoe/projects/adversarial_attacks/attacks.py b/vmoe/projects/adversarial_attacks/attacks.py index bf26cc5..ed3863f 100644 --- a/vmoe/projects/adversarial_attacks/attacks.py +++ b/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: diff --git a/vmoe/projects/adversarial_attacks/lib.py b/vmoe/projects/adversarial_attacks/lib.py index ce0ad97..4682ebf 100644 --- a/vmoe/projects/adversarial_attacks/lib.py +++ b/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) diff --git a/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py b/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py index f160258..f262f20 100644 --- a/vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py +++ b/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,7 +212,7 @@ 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( @@ -220,7 +220,7 @@ def _compute_gates(self, inputs: Array, num_experts: int) -> Array: 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): diff --git a/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py b/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py index 7e34781..ac1f299 100644 --- a/vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py +++ b/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( diff --git a/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py b/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py index a66ba06..f024cab 100644 --- a/vmoe/projects/sparsity_constrained_ot/sparse_projection_routing.py +++ b/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(