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

python/infinicore/__init__.py

@@ -82,6 +82,8 @@
 from infinicore.ops.digamma import digamma
 from infinicore.ops.dist import dist
 from infinicore.ops.equal import equal
+from infinicore.ops.flip import flip
+from infinicore.ops.fliplr import fliplr
 from infinicore.ops.flipud import flipud
 from infinicore.ops.float_power import float_power
 from infinicore.ops.floor import floor
@@ -214,6 +216,8 @@
     "acos",
     "addbmm",
     "floor",
+    "flip",
+    "fliplr",
     "attention",
     "block_diag",
     "kron",

python/infinicore/nn/functional/__init__.py

@@ -8,7 +8,8 @@
 from .binary_cross_entropy_with_logits import binary_cross_entropy_with_logits
 from .causal_softmax import causal_softmax
 from .embedding import embedding
-from .flash_attention import flash_attention
+from .flash_attention import flash_attention 
+from .feature_alpha_dropout import  feature_alpha_dropout
 from .gaussian_nll_loss import gaussian_nll_loss
 from .hardswish import hardswish
 from .hardtanh import hardtanh
@@ -20,7 +21,9 @@
 from .linear_w8a8i8 import linear_w8a8i8
 from .log_softmax import log_softmax
 from .multi_margin_loss import multi_margin_loss
+from .mse_loss import mse_loss
 from .pad import pad
+from .pixel_unshuffle import pixel_unshuffle
 from .prelu import prelu
 from .random_sample import random_sample
 from .relu6 import relu6
@@ -45,6 +48,7 @@
     "adaptive_max_pool1d",
     "causal_softmax",
     "embedding",
+    "feature_alpha_dropout",
     "flash_attention",
     "gaussian_nll_loss",
     "interpolate",
@@ -62,6 +66,7 @@
     "swiglu",
     "interpolate",
     "linear",
+    "mse_loss",
     "triplet_margin_loss",
     "upsample_bilinear",
     "interpolate",
@@ -76,6 +81,7 @@
     "selu",
     "hinge_embedding_loss",
     "pad",
+    "pixel_unshuffle",
     "silu",
     "hardswish",
     "hardtanh",

python/infinicore/nn/functional/feature_alpha_dropout.py

@@ -0,0 +1,63 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def feature_alpha_dropout(
+    input: Tensor,
+    p: float = 0.5,
+    training: bool = False,
+    inplace: bool = False,
+    *,
+    out=None,
+) -> Tensor:
+    r"""Apply feature alpha dropout.
+
+    Equivalent to torch.nn.functional.feature_alpha_dropout.
+    """
+
+    if p < 0.0 or p >= 1.0:
+        raise ValueError(
+            f"dropout probability has to satisfy 0 <= p < 1, but got {p}"
+        )
+
+    if input.ndim < 2:
+        raise RuntimeError("Feature dropout requires at least 2 dimensions in the input")
+
+    if inplace and out is not None:
+        raise RuntimeError("`inplace=True` and `out` cannot be used together.")
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa") and out is None:
+        return infinicore.ntops.torch.feature_alpha_dropout(
+            input,
+            p=p,
+            training=training,
+            inplace=inplace,
+        )
+
+    if inplace:
+        _infinicore.feature_alpha_dropout_(
+            input._underlying,
+            input._underlying,
+            float(p),
+            bool(training),
+        )
+        return input
+
+    if out is None:
+        return Tensor(
+            _infinicore.feature_alpha_dropout(
+                input._underlying,
+                float(p),
+                bool(training),
+            )
+        )
+
+    _infinicore.feature_alpha_dropout_(
+        out._underlying,
+        input._underlying,
+        float(p),
+        bool(training),
+    )
+
+    return out

python/infinicore/nn/functional/mse_loss.py

@@ -0,0 +1,30 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def mse_loss(input: Tensor, target: Tensor, reduction: str = "mean") -> Tensor:
+    r"""Compute mean squared error loss between input and target."""
+
+    assert reduction in (
+        "none",
+        "mean",
+        "sum",
+    ), "`reduction` must be one of 'none', 'mean', or 'sum'."
+
+    assert input.shape == target.shape, "`input` and `target` must have the same shape."
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.mse_loss(
+            input,
+            target,
+            reduction=reduction,
+        )
+
+    return Tensor(
+        _infinicore.mse_loss(
+            input._underlying,
+            target._underlying,
+            reduction,
+        )
+    )

python/infinicore/nn/functional/pixel_unshuffle.py

@@ -0,0 +1,33 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def pixel_unshuffle(input: Tensor, downscale_factor: int, *, out=None) -> Tensor:
+    r"""Rearrange elements in a tensor of shape (*, C, H * r, W * r)
+    to a tensor of shape (*, C * r * r, H, W), where r is downscale_factor.
+    """
+
+    assert isinstance(downscale_factor, int), "`downscale_factor` must be int."
+    assert downscale_factor > 0, "`downscale_factor` must be positive."
+    assert input.ndim == 4, "`pixel_unshuffle` only supports 4D NCHW input."
+
+    n, c, h, w = input.shape
+    r = downscale_factor
+
+    assert h % r == 0, "input height must be divisible by downscale_factor."
+    assert w % r == 0, "input width must be divisible by downscale_factor."
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa") and out is None:
+        return infinicore.ntops.torch.pixel_unshuffle(input, downscale_factor)
+
+    if out is None:
+        return Tensor(_infinicore.pixel_unshuffle(input._underlying, downscale_factor))
+
+    _infinicore.pixel_unshuffle_(
+        out._underlying,
+        input._underlying,
+        downscale_factor,
+    )
+
+    return out

python/infinicore/ops/flip.py

@@ -0,0 +1,40 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+import torch
+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 _is_dims(x):
+    return isinstance(x, int) or (
+        isinstance(x, (tuple, list)) and all(isinstance(i, int) for i in x)
+    )
+def flip(*args) -> Tensor:
+    r"""Reverse the order of an n-D tensor along given dims."""
+    assert len(args) >= 2, "`flip` requires input and dims."
+    dims = args[-1]
+    assert _is_dims(dims), "`dims` must be int, tuple[int], or list[int]."
+    input = None
+    for arg in args[:-1]:
+        if isinstance(arg, Tensor) or hasattr(arg, "_underlying"):
+            input = arg
+            break
+        if isinstance(arg, torch.Tensor):
+            return torch.flip(arg, dims)
+    assert input is not None, "`flip` requires a Tensor input."
+    dims = _normalize_dims(dims, input.ndim)
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.flip(input, dims)
+    return Tensor(
+        _infinicore.flip(
+            input._underlying,
+            dims,
+        )
+    )

python/infinicore/ops/fliplr.py

@@ -0,0 +1,22 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def fliplr(input: Tensor, *, out=None) -> Tensor:
+    r"""Flip tensor in the left/right direction.
+
+    Equivalent to torch.fliplr(input), i.e. torch.flip(input, dims=(1,)).
+    """
+
+    assert input.ndim >= 2, "`fliplr` requires input with ndim >= 2."
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa") and out is None:
+        return infinicore.ntops.torch.fliplr(input)
+
+    if out is None:
+        return Tensor(_infinicore.fliplr(input._underlying))
+
+    _infinicore.fliplr_(out._underlying, input._underlying)
+
+    return out

test/infinicore/ops/feature_alpha_dropout.py

@@ -75,9 +75,9 @@ def get_test_cases(self):
     def torch_operator(self, *args, **kwargs):
         return torch.nn.functional.feature_alpha_dropout(*args, **kwargs)
 
-    # def infinicore_operator(self, *args, **kwargs):
+    def infinicore_operator(self, *args, **kwargs):
     #     """InfiniCore implementation (operator not yet available)."""
-    #     return infinicore.nn.functional.feature_alpha_dropout(*args, **kwargs)
+        return infinicore.nn.functional.feature_alpha_dropout(*args, **kwargs)
 
 
 def main():

test/infinicore/ops/flip.py

@@ -73,17 +73,17 @@ def get_test_cases(self):
         return parse_test_cases()
 
     def torch_operator(self, *args, **kwargs):
-        # dims = kwargs.pop("dims", None)
-        # if dims is not None:
-        #     return infinicore.flip(*args, dims)
+        dims = kwargs.pop("dims", None)
+        if dims is not None:
+            return infinicore.flip(*args, dims)
         return torch.flip(*args, **kwargs)
 
-    # def infinicore_operator(self, *args, **kwargs):
+    def infinicore_operator(self, *args, **kwargs):
     #     """InfiniCore implementation (operator not yet available)."""
-    #     dims = kwargs.pop("dims", None)
-    #     if dims is not None:
-    #         return infinicore.flip(*args, dims)
-    #     return infinicore.flip(*args, **kwargs)
+        dims = kwargs.pop("dims", None)
+        if dims is not None:
+            return infinicore.flip(*args, dims)
+        return infinicore.flip(*args, **kwargs)
 
 
 def main():

test/infinicore/ops/fliplr.py

@@ -72,9 +72,9 @@ def get_test_cases(self):
     def torch_operator(self, *args, **kwargs):
         return torch.fliplr(*args, **kwargs)
 
-    # def infinicore_operator(self, *args, **kwargs):
+    def infinicore_operator(self, *args, **kwargs):
     #     """InfiniCore implementation (operator not yet available)."""
-    #     return infinicore.fliplr(*args, **kwargs)
+        return infinicore.fliplr(*args, **kwargs)
 
 
 def main():

test/infinicore/ops/mse_loss.py

@@ -71,9 +71,9 @@ def get_test_cases(self):
     def torch_operator(self, *args, **kwargs):
         return torch.nn.functional.mse_loss(*args, **kwargs)
 
-    # def infinicore_operator(self, *args, **kwargs):
+    def infinicore_operator(self, *args, **kwargs):
     #     """InfiniCore implementation (operator not yet available)."""
-    #     return infinicore.nn.functional.mse_loss(*args, **kwargs)
+        return infinicore.nn.functional.mse_loss(*args, **kwargs)
 
 
 def main():

test/infinicore/ops/pixel_unshuffle.py

@@ -73,9 +73,9 @@ def get_test_cases(self):
     def torch_operator(self, *args, **kwargs):
         return torch.nn.functional.pixel_unshuffle(*args, **kwargs)
 
-    # def infinicore_operator(self, *args, **kwargs):
+    def infinicore_operator(self, *args, **kwargs):
     #     """InfiniCore implementation (operator not yet available)."""
-    #     return infinicore.nn.functional.pixel_unshuffle(*args, **kwargs)
+        return infinicore.nn.functional.pixel_unshuffle(*args, **kwargs)
 
 
 def main():