formatting

Niccolo-Ajroldi · Niccolo-Ajroldi · commit 816c167810bb · 2025-08-25T15:57:01.000+02:00
diff --git a/algoperf/workloads/wmt/wmt_jax/workload.py b/algoperf/workloads/wmt/wmt_jax/workload.py
@@ -240,7 +240,6 @@ def translate_and_calculate_bleu(
     return bleu_score
 
   def init_model_fn(self, rng: spec.RandomState) -> spec.ModelInitState:
-
     init_fake_batch_size = 8
     input_shape = (init_fake_batch_size, 256)
     target_shape = (init_fake_batch_size, 256)
diff --git a/reference_algorithms/paper_baselines/adafactor/jax/submission.py b/reference_algorithms/paper_baselines/adafactor/jax/submission.py
@@ -169,7 +169,7 @@ def update_params(
     per_device_rngs,
     grad_clip,
     label_smoothing,
-    dropout_rate
+    dropout_rate,
   )
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs
 
diff --git a/reference_algorithms/paper_baselines/adafactor/pytorch/submission.py b/reference_algorithms/paper_baselines/adafactor/pytorch/submission.py
@@ -227,7 +227,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/paper_baselines/adamw/jax/submission.py b/reference_algorithms/paper_baselines/adamw/jax/submission.py
@@ -189,7 +189,7 @@ def update_params(
       rng,
       grad_clip,
       label_smoothing,
-      dropout_rate
+      dropout_rate,
     )
   )
 
diff --git a/reference_algorithms/paper_baselines/adamw/pytorch/submission.py b/reference_algorithms/paper_baselines/adamw/pytorch/submission.py
@@ -84,7 +84,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/paper_baselines/lamb/jax/submission.py b/reference_algorithms/paper_baselines/lamb/jax/submission.py
@@ -176,7 +176,7 @@ def update_params(
     per_device_rngs,
     grad_clip,
     label_smoothing,
-    dropout_rate
+    dropout_rate,
   )
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs
 
diff --git a/reference_algorithms/paper_baselines/lamb/pytorch/submission.py b/reference_algorithms/paper_baselines/lamb/pytorch/submission.py
@@ -225,7 +225,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/paper_baselines/momentum/jax/submission.py b/reference_algorithms/paper_baselines/momentum/jax/submission.py
@@ -227,7 +227,7 @@ def update_params(
     rng,
     grad_clip,
     label_smoothing,
-      dropout_rate
+    dropout_rate,
   )
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs
 
diff --git a/reference_algorithms/paper_baselines/momentum/pytorch/submission.py b/reference_algorithms/paper_baselines/momentum/pytorch/submission.py
@@ -104,7 +104,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/paper_baselines/nadamw/jax/submission.py b/reference_algorithms/paper_baselines/nadamw/jax/submission.py
@@ -340,7 +340,7 @@ def update_params(
       rng,
       grad_clip,
       label_smoothing,
-      dropout_rate
+      dropout_rate,
     )
   )
 
diff --git a/reference_algorithms/paper_baselines/nadamw/pytorch/submission.py b/reference_algorithms/paper_baselines/nadamw/pytorch/submission.py
@@ -264,7 +264,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/paper_baselines/nesterov/pytorch/submission.py b/reference_algorithms/paper_baselines/nesterov/pytorch/submission.py
@@ -104,7 +104,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/paper_baselines/sam/pytorch/submission.py b/reference_algorithms/paper_baselines/sam/pytorch/submission.py
@@ -173,7 +173,7 @@ def _loss_fn(params, update_batch_norm=True):
       mode=spec.ForwardPassMode.TRAIN,
       rng=rng,
       update_batch_norm=update_batch_norm,
-      dropout_rate=hyperparameters.dropout_rate
+      dropout_rate=hyperparameters.dropout_rate,
     )
     label_smoothing = (
       hyperparameters.label_smoothing
diff --git a/reference_algorithms/qualification_baselines/external_tuning/jax_nadamw_full_budget.py b/reference_algorithms/qualification_baselines/external_tuning/jax_nadamw_full_budget.py
@@ -339,7 +339,7 @@ def update_params(
       rng,
       grad_clip,
       label_smoothing,
-      dropout_rate
+      dropout_rate,
     )
   )
 
diff --git a/reference_algorithms/qualification_baselines/external_tuning/jax_nadamw_target_setting.py b/reference_algorithms/qualification_baselines/external_tuning/jax_nadamw_target_setting.py
@@ -335,7 +335,7 @@ def update_params(
     rng,
     grad_clip,
     label_smoothing,
-    dropout_rate
+    dropout_rate,
   )
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs
 
diff --git a/reference_algorithms/qualification_baselines/external_tuning/pytorch_nadamw_full_budget.py b/reference_algorithms/qualification_baselines/external_tuning/pytorch_nadamw_full_budget.py
@@ -264,7 +264,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/qualification_baselines/external_tuning/pytorch_nadamw_target_setting.py b/reference_algorithms/qualification_baselines/external_tuning/pytorch_nadamw_target_setting.py
@@ -264,7 +264,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/qualification_baselines/self_tuning/jax_nadamw_full_budget.py b/reference_algorithms/qualification_baselines/self_tuning/jax_nadamw_full_budget.py
@@ -306,7 +306,7 @@ def update_params(
     grad_clip = None
   dropout_rate = hyperparameters.dropout_rate
 
-    # Create shardings for each argument
+  # Create shardings for each argument
   mesh = jax.sharding.Mesh(jax.devices(), ('batch'))
   replicated = jax_sharding_utils.get_replicate_sharding(
     mesh
@@ -354,7 +354,7 @@ def update_params(
     rng,
     grad_clip,
     label_smoothing,
-    dropout_rate
+    dropout_rate,
   )
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs
 
diff --git a/reference_algorithms/qualification_baselines/self_tuning/jax_nadamw_target_setting.py b/reference_algorithms/qualification_baselines/self_tuning/jax_nadamw_target_setting.py
@@ -312,7 +312,7 @@ def update_params(
     grad_clip = None
   dropout_rate = hyperparameters.dropout_rate
 
-    # Create shardings for each argument
+  # Create shardings for each argument
   mesh = jax.sharding.Mesh(jax.devices(), ('batch'))
   replicated = jax_sharding_utils.get_replicate_sharding(
     mesh
@@ -360,7 +360,7 @@ def update_params(
     rng,
     grad_clip,
     label_smoothing,
-    dropout_rate
+    dropout_rate,
   )
   new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs
 
diff --git a/reference_algorithms/qualification_baselines/self_tuning/pytorch_nadamw_full_budget.py b/reference_algorithms/qualification_baselines/self_tuning/pytorch_nadamw_full_budget.py
@@ -281,7 +281,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/reference_algorithms/qualification_baselines/self_tuning/pytorch_nadamw_target_setting.py b/reference_algorithms/qualification_baselines/self_tuning/pytorch_nadamw_target_setting.py
@@ -281,7 +281,7 @@ def update_params(
     mode=spec.ForwardPassMode.TRAIN,
     rng=rng,
     update_batch_norm=True,
-    dropout_rate=hyperparameters.dropout_rate
+    dropout_rate=hyperparameters.dropout_rate,
   )
 
   label_smoothing = (
diff --git a/tests/modeldiffs/diff.py b/tests/modeldiffs/diff.py
@@ -14,9 +14,7 @@ def torch2jax(
   key_transform=None,
   sd_transform=None,
 ):
-  jax_params, model_state = jax_workload.init_model_fn(
-    jax.random.PRNGKey(0)
-  )
+  jax_params, model_state = jax_workload.init_model_fn(jax.random.PRNGKey(0))
   pytorch_model, _ = pytorch_workload.init_model_fn([0])
   if isinstance(jax_params, dict):
     jax_params = FrozenDict(jax_params)

Original file line number	Diff line number	Diff line change
`@@ -169,7 +169,7 @@ def update_params(`
`169`	`169`	`per_device_rngs,`
`170`	`170`	`grad_clip,`
`171`	`171`	`label_smoothing,`
`172`		`- dropout_rate`
	`172`	`+ dropout_rate,`
`173`	`173`	`)`
`174`	`174`	`new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs`
`175`	`175`
Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,7 @@ def update_params(`
`227`	`227`	`mode=spec.ForwardPassMode.TRAIN,`
`228`	`228`	`rng=rng,`
`229`	`229`	`update_batch_norm=True,`
`230`		`- dropout_rate=hyperparameters.dropout_rate`
	`230`	`+ dropout_rate=hyperparameters.dropout_rate,`
`231`	`231`	`)`
`232`	`232`
`233`	`233`	`label_smoothing = (`
Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,7 @@ def update_params(`
`189`	`189`	`rng,`
`190`	`190`	`grad_clip,`
`191`	`191`	`label_smoothing,`
`192`		`- dropout_rate`
	`192`	`+ dropout_rate,`
`193`	`193`	`)`
`194`	`194`	`)`
`195`	`195`
Original file line number	Diff line number	Diff line change
`@@ -84,7 +84,7 @@ def update_params(`
`84`	`84`	`mode=spec.ForwardPassMode.TRAIN,`
`85`	`85`	`rng=rng,`
`86`	`86`	`update_batch_norm=True,`
`87`		`- dropout_rate=hyperparameters.dropout_rate`
	`87`	`+ dropout_rate=hyperparameters.dropout_rate,`
`88`	`88`	`)`
`89`	`89`
`90`	`90`	`label_smoothing = (`
Original file line number	Diff line number	Diff line change
`@@ -176,7 +176,7 @@ def update_params(`
`176`	`176`	`per_device_rngs,`
`177`	`177`	`grad_clip,`
`178`	`178`	`label_smoothing,`
`179`		`- dropout_rate`
	`179`	`+ dropout_rate,`
`180`	`180`	`)`
`181`	`181`	`new_optimizer_state, new_params, new_model_state, loss, grad_norm = outputs`
`182`	`182`
Original file line number	Diff line number	Diff line change
`@@ -225,7 +225,7 @@ def update_params(`
`225`	`225`	`mode=spec.ForwardPassMode.TRAIN,`
`226`	`226`	`rng=rng,`
`227`	`227`	`update_batch_norm=True,`
`228`		`- dropout_rate=hyperparameters.dropout_rate`
	`228`	`+ dropout_rate=hyperparameters.dropout_rate,`
`229`	`229`	`)`
`230`	`230`
`231`	`231`	`label_smoothing = (`
Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,7 @@ def update_params(`
`104`	`104`	`mode=spec.ForwardPassMode.TRAIN,`
`105`	`105`	`rng=rng,`
`106`	`106`	`update_batch_norm=True,`
`107`		`- dropout_rate=hyperparameters.dropout_rate`
	`107`	`+ dropout_rate=hyperparameters.dropout_rate,`
`108`	`108`	`)`
`109`	`109`
`110`	`110`	`label_smoothing = (`
Original file line number	Diff line number	Diff line change
`@@ -340,7 +340,7 @@ def update_params(`
`340`	`340`	`rng,`
`341`	`341`	`grad_clip,`
`342`	`342`	`label_smoothing,`
`343`		`- dropout_rate`
	`343`	`+ dropout_rate,`
`344`	`344`	`)`
`345`	`345`	`)`
`346`	`346`