Add more tests

[dirmeier]

No description provided.

[codecov]

Codecov Report

Attention: Patch coverage is 81.25000% with 3 lines in your changes missing coverage. Please review.

Project coverage is 80.98%. Comparing base (335bd17) to head (1c8c0de).
Report is 1 commits behind head on main.

Files	Patch %	Lines
sbijax/_src/nn/make_flow.py	57.14%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #49      +/-   ##
==========================================
+ Coverage   80.92%   80.98%   +0.05%     
==========================================
  Files          45       45              
  Lines        1814     1809       -5     
==========================================
- Hits         1468     1465       -3     
+ Misses        346      344       -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.