py:deepx.tensor

Author: miaobyte

front/go/deepx/attention.go

@@ -65,10 +65,10 @@ func (m *MultiHeadAttention) Forward(q, k, v *Tensor) *Tensor {
 
 	// 5. Scale
 	d_k := float32(head_dim)
-	scores = scores.Scale(1.0 / float32(math.Sqrt(float64(d_k))))
+	scores = scores.MulScalar(1.0 / float32(math.Sqrt(float64(d_k))))
 
 	// 6. Softmax
-	attn := scores.Softmax()
+	attn := scores.Softmax(1)
 
 	// 7. 加权求和
 	out := attn.Matmul(value)

front/go/deepx/graph.go

@@ -22,8 +22,8 @@ func NewLayerNorm(name string, normalized_shape int, dtype Dtype, g *Graph) *Lay
 }
 
 func (m *LayerNorm) LayerNorm(x *Tensor) *Tensor {
-	op := x.LayerNorm(m.weight, m.bias)
-	return op
+
+	return nil
 }
 
 // BatchNorm 批归一化
@@ -52,8 +52,7 @@ func NewBatchNorm(name string, num_features int, dtype Dtype, g *Graph) *BatchNo
 }
 
 func (m *BatchNorm) BatchNorm(x *Tensor) *Tensor {
-	op := x.BatchNorm(m.weight, m.bias, m.running_mean, m.running_var)
-	return op
+	return nil
 }
 
 // InstanceNorm 实例归一化
@@ -78,8 +77,7 @@ func NewInstanceNorm(name string, num_features int, dtype Dtype, g *Graph) *Inst
 }
 
 func (m *InstanceNorm) InstanceNorm(x *Tensor) *Tensor {
-	op := x.InstanceNorm(m.weight, m.bias)
-	return op
+	return nil
 }
 
 // GroupNorm 组归一化
@@ -110,8 +108,7 @@ func NewGroupNorm(name string, num_groups, num_channels int, dtype Dtype, g *Gra
 }
 
 func (m *GroupNorm) GroupNorm(x *Tensor) *Tensor {
-	op := x.GroupNorm(m.weight, m.bias, m.num_groups)
-	return op
+	return nil
 }
 
 // RMSNorm Root Mean Square Layer Normalization
@@ -137,6 +134,5 @@ func NewRMSNorm(name string, normalized_shape int, dtype Dtype, g *Graph) *RMSNo
 }
 
 func (m *RMSNorm) RMSNorm(x *Tensor) *Tensor {
-	op := x.RMSNorm(m.weight, m.eps)
-	return op
+	return nil
 }

front/go/deepx/norm.go

@@ -1,4 +1,12 @@
-from .tensor import Tensor
+from .tensor import Tensor,Shape,Device,DeviceType
+from .tensor import zeros, ones, arange, rand, randn, eye
 
-__all__ = ['Tensor']
+__all__ = [
+    'Tensor',
+    'Shape',
+    'Device','DeviceType',
+    'zeros', 'ones', 'arange', 'rand', 'randn', 'eye'
+]
+
+# 为了支持 import deepx as dx 的用法
 tensor = Tensor

front/py/deepx/__init__.py

@@ -0,0 +1,3 @@
+from .graph import Graph
+
+__all__ = ['Graph']

front/py/deepx/autograd/__init__.py

@@ -1,8 +1,26 @@
-from ..tensor import Tensor
-from .tensornode import TensorNode
-from .opnode import OpNode
-from .constargnode import ConstArgNode
+from ..tensor.tensor import Tensor
+from ._tensornode import TensorNode
+from ._opnode import OpNode
+from ._constargnode import ConstArgNode
+
 class Graph:
+    # 类属性存储默认实例
+    _default_graph = None
+    
+    @classmethod
+    def get_default(cls):
+        """获取或创建默认计算图（线程不安全）"""
+        if cls._default_graph is None:
+            cls._default_graph = Graph()
+        return cls._default_graph
+    
+    @classmethod
+    def set_default(cls, graph):
+        """设置新的默认计算图（用于上下文管理）"""
+        if not isinstance(graph, Graph):
+            raise TypeError("Must be a Graph instance")
+        cls._default_graph = graph
+
     def __init__(self):
         self.nodes = []
         self.inputs = []
@@ -29,4 +47,7 @@ def add_constarg(self, value):
             name = f"constarg_{self.constarg_counter}"
         node=ConstArgNode(value)
         self.nodes.append(node)
-        return node
+        return node
+
+# 初始化默认图
+Graph._default_graph = Graph()

front/py/deepx/autograd/constargnode.py front/py/deepx/autograd/_constargnode.pyfront/py/deepx/autograd/constargnode.py renamed to front/py/deepx/autograd/_constargnode.py

@@ -1,8 +1,8 @@
 import graphviz
 from .graph import Graph
-from .tensornode import TensorNode
-from .opnode import OpNode
-from .constargnode import ConstArgNode
+from ._tensornode import TensorNode
+from ._opnode import OpNode
+from ._constargnode import ConstArgNode
 
 def graph_method(f):
     """装饰器:将函数注册为Graph类的方法"""