dropout piping dev JAX workloads

Author: Niccolo-Ajroldi

algoperf/workloads/cifar/cifar_jax/workload.py

@@ -105,9 +105,11 @@ def model_fn(
     rng: spec.RandomState,
     update_batch_norm: bool,
     use_running_average_bn: Optional[bool] = None,
+    dropout_rate: float = 0.0,
   ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     del mode
     del rng
+    del dropout_rate
     variables = {'params': params, **model_state}
     if update_batch_norm:
       logits, new_model_state = self._model.apply(

algoperf/workloads/mnist/mnist_jax/workload.py

@@ -48,10 +48,12 @@ def model_fn(
     mode: spec.ForwardPassMode,
     rng: spec.RandomState,
     update_batch_norm: bool,
+    dropout_rate: float = 0.0,
   ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     del model_state
     del rng
     del update_batch_norm
+    del dropout_rate
     train = mode == spec.ForwardPassMode.TRAIN
     logits_batch = self._model.apply(
       {'params': params},