diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index d84a1ce74..40b54196b 100644 --- a/python/infinicore/__init__.py +++ b/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", diff --git a/python/infinicore/nn/functional/__init__.py b/python/infinicore/nn/functional/__init__.py index f3b9a0a2b..196def0de 100644 --- a/python/infinicore/nn/functional/__init__.py +++ b/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", diff --git a/python/infinicore/nn/functional/fractional_max_pool2d.py b/python/infinicore/nn/functional/fractional_max_pool2d.py new file mode 100644 index 000000000..8d4a09f82 --- /dev/null +++ b/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, + ) + ) diff --git a/python/infinicore/nn/functional/fractional_max_pool3d.py b/python/infinicore/nn/functional/fractional_max_pool3d.py new file mode 100644 index 000000000..147d57596 --- /dev/null +++ b/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, + ) + ) \ No newline at end of file diff --git a/python/infinicore/nn/functional/multilabel_margin_loss.py b/python/infinicore/nn/functional/multilabel_margin_loss.py new file mode 100644 index 000000000..74da7fa0e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/python/infinicore/ops/frac.py b/python/infinicore/ops/frac.py new file mode 100644 index 000000000..12a9d436d --- /dev/null +++ b/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." + ) \ No newline at end of file diff --git a/python/infinicore/ops/scatter_add.py b/python/infinicore/ops/scatter_add.py new file mode 100644 index 000000000..d00a83d79 --- /dev/null +++ b/python/infinicore/ops/scatter_add.py @@ -0,0 +1,88 @@ +import infinicore +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def scatter_add( + input: Tensor, + dim: int, + index: Tensor, + src: Tensor, + *, + out=None, +) -> Tensor: + r"""Add all values from src into input at the indices specified in index along dim.""" + + if infinicore.use_ntops and input.device.type in ("cuda", "musa"): + return infinicore.ntops.torch.scatter_add( + input, + dim, + index, + src, + out=out, + ) + + if out is None: + if not hasattr(_infinicore, "scatter_add"): + raise NotImplementedError( + "scatter_add is not implemented in _infinicore, " + "and ntops path is unavailable." + ) + + return Tensor( + _infinicore.scatter_add( + input._underlying, + dim, + index._underlying, + src._underlying, + ) + ) + + if not hasattr(_infinicore, "scatter_add_"): + raise NotImplementedError( + "scatter_add_ out variant is not implemented in _infinicore." + ) + + _infinicore.scatter_add_( + out._underlying, + input._underlying, + dim, + index._underlying, + src._underlying, + ) + + return out + + +def scatter_add_( + input: Tensor, + dim: int, + index: Tensor, + src: Tensor, +) -> Tensor: + r"""In-place scatter_add.""" + + if infinicore.use_ntops and input.device.type in ("cuda", "musa"): + return infinicore.ntops.torch.scatter_add( + input, + dim, + index, + src, + out=input, + ) + + if not hasattr(_infinicore, "scatter_add_"): + raise NotImplementedError( + "scatter_add_ is not implemented in _infinicore, " + "and ntops path is unavailable." + ) + + _infinicore.scatter_add_( + input._underlying, + input._underlying, + dim, + index._underlying, + src._underlying, + ) + + return input \ No newline at end of file diff --git a/test/infinicore/ops/frac.py b/test/infinicore/ops/frac.py index 40a203c65..7ea62e078 100644 --- a/test/infinicore/ops/frac.py +++ b/test/infinicore/ops/frac.py @@ -87,9 +87,9 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.frac(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): + def infinicore_operator(self, *args, **kwargs): # """InfiniCore implementation (operator not yet available).""" - # return infinicore.frac(*args, **kwargs) + return infinicore.frac(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/fractional_max_pool2d.py b/test/infinicore/ops/fractional_max_pool2d.py index 09a04e8c2..ec64fb65d 100644 --- a/test/infinicore/ops/fractional_max_pool2d.py +++ b/test/infinicore/ops/fractional_max_pool2d.py @@ -65,9 +65,9 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.nn.functional.fractional_max_pool2d(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): + def infinicore_operator(self, *args, **kwargs): # """InfiniCore implementation (operator not yet available).""" - # return infinicore.nn.functional.fractional_max_pool2d(*args, **kwargs) + return infinicore.nn.functional.fractional_max_pool2d(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/fractional_max_pool3d.py b/test/infinicore/ops/fractional_max_pool3d.py index 918874623..0ae77c4a3 100644 --- a/test/infinicore/ops/fractional_max_pool3d.py +++ b/test/infinicore/ops/fractional_max_pool3d.py @@ -3,8 +3,8 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -import infinicore import torch +import infinicore from framework import BaseOperatorTest, TensorSpec, TestCase, GenericTestRunner # Test cases format: (in_shape, in_strides_or_None, kernel_size, output_size_or_None, return_indices) @@ -63,9 +63,9 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.nn.functional.fractional_max_pool3d(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): + def infinicore_operator(self, *args, **kwargs): # """InfiniCore implementation (operator not yet available).""" - # return infinicore.nn.functional.fractional_max_pool3d(*args, **kwargs) + return infinicore.nn.functional.fractional_max_pool3d(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/multilabel_margin_loss.py b/test/infinicore/ops/multilabel_margin_loss.py index 7968433ee..dd099bb6d 100644 --- a/test/infinicore/ops/multilabel_margin_loss.py +++ b/test/infinicore/ops/multilabel_margin_loss.py @@ -72,9 +72,9 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.nn.functional.multilabel_margin_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.multilabel_margin_loss(*args, **kwargs) + return infinicore.nn.functional.multilabel_margin_loss(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/scatter_add.py b/test/infinicore/ops/scatter_add.py index 323e4ae8e..80fb40b2d 100644 --- a/test/infinicore/ops/scatter_add.py +++ b/test/infinicore/ops/scatter_add.py @@ -113,9 +113,9 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.scatter_add(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): + def infinicore_operator(self, *args, **kwargs): # """InfiniCore implementation (operator not yet available).""" - # return infinicore.scatter_add(*args, **kwargs) + return infinicore.scatter_add(*args, **kwargs) def main():