From a284e4258a9108ea88121322f6f2ac0759782527 Mon Sep 17 00:00:00 2001 From: NinaCai Date: Mon, 30 Mar 2026 15:59:41 +0000 Subject: [PATCH 1/2] update model version to gemini 3.1 pro --- .../subagents/gpu_to_jax_agent/constants.py | 15 ++++++++------- .../agents/kernel_gen_agent/constants.py | 8 ++++---- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py b/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py index fe484fb..56cd045 100644 --- a/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py +++ b/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py @@ -1,13 +1,14 @@ """Configuration constants for GPU to JAX conversion agent.""" +from tpu_kernel_gen.agents.kernel_gen_agent import constants -MODEL_NAME = "gemini-3-pro-preview" -MAX_ITERATIONS = 5 -LLM_GEN_RETRY_COUNT = 3 -TEMPERATURE = 0.1 -TOP_P = 0.9 -TOP_K = 5 +MODEL_NAME = constants.MODEL_NAME +MAX_ITERATIONS = constants.MAX_ITERATIONS +LLM_GEN_RETRY_COUNT = constants.LLM_GEN_RETRY_COUNT +TEMPERATURE = constants.TEMPERATURE +TOP_P = constants.TOP_P +TOP_K = constants.TOP_K CONVERSION_TIMEOUT = 180 # Timeout for conversion attempts -EVAL_SERVER_PORT = 1245 +EVAL_SERVER_PORT = constants.EVAL_SERVER_PORT NUMERICAL_TOLERANCE = 1e-5 # Tolerance for numerical correctness checks # Backend selection for evaluation diff --git a/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py b/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py index d900c91..5c401f3 100644 --- a/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py +++ b/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py @@ -1,12 +1,12 @@ -MODEL_NAME = "gemini-3-pro-preview" +MODEL_NAME = "gemini-3.1-pro-preview" MAX_ITERATIONS = 5 LLM_GEN_RETRY_COUNT = 3 TEMPERATURE = 0.1 TOP_P = 0.9 TOP_K = 5 -TPU_TIMEOUT = 120 -REQUEST_TIMEOUT = 1800 +TPU_TIMEOUT=120 +REQUEST_TIMEOUT=1800 TPU_SERVER_PORT = 5463 CPU_SERVER_PORT = 5464 EVAL_SERVER_PORT = 1245 -PERF_THRESHOLD = 1.1 +PERF_THRESHOLD = 1.1 \ No newline at end of file From aae2180c5415fadc3a6e607251d5649ebd647262 Mon Sep 17 00:00:00 2001 From: NinaCai Date: Mon, 30 Mar 2026 16:07:28 +0000 Subject: [PATCH 2/2] format and lint --- .../subagents/gpu_to_jax_agent/constants.py | 1 + .../tpu_kernel_gen/agents/kernel_gen_agent/constants.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py b/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py index 56cd045..282b477 100644 --- a/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py +++ b/MaxKernel/tpu_kernel_gen/agents/hitl_kernel_gen_agent/subagents/gpu_to_jax_agent/constants.py @@ -1,4 +1,5 @@ """Configuration constants for GPU to JAX conversion agent.""" + from tpu_kernel_gen.agents.kernel_gen_agent import constants MODEL_NAME = constants.MODEL_NAME diff --git a/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py b/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py index 5c401f3..ad1aec0 100644 --- a/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py +++ b/MaxKernel/tpu_kernel_gen/agents/kernel_gen_agent/constants.py @@ -4,9 +4,9 @@ TEMPERATURE = 0.1 TOP_P = 0.9 TOP_K = 5 -TPU_TIMEOUT=120 -REQUEST_TIMEOUT=1800 +TPU_TIMEOUT = 120 +REQUEST_TIMEOUT = 1800 TPU_SERVER_PORT = 5463 CPU_SERVER_PORT = 5464 EVAL_SERVER_PORT = 1245 -PERF_THRESHOLD = 1.1 \ No newline at end of file +PERF_THRESHOLD = 1.1