Spaces:
Running
Running
File size: 3,535 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import pytest
from kernel_env.server.static_checks import check_source
from kernels import CUDA_GOOD, GOOD, HEADER, KERNEL, with_forward
def test_good_triton_passes():
assert check_source(GOOD) == []
def test_copied_get_inputs_allowed():
assert check_source(GOOD + "\ndef get_inputs():\n return [torch.rand(4, 4)]\n") == []
def test_layer_named_like_tensor_op_allowed():
# params held in an nn layer whose attribute name matches a tensor op (seen in Qwen3.5 baseline)
src = GOOD.replace("super().__init__()\n", "super().__init__()\n self.conv3d = nn.Conv3d(3, 8, 3)\n", 1)
src = src.replace("x = x.contiguous()", "w = self.conv3d.weight\n x = x.contiguous()")
assert "self.conv3d" in src and check_source(src) == []
called = src.replace("return y", "return self.conv3d(y)")
assert any("nn layer" in e for e in check_source(called))
def test_good_cuda_passes():
assert check_source(CUDA_GOOD, "cuda") == []
@pytest.mark.parametrize(
"src, expect",
[
(with_forward("return torch.relu(x)"), "torch.relu"),
(with_forward("return F.relu(x)"), "torch.nn.functional.relu"),
(with_forward("return self.act(x)"), "nn layer"),
(with_forward("return x.clamp(min=0)"), ".clamp"),
(with_forward("return x @ x.t()"), "'@'"),
(with_forward("mm = torch.matmul\nreturn mm(x, x)"), "reference to PyTorch op"),
(with_forward("f = x.clamp\nreturn f(min=0)"), "reference to tensor op"),
(with_forward("return relu(x)", header=HEADER + "from torch import relu\n"), "torch.relu"),
(with_forward("return getattr(torch, 'relu')(x)"), "getattr"),
(with_forward("return x", extra="import os\n"), "import of 'os'"),
(with_forward("return x", extra="from subprocess import run\n"), "import from 'subprocess'"),
(with_forward("return x", extra="torch.allclose = lambda *a, **k: True\n"), "assigning attributes"),
(with_forward("return x.__class__"), "dunder"),
(with_forward("return torch.compile(lambda t: t)(x)"), "torch.compile"),
(HEADER + "class ModelNew(nn.Module):\n def forward(self, x):\n return x\n", "no @triton.jit"),
(HEADER + KERNEL, "no top-level class ModelNew"),
("def broken(:\n", "SyntaxError"),
],
)
def test_rejects(src, expect):
errors = check_source(src)
assert any(expect in e for e in errors), errors
def test_cuda_requires_load_inline():
assert any("load_inline" in e for e in check_source(GOOD, "cuda"))
def test_prompt_v2_example_is_valid(monkeypatch):
from kernel_env.server.tasks import EXAMPLE_V2, Task, build_prompt
assert check_source(EXAMPLE_V2) == [] # the worked example must itself pass the code checks
task = Task("level1/19_ReLU", 1, "class Model: ...")
assert EXAMPLE_V2 in build_prompt(task)
monkeypatch.setenv("KERNEL_ENV_PROMPT", "1")
assert EXAMPLE_V2 not in build_prompt(task) and "Reply with the complete module source only" in build_prompt(task)
@pytest.mark.parametrize("version", ["1", "2"])
@pytest.mark.parametrize("task_id", ["level1/23_Softmax", "level1/37_FrobeniusNorm_"])
def test_normalization_prompt_discloses_extra_validation(monkeypatch, version, task_id):
from kernel_env.server.tasks import Task, build_prompt
monkeypatch.setenv("KERNEL_ENV_PROMPT", version)
prompt = build_prompt(Task(task_id, 1, "class Model: ..."))
assert "atol=min(configured_atol, max(abs(reference))*1e-5)" in prompt and "signed inputs" in prompt
|