[2026春季][ T1-1-7] zili2004

src/ntops/kernels/__init__.py

@@ -14,6 +14,9 @@
     dropout,
     eq,
     exp,
+    feature_alpha_dropout,
+    flip,
+    fliplr,
     ge,
     gelu,
     gt,
@@ -25,8 +28,10 @@
     max_pool2d,
     mm,
     mul,
+    mse_loss,
     ne,
     neg,
+    pixel_unshuffle,
     pow,
     relu,
     rms_norm,
@@ -57,6 +62,9 @@
     "dropout",
     "eq",
     "exp",
+    "feature_alpha_dropout",
+    "flip",
+    "fliplr",
     "ge",
     "gelu",
     "gt",
@@ -68,8 +76,10 @@
     "max_pool2d",
     "mm",
     "mul",
+    "mse_loss",
     "ne",
     "neg",
+    "pixel_unshuffle",
     "pow",
     "relu",
     "rms_norm",

src/ntops/kernels/feature_alpha_dropout.py

@@ -0,0 +1,36 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, p, seed, a, b, output):
+    alpha_prime = -1.7580993408473766
+
+    keep = ntl.rand(seed, input.offsets()) > p
+
+    dropped = ntl.where(
+        keep,
+        input,
+        alpha_prime,
+    )
+
+    output = dropped * a + b  # noqa: F841
+
+
+def premake(ndim, dtype=None, block_size=None):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
+
+    tensors = (
+        Tensor(ndim, dtype=dtype),                  # input
+        Tensor(0, dtype=ninetoothed.float64),       # p
+        Tensor(0, dtype=ninetoothed.int64),         # seed
+        Tensor(0, dtype=ninetoothed.float64),       # a
+        Tensor(0, dtype=ninetoothed.float64),       # b
+        Tensor(ndim, dtype=dtype),                  # output
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/flip.py

@@ -0,0 +1,78 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+
+def application(input, output):
+    output = input  # noqa: F841
+
+
+def _normalize_dims(dims, ndim):
+    if isinstance(dims, int):
+        dims = (dims,)
+
+    dims = tuple(dim if dim >= 0 else dim + ndim for dim in dims)
+
+    assert all(0 <= dim < ndim for dim in dims), "`dims` out of range."
+
+    result = []
+    for dim in dims:
+        if dim not in result:
+            result.append(dim)
+
+    return tuple(result)
+
+
+def arrangement(
+    input,
+    output,
+    dims,
+    block_size=None,
+):
+    if block_size is None:
+        block_size = ninetoothed.block_size()
+
+    ndim = input.ndim
+    dims = _normalize_dims(dims, ndim)
+    slices = tuple(
+        slice(None, None, -1) if dim in dims else slice(None)
+        for dim in range(ndim)
+    )
+
+    input_arranged = input[slices]
+    input_arranged = input_arranged.flatten().tile((block_size,))
+
+    output_arranged = output.flatten().tile((block_size,))
+
+    return input_arranged, output_arranged
+
+
+def premake(
+    ndim,
+    dims,
+    dtype=None,
+    block_size=None,
+):
+    dims = _normalize_dims(dims, ndim)
+
+    arrangement_ = functools.partial(
+        arrangement,
+        dims=dims,
+        block_size=block_size,
+    )
+
+    tensors = (
+        Tensor(
+            ndim,
+            dtype=dtype,
+            shape_options={"constexpr": True},
+        ),
+        Tensor(
+            ndim,
+            dtype=dtype,
+            shape_options={"constexpr": True},
+        ),
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/fliplr.py

@@ -0,0 +1,59 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+
+def application(input, output):
+    output = input  # noqa: F841
+
+
+def arrangement(
+    input,
+    output,
+    block_size=None,
+):
+    if block_size is None:
+        block_size = ninetoothed.block_size()
+
+    ndim = input.ndim
+    assert ndim >= 2, "`fliplr` requires input with ndim >= 2."
+    slices = tuple(
+        slice(None, None, -1) if dim == 1 else slice(None)
+        for dim in range(ndim)
+    )
+
+    input_arranged = input[slices]
+    input_arranged = input_arranged.flatten().tile((block_size,))
+
+    output_arranged = output.flatten().tile((block_size,))
+
+    return input_arranged, output_arranged
+
+
+def premake(
+    ndim,
+    dtype=None,
+    block_size=None,
+):
+    assert ndim >= 2, "`fliplr` requires input with ndim >= 2."
+
+    arrangement_ = functools.partial(
+        arrangement,
+        block_size=block_size,
+    )
+
+    tensors = (
+        Tensor(
+            ndim,
+            dtype=dtype,
+            shape_options={"constexpr": True},
+        ),
+        Tensor(
+            ndim,
+            dtype=dtype,
+            shape_options={"constexpr": True},
+        ),
+    )
+
+    return arrangement_, application, tensors

src/ntops/kernels/mse_loss.py

@@ -0,0 +1,131 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement as element_wise_arrangement
+from ntops.kernels.reduction import arrangement as reduction_arrangement
+
+
+REDUCTION_NONE = 0
+REDUCTION_MEAN = 1
+REDUCTION_SUM = 2
+
+
+def reduction_all_arrangement(input, target, output, inv_numel=None, block_size=None):
+    input_arranged, target_arranged = reduction_arrangement(
+        input,
+        target,
+        dim=tuple(range(input.ndim)),
+        block_size=block_size,
+    )
+
+    output_arranged = output.tile((1,))
+
+    if inv_numel is None:
+        return input_arranged, target_arranged, output_arranged
+
+    return input_arranged, target_arranged, output_arranged, inv_numel
+
+
+def application_none(input, target, output):
+    diff = input - target
+    output = diff * diff  # noqa: F841
+
+
+def application_sum(input, target, output):
+    dtype = output.dtype
+    acc_dtype = ntl.float32 if dtype == ntl.float16 else dtype
+
+    acc = ntl.cast(0, acc_dtype)
+
+    for i in range(input.shape[0]):
+        diff = ntl.cast(input[i] - target[i], acc_dtype)
+        acc += ntl.sum(diff * diff)
+
+    output = ntl.cast(acc, dtype)  # noqa: F841
+
+
+def application_mean(input, target, output, inv_numel):
+    dtype = output.dtype
+    acc_dtype = ntl.float32 if dtype == ntl.float16 else dtype
+
+    acc = ntl.cast(0, acc_dtype)
+
+    for i in range(input.shape[0]):
+        diff = ntl.cast(input[i] - target[i], acc_dtype)
+        acc += ntl.sum(diff * diff)
+
+    output = ntl.cast(acc * inv_numel, dtype)  # noqa: F841
+
+
+def premake(
+    ndim,
+    reduction=REDUCTION_MEAN,
+    dtype=None,
+    block_size=None,
+):
+    if reduction == REDUCTION_NONE:
+        arrangement_ = functools.partial(
+            element_wise_arrangement,
+            block_size=block_size,
+        )
+
+        tensors = (
+            Tensor(ndim, dtype=dtype),
+            Tensor(ndim, dtype=dtype),
+            Tensor(ndim, dtype=dtype),
+        )
+
+        return arrangement_, application_none, tensors
+
+    assert reduction in (
+        REDUCTION_MEAN,
+        REDUCTION_SUM,
+    ), "`reduction` must be 0, 1, or 2."
+
+    arrangement_ = functools.partial(
+        reduction_all_arrangement,
+        block_size=block_size,
+    )
+
+    input = Tensor(
+        ndim,
+        dtype=dtype,
+        shape_options={"constexpr": True},
+    )
+
+    target = Tensor(
+        ndim,
+        dtype=dtype,
+        shape_options={"constexpr": True},
+    )
+
+    output = Tensor(
+        1,
+        dtype=dtype,
+        shape_options=(
+            {"constexpr": True, "upper_bound": 1},
+        ),
+    )
+    output.shape = (1,)
+
+    if reduction == REDUCTION_SUM:
+        tensors = (
+            input,
+            target,
+            output,
+        )
+        return arrangement_, application_sum, tensors
+
+    inv_numel = Tensor(0, dtype=ninetoothed.float64)
+
+    tensors = (
+        input,
+        target,
+        output,
+        inv_numel,
+    )
+
+    return arrangement_, application_mean, tensors