diff --git a/HONOR_CODE.md b/HONOR_CODE.md new file mode 100644 index 0000000..5d22ad7 --- /dev/null +++ b/HONOR_CODE.md @@ -0,0 +1,71 @@ +# 2026 春季启元人工智能大赛诚信守则（Honor Code） + + +本人作为 2026 春季启元人工智能大赛（以下简称“比赛”）的参赛选手，郑重承诺严格遵守比赛规则及本诚信守则，秉持诚信、公正、廉洁的参赛原则，自觉维护比赛的公平性与严肃性。本人充分理解并认可，违反本准则将导致参赛资格被取消、比赛成绩作废等相应后果，且愿意承担由此产生的一切责任。 + +## 一、参赛诚信承诺 + +1. 本人保证所提交的赛题PR（Pull Request）中包含的算子实现代码及相关文档，均为本人（及参赛团队，如为团队参赛）在比赛期间独立完成或在明确标注参考来源的基础上进行开发，不存在任何欺诈、抄袭、作弊行为。 + +2. 本人承诺主动、全面、真实地披露赛题实现过程中所有参考的外部资源，尤其是开源代码资源，不隐瞒任何可能影响比赛公平性的信息。 + +3. 本人保证不采用任何不正当手段获取比赛优势，包括但不限于窃取其他参赛选手的代码成果、利用非比赛允许的工具或技术、与他人串通作弊等。 + +## 二、参考资源说明 + +本人确认已按比赛要求，将本次赛题实现过程中涉及的参考资源信息单独撰写至`REFERENCE.md`文件中，该文件将与本诚信守则一同作为PR附件提交。`REFERENCE.md`需根据实际参考情况，按以下要求完整填写，信息不完整或虚假填写将视为违反本准则： + +**情况1：无参考外部开源代码及核心实现思路** + +`REFERENCE.md`中需明确声明：“本次赛题提交的算子代码、核心算法逻辑及实现方案均为本人（及参赛团队）独立设计与开发，未参考任何外部开源项目、技术文档中的核心代码片段或实现思路，未接受任何第三方的技术指导或代码支持。” + +**情况2：有参考外部开源代码及相关资源** + +对每个参考资源提供以下信息陈述： +1. 参考开源项目/资源名称 + +2. 参考资源链接（GitHub/Gitee/论文/技术文档等） + +3. 参考的具体内容（请明确说明参考的代码片段、算法逻辑、实现思路等，需标注对应资源的具体位置，如文件路径、代码行数等） + +4. 本人对参考内容的修改与优化说明：（请详细说明在参考基础上，本人所做的独立开发、修改、优化工作，体现自身技术贡献） + +5. 若是开源项目，提供参考资源的开源协议类型：（如MIT、Apache 2.0、GPL等） + +6. 其他需要补充说明的信息 + + +## 三、禁止行为确认 + +本人明确知晓并承诺避免以下违反比赛公平性的行为，若存在以下任一情况，自愿接受比赛组委会的相应处罚： + +1. 未经授权复制、抄袭他人（包括其他参赛选手、开源项目、商业代码）的代码、算法或技术方案，且未进行明确标注； + +2. 隐瞒或虚假披露参考资源信息，包括遗漏重要参考来源、伪造参考内容说明等； + +3. 与其他参赛选手或第三方串通，进行代码共享、成果交换等违规协作； + +4. 利用比赛平台漏洞、技术缺陷或非比赛允许的工具获取不正当利益； + +5. 伪造比赛相关证明材料、提交虚假信息； + +6. 其他违反比赛规则及公序良俗的不诚信行为。 + + +## 四、责任与确认 + +1. 本人充分理解，比赛组委会将对所有提交的PR进行代码溯源、参考信息核查等公平性审查，若发现本人存在违反本准则的行为，有权随时取消本人的参赛资格、作废比赛成绩，情节严重的将在比赛相关平台进行公示。 + +2. 若因本人违反本准则导致比赛争议或第三方权益受损（如开源协议侵权等），本人将独立承担全部法律责任及相关损失，与比赛组委会无关。 + +3. 本人确认已仔细阅读并完全理解本诚信守则的全部内容，自愿签署本准则，接受比赛组委会的监督与审查。 + +## 五、签署信息 + +参赛选手姓名（团队参赛需填写所有成员姓名） + + 李浩坤 + +签署日期 + +___2026___年__6__月__17__日 \ No newline at end of file diff --git a/src/ntops/kernels/__init__.py b/src/ntops/kernels/__init__.py index f6934ef..a91b581 100644 --- a/src/ntops/kernels/__init__.py +++ b/src/ntops/kernels/__init__.py @@ -21,10 +21,14 @@ isnan, layer_norm, le, + linspace, + logit, + logspace, lt, max_pool2d, mm, mul, + nan_to_num, ne, neg, pow, @@ -39,6 +43,7 @@ softmax, sub, tanh, + trapezoid, ) __all__ = [ @@ -64,10 +69,14 @@ "isnan", "layer_norm", "le", + "linspace", + "logit", + "logspace", "lt", "max_pool2d", "mm", "mul", + "nan_to_num", "ne", "neg", "pow", @@ -82,4 +91,5 @@ "softmax", "sub", "tanh", + "trapezoid", ] diff --git a/src/ntops/kernels/linspace.py b/src/ntops/kernels/linspace.py new file mode 100644 index 0000000..509a608 --- /dev/null +++ b/src/ntops/kernels/linspace.py @@ -0,0 +1,23 @@ +import functools + +import ninetoothed.language as ntl +from ninetoothed import Tensor + +from ntops.kernels.element_wise import arrangement + + +def application(start, idx, step, output): + output = start + idx * step # noqa: F841 + + +def premake(ndim, dtype=None, block_size=None): + arrangement_ = functools.partial(arrangement, block_size=block_size) + + tensors = ( + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + ) + + return arrangement_, application, tensors diff --git a/src/ntops/kernels/logit.py b/src/ntops/kernels/logit.py new file mode 100644 index 0000000..918c4c7 --- /dev/null +++ b/src/ntops/kernels/logit.py @@ -0,0 +1,28 @@ +import functools + +import ninetoothed.language as ntl +from ninetoothed import Tensor + +from ntops.kernels.element_wise import arrangement + + +def application(input, output, eps): + one = ntl.cast(1, ntl.float32) + eps_f32 = ntl.cast(eps, ntl.float32) + input_f32 = ntl.cast(input, ntl.float32) + + clamped = ntl.clamp(input_f32, eps_f32, one - eps_f32) + + output = ntl.log(clamped / (one - clamped)) # noqa: F841 + + +def premake(ndim, dtype=None, block_size=None): + arrangement_ = functools.partial(arrangement, block_size=block_size) + + tensors = ( + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + ) + + return arrangement_, application, tensors diff --git a/src/ntops/kernels/logspace.py b/src/ntops/kernels/logspace.py new file mode 100644 index 0000000..db8be14 --- /dev/null +++ b/src/ntops/kernels/logspace.py @@ -0,0 +1,25 @@ +import functools + +import ninetoothed.language as ntl +from ninetoothed import Tensor + +from ntops.kernels.element_wise import arrangement + + +def application(start, idx, step, logbase, output): + linear = start + idx * step + output = ntl.exp(linear * logbase) # noqa: F841 + + +def premake(ndim, dtype=None, block_size=None): + arrangement_ = functools.partial(arrangement, block_size=block_size) + + tensors = ( + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + Tensor(ndim, dtype=dtype), + ) + + return arrangement_, application, tensors diff --git a/src/ntops/kernels/nan_to_num.py b/src/ntops/kernels/nan_to_num.py new file mode 100644 index 0000000..ea2f68b --- /dev/null +++ b/src/ntops/kernels/nan_to_num.py @@ -0,0 +1,30 @@ +import functools + +import ninetoothed.language as ntl +from ninetoothed import Tensor + +from ntops.kernels.element_wise import arrangement + + +def application(input, output): + input_f32 = ntl.cast(input, ntl.float32) + max_finite = ntl.cast(3.4028234663852886e+38, ntl.float32) + min_finite = ntl.cast(-3.4028234663852886e+38, ntl.float32) + + is_nan = input_f32 != input_f32 + is_posinf = input_f32 > max_finite + is_neginf = input_f32 < min_finite + + result = ntl.where(is_nan, ntl.cast(0.0, ntl.float32), input_f32) + result = ntl.where(is_posinf, max_finite, result) + result = ntl.where(is_neginf, min_finite, result) + + output = result # noqa: F841 + + +def premake(ndim, dtype=None, block_size=None): + arrangement_ = functools.partial(arrangement, block_size=block_size) + + tensors = (Tensor(ndim, dtype=dtype), Tensor(ndim, dtype=dtype)) + + return arrangement_, application, tensors diff --git a/src/ntops/kernels/trapezoid.py b/src/ntops/kernels/trapezoid.py new file mode 100644 index 0000000..3afc08b --- /dev/null +++ b/src/ntops/kernels/trapezoid.py @@ -0,0 +1,29 @@ +import functools + +import ninetoothed +import ninetoothed.language as ntl +from ninetoothed import Tensor + +from ntops.kernels.reduction import arrangement + + +def application(areas, output): + dtype = output.dtype.dtype + total = ntl.cast(0, ntl.float32) + + for i in range(areas.shape[0]): + total = total + ntl.cast(ntl.sum(areas[i]), ntl.float32) + + for j in range(output.shape[0]): + output[j] = ntl.cast(total, dtype) + + +def premake(ndim, dtype=None, block_size=None): + arrangement_ = functools.partial(arrangement, dim=-1, block_size=block_size) + + tensors = ( + Tensor(ndim, dtype=dtype, other=0), + Tensor(ndim, dtype=dtype), + ) + + return arrangement_, application, tensors diff --git a/src/ntops/torch/__init__.py b/src/ntops/torch/__init__.py index 82fc596..8a0ab0d 100644 --- a/src/ntops/torch/__init__.py +++ b/src/ntops/torch/__init__.py @@ -20,11 +20,15 @@ from ntops.torch.isnan import isnan from ntops.torch.layer_norm import layer_norm from ntops.torch.le import le +from ntops.torch.linspace import linspace +from ntops.torch.logit import logit +from ntops.torch.logspace import logspace from ntops.torch.lt import lt from ntops.torch.matmul import matmul from ntops.torch.max_pool2d import max_pool2d from ntops.torch.mm import mm from ntops.torch.mul import mul +from ntops.torch.nan_to_num import nan_to_num from ntops.torch.ne import ne from ntops.torch.neg import neg from ntops.torch.pow import pow @@ -39,6 +43,7 @@ from ntops.torch.softmax import softmax from ntops.torch.sub import sub from ntops.torch.tanh import tanh +from ntops.torch.trapezoid import trapezoid __all__ = [ "abs", @@ -63,11 +68,15 @@ "isnan", "layer_norm", "le", + "linspace", + "logit", + "logspace", "lt", "matmul", "max_pool2d", "mm", "mul", + "nan_to_num", "ne", "neg", "pow", @@ -82,4 +91,5 @@ "softmax", "sub", "tanh", + "trapezoid", ] diff --git a/src/ntops/torch/linspace.py b/src/ntops/torch/linspace.py new file mode 100644 index 0000000..cad6f97 --- /dev/null +++ b/src/ntops/torch/linspace.py @@ -0,0 +1,29 @@ +import torch + +import ntops +from ntops.torch.utils import _cached_make + + +def linspace(start, end, steps, *, dtype=None, device=None, layout=None): + if device is None: + device = "cuda" + if dtype is None: + dtype = torch.float32 + + output = torch.empty(steps, dtype=dtype, device=device) + + compute_dtype = torch.float32 if dtype == torch.float16 else dtype + start_t = torch.full((steps,), start, dtype=compute_dtype, device=device) + idx_t = torch.arange(steps, dtype=compute_dtype, device=device) + + if steps > 1: + step_size = (end - start) / (steps - 1) + else: + step_size = 0.0 + + step_t = torch.full((steps,), step_size, dtype=compute_dtype, device=device) + + kernel = _cached_make(ntops.kernels.linspace.premake, output.ndim) + kernel(start_t, idx_t, step_t, output) + + return output diff --git a/src/ntops/torch/logit.py b/src/ntops/torch/logit.py new file mode 100644 index 0000000..157f1fe --- /dev/null +++ b/src/ntops/torch/logit.py @@ -0,0 +1,17 @@ +import torch + +import ntops +from ntops.torch.utils import _cached_make + + +def logit(input, eps=None): + if eps is None: + eps = 0.0 + + output = torch.empty_like(input) + eps_t = torch.full_like(input, eps) + + kernel = _cached_make(ntops.kernels.logit.premake, input.ndim) + kernel(input, output, eps_t) + + return output diff --git a/src/ntops/torch/logspace.py b/src/ntops/torch/logspace.py new file mode 100644 index 0000000..ef1acc4 --- /dev/null +++ b/src/ntops/torch/logspace.py @@ -0,0 +1,33 @@ +import torch + +import ntops +from ntops.torch.utils import _cached_make + + +def logspace(start, end, steps, *, dtype=None, device=None, layout=None): + if device is None: + device = "cuda" + if dtype is None: + dtype = torch.float32 + + output = torch.empty(steps, dtype=dtype, device=device) + + compute_dtype = torch.float32 if dtype == torch.float16 else dtype + start_t = torch.full((steps,), start, dtype=compute_dtype, device=device) + idx_t = torch.arange(steps, dtype=compute_dtype, device=device) + + if steps > 1: + step_size = (end - start) / (steps - 1) + else: + step_size = 0.0 + + step_t = torch.full((steps,), step_size, dtype=compute_dtype, device=device) + + import math + log10 = math.log(10) + logbase_t = torch.full((steps,), log10, dtype=compute_dtype, device=device) + + kernel = _cached_make(ntops.kernels.logspace.premake, output.ndim) + kernel(start_t, idx_t, step_t, logbase_t, output) + + return output diff --git a/src/ntops/torch/nan_to_num.py b/src/ntops/torch/nan_to_num.py new file mode 100644 index 0000000..df34e5d --- /dev/null +++ b/src/ntops/torch/nan_to_num.py @@ -0,0 +1,20 @@ +import torch + +import ntops +from ntops.torch.utils import _cached_make + + +def nan_to_num(input, nan=0.0, posinf=None, neginf=None): + output = torch.empty_like(input) + + kernel = _cached_make(ntops.kernels.nan_to_num.premake, input.ndim) + kernel(input, output) + + # Detect inf/nan from original input, since kernel already replaced them + output = torch.where(torch.isnan(input), nan, output) + if posinf is not None: + output = torch.where(torch.isinf(input) & (input > 0), posinf, output) + if neginf is not None: + output = torch.where(torch.isinf(input) & (input < 0), neginf, output) + + return output diff --git a/src/ntops/torch/trapezoid.py b/src/ntops/torch/trapezoid.py new file mode 100644 index 0000000..fd5a7b1 --- /dev/null +++ b/src/ntops/torch/trapezoid.py @@ -0,0 +1,27 @@ +import torch + +import ntops +from ntops.torch.utils import _cached_make + + +def trapezoid(y, x=None, *, dx=1.0, dim=-1): + if dim != -1 and dim != y.ndim - 1: + y = y.transpose(dim, -1) + if x is not None and x.ndim == y.ndim: + x = x.transpose(dim, -1) + + if x is None: + areas = (y[..., 1:] + y[..., :-1]) * 0.5 * dx + else: + if x.ndim == 1: + dx_vals = x[1:] - x[:-1] + else: + dx_vals = x[..., 1:] - x[..., :-1] + areas = (y[..., 1:] + y[..., :-1]) * 0.5 * dx_vals + + output = torch.zeros(*areas.shape[:-1], 1, dtype=areas.dtype, device=areas.device) + + kernel = _cached_make(ntops.kernels.trapezoid.premake, areas.ndim) + kernel(areas, output) + + return output.squeeze(-1) diff --git a/tests/test_linspace.py b/tests/test_linspace.py new file mode 100644 index 0000000..37b5618 --- /dev/null +++ b/tests/test_linspace.py @@ -0,0 +1,44 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +class TestLinspace: + def test_basic(self): + ninetoothed_output = ntops.torch.linspace(0, 1, 5, device="cuda") + reference_output = torch.linspace(0, 1, 5, device="cuda") + + assert torch.equal(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (5,) + assert ninetoothed_output.device.type == "cuda" + + def test_negative_to_positive(self): + ninetoothed_output = ntops.torch.linspace(-3, 3, 7, device="cuda") + reference_output = torch.linspace(-3, 3, 7, device="cuda") + + assert torch.equal(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (7,) + + def test_single_step(self): + ninetoothed_output = ntops.torch.linspace(5, 5, 1, device="cuda") + reference_output = torch.linspace(5, 5, 1, device="cuda") + + assert torch.equal(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (1,) + + def test_many_steps(self): + ninetoothed_output = ntops.torch.linspace(0, 1, 100, device="cuda") + reference_output = torch.linspace(0, 1, 100, device="cuda") + + assert torch.allclose(ninetoothed_output, reference_output) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_dtype(self, dtype): + ninetoothed_output = ntops.torch.linspace(0, 1, 10, dtype=dtype, device="cuda") + reference_output = torch.linspace(0, 1, 10, dtype=dtype, device="cuda") + + assert torch.allclose(ninetoothed_output, reference_output, atol=1e-3, rtol=1e-3) + assert ninetoothed_output.dtype == dtype diff --git a/tests/test_logit.py b/tests/test_logit.py new file mode 100644 index 0000000..ac9a8e6 --- /dev/null +++ b/tests/test_logit.py @@ -0,0 +1,49 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +class TestLogit: + def test_basic(self): + input = torch.tensor([0.1, 0.5, 0.9], device="cuda") + + ninetoothed_output = ntops.torch.logit(input) + reference_output = torch.logit(input) + + assert torch.allclose(ninetoothed_output, reference_output) + + def test_with_eps(self): + input = torch.tensor([0.0, 0.5, 1.0], device="cuda") + + ninetoothed_output = ntops.torch.logit(input, eps=1e-6) + reference_output = torch.logit(input, eps=1e-6) + + assert torch.allclose(ninetoothed_output, reference_output) + + def test_2d_tensor(self): + input = torch.tensor([[0.2, 0.8], [0.3, 0.7]], device="cuda") + + ninetoothed_output = ntops.torch.logit(input) + reference_output = torch.logit(input) + + assert torch.allclose(ninetoothed_output, reference_output) + + def test_close_to_bounds(self): + input = torch.tensor([1e-6, 1 - 1e-6], device="cuda") + + ninetoothed_output = ntops.torch.logit(input) + reference_output = torch.logit(input) + + assert torch.allclose(ninetoothed_output, reference_output, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_dtype(self, dtype): + input = torch.tensor([0.1, 0.5, 0.9], device="cuda", dtype=dtype) + + ninetoothed_output = ntops.torch.logit(input) + reference_output = torch.logit(input) + + assert torch.allclose(ninetoothed_output, reference_output, atol=0.1, rtol=0.1) diff --git a/tests/test_logspace.py b/tests/test_logspace.py new file mode 100644 index 0000000..e04e70a --- /dev/null +++ b/tests/test_logspace.py @@ -0,0 +1,38 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +class TestLogspace: + def test_basic(self): + ninetoothed_output = ntops.torch.logspace(0, 2, 5, device="cuda") + reference_output = torch.logspace(0, 2, 5, device="cuda") + + assert torch.allclose(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (5,) + assert ninetoothed_output.device.type == "cuda" + + def test_negative_exponents(self): + ninetoothed_output = ntops.torch.logspace(-2, 0, 5, device="cuda") + reference_output = torch.logspace(-2, 0, 5, device="cuda") + + assert torch.allclose(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (5,) + + def test_single_step(self): + ninetoothed_output = ntops.torch.logspace(3, 3, 1, device="cuda") + reference_output = torch.logspace(3, 3, 1, device="cuda") + + assert torch.allclose(ninetoothed_output, reference_output, atol=1e-3, rtol=1e-3) + assert ninetoothed_output.shape == (1,) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_dtype(self, dtype): + ninetoothed_output = ntops.torch.logspace(0, 1, 10, dtype=dtype, device="cuda") + reference_output = torch.logspace(0, 1, 10, dtype=dtype, device="cuda") + + assert torch.allclose(ninetoothed_output, reference_output, atol=1e-2, rtol=1e-2) + assert ninetoothed_output.dtype == dtype diff --git a/tests/test_nan_to_num.py b/tests/test_nan_to_num.py new file mode 100644 index 0000000..22d9da3 --- /dev/null +++ b/tests/test_nan_to_num.py @@ -0,0 +1,75 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +class TestNanToNum: + def test_no_nan_or_inf(self): + input = torch.tensor([1.0, 2.0, 3.0], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input) + reference_output = torch.nan_to_num(input) + + assert torch.equal(ninetoothed_output, reference_output) + + def test_nan_replacement(self): + input = torch.tensor([1.0, float("nan"), 3.0], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input) + reference_output = torch.nan_to_num(input) + + assert torch.allclose(ninetoothed_output, reference_output, equal_nan=True) + assert not torch.isnan(ninetoothed_output).any() + + def test_inf_replacement(self): + input = torch.tensor([1.0, float("inf"), float("-inf")], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input) + reference_output = torch.nan_to_num(input) + + assert torch.allclose(ninetoothed_output, reference_output, equal_nan=True) + assert not torch.isinf(ninetoothed_output).any() + + def test_custom_nan_value(self): + input = torch.tensor([1.0, float("nan"), 3.0], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input, nan=-1.0) + reference_output = torch.nan_to_num(input, nan=-1.0) + + assert torch.allclose(ninetoothed_output, reference_output, equal_nan=True) + + def test_custom_posinf_neginf(self): + input = torch.tensor([1.0, float("inf"), float("-inf")], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input, posinf=1e6, neginf=-1e6) + reference_output = torch.nan_to_num(input, posinf=1e6, neginf=-1e6) + + assert torch.allclose(ninetoothed_output, reference_output) + + def test_mixed_nan_inf(self): + input = torch.tensor([float("nan"), float("inf"), float("-inf"), 42.0], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input) + reference_output = torch.nan_to_num(input) + + assert torch.allclose(ninetoothed_output, reference_output, equal_nan=True) + + def test_2d_tensor(self): + input = torch.tensor([[1.0, float("nan")], [float("inf"), 4.0]], device="cuda") + + ninetoothed_output = ntops.torch.nan_to_num(input) + reference_output = torch.nan_to_num(input) + + assert torch.allclose(ninetoothed_output, reference_output, equal_nan=True) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_dtype(self, dtype): + input = torch.tensor([1.0, float("nan")], device="cuda", dtype=dtype) + + ninetoothed_output = ntops.torch.nan_to_num(input) + reference_output = torch.nan_to_num(input) + + assert torch.allclose(ninetoothed_output, reference_output, equal_nan=True, atol=0.1, rtol=0.1) diff --git a/tests/test_trapezoid.py b/tests/test_trapezoid.py new file mode 100644 index 0000000..c2c7d65 --- /dev/null +++ b/tests/test_trapezoid.py @@ -0,0 +1,70 @@ +import pytest +import torch + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +@skip_if_cuda_not_available +class TestTrapezoid: + def test_1d_basic(self): + y = torch.tensor([1.0, 2.0, 3.0], device="cuda") + + ninetoothed_output = ntops.torch.trapezoid(y) + reference_output = torch.trapezoid(y) + + assert torch.allclose(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == () + + def test_1d_with_dx(self): + y = torch.tensor([1.0, 2.0, 3.0], device="cuda") + + ninetoothed_output = ntops.torch.trapezoid(y, dx=2.0) + reference_output = torch.trapezoid(y, dx=2.0) + + assert torch.allclose(ninetoothed_output, reference_output) + + def test_1d_with_x(self): + y = torch.tensor([1.0, 2.0, 3.0], device="cuda") + x = torch.tensor([0.0, 2.0, 5.0], device="cuda") + + ninetoothed_output = ntops.torch.trapezoid(y, x=x) + reference_output = torch.trapezoid(y, x=x) + + assert torch.allclose(ninetoothed_output, reference_output) + + def test_2d_default_dim(self): + y = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device="cuda") + + ninetoothed_output = ntops.torch.trapezoid(y) + reference_output = torch.trapezoid(y) + + assert torch.allclose(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (2,) + + def test_2d_with_dim(self): + y = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device="cuda") + + ninetoothed_output = ntops.torch.trapezoid(y, dim=0) + reference_output = torch.trapezoid(y, dim=0) + + assert torch.allclose(ninetoothed_output, reference_output) + assert ninetoothed_output.shape == (3,) + + def test_2d_with_x_and_dim(self): + y = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device="cuda") + x = torch.tensor([0.0, 1.0, 3.0], device="cuda") + + ninetoothed_output = ntops.torch.trapezoid(y, x=x, dim=-1) + reference_output = torch.trapezoid(y, x=x, dim=-1) + + assert torch.allclose(ninetoothed_output, reference_output) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) + def test_dtype(self, dtype): + y = torch.tensor([1.0, 2.0, 3.0], device="cuda", dtype=dtype) + + ninetoothed_output = ntops.torch.trapezoid(y) + reference_output = torch.trapezoid(y) + + assert torch.allclose(ninetoothed_output, reference_output, atol=0.1, rtol=0.1)