feat: support customizable and per-input absolute/relative tolerances for kernel evaluation

[shangkunwang01]

This PR enhances the JAX kernel evaluation framework by introducing support for customizable and per-input absolute (atol) and relative (rtol) tolerances. Previously, evaluation tolerances were hardcoded or limited to single scalar float values (defaulting to 1e-3) via string replacement in the harness template.

With this change, tolerances can be specified:

1. Per-task in YAML definitions (kernel_task.yaml).
2. Via CLI arguments (--atol, --rtol) as single scalars or lists of floats (this will override values from the yaml file).
3. Programmatically through JAXKernelEvaluator and benchmark.py.

---

Key Changes

• Core Harness & Per-Input Tolerance Resolution (harness_code.py)

  • Replaced static {atol}/{rtol} template string substitution with JSON-driven configuration loading (task_info["atol"] / task_info["rtol"]).

• Evaluator & Benchmark Pipeline (jax_kernel_evaluator.py, benchmark.py)

  • Updated atol and rtol type annotations across evaluate(), _evaluate_local(), and _evaluate_remote() from float to Optional[Union[float, List[float]]].

• Task Data Model & Code Adapter (kernel_task.py, evaluation_utils.py, code_adapter.py)

  • Extended the KernelTask dataclass to include optional atol and rtol fields (Optional[Union[float, List[float]]] = None).
  • Updated load_kernel_task_from_yaml and CodeAdapter.generate_kernel_task to parse and propagate task-level tolerance configurations.

• Benchmark Dataset Updates (kernel_task.yaml files)

  • Added explicit atol and rtol thresholds across all 50 evaluation task YAMLs in MaxKernel/evaluation/jaxbench_adapted_dataset/ and MaxKernel/evaluation/examples/dsv4/, tailored to each task's numerical stability requirements

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[NinaCai]

We should later add multiple configs later, and have separate tolerance for each set of config.