Finegrain/refiners
1
0
Fork 0
mirror of https://github.com/finegrain-ai/refiners.git synced 2024-09-19 19:05:28 +00:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

add extra tests for Chain

Browse source
This commit is contained in:
Pierre Chapuis 2024-01-29 17:56:25 +01:00 committed by Cédric Deltheil
parent bca50b71f2
commit f43a530254
1 changed files with 31 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
tests/fluxion/layers/test_chain.py
View file

@ -21,6 +21,13 @@ def test_chain_find() -> None:
assert chain.find(fl.Conv2d) is None
def test_chain_getitem_accessor() -> None:
chain = fl.Chain(fl.Linear(1, 1), fl.Linear(1, 1))
assert chain["Linear_2"] == chain.Linear_2
with pytest.raises(KeyError):
chain["Linear_3"]
def test_chain_find_parent():
chain = fl.Chain(fl.Chain(fl.Linear(1, 1)))
@ -48,6 +55,30 @@ def test_chain_slice() -> None:
assert chain[:-1](x).shape == (1, 1)
def test_chain_walk() -> None:
chain = fl.Chain(
fl.Sum(
fl.Chain(fl.Linear(1, 1)),
fl.Linear(1, 1),
),
fl.Chain(),
)
assert list(chain.walk()) == [(chain.Sum, chain), (chain.Chain, chain)]
assert list(chain.walk(fl.Linear)) == [
(chain.Sum.Chain.Linear, chain.Sum.Chain),
(chain.Sum.Linear, chain.Sum),
]
assert list(chain.walk(recurse=True)) == [
(chain.Sum, chain),
(chain.Sum.Chain, chain.Sum),
(chain.Sum.Chain.Linear, chain.Sum.Chain),
(chain.Sum.Linear, chain.Sum),
(chain.Chain, chain),
]
def test_chain_walk_stop_iteration() -> None:
chain = fl.Chain(
fl.Sum(