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

python/infinicore/__init__.py

@@ -88,6 +88,7 @@
 from infinicore.ops.floor_divide import floor_divide
 from infinicore.ops.fmin import fmin
 from infinicore.ops.fmod import fmod
+from infinicore.ops.frac import frac
 from infinicore.ops.hypot import hypot
 from infinicore.ops.index_add import index_add
 from infinicore.ops.index_copy import index_copy
@@ -126,6 +127,7 @@
 from infinicore.ops.rotmg import rotmg
 from infinicore.ops.scal import scal
 from infinicore.ops.scatter import scatter
+from infinicore.ops.scatter_add import scatter_add
 from infinicore.ops.sinh import sinh
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.sum import sum
@@ -211,6 +213,7 @@
     "blas_amin",
     "blas_copy",
     "blas_dot",
+    "scatter_add",
     "acos",
     "addbmm",
     "floor",
@@ -223,6 +226,7 @@
     "baddbmm",
     "bilinear",
     "fmod",
+    "frac",
     "cat",
     "conv2d",
     "inner",

python/infinicore/nn/functional/__init__.py

@@ -9,6 +9,8 @@
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .flash_attention import flash_attention
+from .fractional_max_pool2d import fractional_max_pool2d
+from .fractional_max_pool3d import fractional_max_pool3d
 from .gaussian_nll_loss import gaussian_nll_loss
 from .hardswish import hardswish
 from .hardtanh import hardtanh
@@ -36,6 +38,7 @@
 from .swiglu import swiglu
 from .tanhshrink import tanhshrink
 from .triplet_margin_loss import triplet_margin_loss
+from .multilabel_margin_loss import  multilabel_margin_loss
 from .triplet_margin_with_distance_loss import triplet_margin_with_distance_loss
 from .unfold import unfold
 from .upsample_bilinear import upsample_bilinear
@@ -46,6 +49,8 @@
     "causal_softmax",
     "embedding",
     "flash_attention",
+    "fractional_max_pool2d",
+    "fractional_max_pool3d",
     "gaussian_nll_loss",
     "interpolate",
     "linear",
@@ -80,6 +85,7 @@
     "hardswish",
     "hardtanh",
     "avg_pool1d",
+    "multilabel_margin_loss",
     "swiglu",
     "linear_w8a8i8",
     "silu_and_mul",

python/infinicore/nn/functional/fractional_max_pool2d.py

@@ -0,0 +1,55 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def fractional_max_pool2d(
+    input: Tensor,
+    kernel_size,
+    output_size=None,
+    output_ratio=None,
+    return_indices: bool = False,
+    _random_samples=None,
+) -> Tensor:
+    r"""Apply 2D fractional max pooling over an input signal."""
+
+    assert input.ndim == 4, (
+        "`fractional_max_pool2d` only supports 4D input for now."
+    )
+
+    assert not return_indices, (
+        "`return_indices` is not supported by ntops fractional_max_pool2d yet."
+    )
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.fractional_max_pool2d(
+            input,
+            kernel_size=kernel_size,
+            output_size=output_size,
+            output_ratio=output_ratio,
+            return_indices=return_indices,
+            _random_samples=_random_samples,
+        )
+
+    if hasattr(_infinicore, "fractional_max_pool2d"):
+        if _random_samples is None:
+            return Tensor(
+                _infinicore.fractional_max_pool2d(
+                    input._underlying,
+                    kernel_size,
+                    output_size,
+                    output_ratio,
+                    return_indices,
+                )
+            )
+
+        return Tensor(
+            _infinicore.fractional_max_pool2d(
+                input._underlying,
+                kernel_size,
+                output_size,
+                output_ratio,
+                return_indices,
+                _random_samples._underlying,
+            )
+        )

python/infinicore/nn/functional/fractional_max_pool3d.py

@@ -0,0 +1,55 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def fractional_max_pool3d(
+    input: Tensor,
+    kernel_size,
+    output_size=None,
+    output_ratio=None,
+    return_indices: bool = False,
+    _random_samples=None,
+) -> Tensor:
+    r"""Apply 3D fractional max pooling over an input signal."""
+
+    assert input.ndim == 5, (
+        "`fractional_max_pool3d` only supports 5D input for now."
+    )
+
+    assert not return_indices, (
+        "`return_indices` is not supported by ntops fractional_max_pool3d yet."
+    )
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.fractional_max_pool3d(
+            input,
+            kernel_size=kernel_size,
+            output_size=output_size,
+            output_ratio=output_ratio,
+            return_indices=return_indices,
+            _random_samples=_random_samples,
+        )
+
+    if hasattr(_infinicore, "fractional_max_pool3d"):
+        if _random_samples is None:
+            return Tensor(
+                _infinicore.fractional_max_pool3d(
+                    input._underlying,
+                    kernel_size,
+                    output_size,
+                    output_ratio,
+                    return_indices,
+                )
+            )
+
+        return Tensor(
+            _infinicore.fractional_max_pool3d(
+                input._underlying,
+                kernel_size,
+                output_size,
+                output_ratio,
+                return_indices,
+                _random_samples._underlying,
+            )
+        )

python/infinicore/nn/functional/multilabel_margin_loss.py

@@ -0,0 +1,89 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def _normalize_reduction(size_average=None, reduce=None, reduction="mean"):
+    if size_average is not None or reduce is not None:
+        if reduce is False:
+            return "none"
+
+        if size_average is False:
+            return "sum"
+
+        return "mean"
+
+    return reduction
+
+
+def multilabel_margin_loss(
+    input: Tensor,
+    target: Tensor,
+    size_average=None,
+    reduce=None,
+    reduction: str = "mean",
+    *,
+    out=None,
+) -> Tensor:
+    r"""Compute multilabel margin loss.
+
+    Args:
+        input: Tensor with shape [C], [N, C], or higher dims flattened by ntops wrapper.
+        target: LongTensor with same shape as input, padded by -1.
+        reduction: "none", "mean", or "sum".
+    """
+
+    reduction = _normalize_reduction(
+        size_average=size_average,
+        reduce=reduce,
+        reduction=reduction,
+    )
+
+    assert reduction in (
+        "none",
+        "mean",
+        "sum",
+    ), "`reduction` must be one of 'none', 'mean', or 'sum'."
+    if (
+        infinicore.use_ntops
+        and input.device.type in ("cuda", "musa")
+        and out is None
+    ):
+        return infinicore.ntops.torch.multilabel_margin_loss(
+            input,
+            target,
+            size_average=size_average,
+            reduce=reduce,
+            reduction=reduction,
+        )
+
+    # C++ fallback
+    if not hasattr(_infinicore, "multilabel_margin_loss"):
+        raise NotImplementedError(
+            "multilabel_margin_loss is not implemented in _infinicore, "
+            "and ntops path is unavailable. Enable infinicore.use_ntops "
+            "or add C++ backend implementation."
+        )
+
+    if out is None:
+        return Tensor(
+            _infinicore.multilabel_margin_loss(
+                input._underlying,
+                target._underlying,
+                reduction,
+            )
+        )
+
+    if not hasattr(_infinicore, "multilabel_margin_loss_"):
+        raise NotImplementedError(
+            "multilabel_margin_loss_ out variant is not implemented in _infinicore."
+        )
+
+    _infinicore.multilabel_margin_loss_(
+        out._underlying,
+        input._underlying,
+        target._underlying,
+        reduction,
+    )
+
+    return out

python/infinicore/ops/frac.py

@@ -0,0 +1,65 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def _copy_result(out: Tensor, res: Tensor) -> Tensor:
+    copy_fn = getattr(out._underlying, "copy_", None)
+
+    if callable(copy_fn):
+        copy_fn(res._underlying)
+        return out
+
+    raise NotImplementedError(
+        "frac requires underlying tensor copy_ for inplace/out fallback."
+    )
+
+
+def frac(input: Tensor, inplace: bool = False, *, out=None) -> Tensor:
+    r"""Compute the fractional portion of each element in input."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        res = infinicore.ntops.torch.frac(input)
+
+        if inplace:
+            return _copy_result(input, res)
+
+        if out is not None:
+            return _copy_result(out, res)
+
+        return res
+
+    if inplace:
+        if hasattr(_infinicore, "frac_"):
+            _infinicore.frac_(input._underlying, input._underlying)
+            return input
+
+        if hasattr(_infinicore, "frac"):
+            res = Tensor(_infinicore.frac(input._underlying))
+            return _copy_result(input, res)
+
+        raise NotImplementedError(
+            "frac inplace requires ntops backend, `_infinicore.frac_`, "
+            "or `_infinicore.frac` with copy_ support."
+        )
+
+    if out is None:
+        if hasattr(_infinicore, "frac"):
+            return Tensor(_infinicore.frac(input._underlying))
+
+        raise NotImplementedError(
+            "frac requires ntops backend or `_infinicore.frac`."
+        )
+
+    if hasattr(_infinicore, "frac_"):
+        _infinicore.frac_(out._underlying, input._underlying)
+        return out
+
+    if hasattr(_infinicore, "frac"):
+        res = Tensor(_infinicore.frac(input._underlying))
+        return _copy_result(out, res)
+
+    raise NotImplementedError(
+        "frac out requires ntops backend, `_infinicore.frac_`, "
+        "or `_infinicore.frac` with copy_ support."
+    )