Merge branch 'master' of github.com:InfiniTensor/ntops into develop-softmax

Author: voltjia

src/ntops/__init__.py

@@ -0,0 +1,3 @@
+import torch
+
+__all__ = ["torch"]

src/ntops/kernels/cos.py

@@ -0,0 +1,16 @@
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, output):
+    output = ntl.cos(input)  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    return ninetoothed.make(arrangement, application, (Tensor(ndim), Tensor(ndim)))

src/ntops/kernels/element_wise.py

@@ -1,12 +1,12 @@
 import ninetoothed
 
 
-def arrangement(*tensors):
+def arrangement(*tensors, block_size=ninetoothed.block_size()):
     ndim = max(tensor.ndim for tensor in tensors)
 
     assert all(tensor.ndim == ndim or tensor.ndim == 0 for tensor in tensors)
 
-    block_shape = tuple(1 for _ in range(ndim - 1)) + (ninetoothed.block_size(),)
+    block_shape = tuple(1 for _ in range(ndim - 1)) + (block_size,)
 
     return tuple(
         tensor.tile(block_shape) if tensor.ndim != 0 else tensor for tensor in tensors

src/ntops/kernels/eq.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, other, output):
+    output = input == other  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)

src/ntops/kernels/ge.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, other, output):
+    output = input >= other  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)

src/ntops/kernels/gt.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, other, output):
+    output = input > other  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)

src/ntops/kernels/isinf.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, output):
+    pos_result = input == float("+inf")
+    neg_result = input == float("-inf")
+    output = pos_result or neg_result  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    return ninetoothed.make(arrangement, application, (Tensor(ndim), Tensor(ndim)))

src/ntops/kernels/isnan.py

@@ -0,0 +1,15 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, output):
+    output = input != input  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    return ninetoothed.make(arrangement, application, (Tensor(ndim), Tensor(ndim)))

src/ntops/kernels/le.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, other, output):
+    output = input <= other  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)

src/ntops/kernels/lt.py

@@ -0,0 +1,17 @@
+import functools
+
+import ninetoothed
+from ninetoothed import Tensor
+
+from ntops.kernels.element_wise import arrangement
+
+
+def application(input, other, output):
+    output = input < other  # noqa: F841
+
+
+@functools.cache
+def make(ndim):
+    tensors = (Tensor(ndim), Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)