deep-atomic 0.2.1__tar.gz → 0.2.2__tar.gz

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deep_atomic
-Version: 0.2.1
+Version: 0.2.2
 Summary: A simple deep learning framework built upon numpy only
 Author-email: Elliot Zhang <elliot_zh@proton.me>
 License: MIT
@@ -21,7 +21,7 @@ Description-Content-Type: text/markdown
 
 # Deep Atomic
 
-A simple deep learning framework built upon numpy only. Mainly for practice and learning.
+A simple deep learning framework built on NumPy only. Mainly for practice and learning.
 
 ## Usage
 
@@ -53,7 +53,7 @@ a = da.Tensor(
 
 ### Supported Operations
 
-Most essential deep‑learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
+Most essential deep-learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
 
 ```python
 a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
@@ -70,24 +70,28 @@ c = da.log(a)
 c = da.sin(a)
 # the same for cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh
 
-d = a < b                     # element-wise comparison, create a boolean tensor
+d = a < b                     # element-wise comparison, creates a boolean tensor
 e = a <= b
 c = a > b
 c = a >= b
 c = a == b
 c = a != b
-c = da.fmax(a, b)             # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. same for da.fmin / da.minimum
+c = da.fmax(a, b)             # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. Same for da.fmin / da.minimum
 c = da.maximum(a, b)
 c = da.fmin(a, b)
 c = da.minimum(a, b)
 
 c = da.logical_and(d, e)                     # element-wise and
+c = d & e                                    # equivalence
 c = da.logical_or(d, e)                      # element-wise or
+c = d | e
 c = da.logical_xor(d, e)                     # element-wise xor
+c = d ^ e
 c = da.logical_not(d)                        # element-wise not
-c = d.all(axis=-1, keepdims=False)           # reduction of and operation. axis=None, keepdims=False by default
-c = d.any(axis=-1, keepdims=False)           # reduction of or operation. axis=None, keepdims=False by default
-c = da.where(d, a, b)                        # return elements chosen from a or b depending on condition
+c = ~d
+c = d.all(axis=-1, keepdims=False)           # logical AND reduction. axis=None, keepdims=False by default
+c = d.any(axis=-1, keepdims=False)           # logical OR reduction. axis=None, keepdims=False by default
+c = da.where(d, a, b)                        # returns elements chosen from a or b depending on condition
 
 c = da.topk(a, 2, axis=-1, largest=True)    # same as pytorch. axis=-1, largest=True by default
 
@@ -116,17 +120,17 @@ c = a.tile(2, 2)                              # method‑style alternative
 
 ### Autograd
 
-Autograd is supported via computational graph.
-_Currently only support scalar source points._
+Autograd is supported via a computational graph.
+_Currently only supports scalar source points._
 
 ```python
 x = Tensor(np.random.rand(3, 4)) # requires_grad == True by default
 res = ... # some calculation related to x. res is a **scalar** result
 res.backward()
-print(res.grad) # gradient get!
+print(res.grad) # gradient computed!
 ```
 
-## Todo
+## To Do
 
 - [ ] more basic operations and their autograd
     - [x] topk
@@ -151,8 +155,8 @@ print(res.grad) # gradient get!
 
 ## Development
 
-Recommend manage dependencies using [uv](https://github.com/astral-sh/uv).
-We use [pre-commit](https://github.com/pre-commit/pre-commit) to manage hooks which help to lint and format our code.
+We recommend managing dependencies with [uv](https://github.com/astral-sh/uv).
+We use [pre-commit](https://github.com/pre-commit/pre-commit) to manage hooks that help lint and format our code.
 
 ```bash
 uv sync

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/README.md

@@ -1,6 +1,6 @@
 # Deep Atomic
 
-A simple deep learning framework built upon numpy only. Mainly for practice and learning.
+A simple deep learning framework built on NumPy only. Mainly for practice and learning.
 
 ## Usage
 
@@ -32,7 +32,7 @@ a = da.Tensor(
 
 ### Supported Operations
 
-Most essential deep‑learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
+Most essential deep-learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
 
 ```python
 a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
@@ -49,24 +49,28 @@ c = da.log(a)
 c = da.sin(a)
 # the same for cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh
 
-d = a < b                     # element-wise comparison, create a boolean tensor
+d = a < b                     # element-wise comparison, creates a boolean tensor
 e = a <= b
 c = a > b
 c = a >= b
 c = a == b
 c = a != b
-c = da.fmax(a, b)             # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. same for da.fmin / da.minimum
+c = da.fmax(a, b)             # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. Same for da.fmin / da.minimum
 c = da.maximum(a, b)
 c = da.fmin(a, b)
 c = da.minimum(a, b)
 
 c = da.logical_and(d, e)                     # element-wise and
+c = d & e                                    # equivalence
 c = da.logical_or(d, e)                      # element-wise or
+c = d | e
 c = da.logical_xor(d, e)                     # element-wise xor
+c = d ^ e
 c = da.logical_not(d)                        # element-wise not
-c = d.all(axis=-1, keepdims=False)           # reduction of and operation. axis=None, keepdims=False by default
-c = d.any(axis=-1, keepdims=False)           # reduction of or operation. axis=None, keepdims=False by default
-c = da.where(d, a, b)                        # return elements chosen from a or b depending on condition
+c = ~d
+c = d.all(axis=-1, keepdims=False)           # logical AND reduction. axis=None, keepdims=False by default
+c = d.any(axis=-1, keepdims=False)           # logical OR reduction. axis=None, keepdims=False by default
+c = da.where(d, a, b)                        # returns elements chosen from a or b depending on condition
 
 c = da.topk(a, 2, axis=-1, largest=True)    # same as pytorch. axis=-1, largest=True by default
 
@@ -95,17 +99,17 @@ c = a.tile(2, 2)                              # method‑style alternative
 
 ### Autograd
 
-Autograd is supported via computational graph.
-_Currently only support scalar source points._
+Autograd is supported via a computational graph.
+_Currently only supports scalar source points._
 
 ```python
 x = Tensor(np.random.rand(3, 4)) # requires_grad == True by default
 res = ... # some calculation related to x. res is a **scalar** result
 res.backward()
-print(res.grad) # gradient get!
+print(res.grad) # gradient computed!
 ```
 
-## Todo
+## To Do
 
 - [ ] more basic operations and their autograd
     - [x] topk
@@ -130,8 +134,8 @@ print(res.grad) # gradient get!
 
 ## Development
 
-Recommend manage dependencies using [uv](https://github.com/astral-sh/uv).
-We use [pre-commit](https://github.com/pre-commit/pre-commit) to manage hooks which help to lint and format our code.
+We recommend managing dependencies with [uv](https://github.com/astral-sh/uv).
+We use [pre-commit](https://github.com/pre-commit/pre-commit) to manage hooks that help lint and format our code.
 
 ```bash
 uv sync

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "deep_atomic"
-version = "0.2.1"
+version = "0.2.2"
 description = "A simple deep learning framework built upon numpy only"
 readme = "README.md"
 requires-python = ">=3.10"

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/src/deep_atomic/__init__.py

@@ -1,6 +1,6 @@
 """deep-atomic - A simple deep learning framework built upon numpy only"""
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 __author__ = "elliot_zh@proton.me"
 
 from .op import *

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/src/deep_atomic/nn.py

@@ -46,7 +46,7 @@ class Module(ABC):
         pass
 
     def __call__(self, *args, **kwargs):
-        self.forward(*args, **kwargs)
+        return self.forward(*args, **kwargs)
 
     def __init__(self):
         # TODO: use OrderedDict?
@@ -72,19 +72,19 @@ class Module(ABC):
         for name, buffer in self._buffers.items():
             state_dict_[prefix + name] = buffer
         for name, module in self._modules.items():
-            state_dict_ = {**state_dict_, **module.state_dict(prefix=name)}
+            state_dict_ = {**state_dict_, **module.state_dict(prefix=prefix + name)}
         return state_dict_
 
     def modules(self, recurse=True):
-        for _name, module in self.named_modules(recurse=recurse):
+        for _name, module in self.named_modules():
             yield module
 
-    def named_module(self, prefix="", remove_duplicate=True, memo=None):
+    def named_modules(self, prefix="", remove_duplicate=True, memo=None):
         if memo is None:
             memo = set()
-        if self not in memo:
+        if id(self) not in memo:
             if remove_duplicate:
-                memo.add(self)
+                memo.add(id(self))
             yield prefix, self
             if prefix:
                 prefix += "."
@@ -107,10 +107,10 @@ class Module(ABC):
             if module_prefix:
                 module_prefix += "."
             for k, v in members:
-                if v in memo:
+                if id(v) in memo:
                     continue
                 if remove_duplicate:
-                    memo.add(v)
+                    memo.add(id(v))
                 yield module_prefix + k, v
 
     def parameters(self, recurse=True):

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/src/deep_atomic/op.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from .graph import *
+from . import graph as g
 from .tensor import *
 from .utils import *
 
@@ -89,10 +89,10 @@ def max(x: Tensor, axis=None, keepdims=False):
     if x.requires_grad:
         res.requires_grad = True
         if axis == None:
-            res.dep = MinMax(x, None, keepdims, full_red_value=res.to_np())
+            res.dep = g.MinMax(x, None, keepdims, full_red_value=res.to_np())
         else:
             indices = np.argmax(x.to_np(), axis=axis, keepdims=True)
-            res.dep = MinMax(x, axis, keepdims, indices=indices)
+            res.dep = g.MinMax(x, axis, keepdims, indices=indices)
     return res
 
 
@@ -101,10 +101,10 @@ def min(x: Tensor, axis=None, keepdims=False):
     if x.requires_grad:
         res.requires_grad = True
         if axis == None:
-            res.dep = MinMax(x, None, keepdims, full_red_value=res.to_np())
+            res.dep = g.MinMax(x, None, keepdims, full_red_value=res.to_np())
         else:
             indices = np.argmin(x.to_np(), axis=axis, keepdims=True)
-            res.dep = MinMax(x, axis, keepdims, indices=indices)
+            res.dep = g.MinMax(x, axis, keepdims, indices=indices)
     return res
 
 
@@ -155,7 +155,7 @@ def sigmoid(x: Tensor):
     res[~pos] = exp(x_np[~pos]) / (1 + exp(x_np[~pos]))
     res = Tensor(res, requires_grad=x.requires_grad)
     if res.requires_grad:
-        res.dep = Sigmoid(x, res.to_np())
+        res.dep = g.Sigmoid(x, res.to_np())
     return res
 
 
@@ -247,7 +247,7 @@ def where(condition, x1, x2):
             requires_grad = True
     res = Tensor(np.where(condition, x1, x2), requires_grad=requires_grad)
     if res.requires_grad:
-        res.dep = Where(condition, x1, x2)
+        res.dep = g.Where(condition, x1, x2)
     return res
 
 
@@ -257,7 +257,7 @@ def take_along_axis(x: Tensor, indices, axis=-1):
         indices = indices.to_np()  # cut off gradient
     res = Tensor(np.take_along_axis(x, indices, axis), requires_grad=x.requires_grad)
     if res.requires_grad:
-        res.dep = TakeAlongAxis(x, indices, axis)
+        res.dep = g.TakeAlongAxis(x, indices, axis)
     return res
 
 

deep_atomic-0.2.2/tests/nn_test.py

@@ -0,0 +1,583 @@
+import numpy as np
+import pytest
+
+import deep_atomic as da
+from deep_atomic.nn import (
+    Buffer,
+    BufferList,
+    Linear,
+    Module,
+    ModuleList,
+    Parameter,
+    ParameterList,
+    ReLU,
+    Sequential,
+    Softmax,
+)
+
+from .utils import assert_close, numerical_grad
+
+
+def _make_identity_module():
+    """Factory for a minimal identity Module."""
+
+    class M(Module):
+        def forward(self, x):
+            return x
+
+    return M()
+
+
+class TestBuffer:
+    def test_requires_grad_always_false(self):
+        """Buffer.requires_grad is False even when passed requires_grad=True."""
+        buf = Buffer(np.array([1.0, 2.0, 3.0]), requires_grad=True)
+        assert buf.requires_grad is False
+
+    def test_is_buffer_and_tensor(self):
+        buf = Buffer(np.array([1, 2, 3]))
+        assert isinstance(buf, Buffer)
+        assert isinstance(buf, da.Tensor)
+
+
+MEMBER_PARAMS = [
+    (Parameter, "named_parameters", "parameters"),
+    (Buffer, "named_buffers", "buffers"),
+]  # for TestModuleNamedMembers
+
+
+class TestModule:
+    class TestModuleSetAttr:
+        """__setattr__ routes Parameter/Buffer/Module into their tracking dicts."""
+
+        @pytest.mark.parametrize(
+            "cls, dict_name",
+            [(Parameter, "_parameters"), (Buffer, "_buffers")],
+        )
+        def test_member_routing(self, cls, dict_name):
+            m = _make_identity_module()
+            v = cls(np.array([1.0]))
+            m.w = v
+            assert "w" in getattr(m, dict_name)
+            assert getattr(m, dict_name)["w"] is v
+
+        def test_module_routing(self):
+            child = _make_identity_module()
+            parent = _make_identity_module()
+            parent.child = child
+            assert "child" in parent._modules
+            assert parent._modules["child"] is child
+
+        def test_plain_value_not_tracked(self):
+            m = _make_identity_module()
+            m.x = 3
+            m.label = "hello"
+            assert "x" not in m._parameters
+            assert "x" not in m._buffers
+            assert "x" not in m._modules
+            assert "label" not in m._parameters
+            assert "label" not in m._buffers
+            assert "label" not in m._modules
+
+    class TestModuleStateDict:
+        def test_flat(self):
+            """Flat module: state_dict contains every Parameter + Buffer keyed by name."""
+
+            class M(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = Parameter(np.array([1.0, 2.0]))
+                    self.buf = Buffer(np.array([3.0]))
+
+                def forward(self, x):
+                    return x
+
+            m = M()
+            sd = m.state_dict()
+            assert set(sd.keys()) == {"w", "buf"}
+            assert sd["w"] is m.w
+            assert sd["buf"] is m.buf
+
+        def test_nested(self):
+            """Nested: state_dict flattens recursively with dotted prefix."""
+
+            class Leaf(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = Parameter(np.array([1.0]))
+
+                def forward(self, x):
+                    return x
+
+            class Mid(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.b = Leaf()
+
+                def forward(self, x):
+                    return x
+
+            class Root(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.a = Mid()
+
+                def forward(self, x):
+                    return x
+
+            r = Root()
+            sd = r.state_dict()
+            assert sd == {"a.b.w": r.a.b.w}
+
+    class TestModuleNamedModules:
+        def test_no_submodules(self):
+            """Module with no submodules yields only self."""
+
+            m = _make_identity_module()
+            result = list(m.named_modules())
+            assert len(result) == 1
+            assert result[0][0] == ""
+            assert result[0][1] is m
+
+        def test_recurse(self):
+            """Yields self + all descendants with dotted prefixes."""
+
+            inner = _make_identity_module()
+
+            class Outer(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.a = inner
+
+                def forward(self, x):
+                    return x
+
+            outer = Outer()
+            result = list(outer.named_modules())
+            assert len(result) == 2
+            assert result[0][0] == ""
+            assert result[0][1] is outer
+            assert result[1][0] == "a"
+            assert result[1][1] is outer.a
+
+        def test_remove_duplicate(self):
+            """remove_duplicate=True → shared submodule yielded once."""
+
+            shared = _make_identity_module()
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.left = shared
+                    self.right = shared
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            result = list(p.named_modules(remove_duplicate=True))
+            prefixes = [prefix for prefix, _ in result]
+            assert prefixes.count("") == 1  # self
+            assert prefixes.count("left") == 1
+
+        def test_modules_matches_named_modules(self):
+            """modules() values match named_modules() values (names stripped)."""
+
+            inner = _make_identity_module()
+
+            class Outer(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.a = inner
+
+                def forward(self, x):
+                    return x
+
+            outer = Outer()
+            named_mods = [m for _, m in outer.named_modules()]
+            mods = list(outer.modules())
+            assert mods == named_mods
+
+    class TestModuleNamedMembers:
+        @pytest.mark.parametrize("cls, named, bare", MEMBER_PARAMS)
+        def test_flat(self, cls, named, bare):
+            class M(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.a = cls(np.array([1.0]))
+                    self.b = cls(np.array([2.0]))
+
+                def forward(self, x):
+                    return x
+
+            m = M()
+            result = list(getattr(m, named)())
+            assert len(result) == 2
+            names = {name for name, _ in result}
+            assert names == {"a", "b"}
+
+        @pytest.mark.parametrize("cls, named, bare", MEMBER_PARAMS)
+        def test_nested_recurse(self, cls, named, bare):
+            class Child(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = cls(np.array([1.0]))
+
+                def forward(self, x):
+                    return x
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.child = Child()
+                    self.b = cls(np.array([2.0]))
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            result = list(getattr(p, named)())
+            names = {name for name, _ in result}
+            assert names == {"b", "child.w"}
+
+        @pytest.mark.parametrize("cls, named, bare", MEMBER_PARAMS)
+        def test_no_recurse(self, cls, named, bare):
+            class Child(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = cls(np.array([1.0]))
+
+                def forward(self, x):
+                    return x
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.child = Child()
+                    self.b = cls(np.array([2.0]))
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            result = list(getattr(p, named)(recurse=False))
+            names = {name for name, _ in result}
+            assert names == {"b"}
+
+        @pytest.mark.parametrize("cls, named, bare", MEMBER_PARAMS)
+        def test_remove_duplicate(self, cls, named, bare):
+            shared = cls(np.array([1.0]))
+
+            class Child(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.m = shared
+
+                def forward(self, x):
+                    return x
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.a = Child()
+                    self.b = Child()
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            result = list(getattr(p, named)(remove_duplicate=True))
+            names = [name for name, _ in result]
+            assert names.count("a.m") == 1
+            assert "b.m" not in names
+
+        @pytest.mark.parametrize("cls, named, bare", MEMBER_PARAMS)
+        def test_bare_matches_named(self, cls, named, bare):
+            class Child(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = cls(np.array([1.0]))
+
+                def forward(self, x):
+                    return x
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.child = Child()
+                    self.b = cls(np.array([2.0]))
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            named_vals = [v for _, v in getattr(p, named)()]
+            bare_vals = list(getattr(p, bare)())
+            assert bare_vals == named_vals
+
+        @pytest.mark.parametrize("cls, named, bare", MEMBER_PARAMS)
+        def test_empty(self, cls, named, bare):
+            m = _make_identity_module()
+            assert list(getattr(m, named)()) == []
+            assert list(getattr(m, bare)()) == []
+
+    class TestModuleTrainEval:
+        def test_train(self):
+            """train() sets requires_grad=True on all params including nested."""
+
+            class Child(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = Parameter(np.array([1.0]))
+                    self.w.requires_grad = False
+
+                def forward(self, x):
+                    return x
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.child = Child()
+                    self.b = Parameter(np.array([2.0]))
+                    self.b.requires_grad = False
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            p.train()
+            params = list(p.parameters())
+            assert all(param.requires_grad for param in params)
+
+        def test_eval(self):
+            """eval() sets requires_grad=False on all params including nested."""
+
+            class Child(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.w = Parameter(np.array([1.0]))
+
+                def forward(self, x):
+                    return x
+
+            class Parent(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.child = Child()
+                    self.b = Parameter(np.array([2.0]))
+
+                def forward(self, x):
+                    return x
+
+            p = Parent()
+            p.eval()
+            params = list(p.parameters())
+            assert all(not param.requires_grad for param in params)
+
+        def test_buffers_unaffected(self):
+            """Buffers keep requires_grad=False regardless of train/eval."""
+
+            class M(Module):
+                def __init__(self):
+                    super().__init__()
+                    self.buf = Buffer(np.array([0.0]))
+
+                def forward(self, x):
+                    return x
+
+            m = M()
+            assert m.buf.requires_grad is False
+            m.train()
+            assert m.buf.requires_grad is False
+            m.eval()
+            assert m.buf.requires_grad is False
+
+    class TestModuleCall:
+        def test_call_returns_forward_result(self):
+            """__call__ delegates to forward and returns its result."""
+
+            class M(Module):
+                def forward(self, x):
+                    return x * 2
+
+            m = M()
+            result = m(da.Tensor(np.array([1.0, 2.0])))
+            expected = np.array([2.0, 4.0])
+            assert (result.to_np() == expected).all()
+
+
+LIST_PARAMS = [
+    (
+        ParameterList,
+        lambda: Parameter(np.array([1.0])),
+        "_parameters",
+        "named_parameters",
+        True,
+    ),
+    (ModuleList, _make_identity_module, "_modules", "named_modules", False),
+    (BufferList, lambda: Buffer(np.array([1.0])), "_buffers", "named_buffers", True),
+]
+
+
+class TestListContainers:
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, _c", LIST_PARAMS)
+    def test_len(self, list_cls, make_member, _d, _n, _c):
+        lst = list_cls([make_member(), make_member()])
+        assert len(lst) == 2
+
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, _c", LIST_PARAMS)
+    def test_getitem(self, list_cls, make_member, _d, _n, _c):
+        m0, m1 = make_member(), make_member()
+        lst = list_cls([m0, m1])
+        assert lst[0] is m0
+        assert lst[1] is m1
+        assert lst[-1] is m1
+        assert lst[-2] is m0
+        with pytest.raises(IndexError):
+            _ = lst[2]
+        with pytest.raises(IndexError):
+            _ = lst[-3]
+
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, _c", LIST_PARAMS)
+    def test_setitem(self, list_cls, make_member, _d, _n, _c):
+        m0, m1 = make_member(), make_member()
+        lst = list_cls([m0])
+        lst[0] = m1
+        assert lst[0] is m1
+
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, auto_convert", LIST_PARAMS)
+    def test_setitem_auto_converts(self, list_cls, make_member, _d, _n, auto_convert):
+        if not auto_convert:
+            pytest.skip("auto-convert only applies to ParameterList and BufferList")
+        t = da.Tensor(np.array([1.0]))
+        lst = list_cls([make_member()])
+        lst[0] = t
+        assert isinstance(lst[0], (Parameter, Buffer))
+
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, _c", LIST_PARAMS)
+    def test_iter(self, list_cls, make_member, _d, _n, _c):
+        members = [make_member(), make_member()]
+        lst = list_cls(members)
+        assert list(lst) == members
+
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, _c", LIST_PARAMS)
+    def test_append(self, list_cls, make_member, _d, _n, _c):
+        m = make_member()
+        lst = list_cls([make_member()])
+        initial_len = len(lst)
+        lst.append(m)
+        assert len(lst) == initial_len + 1
+        assert lst[-1] is m
+
+    @pytest.mark.parametrize("list_cls, make_member, _d, _n, _c", LIST_PARAMS)
+    def test_extend(self, list_cls, make_member, _d, _n, _c):
+        m0, m1 = make_member(), make_member()
+        lst = list_cls([make_member()])
+        lst.extend([m0, m1])
+        assert len(lst) == 3
+        assert lst[1] is m0
+        assert lst[2] is m1
+
+    @pytest.mark.parametrize("list_cls, make_member, dict_name, named, _c", LIST_PARAMS)
+    def test_parent_routing(self, list_cls, make_member, dict_name, named, _c):
+        m = make_member()
+
+        class Parent(Module):
+            def __init__(self):
+                super().__init__()
+                self.items = list_cls([m])
+
+            def forward(self, x):
+                return x
+
+        p = Parent()
+        entries = dict(getattr(p, named)())
+        assert "items.0" in entries
+        assert entries["items.0"] is m
+
+
+class TestLinear:
+    def test_weight_shape(self):
+        lin = Linear(3, 5)
+        assert lin.weight.shape == (3, 5)
+
+    def test_bias_shape(self):
+        lin = Linear(3, 5, bias=True)
+        assert lin.bias.shape == (5,)
+
+    def test_bias_is_none_when_disabled(self):
+        lin = Linear(3, 5, bias=False)
+        assert lin.bias is None
+
+    def test_forward_with_bias(self):
+        lin = Linear(3, 5, bias=True)
+        x_np = np.random.randn(2, 3).astype(np.float64)
+        x = da.Tensor(x_np)
+        out = lin(x)
+        expected = x_np @ lin.weight.to_np() + lin.bias.to_np()
+        assert out.shape == (2, 5)
+        assert (out.to_np() == expected).all()
+
+    def test_forward_without_bias(self):
+        lin = Linear(3, 5, bias=False)
+        x_np = np.random.randn(2, 3).astype(np.float64)
+        x = da.Tensor(x_np)
+        out = lin(x)
+        expected = x_np @ lin.weight.to_np()
+        assert out.shape == (2, 5)
+        assert (out.to_np() == expected).all()
+
+
+# --- numpy equivalents for each activation ---
+
+
+def _relu_np(z):
+    return np.maximum(z, 0)
+
+
+def _softmax_np(z, axis=-1):
+    z = z - z.max(axis=axis, keepdims=True)
+    e = np.exp(z)
+    return e / e.sum(axis=axis, keepdims=True)
+
+
+ACTIVATIONS = [
+    (ReLU, _relu_np),
+    (Softmax, _softmax_np),
+]
+
+
+class TestSequentialWithActivations:
+    @pytest.mark.parametrize("act_cls, act_np", ACTIVATIONS)
+    @pytest.mark.parametrize("bias", [True, False])
+    def test_forward(self, act_cls, act_np, bias):
+        lin = Linear(3, 5, bias=bias)
+        seq = Sequential([lin, act_cls()])
+
+        x_np = np.random.randn(2, 3).astype(np.float64)
+        x = da.Tensor(x_np)
+
+        out = seq(x)
+
+        z = x_np @ lin.weight.to_np()
+        if bias:
+            z = z + lin.bias.to_np()
+        expected = act_np(z)
+        assert_close(expected, out.to_np())
+
+    @pytest.mark.parametrize("act_cls, act_np", ACTIVATIONS)
+    @pytest.mark.parametrize("bias", [True, False])
+    def test_grad(self, act_cls, act_np, bias):
+        lin = Linear(3, 5, bias=bias)
+        seq = Sequential([lin, act_cls()])
+
+        x_np = np.random.randn(2, 3).astype(np.float64)
+        x = da.Tensor(x_np, requires_grad=True)
+
+        out = seq(x).sum()
+        out.backward()
+
+        def func(inp):
+            return seq(da.Tensor(inp) if not isinstance(inp, da.Tensor) else inp).sum()
+
+        assert_close(numerical_grad(func, x), x.grad)

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/tests/op_test.py

@@ -285,7 +285,7 @@ class TestShapeOps:
             _test_unary_grad(t, da.squeeze, axis=axis)
 
     @pytest.mark.parametrize(
-        "repeats_axis",
+        "repeats, axis",
         [
             (2, None),
             (2, 0),
@@ -294,12 +294,10 @@ class TestShapeOps:
         ],
     )
     class TestRepeat:
-        def test_forward(self, unary, repeats_axis):
-            repeats, axis = repeats_axis
+        def test_forward(self, unary, repeats, axis):
             _test_unary_forward(unary, da.repeat, np.repeat, repeats=repeats, axis=axis)
 
-        def test_grad(self, unary, repeats_axis):
-            repeats, axis = repeats_axis
+        def test_grad(self, unary, repeats, axis):
             log_softmax_axis = -1 if axis is None else axis
             _test_unary_grad(
                 unary,

{deep_atomic-0.2.1 → deep_atomic-0.2.2}/uv.lock

@@ -186,7 +186,7 @@ toml = [
 
 [[package]]
 name = "deep-atomic"
-version = "0.2.1"
+version = "0.2.2"
 source = { editable = "." }
 dependencies = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/" }, marker = "python_full_version < '3.11'" },