Automated Code Change

h-joo · copybara-github · commit 948deddc9daf · 2026-07-07T17:14:46.000-07:00
PiperOrigin-RevId: 944173080
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_GRU.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_GRU.py
@@ -78,9 +78,9 @@ def test_jax_independent():
 
   expected_shape = (CONFIG['batch_size'], CONFIG['output_size'])
   assert (
-      y.shape == expected_shape
+      y.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f'JAX Shape Mismatch: Expected {expected_shape}, got {y.shape}'
-  assert not jnp.isnan(y).any(), 'JAX Output contains NaNs'
+  assert not jnp.isnan(y).any(), 'JAX Output contains NaNs'  # pyrefly: ignore[bad-argument-type]
 
 
 def main(argv: Sequence[str]):
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_GenerativeAdversarialNetwork.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_GenerativeAdversarialNetwork.py
@@ -72,23 +72,23 @@ def test_jax_validity():
   gen_params = generator.init(key, dummy_z)
   gen_out = generator.apply(gen_params, dummy_z, train=False)
   expected_gen_shape = (batch_size, int(np.prod(img_shape)))
-  assert gen_out.shape == expected_gen_shape, (
+  assert gen_out.shape == expected_gen_shape, (  # pyrefly: ignore[missing-attribute]
       f'JAX Generator Shape Mismatch: Expected {expected_gen_shape}, got'
       f' {gen_out.shape}'
   )
-  assert not jnp.isnan(gen_out).any(), 'JAX Generator Output contains NaNs'
+  assert not jnp.isnan(gen_out).any(), 'JAX Generator Output contains NaNs'  # pyrefly: ignore[bad-argument-type]
 
   # Test Discriminator
   discriminator = model_jax.Discriminator()
   dummy_img = jnp.zeros((batch_size, *img_shape))
   disc_params = discriminator.init(key, dummy_img)
   disc_out = discriminator.apply(disc_params, dummy_img)
   expected_disc_shape = (batch_size, 1)
-  assert disc_out.shape == expected_disc_shape, (
+  assert disc_out.shape == expected_disc_shape, (  # pyrefly: ignore[missing-attribute]
       f'JAX Discriminator Shape Mismatch: Expected {expected_disc_shape}, got'
       f' {disc_out.shape}'
   )
-  assert not jnp.isnan(disc_out).any(), 'JAX Discriminator Output contains NaNs'
+  assert not jnp.isnan(disc_out).any(), 'JAX Discriminator Output contains NaNs'  # pyrefly: ignore[bad-argument-type]
 
 
 def test_generator_equivalence():
@@ -148,7 +148,7 @@ def copy_bn(torch_layer, jax_layer_name):
   out_jax = jax_model.apply(params, z_jax, train=False)
 
   print(f'Generator Output Max Diff: {np.max(np.abs(out_torch - out_jax))}')
-  np.testing.assert_allclose(out_torch, out_jax, rtol=1e-5, atol=1e-5)
+  np.testing.assert_allclose(out_torch, out_jax, rtol=1e-5, atol=1e-5)  # pyrefly: ignore[no-matching-overload]
   print('Generator Equivalence: PASS')
 
 
@@ -196,7 +196,7 @@ def copy_linear(torch_layer, jax_layer_name):
   out_jax = jax_model.apply(params, img_jax)
 
   print(f'Discriminator Output Max Diff: {np.max(np.abs(out_torch - out_jax))}')
-  np.testing.assert_allclose(out_torch, out_jax, rtol=1e-5, atol=1e-5)
+  np.testing.assert_allclose(out_torch, out_jax, rtol=1e-5, atol=1e-5)  # pyrefly: ignore[no-matching-overload]
   print('Discriminator Equivalence: PASS')
 
 
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_GraphConvolution.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_GraphConvolution.py
@@ -65,9 +65,9 @@ def test_jax_independent():
   expected_shape = (CONFIG["n_nodes"], CONFIG["out_features"])
 
   assert (
-      out.shape == expected_shape
+      out.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f"JAX Shape Mismatch: Expected {expected_shape}, got {out.shape}"
-  assert not jnp.isnan(out).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(out).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_LSTMAttentionForecaster.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_LSTMAttentionForecaster.py
@@ -74,9 +74,9 @@ def test_jax_independent():
 
   expected_shape = (CONFIG["batch_size"], CONFIG["output_size"])
   assert (
-      y.shape == expected_shape
+      y.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f"JAX Shape Mismatch: Expected {expected_shape}, got {y.shape}"
-  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_LSTMForecaster.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_LSTMForecaster.py
@@ -77,9 +77,9 @@ def test_jax_independent():
 
   expected_shape = (CONFIG["batch_size"], CONFIG["output_size"])
   assert (
-      y.shape == expected_shape
+      y.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f"JAX Shape Mismatch: Expected {expected_shape}, got {y.shape}"
-  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_MLP.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_MLP.py
@@ -73,9 +73,9 @@ def test_jax_independent():
   # Checks
   expected_shape = (TEST_CONFIG["batch_size"], 1)
   assert (
-      y.shape == expected_shape
+      y.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f"JAX Shape Mismatch: Expected {expected_shape}, got {y.shape}"
-  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_SelfAttention.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_SelfAttention.py
@@ -84,9 +84,9 @@ def test_jax_independent():
       config["embed_size"],
   )
   assert (
-      out.shape == expected_shape
+      out.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f"JAX Shape Mismatch: Expected {expected_shape}, got {out.shape}"
-  assert not jnp.isnan(out).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(out).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
   return True
 
 
@@ -143,7 +143,7 @@ def copy_layer(pt_layer, jax_name):
   print(f"   >> Head Dim: {head_dim}, Scaling Factor: {scale}")
 
   # Final Assertion
-  np.testing.assert_allclose(pt_out, jax_out, rtol=1e-4, atol=1e-5)
+  np.testing.assert_allclose(pt_out, jax_out, rtol=1e-4, atol=1e-5)  # pyrefly: ignore[no-matching-overload]
   return True
 
 
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_SimpleCNN.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_SimpleCNN.py
@@ -52,11 +52,11 @@ def test_jax_model_validity():
   params = model.init(key, dummy_input)
   output = model.apply(params, dummy_input)
 
-  assert output.shape == (
+  assert output.shape == (  # pyrefly: ignore[missing-attribute]
       4,
       10,
   ), f'Expected output shape (4, 10), got {output.shape}'
-  assert not jnp.isnan(output).any(), 'Output contains NaNs'
+  assert not jnp.isnan(output).any(), 'Output contains NaNs'  # pyrefly: ignore[bad-argument-type]
 
 
 def test_pytorch_jax_equivalence():
@@ -130,7 +130,7 @@ def get_b(layer):
   diff = np.abs(out_torch - jax_numpy)
   max_diff = np.max(diff)
   print(f'\n   >> Max Absolute Difference: {max_diff:.2e}')
-  np.testing.assert_allclose(out_torch, out_jax, rtol=1e-5, atol=1e-5)
+  np.testing.assert_allclose(out_torch, out_jax, rtol=1e-5, atol=1e-5)  # pyrefly: ignore[no-matching-overload]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_SimpleRNN.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_SimpleRNN.py
@@ -45,8 +45,8 @@ def test_jax_independent():
   x = jnp.ones((2, 5, 4))
   variables = model.init(jax.random.PRNGKey(42), x)
   y = model.apply(variables, x)
-  assert y.shape == (2, 3)
-  assert not jnp.isnan(y).any()
+  assert y.shape == (2, 3)  # pyrefly: ignore[missing-attribute]
+  assert not jnp.isnan(y).any()  # pyrefly: ignore[bad-argument-type]
 
 
 def test_rnn_equivalence(
@@ -118,7 +118,7 @@ def test_rnn_equivalence(
   diff = np.abs(pt_out - jax_numpy)
   max_diff = np.max(diff)
   print(f'\n   >> Max Absolute Difference: {max_diff:.2e}')
-  np.testing.assert_allclose(pt_out, jax_out, rtol=1e-5, atol=1e-5)
+  np.testing.assert_allclose(pt_out, jax_out, rtol=1e-5, atol=1e-5)  # pyrefly: ignore[no-matching-overload]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_TCU.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_TCU.py
@@ -72,9 +72,9 @@ def test_jax_independent():
   )
 
   assert (
-      y.shape == expected_shape
+      y.shape == expected_shape  # pyrefly: ignore[missing-attribute]
   ), f"JAX Shape Mismatch: Expected {expected_shape}, got {y.shape}"
-  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(y).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
 
 
 # pylint: disable=unused-argument
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_VAE.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_VAE.py
@@ -79,7 +79,7 @@ def test_jax_independent():
   variables = model.init(keys, dummy_input)
 
   # The .apply() call was already correct using rngs=...
-  recon, mu, log_var = model.apply(
+  recon, mu, log_var = model.apply(  # pyrefly: ignore[bad-unpacking]
       variables, dummy_input, rngs={"sampling": jax.random.PRNGKey(0)}
   )
   # Shape and NaN checks
@@ -96,7 +96,7 @@ def test_jax_independent():
       f" {(config['batch_size'], config['latent_dim'])}, got {log_var.shape}"
   )
   assert not jnp.isnan(recon).any(), "JAX Reconstruction contains NaNs"
-  assert not jnp.isnan(mu).any(), "JAX Mu contains NaNs"
+  assert not jnp.isnan(mu).any(), "JAX Mu contains NaNs"  # pyrefly: ignore[bad-argument-type]
   assert not jnp.isnan(log_var).any(), "JAX LogVar contains NaNs"
 
 
diff --git a/MaxCode/evaluation/code_agent/test_cases/level2/test_encoder.py b/MaxCode/evaluation/code_agent/test_cases/level2/test_encoder.py
@@ -107,10 +107,10 @@ def test_jax_independent():
   out = model.apply(variables, x, src_mask=None, deterministic=True)
 
   expected_shape = (CONFIG["batch_size"], CONFIG["seq_len"], CONFIG["ntoken"])
-  assert out.shape == expected_shape, (
+  assert out.shape == expected_shape, (  # pyrefly: ignore[missing-attribute]
       f"JAX Shape Mismatch: Expected {expected_shape}, got {out.shape}"
   )
-  assert not jnp.isnan(out).any(), "JAX Output contains NaNs"
+  assert not jnp.isnan(out).any(), "JAX Output contains NaNs"  # pyrefly: ignore[bad-argument-type]
 
 
 def transfer_and_align_weights(pt_model, jax_variables):
diff --git a/MaxCode/mcp_server/primary_agent_server.py b/MaxCode/mcp_server/primary_agent_server.py
@@ -43,10 +43,10 @@ async def _execute_adk_agent(
       break
   if handler_to_reset:
     handler_to_reset.close()
-    handler_to_reset.mode = "w"
+    handler_to_reset.mode = "w"  # pyrefly: ignore[missing-attribute]
   logging.info("Job started, clearing agent server log.")
   if handler_to_reset:
-    handler_to_reset.mode = "a"
+    handler_to_reset.mode = "a"  # pyrefly: ignore[missing-attribute]
   logging.info("Received agent call. Prompt: %s, API Key: %s", prompt, api_key)
   await ctx.info("Starting ADK Agent task...")
   logging.info("Starting ADK Agent task...")
diff --git a/MaxCode/models.py b/MaxCode/models.py
@@ -71,7 +71,7 @@ def __call__(self, user_prompt: str):
     headers = {"Content-Type": "application/json"}
     payload = {"contents": [{"parts": [{"text": user_prompt}], "role": "user"}]}
     if self.system_instruction:
-      payload["system_instruction"] = {
+      payload["system_instruction"] = {  # pyrefly: ignore[bad-assignment]
           "parts": [{"text": self.system_instruction}]
       }
 
diff --git a/MaxCode/rag/embedding.py b/MaxCode/rag/embedding.py
@@ -38,9 +38,9 @@ def embed(self, text: str) -> list[float]:
 
     time.sleep(2)
     embedding = (
-        self.client.models.embed_content(model=self.model, contents=text)
+        self.client.models.embed_content(model=self.model, contents=text)  # pyrefly: ignore[unsupported-operation]
         .embeddings[0]
         .values
     )
     self._cache[text] = embedding
-    return embedding
+    return embedding  # pyrefly: ignore[bad-return]

Original file line number	Diff line number	Diff line change
`@@ -72,9 +72,9 @@ def test_jax_independent():`
`72`	`72`	`)`
`73`	`73`
`74`	`74`	`assert (`
`75`		`- y.shape == expected_shape`
	`75`	`+ y.shape == expected_shape # pyrefly: ignore[missing-attribute]`
`76`	`76`	`), f"JAX Shape Mismatch: Expected {expected_shape}, got {y.shape}"`
`77`		`- assert not jnp.isnan(y).any(), "JAX Output contains NaNs"`
	`77`	`+ assert not jnp.isnan(y).any(), "JAX Output contains NaNs" # pyrefly: ignore[bad-argument-type]`
`78`	`78`
`79`	`79`
`80`	`80`	`# pylint: disable=unused-argument`