feat: parse atol and rtol from kernel_task.yaml to pass them to search agents

[shangkunwang01]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces changes to MaxKernel/auto_search/run_batch_search.py to parse atol and rtol values from a kernel_task.yaml file and inject them into the agent_config of problem_kwargs. Feedback on these changes highlights three key issues: potential TypeError exceptions when atol or rtol are lists rather than single floats, a lack of robustness if the YAML file does not parse into a dictionary, and unnecessary dictionary copying that could inject an empty agent_config when no values are parsed. A refactored code block was suggested to resolve these issues.

[gemini-code-assist]

There are three issues with the current parsing and configuration update logic:

1. Type Safety / Support for Lists: As defined in KernelTask (in MaxKernel/evaluation/custom_types/kernel_task.py), atol and rtol can be either a single float or a List[float]. Attempting to unconditionally cast them using float(...) will raise a TypeError if they are lists.
2. Robustness against Non-Dictionary YAML: If kernel_task.yaml is empty or contains a non-dictionary structure (e.g., a list or string), yaml.safe_load(f) will return None or a non-dict object. Checking "atol" in task_data will then raise a TypeError (e.g., 'NoneType' is not iterable).
3. Unnecessary Dict Copying & Empty Config Injection: problem_kwargs is copied and agent_config is initialized to {} even when both atol and rtol are None. This can inject an empty agent_config dictionary into problem_kwargs when it was originally None or not provided, potentially altering downstream behavior.

We can resolve all of these issues by checking if task_data is a dictionary, supporting list parsing, and only copying/modifying problem_kwargs if atol or rtol is actually parsed.

      atol = None
      rtol = None
      kernel_task_file = os.path.join(problem_dir, "kernel_task.yaml")
      if os.path.exists(kernel_task_file):
        with open(kernel_task_file, "r") as f:
          try:
            task_data = yaml.safe_load(f)
            if isinstance(task_data, dict):
              if "atol" in task_data and task_data["atol"] is not None:
                atol_val = task_data["atol"]
                atol = [float(x) for x in atol_val] if isinstance(atol_val, list) else float(atol_val)
              if "rtol" in task_data and task_data["rtol"] is not None:
                rtol_val = task_data["rtol"]
                rtol = [float(x) for x in rtol_val] if isinstance(rtol_val, list) else float(rtol_val)
          except Exception as e:
            logger.warning(
              f"Failed to parse kernel_task.yaml for {problem_id}: {e}"
            )

      problem_kwargs = kwargs
      if atol is not None or rtol is not None:
        problem_kwargs = dict(kwargs)
        agent_config = dict(problem_kwargs.get("agent_config") or {})
        if atol is not None:
          agent_config["atol"] = atol
        if rtol is not None:
          agent_config["rtol"] = rtol
        problem_kwargs["agent_config"] = agent_config

[shangkunwang01]

1. fixed. However, the agent does not support tol as list in the test harness. If a list is passed in, the first tol will be used.
2. fixed
3. fixed

[NinaCai]

1. I think it makes sense to have different tols for different input, especially for different dtypes?
2. Does edge cases also use the tol you provides?

[shangkunwang01]

1. Yet, it makes sense to have different tol for different input. I can have another PR for this because several parts (test, prompts and tool) need to be updated. For now, the task yaml file is only parsed for our benchmark and they only contain one tolerance. For general tasks without yaml file, the agent will decide a tolerance as before.
2. Yes but since this tolerance will be consumed only if there is a task yaml file, it will not change the behavior of the agent outside the benchmark. Ideally the test generation should have different tol for different inputs.

[NinaCai]

Sounds good. If we are not going to support list of float anytime soon, could we change

accelerator-agents/MaxKernel/evaluation/jax_kernel_evaluator.py

Lines 77 to 78 in 9b6f768

	atol: Optional[Union[float, List[float]]] = None,
	rtol: Optional[Union[float, List[float]]] = None,

,

accelerator-agents/MaxKernel/evaluation/benchmark.py

Lines 31 to 32 in 9b6f768

	atol: Optional[Union[float, List[float]]] = None,
	rtol: Optional[Union[float, List[float]]] = None,

,

accelerator-agents/MaxKernel/evaluation/code_adapter/code_adapter.py

Lines 100 to 101 in 9b6f768

	atol: Optional[Union[float, List[float]]] = None,
	rtol: Optional[Union[float, List[float]]] = None,

so that only float is supported?

[shangkunwang01]

The eval is seperate from the agent so I think eval can still keep this feature because even the agent only support one tol, we can test different tol with the eval.

[NinaCai]

I see. but looks like the evaluation only accepts float type

accelerator-agents/MaxKernel/evaluation/benchmark.py

Line 258 in 9b6f768

type=float,

[shangkunwang01]

The nargs="+" supports it to have multiple inputs.

[NinaCai]

I see. I didn't know we support list of atol in evaluation.
I find that usually the test agent will pick a threshold smaller than the one we set in the yaml file. Why do we want to use a larger threshold instead?

[shangkunwang01]

This change is mainly for the benchmark experiment. In real use case, no yaml file will be provide so the change will not matter. In benchmark, I have seen the agent picking looser threshhold in the experiment and that node was chosen. This makes the final result a bit unreliable.