change var name to enhance readibility

Naeemkh · Naeemkh · commit 5efdfad14058 · 2025-03-28T01:29:07.000Z
pyt --&gt; pytorch
k --&gt; eval_metric_key
diff --git a/tests/modeldiffs/criteo1tb/compare.py b/tests/modeldiffs/criteo1tb/compare.py
@@ -53,16 +53,16 @@ def sd_transform(sd):
   jax_workload = JaxWorkload()
   pytorch_workload = PyTorchWorkload()
 
-  pyt_batch = {
+  pytorch_batch = {
       'inputs': torch.ones((2, 13 + 26)),
       'targets': torch.randint(low=0, high=1, size=(2,)),
       'weights': torch.ones(2),
   }
-  jax_batch = {k: np.array(v) for k, v in pyt_batch.items()}
+  jax_batch = {k: np.array(v) for k, v in pytorch_batch.items()}
 
   # Test outputs for identical weights and inputs.
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/criteo1tb_embed_init/compare.py b/tests/modeldiffs/criteo1tb_embed_init/compare.py
@@ -52,16 +52,16 @@ def sd_transform(sd):
   jax_workload = JaxWorkload()
   pytorch_workload = PyTorchWorkload()
 
-  pyt_batch = {
+  pytorch_batch = {
       'inputs': torch.ones((2, 13 + 26)),
       'targets': torch.randint(low=0, high=1, size=(2,)),
       'weights': torch.ones(2),
   }
-  jax_batch = {k: np.array(v) for k, v in pyt_batch.items()}
+  jax_batch = {k: np.array(v) for k, v in pytorch_batch.items()}
 
   # Test outputs for identical weights and inputs.
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/criteo1tb_layernorm/compare.py b/tests/modeldiffs/criteo1tb_layernorm/compare.py
@@ -64,16 +64,16 @@ def sd_transform(sd):
   jax_workload = JaxWorkload()
   pytorch_workload = PyTorchWorkload()
 
-  pyt_batch = {
+  pytorch_batch = {
       'inputs': torch.ones((2, 13 + 26)),
       'targets': torch.randint(low=0, high=1, size=(2,)),
       'weights': torch.ones(2),
   }
-  jax_batch = {k: np.array(v) for k, v in pyt_batch.items()}
+  jax_batch = {k: np.array(v) for k, v in pytorch_batch.items()}
 
   # Test outputs for identical weights and inputs.
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/criteo1tb_resnet/compare.py b/tests/modeldiffs/criteo1tb_resnet/compare.py
@@ -64,7 +64,7 @@ def sd_transform(sd):
   jax_workload = JaxWorkload()
   pytorch_workload = PyTorchWorkload()
 
-  pyt_batch = {
+  pytorch_batch = {
       'inputs': torch.ones((2, 13 + 26)),
       'targets': torch.randint(low=0, high=1, size=(2,)),
       'weights': torch.ones(2),
@@ -75,12 +75,12 @@ def sd_transform(sd):
   input_size = 13 + num_categorical_features
   input_shape = (init_fake_batch_size, input_size)
   fake_inputs = jnp.ones(input_shape, jnp.float32)
-  jax_batch = {k: np.array(v) for k, v in pyt_batch.items()}
+  jax_batch = {k: np.array(v) for k, v in pytorch_batch.items()}
   jax_batch['inputs'] = fake_inputs
 
   # Test outputs for identical weights and inputs.
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/fastmri/compare.py b/tests/modeldiffs/fastmri/compare.py
@@ -61,10 +61,10 @@ def sort_key(k):
   image = torch.randn(2, 320, 320)
 
   jax_batch = {'inputs': image.detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/fastmri_layernorm/compare.py b/tests/modeldiffs/fastmri_layernorm/compare.py
@@ -68,10 +68,10 @@ def sort_key(k):
   image = torch.randn(2, 320, 320)
 
   jax_batch = {'inputs': image.detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/fastmri_model_size/compare.py b/tests/modeldiffs/fastmri_model_size/compare.py
@@ -61,10 +61,10 @@ def sort_key(k):
   image = torch.randn(2, 320, 320)
 
   jax_batch = {'inputs': image.detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/fastmri_tanh/compare.py b/tests/modeldiffs/fastmri_tanh/compare.py
@@ -61,10 +61,10 @@ def sort_key(k):
   image = torch.randn(2, 320, 320)
 
   jax_batch = {'inputs': image.detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_resnet/compare.py b/tests/modeldiffs/imagenet_resnet/compare.py
@@ -78,10 +78,10 @@ def sd_transform(sd):
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_resnet/gelu_compare.py b/tests/modeldiffs/imagenet_resnet/gelu_compare.py
@@ -25,10 +25,10 @@
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_resnet/silu_compare.py b/tests/modeldiffs/imagenet_resnet/silu_compare.py
@@ -25,10 +25,10 @@
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_vit/compare.py b/tests/modeldiffs/imagenet_vit/compare.py
@@ -91,10 +91,10 @@ def key_transform(k):
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_vit_glu/compare.py b/tests/modeldiffs/imagenet_vit_glu/compare.py
@@ -27,10 +27,10 @@
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_vit_map/compare.py b/tests/modeldiffs/imagenet_vit_map/compare.py
@@ -38,10 +38,10 @@ def sd_transform(sd):
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/imagenet_vit_postln/compare.py b/tests/modeldiffs/imagenet_vit_postln/compare.py
@@ -27,10 +27,10 @@
   image = torch.randn(2, 3, 224, 224)
 
   jax_batch = {'inputs': image.permute(0, 2, 3, 1).detach().numpy()}
-  pyt_batch = {'inputs': image}
+  pytorch_batch = {'inputs': image}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_conformer/compare.py b/tests/modeldiffs/librispeech_conformer/compare.py
@@ -66,10 +66,10 @@ def sd_transform(sd):
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_conformer_attention_temperature/compare.py b/tests/modeldiffs/librispeech_conformer_attention_temperature/compare.py
@@ -66,10 +66,10 @@ def sd_transform(sd):
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_conformer_gelu/compare.py b/tests/modeldiffs/librispeech_conformer_gelu/compare.py
@@ -66,10 +66,10 @@ def sd_transform(sd):
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_conformer_layernorm/compare.py b/tests/modeldiffs/librispeech_conformer_layernorm/compare.py
@@ -66,10 +66,10 @@ def sd_transform(sd):
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_deepspeech/compare.py b/tests/modeldiffs/librispeech_deepspeech/compare.py
@@ -91,10 +91,10 @@ def sd_transform(sd):
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_deepspeech_noresnet/compare.py b/tests/modeldiffs/librispeech_deepspeech_noresnet/compare.py
@@ -27,10 +27,10 @@
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_deepspeech_normaug/compare.py b/tests/modeldiffs/librispeech_deepspeech_normaug/compare.py
@@ -27,10 +27,10 @@
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/librispeech_deepspeech_tanh/compare.py b/tests/modeldiffs/librispeech_deepspeech_tanh/compare.py
@@ -27,10 +27,10 @@
   pad[0, 200000:] = 1
 
   jax_batch = {'inputs': (wave.detach().numpy(), pad.detach().numpy())}
-  pyt_batch = {'inputs': (wave, pad)}
+  pytorch_batch = {'inputs': (wave, pad)}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/ogbg/compare.py b/tests/modeldiffs/ogbg/compare.py
@@ -79,7 +79,7 @@ def sd_transform(sd):
   jax_workload = JaxWorkload()
   pytorch_workload = PyTorchWorkload()
 
-  pyt_batch = dict(
+  pytorch_batch = dict(
       n_node=torch.LongTensor([5]),
       n_edge=torch.LongTensor([5]),
       nodes=torch.randn(5, 9),
@@ -88,17 +88,17 @@ def sd_transform(sd):
       senders=torch.LongTensor(list(range(5))),
       receivers=torch.LongTensor([(i + 1) % 5 for i in range(5)]))
 
-  jax_batch = {k: np.array(v) for k, v in pyt_batch.items()}
+  jax_batch = {k: np.array(v) for k, v in pytorch_batch.items()}
 
   # Test outputs for identical weights and inputs.
   graph_j = jraph.GraphsTuple(**jax_batch)
-  graph_p = jraph.GraphsTuple(**pyt_batch)
+  graph_p = jraph.GraphsTuple(**pytorch_batch)
 
   jax_batch = {'inputs': graph_j}
-  pyt_batch = {'inputs': graph_p}
+  pytorch_batch = {'inputs': graph_p}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/ogbg_gelu/compare.py b/tests/modeldiffs/ogbg_gelu/compare.py
@@ -79,7 +79,7 @@ def sd_transform(sd):
   jax_workload = JaxWorkload()
   pytorch_workload = PyTorchWorkload()
 
-  pyt_batch = dict(
+  pytorch_batch = dict(
       n_node=torch.LongTensor([5]),
       n_edge=torch.LongTensor([5]),
       nodes=torch.randn(5, 9),
@@ -88,17 +88,17 @@ def sd_transform(sd):
       senders=torch.LongTensor(list(range(5))),
       receivers=torch.LongTensor([(i + 1) % 5 for i in range(5)]))
 
-  jax_batch = {k: np.array(v) for k, v in pyt_batch.items()}
+  jax_batch = {k: np.array(v) for k, v in pytorch_batch.items()}
 
   # Test outputs for identical weights and inputs.
   graph_j = jraph.GraphsTuple(**jax_batch)
-  graph_p = jraph.GraphsTuple(**pyt_batch)
+  graph_p = jraph.GraphsTuple(**pytorch_batch)
 
   jax_batch = {'inputs': graph_j}
-  pyt_batch = {'inputs': graph_p}
+  pytorch_batch = {'inputs': graph_p}
 
   pytorch_model_kwargs = dict(
-      augmented_and_preprocessed_input_batch=pyt_batch,
+      augmented_and_preprocessed_input_batch=pytorch_batch,
       model_state=None,
       mode=spec.ForwardPassMode.EVAL,
       rng=None,
diff --git a/tests/modeldiffs/ogbg_model_size/compare.py b/tests/modeldiffs/ogbg_model_size/compare.py
diff --git a/tests/modeldiffs/ogbg_silu/compare.py b/tests/modeldiffs/ogbg_silu/compare.py
diff --git a/tests/modeldiffs/wmt/compare.py b/tests/modeldiffs/wmt/compare.py
diff --git a/tests/modeldiffs/wmt_attention_temp/compare.py b/tests/modeldiffs/wmt_attention_temp/compare.py
diff --git a/tests/modeldiffs/wmt_glu_tanh/compare.py b/tests/modeldiffs/wmt_glu_tanh/compare.py
diff --git a/tests/modeldiffs/wmt_post_ln/compare.py b/tests/modeldiffs/wmt_post_ln/compare.py
diff --git a/tests/test_traindiffs.py b/tests/test_traindiffs.py
