No public description

PiperOrigin-RevId: 557378061
google-research · Aug 16, 2023 · c2b4ce0 · c2b4ce0
1 parent 662341d
commit c2b4ce0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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(
+    eval_state_dtype_struct = EvalState(  # pytype: disable=wrong-arg-types  # dataclass_transform
         num=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),
         sum_correct=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),
         sum_loss=jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),

diff --git 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):
+class AttackState(flax.struct.PyTreeNode):  # pytype: disable=invalid-function-definition  # dataclass_transform
   """State used to aggregate attack statistics over several batches of images.
 
   Attributes: