No public description

PiperOrigin-RevId: 624877716

vmoe/nn/routing.py

@@ -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
+    return dispatcher, metrics  # pytype: disable=bad-return-type
 
   @nn.nowrap
   def _compute_gates_softmax(self, inputs: Array, num_experts: int) -> Array:

vmoe/projects/sparsity_constrained_ot/kl_projection_routing.py

@@ -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
+    return dispatcher, metrics  # pytype: disable=bad-return-type
 
   @nn.nowrap
   def _compute_gates(self, inputs: Array, num_experts: int) -> Array:

vmoe/projects/sparsity_constrained_ot/ksparse_projection_routing.py

@@ -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
+    return dispatcher, metrics  # pytype: disable=bad-return-type
 
   @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
+    return dispatcher, metrics  # pytype: disable=bad-return-type
 
   @nn.nowrap
   def _compute_gates(