nn.functional:增加

Author: miaobyte

excuter/op-mem-ompsimd/src/client/main.cpp

@@ -19,7 +19,7 @@ int main()
 
     client::udpserver server(8080);
     deepx::op::OpFactory opfactory;
-    deepx::op::register_all(opfactory);
+    register_all(opfactory);
 
     server.func = [&mem, &opfactory, &memmutex](const char *buffer)
     {

front/py/deepx/__init__.py

@@ -1,17 +1,20 @@
 from .tensor import Tensor,Shape,Device,DeviceType
 from deepx.nn.functional import full,zeros,ones,arange,rand,randn,eye
-from deepx.nn.functional import add,sub,mul,div
+from deepx.nn.functional import add,sub,mul,div,clamp
 from deepx.nn.functional import matmul
+from deepx.nn.functional import max,min,sum,prod,mean
 __all__ = [
     'Tensor',
     'Shape',
     'Device','DeviceType',
     #init
     'full','zeros', 'ones', 'arange', 'rand', 'randn', 'eye',
     #elementwise
-    "add","sub","mul","div",
+    "add","sub","mul","div","clamp",
     #matmul
     "matmul",
+    #reduce
+    "max","min","sum","prod","mean",
 ]
 
 # 为了支持 import deepx as dx 的用法

front/py/deepx/nn/functional/__init__.py

@@ -1,12 +1,15 @@
-from .elementwise import add,sub,mul,div
+from .elementwise import add,sub,mul,div,clamp
 from .new import newtensor
 from .print import printtensor
 from .matmul import matmul
 from .init import full,zeros,ones,arange,rand,randn,eye
+from .reduce import max,min,sum,prod,mean
+
 __all__ = [
     "newtensor",
     "printtensor",
     "full","zeros","ones","arange","rand","randn","eye",
-    "add","sub","mul","div",
+    "add","sub","mul","div","clamp",
     "matmul",
+    "max","min","sum","prod","mean",
 ]

front/py/deepx/nn/functional/elementwise.py

@@ -29,7 +29,7 @@ def _A_b_elementwiseop_C(
     if a.graph.eager:
         varir=DeepxIR("argset", a.dtype, [b], [varnode.name])
         send(str(varir))
-        ir=DeepxIR(op+"_scalar", a.dtype, [a.node.name,varnode.name], [out.node.name])
+        ir=DeepxIR(op, a.dtype, [a.node.name,varnode.name], [out.node.name])
         send(str(ir))
 #add
 OpNode.register("add")
@@ -42,7 +42,7 @@ def add(
     if isinstance(b,Tensor):
         _A_B_elementwiseop_C(a,b,"add",out)
     else:
-        _A_b_elementwiseop_C(a,b,"add",out)
+        _A_b_elementwiseop_C(a,b,"add_scalar",out)
 
 
 #sub
@@ -56,7 +56,7 @@ def sub(
     if isinstance(b,Tensor):
         _A_B_elementwiseop_C(a,b,"sub",out)
     else:
-        _A_b_elementwiseop_C(a,b,"sub",out)
+        _A_b_elementwiseop_C(a,b,"sub_scalar",out)
 
 
 #mul
@@ -70,7 +70,7 @@ def mul(
     if isinstance(b,Tensor):
         _A_B_elementwiseop_C(a,b,"mul",out)
     else:
-        _A_b_elementwiseop_C(a,b,"mul",out)
+        _A_b_elementwiseop_C(a,b,"mul_scalar",out)
 
 
 #div
@@ -84,10 +84,28 @@ def div(
     if isinstance(b,Tensor):
         _A_B_elementwiseop_C(a,b,"div",out)
     else:
-        _A_b_elementwiseop_C(a,b,"div",out)
+        _A_b_elementwiseop_C(a,b,"div_scalar",out)
 
-
-
+ 
+#clamp
+OpNode.register("clamp")
+def clamp(
+        a:Tensor,
+        min: Optional[Union[ float, int]] = None, 
+        max: Optional[Union[ float, int]] = None, 
+        out:Tensor=None):   
+    opnode = a.graph.add_op("clamp")
+    opnode.add_input(a.node)
+    if min is not None:
+        min_node = a.graph.add_var("", min)
+        opnode.add_input(min_node)
+    if max is not None:
+        max_node = a.graph.add_var("", max)
+        opnode.add_input(max_node)
+    out.node.add_input(opnode)
+    if a.graph.eager:
+        varir=DeepxIR("clamp", a.dtype, [a.node.name,min,max], [out.node.name])
+        send(str(varir))
 
 # OpNode.register("ReLU", 101)
 # OpNode.register("Placeholder", 102)
@@ -98,22 +116,7 @@ def div(
 # NodeType.register("Tanh", 107)
 # NodeType.register("Reshape", 108)
 # NodeType.register("Transpose", 109)
-# NodeType.register("Sum", 110)
-# NodeType.register("Mean", 111)
-
-# # 操作节点创建函数
-# def matmul(a, b, name=None):
-#     node = OpNode("MatMul", name)
-#     node.add_input("a", a)
-#     node.add_input("b", b)
-#     return node
-
-# def add(a, b, name=None):
-#     node = OpNode("Add", name)
-#     node.add_input("a", a)
-#     node.add_input("b", b)
-#     return node
-
+  
 # def relu(x, name=None):
 #     node = OpNode("ReLU", name)
 #     node.add_input("x", x)
@@ -129,25 +132,7 @@ def div(
 #     node = OpNode("Neg")
 #     node.add_input("x", x)
 #     return node
-
-# def mul(a, b):
-#     node = OpNode("Mul")
-#     node.add_input("a", a)
-#     node.add_input("b", b)
-#     return node
-
-# def div(a, b):
-#     node = OpNode("Div")
-#     node.add_input("a", a)
-#     node.add_input("b", b)
-#     return node
-
-# def sub(a, b):
-#     node = OpNode("Sub")
-#     node.add_input("a", a)
-#     node.add_input("b", b)
-#     return node
-
+ 
 # def less(a, b):
 #     node = OpNode("Less")
 #     node.add_input("a", a)
@@ -182,17 +167,4 @@ def div(
 #     node.set_attr("dim0", dim0)
 #     node.set_attr("dim1", dim1)
 #     return node
-
-# def sum(x, dim=None, keepdim=False):
-#     node = OpNode("Sum")
-#     node.add_input("x", x)
-#     node.set_attr("dim", dim)
-#     node.set_attr("keepdim", keepdim)
-#     return node
-
-# def mean(x, dim=None, keepdim=False):
-#     node = OpNode("Mean")
-#     node.add_input("x", x)
-#     node.set_attr("dim", dim)
-#     node.set_attr("keepdim", keepdim)
-#     return node 
+ 

front/py/deepx/nn/functional/reduce.py

@@ -1,14 +1,14 @@
 from typing import Optional, Union
 
-from .tensor import Tensor,tensor_method
+from deepx.tensor import Tensor
 from deepx.autograd.graph import OpNode
 from deepx.nn.deepxir import DeepxIR    
 from deepx.scheduler import send
 from .elementwise import _A_b_elementwiseop_C
 
 def _A_v_reduceop_C(
         a:Tensor,
-        v: Optional[Union[Tensor, float, int]] = None, 
+        v: Optional[Union[list[int],tuple[int]]] = None, 
         op:str=None,
         out:Tensor=None):
     opnode = a.graph.add_op(op)
@@ -20,56 +20,84 @@ def _A_v_reduceop_C(
     if a.graph.eager:
         varir=DeepxIR("argset", a.dtype, v, [vector_node.name])
         send(str(varir))
-        ir=DeepxIR(op+"_scalar", a.dtype, [a.node.name,vector_node.name], [out.node.name])
+        ir=DeepxIR(op, a.dtype, [a.node.name,vector_node.name], [out.node.name])
         send(str(ir))
 
 
 #max
 OpNode.register("max")
-OpNode.register("max_scalar")
-
 def max(
         a:Tensor,
-        b:Optional[Union[float,int],Union[Tensor,float,int]]=None,
+        b: Optional[Union[
+            int,float,
+            Tensor,
+            ]] = None, 
+        dims:Optional[Union[list[int],tuple[int]]]=None,
         out:Tensor=None):
-    if isinstance(b,list):
-        _A_v_reduceop_C(a,b,"max",out)
-    else:
+    if b is not None and isinstance(b,int,float):
         _A_b_elementwiseop_C(a,b,"max_scalar",out)
-
-@tensor_method
-def max_(self, other):
-    result = Tensor(dtype=self.dtype,shape=self.shape)
-    max(self,other,result)
-    return result
+    elif b is not None and isinstance(b,Tensor):
+        _A_b_elementwiseop_C(a,b,"max_tensor",out)
+    else:
+        if dims is None:
+            dims=list(range(a.ndim))
+        _A_v_reduceop_C(a,dims,"max",out)
 
 #min    
 OpNode.register("min")
-OpNode.register("min_scalar")
-
-def min(a:Tensor,b:Tensor,out:Tensor):
-    if isinstance(b,list):
-        _A_v_reduceop_C(a,b,"min",out)
-    else:
+def min(
+        a:Tensor,
+        b: Optional[Union[
+            int,float,
+            Tensor,
+            ]] = None, 
+        dims:Optional[Union[list[int],tuple[int]]]=None,
+        out:Tensor=None):
+    if b is not None and isinstance(b,int,float):
         _A_b_elementwiseop_C(a,b,"min_scalar",out)
-
-@tensor_method
-def min_(self, other):
-    result = Tensor(dtype=self.dtype,shape=self.shape)
-    min(self,other,result)
-    return result
-
-
+    elif b is not None and isinstance(b,Tensor):
+        _A_b_elementwiseop_C(a,b,"min_tensor",out)
+    else:
+        if dims is None:
+            dims=list(range(a.ndim))
+        _A_v_reduceop_C(a,dims,"min",out)
+ 
 #sum    
 OpNode.register("sum")
 def sum(
         a:Tensor,
-        b:list[int],
-        out:Tensor):
-    _A_v_reduceop_C(a,b,"sum",out)
+        dims:Optional[Union[
+            list[int],
+            tuple[int],
+            ]]=None,
+        out:Tensor=None):
+    if dims is None:
+        dims=list(range(a.ndim))
+    _A_v_reduceop_C(a,dims,"sum",out)
+
+#prod
+OpNode.register("prod")
+def prod(
+        a:Tensor,
+        dims:Optional[Union[
+            list[int],
+            tuple[int],
+            ]]=None,
+        out:Tensor=None):
+    if dims is None:
+        dims=list(range(a.ndim))
+    _A_v_reduceop_C(a,dims,"prod",out)
+
+#mean
+OpNode.register("mean")
+def mean(
+        a:Tensor,
+        dims:Optional[Union[list[int],tuple[int]]]=None,
+        out:Tensor=None):
+    if dims is None:
+        dims=list(range(a.ndim))
+    _A_v_reduceop_C(a,dims,"mean",out)
+
 
-@tensor_method
-def sum_(self, other):
-    result = Tensor(dtype=self.dtype,shape=self.shape)
-    sum(self,other,result)
-    return result
+# #var
+# OpNode.register("var")

front/py/deepx/nn/functional/reduction.py

@@ -28,4 +28,23 @@ def div_(self, other):
     div(self,other,result)
     return result
 
- 
+@tensor_method
+def min_scalar_(self, other):
+    result = Tensor(dtype=self.dtype,shape=self.shape)
+    from deepx.nn.functional import min_scalar
+    min_scalar(self,other,result)
+    return result
+
+@tensor_method
+def max_scalar_(self, other):
+    result = Tensor(dtype=self.dtype,shape=self.shape)
+    from deepx.nn.functional import max_scalar
+    max_scalar(self,other,result)
+    return result   
+
+@tensor_method
+def clamp_(self, min, max):
+    result = Tensor(dtype=self.dtype,shape=self.shape)
+    from deepx.nn.functional import clamp
+    clamp(self,min,max,result)
+    return result