Repository: ivannz/cplxmodule (language: python)

Issue:
Nit picks / Bugs
1. Module should have a `cplxmodule.__version__` defined.
2. In your `setup.py` set the `version=version` and use this up top somewhere. This might look dirty but believe me it's the best way. We use this for [gnuradio/sigmf](https://github.com/gnuradio/SigMF/blob/sigmf-v1.x/setup.py#L15) and other projects:
```python
import os

with open(os.path.join('cplxmodule', '__init__.py')) as derp:
    version = re.search(r'__version__\s*=\s*[\'"]([^\'"]*)[\'"]', derp.read()).group(1)
```
3. I'm fairly certain that [this line](https://github.com/ivannz/cplxmodule/blob/ff051d96a7465ce7f155850a90a3664101f2b032/cplxmodule/nn/modules/conv.py#L108) should be changed to allow passing a tuple of padding parameters like (5,7,) or whatever to `padding`:
```python
self.stride[0], self.padding[0], self.dilation[0], # from this
self.stride, self.padding, self.dilation, # to this
```
4. ~~Due to the way that you are using relative imports (correctly) this project is compatible with Python 3.7+ and NOT prior versions. You can add an indicator for this in the setup.~~ This was actually a bug related to your pip package reporting the same version as the version installed from git, but they were different and had different submodules. More reason to properly version.

I probably have a PR for you for a different feature, but I have to get it approved for release first.
Nit picks / Bugs
1. Module should have a `cplxmodule.__version__` defined.
2. In your `setup.py` set the `version=version` and use this up top somewhere. This might look dirty but believe me it's the best way. We use this for [gnuradio/sigmf](https://github.com/gnuradio/SigMF/blob/sigmf-v1.x/setup.py#L15) and other projects:
```python
import os

with open(os.path.join('cplxmodule', '__init__.py')) as derp:
    version = re.search(r'__version__\s*=\s*[\'"]([^\'"]*)[\'"]', derp.read()).group(1)
```
3. I'm fairly certain that [this line](https://github.com/ivannz/cplxmodule/blob/ff051d96a7465ce7f155850a90a3664101f2b032/cplxmodule/nn/modules/conv.py#L108) should be changed to allow passing a tuple of padding parameters like (5,7,) or whatever to `padding`:
```python
self.stride[0], self.padding[0], self.dilation[0], # from this
self.stride, self.padding, self.dilation, # to this
```
4. ~~Due to the way that you are using relative imports (correctly) this project is compatible with Python 3.7+ and NOT prior versions. You can add an indicator for this in the setup.~~ This was actually a bug related to your pip package reporting the same version as the version installed from git, but they were different and had different submodules. More reason to properly version.

I probably have a PR for you for a different feature, but I have to get it approved for release first.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 1b2efea..b6abad8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,5 @@
+cplxmodule/__version__.py
+
 # General
 .DS_Store
 .AppleDouble
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..4e68853
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,17 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+-   repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+    -   id: black
+-   repo: https://gitlab.com/pycqa/flake8
+    rev: 3.9.2
+    hooks:
+    -   id: flake8
+        types: [file, python]
+        args: [--max-line-length=88, --extend-ignore=E203]
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b20e00e..2255d18 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,50 +1,71 @@
-# Version 2020.08.17
+# Changelog
 
-* FIX: Fixed shape mismatch in `nn.init.cplx_trabelsi_independent_`, which prevented it from working properly [# 11](https://github.com/ivannz/cplxmodule/issues/11)
-* ENH: [Hendrik Schröter](https://github.com/Rikorose) implemented Complex Transposed Convolutions [# 8](https://github.com/ivannz/cplxmodule/pull/8), squeeze/unsqueeze methods for `Cplx` [# 7](https://github.com/ivannz/cplxmodule/pull/7), and added support for `view` and `view_as` methods for `Cplx` [# 6](https://github.com/ivannz/cplxmodule/pull/6)
-* ENH: Introduce converters for special torch format of complex tensors (last dim is exactly 2) see [torch.fft](https://pytorch.org/docs/stable/generated/torch.fft.html#torch.fft)
-* ENH: `Cplx` now also has `.size()` method, which mimics `torch.Tensor.size()`
-* DOC: Improved documentation of `nn.casting` modules
+## Version 2022.06
 
+- COSMIT: accepted [Black](https://black.readthedocs.io/en/stable/index.html) a the code style of choice, introduced pre-commit hooks for developers
+- FIX: having a dunder-version in the root of the package is a the standard (issue #24)
+- FIX: set the minimal python to `3.7` as pointed out in issue #24
+- UPD: bumped the base version of torch to at least `1.8`
+- FIX: upgraded `.utils.spectrum` to new native torch complex backend (`torch>=1.8`)
+- FIX: ensured ONNX support in PR #14
+- ENH: implemented modulus-based maxpooling, requested in issue #17
+- FIX: made `.Cplx` instances `deepcopy`-able, fixing issue #18
+- DOC: improved docs for `.nn.ModReLU` indicating the sign-deviation from the original paper proposing it (issue #22)
+- DOC: added a basic TOC to the main README docs
 
-# Version 2020.08
+### Completed DEPRECATION cycles
 
-* structure of the `.nn.relevance` was simplified
-    - importing from `ard` has been deprecated, and ARD layers have been moved
-    to `.real` or `.complex` depending on their type
-* changed relevance layers class hierarchy in `.relevance.real` and `.relevance.complex`:
-    - factored out Gaussian Local Reparameterization into pure `*Gaussian` layers,
-    that reside in `real.base` and `complex.base`
-    - subclassed Variational Dropout layers (`*VD`) from `*Gaussian` with improper prior KL mixin
-    - subclassed ARD layers (`*ARD`) from Variational Dropout layers `*VD` with ARD Gaussian prior KL mixin
+- misnamed VD and misplaced ARD layers in `.nn.relevance`
+- sparsity stats badly placed in `.utils.stats`
+- misnamed $\ell_0$ probabilistic pruning layer in `.nn.relevance.extensions.real`, since it had nothing to do with the Automatic Relevance Determination Bayesian approach
 
+## Version 2020.08.17
 
-# Version 2020.03
+- FIX: Fixed shape mismatch in `.nn.init.cplx_trabelsi_independent_`, which prevented it from working properly [# 11](https://github.com/ivannz/cplxmodule/issues/11)
+- ENH: [Hendrik Schröter](https://github.com/Rikorose) implemented Complex Transposed Convolutions [# 8](https://github.com/ivannz/cplxmodule/pull/8), squeeze/unsqueeze methods for `Cplx` [# 7](https://github.com/ivannz/cplxmodule/pull/7), and added support for `.view` and `.view_as` methods for `Cplx` [# 6](https://github.com/ivannz/cplxmodule/pull/6)
+- ENH: Introduce converters for special torch format of complex tensors (last dim is exactly 2) see [torch.fft](https://pytorch.org/docs/stable/generated/torch.fft.html#torch.fft)
+- ENH: `Cplx` now also has `.size()` method, which mimics `torch.Tensor.size()`
+- DOC: Improved documentation of `.nn.casting` modules
 
-## Major changes in `.nn`
-* The structure of the `.nn` sub-module now more closely resembles that of `torch`
-    _ `.base` : `CplxToCplx` and parameter type `CplxParameter`
-    _ `.casting` : real-Cplx tensor conversion layers
-    _ `.linear`, `.conv`, `.activation` : essential layers and activations
-    _ `.container` : sequential container which explicitly checks types of internal layers
-    _ `.extra` : 1-dim Bernoulli Dropout for complex-valued tensors (Cplx)
-* `CplxToCplx` can now promote torch's univariate functions to split-complex activations, e.g. use `CplxToCplx[AvgPoool1d]` instead of `CplxAvgPool1d`
-* Niche complex-valued containers were removed, dropped dedicated activations, like `CplxLog` and `CplxExp`
+## Version 2020.08
 
+- structure of the `.nn.relevance` was simplified
+  - importing from `nn.relevance.ard` has been deprecated, and ARD layers have been moved
+  to `.real` or `.complex` depending on their type
+- changed relevance layers class hierarchy in `.relevance.real` and `.relevance.complex`:
+  - factored out Gaussian Local Reparameterization into pure `*Gaussian` layers,
+  that reside in `.real.base` and `.complex.base`
+  - subclassed Variational Dropout layers (`*VD`) from `*Gaussian` with improper prior KL mixin
+  - subclassed ARD layers (`*ARD`) from Variational Dropout layers `*VD` with ARD Gaussian prior KL mixin
 
-## Major changes in `.nn.relevance`
-* misnamed Bayesian layers in `.nn.relevance` were moved around and corrected
-    - layers in `.real` and `.complex` were renamed to Var Dropout, with deprecation warnings for old names
-    - `.ard` implements the Bayesian Dropout methods with Automatic Relevance Determination priors
-* `.extensions` submodule contains relaxations, approximations, and related but non-Bayesian layers
-    - `\ell_0` stochastic regularization layer was moved to `.real`
-    - `Lasso` was kept to illustrate extensibility, but similarly moved to `.real`
-    - Variational Dropout approximations and speeds ups were moved to `.complex`
+## Version 2020.03
 
+### Major changes in `.nn`
 
-## Enhancements
-* `CplxParameter` now supports real-to-complex promotion during `load_state_dict`
-* added submodule-specific readme's, explaining typical use cases and peculiarities
+- The structure of the `.nn` sub-module now more closely resembles that of `torch`
+  - `.base` : `CplxToCplx` and parameter type `CplxParameter`
+  - `.casting` : real-Cplx tensor conversion layers
+  - `.linear`, `.conv`, `.activation` : essential layers and activations
+  - `.container` : sequential container which explicitly checks types of internal layers
+  - `.extra` : 1-dim Bernoulli Dropout for complex-valued tensors (Cplx)
+- `CplxToCplx` can now promote torch's univariate functions to split-complex activations, e.g. use `CplxToCplx[AvgPoool1d]` instead of `CplxAvgPool1d`
+- Niche complex-valued containers were removed, dropped dedicated activations, like `CplxLog` and `CplxExp`
 
-# Prior to 2020.03
-Prior version used different version numbering and although the layers are backwards compatible, their location within the library was much different.
\ No newline at end of file
+### Major changes in `.nn.relevance`
+
+- misnamed Bayesian layers in `.nn.relevance` were moved around and corrected
+  - layers in `.real` and `.complex` were renamed to Var Dropout, with deprecation warnings for old names
+  - `.ard` implements the Bayesian Dropout methods with Automatic Relevance Determination priors
+- `.extensions` submodule contains relaxations, approximations, and related but non-Bayesian layers
+  - `\ell_0` stochastic regularization layer was moved to `.real`
+  - `Lasso` was kept to illustrate extensibility, but similarly moved to `.real`
+  - Variational Dropout approximations and speeds ups were moved to `.complex`
+
+### Enhancements
+
+- `CplxParameter` now supports real-to-complex promotion during `.load_state_dict`
+- added submodule-specific README's, explaining typical use cases and peculiarities
+
+## Prior to 2020.03
+
+Prior version used different version numbering and although the layers are backwards compatible, their location within the library was much different.
diff --git a/LICENSE b/LICENSE
index d2b2996..8e6650a 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2019 Ivan Nazarov
+Copyright (c) 2019-present Ivan Nazarov
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/PLAN.md b/PLAN.md
index 7682ad6..5ec825f 100644
--- a/PLAN.md
+++ b/PLAN.md
@@ -1,108 +1,101 @@
-# Plan for the 2020.09 release
-
-[ ] update the complex layer initialization from kaiming to independent by default (check Trabelsi et al. 2018)
-
-[+] deal with the `torch.nonzero(..., as_tuple=True)` deprecation warning in `utils.spectrum`
-
-[ ] Improve implementation
-* (Bernoulli Dropout) need 1d (exists), 2d and 3d
-* ~(Convolutions) implement 3d convolutions and 3d VarDropout convolutions both real and complex~
-* ~(Transposed Convolutions) figure out the math and implement var dropout for transposed convos~
-
-[+] FIgure out the issues with ONNX support
-
-[ ] begin migration to `complex` tensors in `pytorch>=1.6`
-* For C->R real-valued loss functions [grad.conj()](https://pytorch.org/docs/master/complex_numbers.html#autograd) gives a descent direction.
-* complex [autograd](https://pytorch.org/docs/master/notes/autograd.html#complex-autograd-doc)
-
-[ ] Consider replacing `Real` with `Tensor` in format-conversion layers, like `RealToCplx`, `CplxToReal`
-* the term Real has connotations with real numbers, making it very unintuitive to convert between Cplx, which is perceived as a complex number, to a torch Tensor, which serves merely as a storage format.
-* need a deprecation cycle for these and related functions
-  * in `cplx`: `from_interleaved_real`, `from_concatenated_real`, `to_interleaved_real`, `to_concatenated_real`, aliases `from_real` and `to_real` (affects `__init__.py`)
-  * `nn.modules.casting`: `InterleavedRealToCplx`, `ConcatenatedRealToCplx`, `CplxToInterleavedReal`, `CplxToConcatenatedReal`, also base classes `BaseRealToCplx` and `BaseCplxToReal`.
-    * three basic types? Tensor -- aka Storage, Real -- real-valued tensor, Cplx -- complex-valued tensor
-
-[ ] Implement scheduled mag-pruning of [Zhu and Gupta (2017)](http://arxiv.org/abs/1710.01878) or thresholded of [Wu et al. (2019)](https://arxiv.org/abs/1903.02358).
-* use `nn.masked` as a `backend` -- this will automatically support real and Cplx layers!!!!
-* implement as either wrapper around optimizer (bad), or as a separate entity (better)
-  * settings of the target sparsity per eligible layer (`dict`)
-  * method `.step()` which updates the masks according to the schedule and the current sorted magnitudes of the parameters
-
-
-# Plan for the first release
-
-[+] make load_state_dict respect components of CplxParameter and allow promoting real-tensors to complex-tensors provided the state dict has no .real or .imag, but a correct key referring to the parameter.
-
-[+] fix the incorrect naming of bayesain methods in `nn.relevance`
-* rename `*ARD` named layers in `.real` and `.complex` to `*VD` layers, since they use log-uniform prior and thus are in fact Variational Dropout layers
-* start deprecating importing `*ARD` named layers from `.real` and `.complex`
-* fix aliases of imported layers in `.extensions`
-* expose all base VD/ARD layers in `\_\_init\_\_.py` and require importing modifications from `.extensions`
-* fix the text in nn/relevance/README.md
-
-[+] fix the names for L0 regularized layer whith in fact performs probabilistic sparsification, and is not related to Variational inference
-
-[+] check if `setup.py` has correct requirements and specifiy them explicitly
-* `requires` is not a keyword, use `install_requires` and `tests_require`
-
-[+] investigate reordering base classes in `LinearMasked(MaskedWeightMixin, Linear, _BaseRealMixin)` and similar in `nn.masked`.
-* could moving it further into the bases result in a slower property lookup? It seems no:
-  * from python decsriptors [doc](https://docs.python.org/3/howto/descriptor.html)
-  > The implementation works through a precedence chain that gives data descriptors priority over instance variables, instance variables priority over non-data descriptors, and assigns lowest priority to \_\_getattr\_\_
-  * lookup order is thus by \_\_getattribute\_\_: descriptors (aka @property), instance \_\_dict\_\_, class attributes \_\_dict\_\_, and lastly \_\_getattr\_\_.
-* moved MaskedWeightMixin into \_BaseMixin
-
-[+] get rid of `torch_module` from `.utils` and declare `activations` explicitly
-
-[+] clean up the `nn` module itself
-* remove crap from `.sequential`: `CplxResidualBottleneck`, `CplxResidualSequential` and CplxBusResidualSequential must go, and move CplxSequential to base layers
-* split `.layers`, `.activation`, and `.sequential`
-  * `.modules.base` : base classes (CplxToCplx, BaseRealToCplx, BaseCplxToReal), and parameter type (CplxParameter, CplxParameterAccessor)
-  * `.modules.casting` : converting real tensors in various formats to and from Cplx (InterleavedRealToCplx, ConcatenatedRealToCplx, CplxToInterleavedReal, CplxToConcatenatedReal, AsTypeCplx)
-  * `.modules.linear` : Linear, Bilinear, Identity, PhaseShift
-  * `.modules.conv` : everything convolutional
-  * `.modules.activation` : activations (CplxModReLU, CplxAdaptiveModReLU, CplxModulus, CplxAngle) and layers (CplxReal, CplxImag)
-  * `.modules.container` : CplxSequential
-  * `.modules.extra` : Dropout, AvgPool1d
-* move `.batchnorm` to modules, keep `.init` in `.nn`
-* fix imports from adjacent modules: `nn.masked` and `nn.relevance`.
-
-[+] in `nn.relevance.complex` : drop `Cplx(*map(torch.randn_like, (s2, s2)))` and write `Cplx(torch.randn_like(s2), torch.randn_like(s2))` explicitly
-* implemented `cplx.randn` and `cplx.randn_like`
-
-[+] residual clean up in `nn` module
-* `.activation` : `CplxActivation` is the same as CplxToCplx[...]
-  * CplxActivation promotes classic (real) torch functions to split activations, so yes.
-  * See if it is possible to implement function promotion through CplxToCplx[...]
-    * it is possbile: just reuse CplxActivation
-  * Currently CplxToCplx promotes layers and real functions to inpdependently applied layers/functions (split)
-    * how should we proceed with cplx. trig functions? a wrapper, or hardcoded activations?
-      * the latter seems more natural, as the trig functions are vendored by this module
-      * since torch is the base, and implements a great number of univariate tensor functions and could potentially be extended, it is more natural to use a wrapper (rationale behind CplxToCplx[...]).
-* `.modules.extra` : this needs thorough cleaning
-  * drop CplxResidualBottleneck, CplxResidualSequential and CplxBusResidualSequential
-  * abandon `torch_module` and code the trig activations by hand.
-  * remove alias CplxDropout1d : use torch.nn names as much as possible
-  * deprecate CplxAvgPool1d: it can be created in runtime with CplxToCplx\[torch.nn.AvgPool1d\]
-
-[+] documentation for bayesian and maskable layers
-* in `nn.relevance.base`, making it like in `nn.masked`
-* classes in `nn.relevance`  `.real` and `.complex` should be also documented properly, the same goes for `.extensions`
-
-[+] restrucure the extensions and non-bayesian layers
-* new folder structure
-  * take ard-related declarations and move them to `relevance/ard.py`, everythin else to a submodule
-  * `.extensions` submodule:
-    * `complex` for cplx specific etended layers: bogus penalties, approximations and other stuff, -- not directly related to variational dropout or automatic relevance determination
-    * `real` for supplementary real-valued layers
-* decide the fate of `lasso` class in `nn.relevance`:
-  * it is irrelevant to Bayesian methods: move it to `extensions/real`
-
-[+] documentation
-* go through README-s in each submodule to make sure that info there is correct and typical use cases described
-* `nn.init` : document the initializations according to Trabelsi et al. (2018)
-  * seems to be automatically documented using `functools.wraps` from the original `torch.nn.init` procedures.
-
-[+] add missing tests to the unit test suite
-* tests for `*state_dict` api compliance of `nn.masked` and `nn.base.CplxParameter`
-  * implementing these test helped figure out and fix edge cases and fix them, so yay for TDD!
+# Plan
+
+## Plans for the 2020.09 release
+
+- [ ] Extend the implementation
+  - [ ] (Bernoulli Dropout) need 1d (exists), 2d and 3d
+  - [x] (Convolutions) implement 3d convolutions and 3d Variational Dropout convolutions both real and complex
+  - [x] (Transposed Convolutions) figure out the math and implement var dropout for transposed convos
+- [ ] update the complex layer initialization from Kaiming to independent by default (check Trabelsi et al. 2018)
+  - [ ] this may break older experiments in third-party repos, so need to issue a warning and a patch
+- [x] deal with the `torch.nonzero(..., as_tuple=True)` deprecation warning in `utils.spectrum`
+- [x] figure out the issues with ONNX support
+
+## Plans for the first release
+
+- [x] make `.load_state_dict` respect components of CplxParameter and allow promoting real-tensors to complex-tensors provided the state dict has no `.real` or `.imag`, but a correct key referring to the parameter.
+- [x] fix the incorrect naming of Bayesain methods in `nn.relevance`
+  - rename `*ARD` named layers in `.real` and `.complex` to `*VD` layers, since they use log-uniform prior and thus are in fact Variational Dropout layers
+  - start deprecating importing `*ARD` named layers from `.real` and `.complex`
+  - fix aliases of imported layers in `.extensions`
+  - expose all base VD/ARD layers in `\_\_init\_\_.py` and require importing modifications from `.extensions`
+  - fix the text in `nn/relevance/README.md`
+- [x] fix the names for `L0` regularized layer which in fact performs *probabilistic sparsification*, and is not related to Variational inference
+- [x] check if `setup.py` has correct requirements and specify them explicitly
+  - `requires` is not a keyword, use `install_requires` and `tests_require`
+- [x] investigate reordering base classes in `LinearMasked(MaskedWeightMixin, Linear, _BaseRealMixin)` and similar in `nn.masked`.
+  - could moving it further into the bases result in a slower property lookup? It seems no:
+    - from python descriptors [doc](https://docs.python.org/3/howto/descriptor.html)
+    > The implementation works through a precedence chain that gives data descriptors priority over instance variables, instance variables priority over non-data descriptors, and assigns lowest priority to \_\_getattr\_\_
+    - lookup order is thus by \_\_getattribute\_\_: descriptors (aka @property), instance \_\_dict\_\_, class attributes \_\_dict\_\_, and lastly \_\_getattr\_\_.
+  - moved MaskedWeightMixin into \_BaseMixin
+- [x] get rid of `torch_module` from `.utils` and declare `activations` explicitly
+- [x] clean up the `nn` module itself
+  - remove crap from `.sequential`: `CplxResidualBottleneck`, `CplxResidualSequential` and CplxBusResidualSequential must go, and move CplxSequential to base layers
+  - split `.layers`, `.activation`, and `.sequential`
+    - `.modules.base` : base classes (`CplxToCplx`, `BaseRealToCplx`, `BaseCplxToReal`), and parameter type (`CplxParameter`, `CplxParameterAccessor`)
+    - `.modules.casting` : converting real tensors in various formats to and from `Cplx` (`InterleavedRealToCplx`, `ConcatenatedRealToCplx`, `CplxToInterleavedReal`, `CplxToConcatenatedReal`, `AsTypeCplx`)
+    - `.modules.linear` : `Linear`, `Bilinear`, `Identity`, `PhaseShift`
+    - `.modules.conv` : everything convolutional
+    - `.modules.activation` : activations (`CplxModReLU`, `CplxAdaptiveModReLU`, `CplxModulus`, `CplxAngle`) and layers (`CplxReal`, `CplxImag`)
+    - `.modules.container` : `CplxSequential`
+    - `.modules.extra` : `Dropout`, `AvgPool1d`
+  - move `.batchnorm` to modules, keep `.init` in `.nn`
+  - fix imports from adjacent modules: `nn.masked` and `nn.relevance`.
+- [x] in `nn.relevance.complex` : drop `Cplx(*map(torch.randn_like, (s2, s2)))` and write `Cplx(torch.randn_like(s2), torch.randn_like(s2))` explicitly
+  - implemented `cplx.randn` and `cplx.randn_like`
+- [x] residual clean up in `nn` module
+  - `.activation` : `CplxActivation` is the same as `CplxToCplx[...]`
+    - `CplxActivation` promotes classic (real) torch functions to split activations, so yes.
+    - See if it is possible to implement function promotion through `CplxToCplx[...]`
+      - it is possible: just reuse `CplxActivation`
+    - Currently `CplxToCplx` promotes layers and real functions to independently applied layers/functions (split)
+      - how should we proceed with `Cplx` trig. functions? a wrapper, or hardcoded activations?
+        - the latter seems more natural, as the trig. functions are vendored by this module
+        - since torch is the base, and implements a great number of univariate tensor functions and could potentially be extended, it is more natural to use a wrapper (rationale behind `CplxToCplx[...]`).
+  - `.modules.extra`: this needs thorough cleaning
+    - drop `CplxResidualBottleneck`, `CplxResidualSequential` and `CplxBusResidualSequential`
+    - abandon `torch_module` and code the trig. activations by hand.
+    - remove alias `CplxDropout1d`: use torch.nn names as much as possible
+    - deprecate `CplxAvgPool1d`: it can be created in runtime with `CplxToCplx\[torch.nn.AvgPool1d\]`
+- [x] documentation for Bayesian and maskable layers
+  - in `nn.relevance.base`, making it like in `nn.masked`
+  - classes in `nn.relevance`  `.real` and `.complex` should be also documented properly, the same goes for `.extensions`
+- [x] restructure the extensions and non-Bayesian layers
+  - new folder structure
+    - take ARD-related declarations and move them to `relevance/ard.py`, everything else to a submodule
+    - `.extensions` submodule:
+      - `complex` for `Cplx`-specific extended layers: bogus penalties, approximations and other stuff, -- not directly related to Variational Dropout or automatic relevance determination
+      - `real` for supplementary real-valued layers
+  - decide the fate of `lasso` class in `nn.relevance`:
+    - it is irrelevant to Bayesian methods: move it to `extensions/real`
+- [x] documentation
+  - go through README-s in each submodule to make sure that info there is correct and typical use cases are described
+  - `nn.init` : document the initializations according to Trabelsi et al. (2018)
+    - seems to be automatically documented using `functools.wraps` from the original `torch.nn.init` procedures.
+- [x] add missing tests to the unit test suite
+  - tests for `*state_dict` api compliance of `nn.masked` and `nn.base.CplxParameter`
+    - implementing these test helped figure out and fix edge cases and fix them, so yay for TDD!
+
+## Abandoned plans
+
+The following ideas were at some point planned out, but ultimately were not pursued due to various reasons.
+
+- [ ] begin migration to `complex` tensors in `pytorch>=1.6`
+  - #20220608 consider the discussion in [issue #21](https://github.com/ivannz/cplxmodule/issues/21)
+  - For C->R real-valued loss functions [grad.conj()](https://pytorch.org/docs/master/complex_numbers.html#autograd) gives a descent direction.
+  - complex [autograd](https://pytorch.org/docs/master/notes/autograd.html#complex-autograd-doc)
+- [ ] Consider replacing `Real` with `Tensor` in format-conversion layers, like `RealToCplx`, `CplxToReal`
+  - the term `Real` has connotations with real numbers, making it very *unintuitive* to convert between `Cplx`, which is perceived as a complex number, to a torch `Tensor`, which serves merely as a storage format.
+  - need a deprecation cycle for these and related functions
+    - in `cplx`: `from_interleaved_real`, `from_concatenated_real`, `to_interleaved_real`, `to_concatenated_real`, aliases `from_real` and `to_real` (affects `__init__.py`)
+    - `nn.modules.casting`: `InterleavedRealToCplx`, `ConcatenatedRealToCplx`, `CplxToInterleavedReal`, `CplxToConcatenatedReal`, also base classes `BaseRealToCplx` and `BaseCplxToReal`.
+    - three basic types?
+      - Tensor -- aka Storage
+      - Real -- real-valued tensor
+      - Cplx -- complex-valued tensor
+- [ ] Implement scheduled mag-pruning of [Zhu and Gupta (2017)](http://arxiv.org/abs/1710.01878) or thresholded of [Wu et al. (2019)](https://arxiv.org/abs/1903.02358).
+  - use `nn.masked` as a `backend` -- this will automatically support real and `Cplx` layers!!!!
+  - implement as either wrapper around optimizer (bad), or as a separate entity (better)
+    - settings of the target sparsity per eligible layer (`dict`)
+    - method `.step()` which updates the masks according to the schedule and the current sorted magnitudes of the parameters
diff --git a/README.md b/README.md
index 8b97707..50c66a0 100644
--- a/README.md
+++ b/README.md
@@ -1,12 +1,14 @@
 # CplxModule
 
-A lightweight extension for `torch.nn` that adds layers and activations, which respect algebraic operations over the field of complex numbers, and implements real- and complex-valued Variational Dropout methods for weight sparsification.
-
+A lightweight extension for `torch.nn` that adds layers and activations, which respect algebraic operations over the *field of complex numbers*, and implements *real- and complex-valued Variational Dropout* methods for weight sparsification.
 
 ## Documentation
 
-Please refer to README files located in `cplxmodule.nn`, `cplxmodule.nn.relevance`, and `cplxmodule.nn.masked` for a high-level description of the implementation, functionality and useful code patterns.
+For a high-level description of the implementation, functionality and useful code patterns, please refer to the following READMEs
 
+- [cplxmodule.nn](./cplxmodule/nn) the implemented complex-valued layers and their basic use
+- [cplxmodule.nn.relevance](./cplxmodule/nn/relevance) the *plug-and-play* layers for Variational Dropout and how to use them ([[3]](#user-content-ref3), [[4]](#user-content-ref4), [[5]](#user-content-ref5)).
+- [cplxmodule.nn.masked](./cplxmodule/nn/masked) supported masked layers for fine-tuning pruned networks and how to migrate parameters between classic `torch.nn` layers
 
 ## Implementation
 
@@ -18,22 +20,76 @@ The implementations of the real-valued Variational Dropout and Automatic Relevan
 
 Complex-valued Bayesian sparsification layers are based on the research by [Nazarov and Burnaev (2020)](http://proceedings.mlr.press/v119/nazarov20a.html) [[5]](#user-content-ref5).
 
-
 ## Installation
 
-You can install this package with `pip`:
+The essential dependencies of `cplxmodule` are `numpy`, `torch` and `scipy`, which can be installed via
+
+```bash
+# essential dependencies
+# conda update -n base -c defaults conda
+conda create -n cplxmodule "python>=3.7" pip numpy scipy "pytorch::pytorch" \
+  && conda activate cplxmodule
+```
+
+Extra dependencies, that are used in tests and needed for development, can be added on top of the essentials. Check [ONNX Runtime](https://onnxruntime.ai/) to see of your system is compatible.
+
 ```bash
+conda activate cplxmodule
+
+# extra deps for development
+conda install -n cplxmodule matplotlib scikit-learn tqdm pytest "pytorch::torchvision" \
+  && pip install black pre-commit
+
+# ONNX (for compatible systems)
+conda install -n cplxmodule onnx && pip install onnxruntime
+```
+
+The package itself can be installed this package with `pip`:
+
+```bash
+conda activate cplxmodule
+
 pip install cplxmodule
 ```
+
 or from the git repo to get the latest version:
+
 ```bash
+conda activate cplxmodule
+
 pip install --upgrade git+https://github.com/ivannz/cplxmodule.git
 ```
-If you prefer a developer install (editable), then run the following from the root of the locally cloned repo
+
+or locally from *the root of the locally cloned repo*, if you prefer an editable developer install:
+
 ```bash
+conda activate cplxmodule
+
+# enable basic checks (codestyle, stray whitespace, eof newline)
+pre-commit install
+
+# editable install
 pip install -e .
+
+# run tests to verify installation (batchnorm test )
+# XXX `test_batchnorm.py` depends on the precision of the outcome of SGD, hence
+#  may occasionally fail
+# XXX A user warning concerning non-writable numpy array is expected
+pytest
 ```
 
+Additionally, you may want to study the following examples and test [Variational Dropout](./cplxmodule/nn/relevance):
+
+```bash
+conda activate cplxmodule
+
+# test real- and complex-valued Bayesian sparsification layers
+python tests/test_relevance.py
+
+# showcase the train-sparisify-fine-tune staged pipeline on a basic
+#  real-valued CNN on MNIST
+python tests/test_mnist.py
+```
 
 ## Citation
 
@@ -56,7 +112,6 @@ The proper citation for the real-valued Bayesian Sparsification layers from `cpl
 }
 ```
 
-
 ## References
 
 <a id="user-content-ref1">[1]</a>
diff --git a/VERSION b/VERSION
new file mode 100644
index 0000000..97b50db
--- /dev/null
+++ b/VERSION
@@ -0,0 +1,1 @@
+2022.06
\ No newline at end of file
diff --git a/cplxmodule/__init__.py b/cplxmodule/__init__.py
index e4eea82..38249a3 100644
--- a/cplxmodule/__init__.py
+++ b/cplxmodule/__init__.py
@@ -1,4 +1,5 @@
 """Complex-valued modules for pytorch."""
+from .__version__ import __version__
 
 from .cplx import Cplx, from_real, to_real
 
diff --git a/cplxmodule/cplx.py b/cplxmodule/cplx.py
index 147b21c..d67e873 100644
--- a/cplxmodule/cplx.py
+++ b/cplxmodule/cplx.py
@@ -1,4 +1,3 @@
-import warnings
 from copy import deepcopy
 
 import torch
@@ -32,22 +31,21 @@ def __new__(cls, real, imag=None):
                 imag = 0.0
 
             elif not isinstance(imag, float):
-                raise TypeError("""Imaginary part must be float.""")
+                raise TypeError("Imaginary part must be float.")
 
             real, imag = torch.tensor(real), torch.tensor(imag)
 
         elif not isinstance(real, torch.Tensor):
-            raise TypeError("""Real part must be torch.Tensor.""")
+            raise TypeError("Real part must be torch.Tensor.")
 
         if imag is None:
             imag = torch.zeros_like(real)
 
         elif not isinstance(imag, torch.Tensor):
-            raise TypeError("""Imaginary part must be torch.Tensor.""")
+            raise TypeError("Imaginary part must be torch.Tensor.")
 
         if real.shape != imag.shape:
-            raise ValueError("""Real and imaginary parts have """
-                             """mistmatching shape.""")
+            raise ValueError("Real and imaginary parts have mistmatching shape.")
 
         self = super().__new__(cls)
         self.__real, self.__imag = real, imag
@@ -142,8 +140,10 @@ def __mul__(u, v):
         # (a + ib) (u + iv) = au - bv + i(av + bu)
         # (a+u)(b+v) = ab + uv + (av + ub)
         # (a-v)(b+u) = ab - uv + (au - vb)
-        return type(u)(u.__real * v.real - u.__imag * v.imag,
-                       u.__imag * v.real + u.__real * v.imag)
+        return type(u)(
+            u.__real * v.real - u.__imag * v.imag,
+            u.__imag * v.real + u.__real * v.imag,
+        )
 
     __rmul__ = __mul__
     __imul__ = __mul__
@@ -225,8 +225,10 @@ def h(self):
         return self.conj.t()  # Cplx(self.__real.t(), -self.__imag.t())
 
     def flatten(self, start_dim=0, end_dim=-1):
-        return type(self)(self.__real.flatten(start_dim, end_dim),
-                          self.__imag.flatten(start_dim, end_dim))
+        return type(self)(
+            self.__real.flatten(start_dim, end_dim),
+            self.__imag.flatten(start_dim, end_dim),
+        )
 
     def view(self, *shape):
         r"""Return a view of the complex tensor."""
@@ -248,7 +250,9 @@ def size(self, *dim):
         return self.__real.size(*dim)
 
     def squeeze(self, dim=None):
-        r"""Returns the complex tensor with all the dimensions of input of size 1 removed."""
+        r"""Returns the complex tensor with all the dimensions of input of
+        size one removed.
+        """
         if dim is None:
             return type(self)(self.__real.squeeze(), self.__imag.squeeze())
         else:
@@ -257,7 +261,9 @@ def squeeze(self, dim=None):
             )
 
     def unsqueeze(self, dim=None):
-        r"""Returns a new complex tensor with a dimension of size one inserted at the specified position."""
+        r"""Returns a new complex tensor with a dimension of size one inserted
+        at the specified position.
+        """
         if dim is None:
             return type(self)(self.__real.unsqueeze(), self.__imag.unsqueeze())
         else:
@@ -281,8 +287,10 @@ def numpy(self):
         return self.__real.numpy() + 1j * self.__imag.numpy()
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(\n" \
-               f"  real={self.__real},\n  imag={self.__imag}\n)"
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  real={self.__real},\n  imag={self.__imag}\n)"
+        )
 
     def detach(self):
         r"""Return a copy of the complex tensor detached from autograd graph."""
@@ -290,8 +298,10 @@ def detach(self):
 
     def requires_grad_(self, requires_grad=True):
         r"""Toggle the gradient of real and imaginary parts."""
-        return type(self)(self.__real.requires_grad_(requires_grad),
-                          self.__imag.requires_grad_(requires_grad))
+        return type(self)(
+            self.__real.requires_grad_(requires_grad),
+            self.__imag.requires_grad_(requires_grad),
+        )
 
     @property
     def grad(self):
@@ -311,8 +321,9 @@ def cpu(self):
 
     def to(self, *args, **kwargs):
         r"""Move / typecast the complex tensor."""
-        return type(self)(self.__real.to(*args, **kwargs),
-                          self.__imag.to(*args, **kwargs))
+        return type(self)(
+            self.__real.to(*args, **kwargs), self.__imag.to(*args, **kwargs)
+        )
 
     @property
     def device(self):
@@ -330,13 +341,13 @@ def dim(self):
 
     def permute(self, *dims):
         r"""Shuffle the dimensions of the complex tensor."""
-        return type(self)(self.__real.permute(*dims),
-                          self.__imag.permute(*dims))
+        return type(self)(self.__real.permute(*dims), self.__imag.permute(*dims))
 
     def transpose(self, dim0, dim1):
         r"""Transpose the specified dimensions of the complex tensor."""
-        return type(self)(self.__real.transpose(dim0, dim1),
-                          self.__imag.transpose(dim0, dim1))
+        return type(self)(
+            self.__real.transpose(dim0, dim1), self.__imag.transpose(dim0, dim1)
+        )
 
     def is_complex(self):
         r"""Test if the tensor indeed represents a complex number."""
@@ -345,83 +356,96 @@ def is_complex(self):
     @classmethod
     def empty(cls, *sizes, dtype=None, device=None, requires_grad=False):
         r"""Create an empty complex tensor."""
-        re = torch.empty(*sizes, dtype=dtype, device=device,
-                         requires_grad=requires_grad)
+        re = torch.empty(
+            *sizes, dtype=dtype, device=device, requires_grad=requires_grad
+        )
         return cls(re, torch.empty_like(re, requires_grad=requires_grad))
 
     @classmethod
     def zeros(cls, *sizes, dtype=None, device=None, requires_grad=False):
         r"""Create an empty complex tensor."""
-        re = torch.zeros(*sizes, dtype=dtype, device=device,
-                         requires_grad=requires_grad)
+        re = torch.zeros(
+            *sizes, dtype=dtype, device=device, requires_grad=requires_grad
+        )
         return cls(re, torch.zeros_like(re, requires_grad=requires_grad))
 
     @classmethod
     def ones(cls, *sizes, dtype=None, device=None, requires_grad=False):
         r"""Create an empty complex tensor."""
-        re = torch.ones(*sizes, dtype=dtype, device=device,
-                        requires_grad=requires_grad)
+        re = torch.ones(*sizes, dtype=dtype, device=device, requires_grad=requires_grad)
         return cls(re, torch.zeros_like(re, requires_grad=requires_grad))
 
 
 def cat(tensors, dim):
     tensors = [*map(Cplx, tensors)]
-    return Cplx(torch.cat([z.real for z in tensors], dim=dim),
-                torch.cat([z.imag for z in tensors], dim=dim))
+    return Cplx(
+        torch.cat([z.real for z in tensors], dim=dim),
+        torch.cat([z.imag for z in tensors], dim=dim),
+    )
 
 
 def split(input, split_size_or_sections, dim=0):
     """see documentation for `torch.split`"""
-    return tuple(Cplx(re, im) for re, im in zip(
-        torch.split(input.real, split_size_or_sections, dim),
-        torch.split(input.imag, split_size_or_sections, dim),
-    ))
+    return tuple(
+        Cplx(re, im)
+        for re, im in zip(
+            torch.split(input.real, split_size_or_sections, dim),
+            torch.split(input.imag, split_size_or_sections, dim),
+        )
+    )
 
 
 def chunk(input, chunks, dim=0):
     """see documentation for `torch.chunk`"""
-    return tuple(Cplx(re, im) for re, im in zip(
-        torch.chunk(input.real, chunks, dim),
-        torch.chunk(input.imag, chunks, dim),
-    ))
+    return tuple(
+        Cplx(re, im)
+        for re, im in zip(
+            torch.chunk(input.real, chunks, dim),
+            torch.chunk(input.imag, chunks, dim),
+        )
+    )
 
 
 def stack(tensors, dim):
     tensors = [*map(Cplx, tensors)]
-    return Cplx(torch.stack([z.real for z in tensors], dim=dim),
-                torch.stack([z.imag for z in tensors], dim=dim))
+    return Cplx(
+        torch.stack([z.real for z in tensors], dim=dim),
+        torch.stack([z.imag for z in tensors], dim=dim),
+    )
 
 
 def unbind(input, dim=0):
     """see documentation for `torch.unbind`"""
-    return tuple(Cplx(re, im) for re, im in zip(
-        torch.unbind(input.real, dim),
-        torch.unbind(input.imag, dim),
-    ))
+    return tuple(
+        Cplx(re, im)
+        for re, im in zip(
+            torch.unbind(input.real, dim),
+            torch.unbind(input.imag, dim),
+        )
+    )
 
 
 def take(input, index):
     """see documentation for `torch.take`"""
-    return Cplx(torch.take(input.real, index),
-                torch.take(input.imag, index))
+    return Cplx(torch.take(input.real, index), torch.take(input.imag, index))
 
 
 def narrow(input, dim, start, length):
     """see documentation for `torch.narrow`"""
-    return Cplx(torch.narrow(input.real, dim, start, length),
-                torch.narrow(input.imag, dim, start, length))
+    return Cplx(
+        torch.narrow(input.real, dim, start, length),
+        torch.narrow(input.imag, dim, start, length),
+    )
 
 
 def squeeze(input, dim=None):
     """see documentation for `torch.squeeze`"""
-    return Cplx(torch.squeeze(input.real, dim),
-                torch.squeeze(input.imag, dim))
+    return Cplx(torch.squeeze(input.real, dim), torch.squeeze(input.imag, dim))
 
 
 def unsqueeze(input, dim):
     """see documentation for `torch.unsqueeze`"""
-    return Cplx(torch.unsqueeze(input.real, dim),
-                torch.unsqueeze(input.imag, dim))
+    return Cplx(torch.unsqueeze(input.real, dim), torch.unsqueeze(input.imag, dim))
 
 
 def from_interleaved_real(input, copy=True, dim=-1):
@@ -443,7 +467,7 @@ def to_interleaved_real(input, flatten=True, dim=-1):
     """Interleave the complex re-im pair into a real tensor."""
     dim = 1 + fix_dim(dim, input.dim())
     input = torch.stack([input.real, input.imag], dim=dim)
-    return input.flatten(dim-1, dim) if flatten else input
+    return input.flatten(dim - 1, dim) if flatten else input
 
 
 to_real = to_interleaved_real
@@ -458,8 +482,7 @@ def to_concatenated_real(input, flatten=None, dim=-1):
 def exp(input):
     r"""Compute the exponential of the complex tensor in re-im pair."""
     scale = torch.exp(input.real)
-    return Cplx(scale * torch.cos(input.imag),
-                scale * torch.sin(input.imag))
+    return Cplx(scale * torch.cos(input.imag), scale * torch.sin(input.imag))
 
 
 def log(input):
@@ -469,14 +492,18 @@ def log(input):
 
 def sin(input):
     r"""Compute the sine of the complex tensor in re-im pair."""
-    return Cplx(torch.sin(input.real) * torch.cosh(input.imag),
-                torch.cos(input.real) * torch.sinh(input.imag))
+    return Cplx(
+        torch.sin(input.real) * torch.cosh(input.imag),
+        torch.cos(input.real) * torch.sinh(input.imag),
+    )
 
 
 def cos(input):
     r"""Compute the cosine of the complex tensor in re-im pair."""
-    return Cplx(torch.cos(input.real) * torch.cosh(input.imag),
-                - torch.sin(input.real) * torch.sinh(input.imag))
+    return Cplx(
+        torch.cos(input.real) * torch.cosh(input.imag),
+        -torch.sin(input.real) * torch.sinh(input.imag),
+    )
 
 
 def tan(input):
@@ -489,8 +516,10 @@ def sinh(input):
 
     sinh(z) = - j sin(j z)
     """
-    return Cplx(torch.sinh(input.real) * torch.cos(input.imag),
-                torch.cosh(input.real) * torch.sin(input.imag))
+    return Cplx(
+        torch.sinh(input.real) * torch.cos(input.imag),
+        torch.cosh(input.real) * torch.sin(input.imag),
+    )
 
 
 def cosh(input):
@@ -498,8 +527,10 @@ def cosh(input):
 
     cosh(z) = cos(j z)
     """
-    return Cplx(torch.cosh(input.real) * torch.cos(input.imag),
-                torch.sinh(input.real) * torch.sin(input.imag))
+    return Cplx(
+        torch.cosh(input.real) * torch.cos(input.imag),
+        torch.sinh(input.real) * torch.sin(input.imag),
+    )
 
 
 def tanh(input):
@@ -512,8 +543,9 @@ def tanh(input):
 
 def randn(*size, dtype=None, device=None, requires_grad=False):
     """Generate standard complex Gaussian noise."""
-    normal = torch.randn(2, *size, dtype=dtype, layout=torch.strided,
-                         device=device, requires_grad=False) / sqrt(2)
+    normal = torch.randn(
+        2, *size, dtype=dtype, layout=torch.strided, device=device, requires_grad=False
+    ) / sqrt(2)
     z = Cplx(normal[0], normal[1])
     return z.requires_grad_(True) if requires_grad else z
 
@@ -522,10 +554,12 @@ def randn_like(input, dtype=None, device=None, requires_grad=False):
     """Returns a tensor with the same size as `input` that is filled with
     standard comlpex Gaussian random numbers.
     """
-    return randn(*input.size(),
-                 dtype=input.dtype if dtype is None else dtype,
-                 device=input.device if device is None else device,
-                 requires_grad=requires_grad)
+    return randn(
+        *input.size(),
+        dtype=input.dtype if dtype is None else dtype,
+        device=input.device if device is None else device,
+        requires_grad=requires_grad,
+    )
 
 
 def modrelu(input, threshold=0.5):
@@ -579,7 +613,7 @@ def modrelu(input, threshold=0.5):
     """
     # scale = (1 - \trfac{b}{|z|})_+
     modulus = torch.clamp(abs(input), min=1e-5)
-    return input * torch.relu(1. - threshold / modulus)
+    return input * torch.relu(1.0 - threshold / modulus)
 
 
 def phaseshift(input, phi=0.0):
@@ -604,10 +638,8 @@ def linear_naive(input, weight, bias=None):
     # W = U + i V,  z = u + i v, c = \Re c + i \Im c
     #  W z + c = (U + i V) (u + i v) + \Re c + i \Im c
     #          = (U u + \Re c - V v) + i (V u + \Im c + U v)
-    re = F.linear(input.real, weight.real) \
-        - F.linear(input.imag, weight.imag)
-    im = F.linear(input.real, weight.imag) \
-        + F.linear(input.imag, weight.real)
+    re = F.linear(input.real, weight.real) - F.linear(input.imag, weight.imag)
+    im = F.linear(input.real, weight.imag) + F.linear(input.imag, weight.real)
 
     output = Cplx(re, im)
     if bias is not None:
@@ -618,10 +650,13 @@ def linear_naive(input, weight, bias=None):
 
 def linear_cat(input, weight, bias=None):
     # [n_out, n_in] -> [2 * n_out, 2 * n_in] : [[U, V], [-V, U]]
-    ww = torch.cat([
-        torch.cat([ weight.real, weight.imag], dim=0),
-        torch.cat([-weight.imag, weight.real], dim=0)
-    ], dim=1)
+    ww = torch.cat(
+        [
+            torch.cat([weight.real, weight.imag], dim=0),
+            torch.cat([-weight.imag, weight.real], dim=0),
+        ],
+        dim=1,
+    )
 
     xx = to_concatenated_real(input, dim=-1)  # [..., 2 * n_in]
     output = from_concatenated_real(F.linear(xx, ww, None))
@@ -648,7 +683,7 @@ def linear_3m(input, weight, bias=None):
     #          = [± U (v + j u)]  # Gauss trick
     #          =     (U (u + v) - (U + V) v + \Re c)
     #            + j (U (u + v) + (V - U) u + \Im c)
-    K1 = F.linear(input.real + input.imag,  weight.real)
+    K1 = F.linear(input.real + input.imag, weight.real)
     K2 = F.linear(input.real, weight.imag - weight.real)
     K3 = F.linear(input.imag, weight.real + weight.imag)
 
@@ -679,23 +714,19 @@ def symmetric_circular_padding(input, padding):
     return input.apply(F.pad, tuple(expanded_padding), mode="circular")
 
 
-def convnd_naive(conv, input, weight, stride=1,
-                 padding=0, dilation=1, groups=1):
+def convnd_naive(conv, input, weight, stride=1, padding=0, dilation=1, groups=1):
 
-    re = conv(input.real, weight.real, None,
-              stride, padding, dilation, groups) \
-        - conv(input.imag, weight.imag, None,
-               stride, padding, dilation, groups)
-    im = conv(input.real, weight.imag, None,
-              stride, padding, dilation, groups) \
-        + conv(input.imag, weight.real, None,
-               stride, padding, dilation, groups)
+    re = conv(input.real, weight.real, None, stride, padding, dilation, groups) - conv(
+        input.imag, weight.imag, None, stride, padding, dilation, groups
+    )
+    im = conv(input.real, weight.imag, None, stride, padding, dilation, groups) + conv(
+        input.imag, weight.real, None, stride, padding, dilation, groups
+    )
 
     return Cplx(re, im)
 
 
-def convnd_quick(conv, input, weight, stride=1,
-                 padding=0, dilation=1):
+def convnd_quick(conv, input, weight, stride=1, padding=0, dilation=1):
     r"""Applies a complex convolution transformation to the complex data
     :math:`y = x \ast W + b` using two calls to `conv` at the cost of extra
     concatenation and slicing.
@@ -711,8 +742,7 @@ def convnd_quick(conv, input, weight, stride=1,
     return Cplx(rwr - iwi, iwr + rwi)
 
 
-def convnd_3m(conv, input, weight, stride=1,
-              padding=0, dilation=1, groups=1):
+def convnd_3m(conv, input, weight, stride=1, padding=0, dilation=1, groups=1):
     r"""Applies a complex convolution transformation to the complex data
     :math:`y = x \ast W + b` using 3 calls to `conv`.
     """
@@ -722,36 +752,46 @@ def convnd_3m(conv, input, weight, stride=1,
     #          = [± U (v + j u)]  # Gauss trick
     #          =     (U (u + v) - (U + V) v)
     #            + j (U (u + v) + (V - U) u)
-    K1 = conv(input.real + input.imag,  weight.real, None,
-              stride, padding, dilation, groups)
+    K1 = conv(
+        input.real + input.imag, weight.real, None, stride, padding, dilation, groups
+    )
 
-    K2 = conv(input.real, weight.imag - weight.real, None,
-              stride, padding, dilation, groups)
+    K2 = conv(
+        input.real, weight.imag - weight.real, None, stride, padding, dilation, groups
+    )
 
-    K3 = conv(input.imag, weight.real + weight.imag, None,
-              stride, padding, dilation, groups)
+    K3 = conv(
+        input.imag, weight.real + weight.imag, None, stride, padding, dilation, groups
+    )
 
     return Cplx(K1 - K3, K1 + K2)
 
 
-def convnd(conv, input, weight, bias=None, stride=1,
-           padding=0, dilation=1, groups=1, padding_mode="zeros"):
+def convnd(
+    conv,
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex n-d convolution to the incoming complex
     tensor `B x c_in x L_1 x ... L_n`: :math:`y = x \star W + b`.
     """
-    if padding_mode == 'circular':
+    if padding_mode == "circular":
         input = symmetric_circular_padding(input, padding)
         padding = 0
-    elif padding_mode != 'zeros':
+    elif padding_mode != "zeros":
         raise ValueError("padding_mode must be 'zeros' or 'circular'.")
 
     if groups == 1:
         # ungroupped convolution can be done a little bit faster
-        output = convnd_quick(conv, input, weight, stride,
-                              padding, dilation)
+        output = convnd_quick(conv, input, weight, stride, padding, dilation)
     else:
-        output = convnd_naive(conv, input, weight, stride,
-                              padding, dilation, groups)
+        output = convnd_naive(conv, input, weight, stride, padding, dilation, groups)
 
     if bias is not None:
         broadcast = (input.dim() - 2) * [1]
@@ -760,81 +800,143 @@ def convnd(conv, input, weight, bias=None, stride=1,
     return output
 
 
-def conv1d(input, weight, bias=None, stride=1, padding=0,
-           dilation=1, groups=1, padding_mode="zeros"):
+def conv1d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex 1d convolution to the incoming complex
     tensor `B x c_in x L`: :math:`y = x \star W + b`.
     """
 
-    return convnd(F.conv1d, input, weight, bias, stride,
-                  padding, dilation, groups, padding_mode)
-
-
-def conv2d(input, weight, bias=None, stride=1, padding=0,
-           dilation=1, groups=1, padding_mode="zeros"):
+    return convnd(
+        F.conv1d, input, weight, bias, stride, padding, dilation, groups, padding_mode
+    )
+
+
+def conv2d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex 2d convolution to the incoming complex
     tensor `B x c_in x H x W`: :math:`y = x \star W + b`.
     """
 
-    return convnd(F.conv2d, input, weight, bias, stride,
-                  padding, dilation, groups, padding_mode)
-
-
-def conv3d(input, weight, bias=None, stride=1, padding=0,
-           dilation=1, groups=1, padding_mode="zeros"):
+    return convnd(
+        F.conv2d, input, weight, bias, stride, padding, dilation, groups, padding_mode
+    )
+
+
+def conv3d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    dilation=1,
+    groups=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex 3d convolution to the incoming complex
     tensor `B x c_in x H x W x D`: :math:`y = x \star W + b`.
     """
 
-    return convnd(F.conv3d, input, weight, bias, stride,
-                  padding, dilation, groups, padding_mode)
-
-
-def conv_transposend_naive(conv_t, input, weight, stride=1,
-                           padding=0, output_padding=0,
-                           groups=1, dilation=1):
-    re = conv_t(input.real, weight.real, None, stride,
-                padding, output_padding, groups, dilation) \
-        - conv_t(input.imag, weight.imag, None, stride,
-                 padding, output_padding, groups, dilation)
-    im = conv_t(input.real, weight.imag, None, stride,
-                padding, output_padding, groups, dilation) \
-        + conv_t(input.imag, weight.real, None, stride,
-                 padding, output_padding, groups, dilation)
+    return convnd(
+        F.conv3d, input, weight, bias, stride, padding, dilation, groups, padding_mode
+    )
+
+
+def conv_transposend_naive(
+    conv_t, input, weight, stride=1, padding=0, output_padding=0, groups=1, dilation=1
+):
+    re = conv_t(
+        input.real, weight.real, None, stride, padding, output_padding, groups, dilation
+    ) - conv_t(
+        input.imag, weight.imag, None, stride, padding, output_padding, groups, dilation
+    )
+    im = conv_t(
+        input.real, weight.imag, None, stride, padding, output_padding, groups, dilation
+    ) + conv_t(
+        input.imag, weight.real, None, stride, padding, output_padding, groups, dilation
+    )
     return Cplx(re, im)
 
 
-def conv_transposend_3m(conv_t, input, weight, stride=1,
-                        padding=0, output_padding=0,
-                        groups=1, dilation=1):
+def conv_transposend_3m(
+    conv_t, input, weight, stride=1, padding=0, output_padding=0, groups=1, dilation=1
+):
     """3m interface for transposed convolutions."""
-    K1 = conv_t(input.real + input.imag,  weight.real, None, stride,
-                padding, output_padding, groups, dilation)
-
-    K2 = conv_t(input.real, weight.imag - weight.real, None, stride,
-                padding, output_padding, groups, dilation)
-
-    K3 = conv_t(input.imag, weight.real + weight.imag, None, stride,
-                padding, output_padding, groups, dilation)
+    K1 = conv_t(
+        input.real + input.imag,
+        weight.real,
+        None,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+    )
+
+    K2 = conv_t(
+        input.real,
+        weight.imag - weight.real,
+        None,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+    )
+
+    K3 = conv_t(
+        input.imag,
+        weight.real + weight.imag,
+        None,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+    )
 
     return Cplx(K1 - K3, K1 + K2)
 
 
-def conv_transposend(conv, input, weight, bias=None, stride=1,
-                     padding=0, output_padding=1, groups=1,
-                     dilation=1, padding_mode="zeros"):
+def conv_transposend(
+    conv,
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    output_padding=1,
+    groups=1,
+    dilation=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex n-d transposed convolution to the
     incoming complex tensor. See torch.nn.ConvTranspose2d
     for documentation."""
-    if padding_mode == 'circular':
+    if padding_mode == "circular":
         input = symmetric_circular_padding(input, padding)
         padding = 0
-    elif padding_mode != 'zeros':
+    elif padding_mode != "zeros":
         raise ValueError("padding_mode must be 'zeros' or 'circular'.")
 
-    output = conv_transposend_naive(conv, input, weight, stride,
-                                    padding, output_padding,
-                                    groups, dilation)
+    output = conv_transposend_naive(
+        conv, input, weight, stride, padding, output_padding, groups, dilation
+    )
 
     if bias is not None:
         broadcast = (input.dim() - 2) * [1]
@@ -843,64 +945,118 @@ def conv_transposend(conv, input, weight, bias=None, stride=1,
     return output
 
 
-def conv_transpose1d(input, weight, bias=None, stride=1,
-                     padding=0, output_padding=0, groups=0,
-                     dilation=1, padding_mode="zeros"):
+def conv_transpose1d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    output_padding=0,
+    groups=0,
+    dilation=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex 1d transposed convolution to the
     incoming complex tensor. See torch.nn.ConvTranspose1d
     for documentation."""
-    return conv_transposend(F.conv_transpose1d, input, weight,
-                            bias, stride, padding, output_padding,
-                            groups, dilation, padding_mode)
-
-
-def conv_transpose2d(input, weight, bias=None, stride=1,
-                     padding=0, output_padding=0, groups=0,
-                     dilation=1, padding_mode="zeros"):
+    return conv_transposend(
+        F.conv_transpose1d,
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+        padding_mode,
+    )
+
+
+def conv_transpose2d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    output_padding=0,
+    groups=0,
+    dilation=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex 2d transposed convolution to the
     incoming complex tensor. See torch.nn.ConvTranspose2d
     for documentation."""
-    return conv_transposend(F.conv_transpose2d, input, weight,
-                            bias, stride, padding, output_padding,
-                            groups, dilation, padding_mode)
-
-
-def conv_transpose3d(input, weight, bias=None, stride=1,
-                     padding=0, output_padding=0, groups=0,
-                     dilation=1, padding_mode="zeros"):
+    return conv_transposend(
+        F.conv_transpose2d,
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+        padding_mode,
+    )
+
+
+def conv_transpose3d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    output_padding=0,
+    groups=0,
+    dilation=1,
+    padding_mode="zeros",
+):
     r"""Applies a complex 3d transposed convolution to the
     incoming complex tensor. See torch.nn.ConvTranspose3d
     for documentation."""
-    return conv_transposend(F.conv_transpose3d, input, weight,
-                            bias, stride, padding, output_padding,
-                            groups, dilation, padding_mode)
+    return conv_transposend(
+        F.conv_transpose3d,
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+        padding_mode,
+    )
 
 
 def einsum(equation, *tensors):
     """2-tensor einstein summation."""
     if not tensors:
-        raise RuntimeError("""`einsum()` requires """
-                           """at least one tensor.""")
+        raise RuntimeError("`einsum()` requires at least one tensor.")
 
     cplx1, *tensors = tensors
     if not tensors:
         # no complex multiplication with only one tensor
-        return Cplx(torch.einsum(equation, cplx1.real),
-                    torch.einsum(equation, cplx1.imag))
+        return Cplx(
+            torch.einsum(equation, cplx1.real), torch.einsum(equation, cplx1.imag)
+        )
 
     cplx2, *tensors = tensors
     if not tensors:
         # Einsum is complex bilinear -- the logic is same here.
-        re = torch.einsum(equation, cplx1.real, cplx2.real) \
-            - torch.einsum(equation, cplx1.imag, cplx2.imag)
+        re = torch.einsum(equation, cplx1.real, cplx2.real) - torch.einsum(
+            equation, cplx1.imag, cplx2.imag
+        )
 
-        im = torch.einsum(equation, cplx1.real, cplx2.imag) \
-            + torch.einsum(equation, cplx1.imag, cplx2.real)
+        im = torch.einsum(equation, cplx1.real, cplx2.imag) + torch.einsum(
+            equation, cplx1.imag, cplx2.real
+        )
 
         return Cplx(re, im)
 
-    raise RuntimeError(f"""`einsum()` does not support more """
-                       f"""than 2 tensors. Got {2 + len(tensors)}.""")
+    raise RuntimeError(
+        f"`Cplx.einsum` does not support more than 2 tensors. Got {2 + len(tensors)}."
+    )
 
 
 def bilinear_naive(input1, input2, weight, bias=None, conjugate=True):
@@ -939,17 +1095,11 @@ def bilinear_cat(input1, input2, weight, bias=None, conjugate=True):
     U, V = weight.real, weight.imag
 
     UV = torch.cat([U, -V], dim=2)
-    VU = torch.cat([V,  U], dim=2)
+    VU = torch.cat([V, U], dim=2)
     if conjugate:
-        ww = torch.cat([
-            torch.cat([UV,  VU], dim=1),
-            torch.cat([VU, -UV], dim=1)
-        ], dim=0)
+        ww = torch.cat([torch.cat([UV, VU], dim=1), torch.cat([VU, -UV], dim=1)], dim=0)
     else:
-        ww = torch.cat([
-            torch.cat([UV, -VU], dim=1),
-            torch.cat([VU,  UV], dim=1)
-        ], dim=0)
+        ww = torch.cat([torch.cat([UV, -VU], dim=1), torch.cat([VU, UV], dim=1)], dim=0)
 
     x1 = to_concatenated_real(input1, dim=-1)
     x2 = to_concatenated_real(input2, dim=-1)
@@ -961,8 +1111,9 @@ def bilinear_cat(input1, input2, weight, bias=None, conjugate=True):
     return output
 
 
-def max_poolnd(pool, input, kernel_size, stride=None,
-               padding=0, dilation=1, ceil_mode=False):
+def max_poolnd(
+    pool, input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False
+):
     r"""Applies complex version of n-d max-pooling to the incoming complex
     tensor `B x c_in x L_1 x ... L_n`.
 
@@ -1001,8 +1152,15 @@ def max_poolnd(pool, input, kernel_size, stride=None,
 
     # 1. get the abs max-pooling indices
     #  XXX should we care about the extra sqrt in abs?
-    _, indices = pool(abs(input), kernel_size, stride, padding,
-                      dilation, ceil_mode, return_indices=True)
+    _, indices = pool(
+        abs(input),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        ceil_mode,
+        return_indices=True,
+    )
 
     # 2. the indices refer to the positions in the trailing flattened
     #    magnitudes, so we flatten the trailing dims after the channels.
@@ -1017,25 +1175,25 @@ def max_poolnd(pool, input, kernel_size, stride=None,
     return Cplx(re, im).reshape(input.shape[:2] + indices.shape[2:])
 
 
-def max_pool1d(input, kernel_size, stride=None, padding=0,
-               dilation=1, ceil_mode=False):
+def max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False):
     """Complex 1d max pooling on the complex tensor `B x c_in x L`."""
 
-    return max_poolnd(F.max_pool1d, input, kernel_size,
-                      stride, padding, dilation, ceil_mode)
+    return max_poolnd(
+        F.max_pool1d, input, kernel_size, stride, padding, dilation, ceil_mode
+    )
 
 
-def max_pool2d(input, kernel_size, stride=None, padding=0,
-               dilation=1, ceil_mode=False):
+def max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False):
     """Complex 2d max pooling on the complex tensor `B x c_in x H x W`."""
 
-    return max_poolnd(F.max_pool2d, input, kernel_size,
-                      stride, padding, dilation, ceil_mode)
+    return max_poolnd(
+        F.max_pool2d, input, kernel_size, stride, padding, dilation, ceil_mode
+    )
 
 
-def max_pool3d(input, kernel_size, stride=None, padding=0,
-               dilation=1, ceil_mode=False):
+def max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False):
     """Complex 3d max pooling on the complex tensor `B x c_in x H x W x D`."""
 
-    return max_poolnd(F.max_pool3d, input, kernel_size,
-                      stride, padding, dilation, ceil_mode)
+    return max_poolnd(
+        F.max_pool3d, input, kernel_size, stride, padding, dilation, ceil_mode
+    )
diff --git a/cplxmodule/nn/init.py b/cplxmodule/nn/init.py
index b3b5775..1e4cccd 100644
--- a/cplxmodule/nn/init.py
+++ b/cplxmodule/nn/init.py
@@ -13,8 +13,10 @@ def get_fans(cplxtensor):
     """Almost verbatim copy of `init._calculate_fan_in_and_fan_out`"""
     ndim = cplxtensor.dim()
     if ndim < 2:
-        raise ValueError("Fan in and fan out can not be computed "
-                         "for tensor with fewer than 2 dimensions.")
+        raise ValueError(
+            "Fan in and fan out can not be computed for tensor with "
+            "fewer than 2 dimensions."
+        )
 
     n_fmaps_output, n_fmaps_input, *rest = cplxtensor.shape
     if ndim == 2:
@@ -29,7 +31,7 @@ def get_fans(cplxtensor):
 
 
 @wraps(init.kaiming_normal_, assigned=("__name__", "__doc__", "__annotations__"))
-def cplx_kaiming_normal_(tensor, a=0., mode='fan_in', nonlinearity='leaky_relu'):
+def cplx_kaiming_normal_(tensor, a=0.0, mode="fan_in", nonlinearity="leaky_relu"):
     assert isinstance(tensor, Cplx)
 
     a = math.sqrt(1 + 2 * a * a)
@@ -41,12 +43,12 @@ def cplx_kaiming_normal_(tensor, a=0., mode='fan_in', nonlinearity='leaky_relu')
 def cplx_xavier_normal_(tensor, gain=1.0):
     assert isinstance(tensor, Cplx)
 
-    init.xavier_normal_(tensor.real, gain=gain/math.sqrt(2))
-    init.xavier_normal_(tensor.imag, gain=gain/math.sqrt(2))
+    init.xavier_normal_(tensor.real, gain=gain / math.sqrt(2))
+    init.xavier_normal_(tensor.imag, gain=gain / math.sqrt(2))
 
 
 @wraps(init.kaiming_uniform_, assigned=("__name__", "__doc__", "__annotations__"))
-def cplx_kaiming_uniform_(tensor, a=0., mode='fan_in', nonlinearity='leaky_relu'):
+def cplx_kaiming_uniform_(tensor, a=0.0, mode="fan_in", nonlinearity="leaky_relu"):
     assert isinstance(tensor, Cplx)
 
     a = math.sqrt(1 + 2 * a * a)
@@ -58,8 +60,8 @@ def cplx_kaiming_uniform_(tensor, a=0., mode='fan_in', nonlinearity='leaky_relu'
 def cplx_xavier_uniform_(tensor, gain=1.0):
     assert isinstance(tensor, Cplx)
 
-    init.xavier_uniform_(tensor.real, gain=gain/math.sqrt(2))
-    init.xavier_uniform_(tensor.imag, gain=gain/math.sqrt(2))
+    init.xavier_uniform_(tensor.real, gain=gain / math.sqrt(2))
+    init.xavier_uniform_(tensor.imag, gain=gain / math.sqrt(2))
 
 
 def cplx_trabelsi_standard_(cplx, kind="glorot"):
@@ -102,8 +104,7 @@ def cplx_trabelsi_independent_(cplx, kind="glorot"):
     Z = np.random.rand(*shape) + 1j * np.random.rand(*shape)
 
     # Z is n x m, so u is n x n and vh is m x m
-    u, _, vh = np.linalg.svd(Z, compute_uv=True, full_matrices=True,
-                             hermitian=False)
+    u, _, vh = np.linalg.svd(Z, compute_uv=True, full_matrices=True, hermitian=False)
     k = min(*shape)
     M = np.dot(u[:, :k], vh[:, :k].conjugate().T)
 
@@ -122,7 +123,7 @@ def cplx_trabelsi_independent_(cplx, kind="glorot"):
     return cplx
 
 
-def cplx_uniform_independent_(cplx, a=0., b=1.):
+def cplx_uniform_independent_(cplx, a=0.0, b=1.0):
     torch.nn.init.uniform_(cplx.real, a, b)
     torch.nn.init.uniform_(cplx.imag, a, b)
 
diff --git a/cplxmodule/nn/masked/base.py b/cplxmodule/nn/masked/base.py
index b87e7b9..6d04fcf 100644
--- a/cplxmodule/nn/masked/base.py
+++ b/cplxmodule/nn/masked/base.py
@@ -29,6 +29,7 @@ class BaseMasked(torch.nn.Module):
     The masks could be set either manually through setting `.mask`, or loaded in
     bulk with `deploy_masks`, or via `model.load_state_dict(..., strict=False)`.
     """
+
     def __init__(self):
         super().__init__()
         self.register_buffer("mask", None)
@@ -52,8 +53,9 @@ def mask_(self, mask):
         Effectively switches on / off masking of the weights.
         """
         if mask is not None and not isinstance(mask, torch.Tensor):
-            raise TypeError(f"`mask` must be either a Tensor or "
-                            f"`None`. Got {type(mask).__name__}.")
+            raise TypeError(
+                f"`mask` must be either a Tensor or `None`. Got {type(mask).__name__}."
+            )
 
         if mask is not None:
             # None -> sparse : register mask, turning on sparsity
@@ -86,15 +88,28 @@ def __setattr__(self, name, value):
 
         self.mask_(value)
 
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata,
-                              strict, missing_keys, unexpected_keys,
-                              error_msgs):
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
         """Surgically load the state with the runtime masks from a dict."""
         # this next call loads everything, expect mask into this module only!
         mask = prefix + "mask"
         super()._load_from_state_dict(
-            {k: v for k, v in state_dict.items() if k != mask}, prefix,
-            local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+            {k: v for k, v in state_dict.items() if k != mask},
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
 
         # here: mask in missing <=> buffer exists and not None <=> is_sparse
         # https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/module.py#L758
@@ -117,14 +132,17 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
             missing_keys.remove(mask)
 
 
-class MaskedWeightMixin():
+class MaskedWeightMixin:
     """A mixin for accessing read-only masked weight,"""
+
     @property
     def weight_masked(self):
         """Return a sparsified weight of the parent *Linear."""
         if not self.is_sparse:
-            msg = f"`{type(self).__name__}` has no sparsity mask. Please, " \
-                  f"either set a mask attribute, or call `deploy_masks()`."
+            msg = (
+                f"`{type(self).__name__}` has no sparsity mask. Please, "
+                f"either set a mask attribute, or call `deploy_masks()`."
+            )
             raise RuntimeError(msg)
 
         # __rmul__ by the mask in case of weird weight types (like Cplx)
@@ -188,8 +206,7 @@ def deploy_masks(module, *, state_dict=None, prefix="", reset=False):
         Whether to forcefully reset the sparsity of the layers, masks for
         which were NOT provided in `state_dict`.
     """
-    if not isinstance(state_dict, dict) \
-       or not isinstance(module, torch.nn.Module):
+    if not isinstance(state_dict, dict) or not isinstance(module, torch.nn.Module):
         return module
 
     for name, mod in module.named_modules(prefix=prefix):
@@ -242,7 +259,6 @@ def binarize_masks(state_dict, masks):
 
             new_state_dict[name] = par
 
-        new_masks = {name: torch.ne(mask, 0).to(mask)
-                     for name, mask in masks.items()}
+        new_masks = {name: torch.ne(mask, 0).to(mask) for name, mask in masks.items()}
 
     return new_state_dict, new_masks
diff --git a/cplxmodule/nn/masked/complex.py b/cplxmodule/nn/masked/complex.py
index 34e61d2..347ef05 100644
--- a/cplxmodule/nn/masked/complex.py
+++ b/cplxmodule/nn/masked/complex.py
@@ -22,9 +22,12 @@ def sparsity(self, *, hard=True, **kwargs):
             n_dropped -= float(mask.sum().item())
 
         else:
-            n_dropped = 0.
+            n_dropped = 0.0
 
-        return [(id(weight.real), n_dropped), (id(weight.imag), n_dropped), ]
+        return [
+            (id(weight.real), n_dropped),
+            (id(weight.imag), n_dropped),
+        ]
 
 
 class CplxLinearMasked(CplxLinear, _BaseCplxMixin):
@@ -39,20 +42,41 @@ def forward(self, input1, input2):
 
 class CplxConv1dMasked(CplxConv1d, _BaseCplxMixin):
     def forward(self, input):
-        return cplx.conv1d(input, self.weight_masked, self.bias,
-                           self.stride, self.padding, self.dilation,
-                           self.groups, self.padding_mode)
+        return cplx.conv1d(
+            input,
+            self.weight_masked,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.padding_mode,
+        )
 
 
 class CplxConv2dMasked(CplxConv2d, _BaseCplxMixin):
     def forward(self, input):
-        return cplx.conv2d(input, self.weight_masked, self.bias,
-                           self.stride, self.padding, self.dilation,
-                           self.groups, self.padding_mode)
+        return cplx.conv2d(
+            input,
+            self.weight_masked,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.padding_mode,
+        )
 
 
 class CplxConv3dMasked(CplxConv3d, _BaseCplxMixin):
     def forward(self, input):
-        return cplx.conv3d(input, self.weight_masked, self.bias,
-                           self.stride, self.padding, self.dilation,
-                           self.groups, self.padding_mode)
+        return cplx.conv3d(
+            input,
+            self.weight_masked,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.padding_mode,
+        )
diff --git a/cplxmodule/nn/masked/real.py b/cplxmodule/nn/masked/real.py
index 9612e79..a2bcf9d 100644
--- a/cplxmodule/nn/masked/real.py
+++ b/cplxmodule/nn/masked/real.py
@@ -17,7 +17,7 @@ def sparsity(self, *, hard=True, **kwargs):
             n_dropped -= float(mask.sum().item())
 
         else:
-            n_dropped = 0.
+            n_dropped = 0.0
 
         return [(id(self.weight), n_dropped)]
 
@@ -29,20 +29,41 @@ def forward(self, input):
 
 class Conv1dMasked(Conv1d, _BaseRealMixin):
     def forward(self, input):
-        return F.conv1d(input, self.weight_masked, self.bias,
-                        self.stride, self.padding, self.dilation, self.groups)
+        return F.conv1d(
+            input,
+            self.weight_masked,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
 
 
 class Conv2dMasked(Conv2d, _BaseRealMixin):
     def forward(self, input):
-        return F.conv2d(input, self.weight_masked, self.bias,
-                        self.stride, self.padding, self.dilation, self.groups)
+        return F.conv2d(
+            input,
+            self.weight_masked,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
 
 
 class Conv3dMasked(Conv3d, _BaseRealMixin):
     def forward(self, input):
-        return F.conv3d(input, self.weight_masked, self.bias,
-                        self.stride, self.padding, self.dilation, self.groups)
+        return F.conv3d(
+            input,
+            self.weight_masked,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
 
 
 class BilinearMasked(Bilinear, _BaseRealMixin):
diff --git a/cplxmodule/nn/modules/activation.py b/cplxmodule/nn/modules/activation.py
index b7154a4..c62246a 100644
--- a/cplxmodule/nn/modules/activation.py
+++ b/cplxmodule/nn/modules/activation.py
@@ -16,6 +16,7 @@ class CplxModReLU(CplxToCplx):
     with $\tau \in \mathbb{R}$. The if threshold=None then it
     becomes a learnable parameter.
     """
+
     def __init__(self, threshold=0.5):
         super().__init__()
         if not isinstance(threshold, float):
@@ -40,6 +41,7 @@ class CplxAdaptiveModReLU(CplxToCplx):
     features of the $d$-dim complex vector, and `CplxChanneledModReLU(d)` lets
     each dimension have its own threshold.
     """
+
     def __init__(self, *dim):
         super().__init__()
         self.dim = dim if dim else (1,)
diff --git a/cplxmodule/nn/modules/base.py b/cplxmodule/nn/modules/base.py
index e402a6a..b39fa7e 100644
--- a/cplxmodule/nn/modules/base.py
+++ b/cplxmodule/nn/modules/base.py
@@ -7,26 +7,36 @@
 
 class CplxParameter(torch.nn.ParameterDict):
     """Torch-friendly container for complex-valued parameter."""
+
     def __init__(self, cplx):
         if not isinstance(cplx, Cplx):
-            raise TypeError(f"""`{type(self).__name__}` accepts only """
-                            f"""Cplx tensors.""")
+            raise TypeError(f"`{type(self).__name__}` accepts only Cplx tensors.")
 
-        super().__init__({
-            "real": torch.nn.Parameter(cplx.real),
-            "imag": torch.nn.Parameter(cplx.imag),
-        })
+        super().__init__(
+            {
+                "real": torch.nn.Parameter(cplx.real),
+                "imag": torch.nn.Parameter(cplx.imag),
+            }
+        )
 
         # save a reference to the underlying cplx data (bypass ParameterDict)
-        torch.nn.Module.__setattr__(self, '_cplx', cplx)
-
-    def _load_from_state_dict(self, state_dict, prefix, local_metadata,
-                              strict, missing_keys, unexpected_keys,
-                              error_msgs):
+        torch.nn.Module.__setattr__(self, "_cplx", cplx)
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
 
         missing, unexpected = [], []
-        super()._load_from_state_dict(state_dict, prefix, local_metadata,
-                                      strict, missing, unexpected, error_msgs)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing, unexpected, error_msgs
+        )
 
         if len(missing) == 2:
             # By desgin of `__init__` we get here if the last call to
@@ -56,21 +66,30 @@ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
                 # recursively call on a special state_dict to promote R to C
                 # NO missing or unexpected items are possible by design, only
                 # shape mismatch or general exceptions are now possible.
-                self._load_from_state_dict({
-                        f"{prefix}real": par,
-                        f"{prefix}imag": torch.zeros_like(par)
-                    }, prefix, local_metadata, strict, [], [], error_msgs)
+                self._load_from_state_dict(
+                    {f"{prefix}real": par, f"{prefix}imag": torch.zeros_like(par)},
+                    prefix,
+                    local_metadata,
+                    strict,
+                    [],
+                    [],
+                    error_msgs,
+                )
 
         elif len(missing) == 1:
             # Although loaded either `.real` or `.imag`, the state_dict
             # contains only one of ".real" or ".imag", so the other part is
             # missing. Therefore append to `error_msg`.
-            error_msgs.append(f"Complex parameter requires both `.imag`"
-                              f" and `.imag` parts. Missing `{missing[0]}`.")
+            error_msgs.append(
+                f"Complex parameter requires both `.real`"
+                f" and `.imag` parts. Missing `{missing[0]}`."
+            )
 
         if strict and unexpected:
-            error_msgs.append(f"Complex parameter disallows redundant key(s)"
-                              f" in state_dict: {unexpected}.")
+            error_msgs.append(
+                f"Complex parameter disallows redundant key(s)"
+                f" in state_dict: {unexpected}."
+            )
 
         unexpected_keys.extend(unexpected)
         missing_keys.extend(missing)
@@ -83,7 +102,7 @@ def data(self):
         return self._cplx
 
 
-class CplxParameterAccessor():
+class CplxParameterAccessor:
     """Cosmetic complex parameter accessor.
 
     Details
@@ -99,6 +118,7 @@ class CplxParameterAccessor():
     is why `SparseWeightMixin` in .masked couldn't work with 'weight'
     as a read-only @property.
     """
+
     def __getattr__(self, name):
         # default attr lookup straight to parent's __getattr__
         attr = super().__getattr__(name)
@@ -120,6 +140,7 @@ class BaseCplxToReal(torch.nn.Module):
 
 def _promote_callable_to_split(fn):
     """Create a runtime class promoting a real function to split activation."""
+
     class template(CplxToCplx):
         @wraps(fn)
         def __init__(self, *args, **kwargs):
@@ -137,13 +158,15 @@ def forward(self, input):
     template.__name__ = f"CplxSplitFunc{fn.__name__.title()}"
     template.__qualname__ = f"<runtime type `{template.__name__}`>"
     template.__doc__ = f"Split activation based on `{fn.__name__}`"
-    template.forward.__doc__ = (f"Call `{fn.__name__}` on real and "
-                                "imaginary components independently.")
+    template.forward.__doc__ = (
+        f"Call `{fn.__name__}` on real and " "imaginary components independently."
+    )
     return template
 
 
 def _promote_module_to_split(Module):
     """Make a runtime class promoting a Module subclass to split activation."""
+
     class template(Module, CplxToCplx):
         def forward(self, input):
             """Apply to real and imaginary parts independently."""
@@ -157,6 +180,7 @@ def forward(self, input):
 
 class _CplxToCplxMeta(type):
     """Meta class for promoting real activations to split complex ones."""
+
     @lru_cache(maxsize=None)
     def __getitem__(self, Base):
         if not isinstance(Base, type) and callable(Base):
@@ -174,12 +198,13 @@ def __getitem__(self, Base):
 
             return _promote_module_to_split(Base)
 
-        raise TypeError("Expecting either a torch.nn.Module subclass, or"
-                        f" a callable for promotion. Got `{type(Base)}`.")
+        raise TypeError(
+            "Expecting either a torch.nn.Module subclass, or"
+            f" a callable for promotion. Got `{type(Base)}`."
+        )
 
 
-class CplxToCplx(CplxParameterAccessor, torch.nn.Module,
-                 metaclass=_CplxToCplxMeta):
+class CplxToCplx(CplxParameterAccessor, torch.nn.Module, metaclass=_CplxToCplxMeta):
     pass
 
 
diff --git a/cplxmodule/nn/modules/batchnorm.py b/cplxmodule/nn/modules/batchnorm.py
index 95d5a64..259d715 100644
--- a/cplxmodule/nn/modules/batchnorm.py
+++ b/cplxmodule/nn/modules/batchnorm.py
@@ -6,8 +6,14 @@
 from ... import cplx
 
 
-def whiten2x2(tensor, training=True, running_mean=None, running_cov=None,
-              momentum=0.1, nugget=1e-5):
+def whiten2x2(
+    tensor,
+    training=True,
+    running_mean=None,
+    running_cov=None,
+    momentum=0.1,
+    nugget=1e-5,
+):
     r"""Solve R M R = I for R and a given 2x2 matrix M = [[a, b], [c, d]].
 
     Arguments
@@ -81,10 +87,15 @@ def whiten2x2(tensor, training=True, running_mean=None, running_cov=None,
         # has to mul-mean here anyway (naïve) : reduction axes shifted left.
         cov_vu = cov_uv = (tensor[0] * tensor[1]).mean([a - 1 for a in axes])
         if running_cov is not None:
-            cov = torch.stack([
-                cov_uu.data, cov_uv.data,
-                cov_vu.data, cov_vv.data,
-            ], dim=0).reshape(2, 2, -1)
+            cov = torch.stack(
+                [
+                    cov_uu.data,
+                    cov_uv.data,
+                    cov_vu.data,
+                    cov_vv.data,
+                ],
+                dim=0,
+            ).reshape(2, 2, -1)
             running_cov += momentum * (cov - running_cov)
 
     else:
@@ -102,15 +113,24 @@ def whiten2x2(tensor, training=True, running_mean=None, running_cov=None,
     r, s = -cov_vu / denom, (cov_uu + sqrdet) / denom
 
     # 4. apply Q to x (manually)
-    out = torch.stack([
-        tensor[0] * p.reshape(tail) + tensor[1] * r.reshape(tail),
-        tensor[0] * q.reshape(tail) + tensor[1] * s.reshape(tail),
-    ], dim=0)
+    out = torch.stack(
+        [
+            tensor[0] * p.reshape(tail) + tensor[1] * r.reshape(tail),
+            tensor[0] * q.reshape(tail) + tensor[1] * s.reshape(tail),
+        ],
+        dim=0,
+    )
     return out  # , torch.cat([p, q, r, s], dim=0).reshape(2, 2, -1)
 
 
-def whitendxd(tensor, training=True, running_mean=None, running_cov=None,
-              momentum=0.1, nugget=1e-5):
+def whitendxd(
+    tensor,
+    training=True,
+    running_mean=None,
+    running_cov=None,
+    momentum=0.1,
+    nugget=1e-5,
+):
     """Jointly whiten features in tensors [P x B x F x ...]: take dim-P vectors
     and whiten individually for each F over [B x ...].
 
@@ -159,7 +179,8 @@ def whitendxd(tensor, training=True, running_mean=None, running_cov=None,
     ell = torch.cholesky(cov + eye, upper=True)
     soln = torch.triangular_solve(
         tensor.unsqueeze(-1).permute(*range(1, tensor.dim()), 0, -1),
-        ell.reshape(*shape, d, d))
+        ell.reshape(*shape, d, d),
+    )
 
     soln = soln.solution.squeeze(-1)
     return torch.stack(torch.unbind(soln, dim=-1), dim=0)
@@ -223,25 +244,36 @@ def cplx_batch_norm(
     speedup both on host and device computations).
     """
     # check arguments
-    assert ((running_mean is None and running_var is None)
-            or (running_mean is not None and running_var is not None))
-    assert ((weight is None and bias is None)
-            or (weight is not None and bias is not None))
+    assert (running_mean is None and running_var is None) or (
+        running_mean is not None and running_var is not None
+    )
+    assert (weight is None and bias is None) or (
+        weight is not None and bias is not None
+    )
 
     # stack along the first axis
     x = torch.stack([input.real, input.imag], dim=0)
 
     # whiten and apply affine transformation
-    z = whiten2x2(x, training=training, running_mean=running_mean,
-                  running_cov=running_var, momentum=momentum, nugget=eps)
+    z = whiten2x2(
+        x,
+        training=training,
+        running_mean=running_mean,
+        running_cov=running_var,
+        momentum=momentum,
+        nugget=eps,
+    )
 
     if weight is not None and bias is not None:
         shape = 1, x.shape[2], *([1] * (x.dim() - 3))
         weight = weight.reshape(2, 2, *shape)
-        z = torch.stack([
-            z[0] * weight[0, 0] + z[1] * weight[0, 1],
-            z[0] * weight[1, 0] + z[1] * weight[1, 1],
-        ], dim=0) + bias.reshape(2, *shape)
+        z = torch.stack(
+            [
+                z[0] * weight[0, 0] + z[1] * weight[0, 1],
+                z[0] * weight[1, 0] + z[1] * weight[1, 1],
+            ],
+            dim=0,
+        ) + bias.reshape(2, *shape)
 
     return cplx.Cplx(z[0], z[1])
 
@@ -251,6 +283,7 @@ class _CplxBatchNorm(CplxToCplx):
 
     Taken from `torch.nn.modules.batchnorm` verbatim.
     """
+
     def __init__(
         self,
         num_features,
@@ -271,18 +304,20 @@ def __init__(
             self.bias = torch.nn.Parameter(torch.empty(2, num_features))
 
         else:
-            self.register_parameter('weight', None)
-            self.register_parameter('bias', None)
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
 
         if self.track_running_stats:
-            self.register_buffer('running_mean', torch.empty(2, num_features))
-            self.register_buffer('running_var', torch.empty(2, 2, num_features))
-            self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
+            self.register_buffer("running_mean", torch.empty(2, num_features))
+            self.register_buffer("running_var", torch.empty(2, 2, num_features))
+            self.register_buffer(
+                "num_batches_tracked", torch.tensor(0, dtype=torch.long)
+            )
 
         else:
-            self.register_parameter('running_mean', None)
-            self.register_parameter('running_var', None)
-            self.register_parameter('num_batches_tracked', None)
+            self.register_parameter("running_mean", None)
+            self.register_parameter("running_var", None)
+            self.register_parameter("num_batches_tracked", None)
 
         self.reset_running_stats()
         self.reset_parameters()
@@ -294,13 +329,13 @@ def reset_running_stats(self):
         self.num_batches_tracked.zero_()
 
         self.running_mean.zero_()
-        self.running_var.copy_(torch.eye(2,  2).unsqueeze(-1))
+        self.running_var.copy_(torch.eye(2, 2).unsqueeze(-1))
 
     def reset_parameters(self):
         if not self.affine:
             return
 
-        self.weight.data.copy_(torch.eye(2,  2).unsqueeze(-1))
+        self.weight.data.copy_(torch.eye(2, 2).unsqueeze(-1))
         init.zeros_(self.bias)
 
     def _check_input_dim(self, input):
@@ -323,40 +358,50 @@ def forward(self, input):
                     exponential_average_factor = self.momentum
 
         return cplx_batch_norm(
-            input, self.running_mean, self.running_var, self.weight, self.bias,
+            input,
+            self.running_mean,
+            self.running_var,
+            self.weight,
+            self.bias,
             self.training or not self.track_running_stats,
-            exponential_average_factor, self.eps)
+            exponential_average_factor,
+            self.eps,
+        )
 
     def extra_repr(self):
-        return '{num_features}, eps={eps}, momentum={momentum}, affine={affine}, ' \
-               'track_running_stats={track_running_stats}'.format(**vars(self))
+        return (
+            "{num_features}, eps={eps}, momentum={momentum}, affine={affine}, "
+            "track_running_stats={track_running_stats}".format(**vars(self))
+        )
 
 
 class CplxBatchNorm1d(_CplxBatchNorm):
     """Complex-valued batch normalization for 2D or 3D data.
     See torch.nn.BatchNorm1d for details.
     """
+
     def _check_input_dim(self, input):
         if input.dim() != 2 and input.dim() != 3:
-            raise ValueError('expected 2D or 3D input (got {}D input)'
-                             .format(input.dim()))
+            raise ValueError(
+                "expected 2D or 3D input (got {}D input)".format(input.dim())
+            )
 
 
 class CplxBatchNorm2d(_CplxBatchNorm):
     """Complex-valued batch normalization for 4D data.
     See torch.nn.BatchNorm2d for details.
     """
+
     def _check_input_dim(self, input):
         if input.dim() != 4:
-            raise ValueError('expected 4D input (got {}D input)'
-                             .format(input.dim()))
+            raise ValueError("expected 4D input (got {}D input)".format(input.dim()))
 
 
 class CplxBatchNorm3d(_CplxBatchNorm):
     """Complex-valued batch normalization for 5D data.
     See torch.nn.BatchNorm3d for details.
     """
+
     def _check_input_dim(self, input):
         if input.dim() != 5:
-            raise ValueError('expected 5D input (got {}D input)'
-                             .format(input.dim()))
+            raise ValueError("expected 5D input (got {}D input)".format(input.dim()))
diff --git a/cplxmodule/nn/modules/casting.py b/cplxmodule/nn/modules/casting.py
index 78e6993..6a81af2 100644
--- a/cplxmodule/nn/modules/casting.py
+++ b/cplxmodule/nn/modules/casting.py
@@ -20,6 +20,7 @@ class InterleavedRealToCplx(BaseRealToCplx):
 
     Inverts `CplxToInterleavedReal`.
     """
+
     def __init__(self, copy=False, dim=-1):
         super().__init__()
         self.copy, self.dim = copy, dim
@@ -43,6 +44,7 @@ class ConcatenatedRealToCplx(BaseRealToCplx):
 
     Inverts `CplxToConcatenatedReal`.
     """
+
     def __init__(self, copy=False, dim=-1):
         super().__init__()
         self.copy, self.dim = copy, dim
@@ -63,6 +65,7 @@ class CplxToInterleavedReal(BaseCplxToReal):
 
     Inverts `InterleavedRealToCplx`.
     """
+
     def __init__(self, dim=-1):
         super().__init__()
         self.dim = dim
@@ -83,6 +86,7 @@ class CplxToConcatenatedReal(BaseCplxToReal):
 
     Inverts `ConcatenatedRealToCplx`.
     """
+
     def __init__(self, dim=-1):
         super().__init__()
         self.dim = dim
@@ -103,6 +107,7 @@ class AsTypeCplx(BaseRealToCplx):
 
     Inverts `nn.linear.CplxReal`.
     """
+
     def forward(self, input):
         return cplx.Cplx(input)
 
@@ -120,6 +125,7 @@ class TensorToCplx(BaseRealToCplx):
 
     Inverts `CplxToTensor`.
     """
+
     def forward(self, input):
         """input must be a , and may have
         arbitrary number of.
@@ -141,5 +147,6 @@ class CplxToTensor(BaseCplxToReal):
 
     Inverts `TensorToCplx`.
     """
+
     def forward(self, input):
         return cplx.to_interleaved_real(input, False, -1)
diff --git a/cplxmodule/nn/modules/container.py b/cplxmodule/nn/modules/container.py
index 608b885..03d3b99 100644
--- a/cplxmodule/nn/modules/container.py
+++ b/cplxmodule/nn/modules/container.py
@@ -13,6 +13,7 @@ class CplxSequential(torch.nn.Sequential, CplxToCplx):
         \,, $$
     for $l=1..L$ and the complex input $z_0$.
     """
+
     def __init__(self, *args):
         # make a simple typecheck on the passed arguments
         if len(args) == 1 and isinstance(args[0], OrderedDict):
@@ -22,8 +23,10 @@ def __init__(self, *args):
 
         bad_modules = [str(n) for n, m in modules if not is_cplx_to_cplx(m)]
         if bad_modules:
-            raise TypeError(f"""Only complex-to-complex modules can be used """
-                            f"""in {self.__class__.__name__}. The following """
-                            f"""modules failed: {bad_modules}.""")
+            raise TypeError(
+                f"Only complex-to-complex modules can be used "
+                f"in {self.__class__.__name__}. The following "
+                f"modules failed: {bad_modules}."
+            )
 
         super().__init__(*args)
diff --git a/cplxmodule/nn/modules/conv.py b/cplxmodule/nn/modules/conv.py
index 26e8b8c..8b18e3d 100644
--- a/cplxmodule/nn/modules/conv.py
+++ b/cplxmodule/nn/modules/conv.py
@@ -10,9 +10,21 @@
 
 class CplxConvNd(CplxToCplx):
     r"""An almost verbatim copy of `_ConvNd` from torch/nn/modules/conv.py"""
-    def __init__(self, in_channels, out_channels, kernel_size, stride,
-                 padding, dilation, transposed, output_padding, groups,
-                 bias, padding_mode):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        transposed,
+        output_padding,
+        groups,
+        bias,
+        padding_mode,
+    ):
         super().__init__()
 
         if in_channels % groups != 0:
@@ -31,11 +43,13 @@ def __init__(self, in_channels, out_channels, kernel_size, stride,
         self.groups, self.padding_mode = groups, padding_mode
 
         if transposed:
-            self.weight = CplxParameter(cplx.Cplx.empty(
-                in_channels, out_channels // groups, *kernel_size))
+            self.weight = CplxParameter(
+                cplx.Cplx.empty(in_channels, out_channels // groups, *kernel_size)
+            )
         else:
-            self.weight = CplxParameter(cplx.Cplx.empty(
-                out_channels, in_channels // groups, *kernel_size))
+            self.weight = CplxParameter(
+                cplx.Cplx.empty(out_channels, in_channels // groups, *kernel_size)
+            )
 
         if bias:
             self.bias = CplxParameter(cplx.Cplx.empty(out_channels))
@@ -52,8 +66,10 @@ def reset_parameters(self):
             init.cplx_uniform_independent_(self.bias, -bound, bound)
 
     def extra_repr(self):
-        s = ("{in_channels}, {out_channels}, kernel_size={kernel_size}"
-             ", stride={stride}")
+        s = (
+            "{in_channels}, {out_channels}, kernel_size={kernel_size}"
+            ", stride={stride}"
+        )
 
         if self.padding != (0,) * len(self.padding):
             s += ", padding={padding}"
@@ -88,25 +104,44 @@ class CplxConv1d(CplxConvNd):
 
     See torch.nn.Conv1d for reference on the input dimensions.
     """
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 dilation=1,
-                 groups=1,
-                 bias=True,
-                 padding_mode="zeros"):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
         super().__init__(
-            in_channels, out_channels, _single(kernel_size), _single(stride),
-            _single(padding), _single(dilation), False, _single(0), groups,
-            bias, padding_mode)
+            in_channels,
+            out_channels,
+            _single(kernel_size),
+            _single(stride),
+            _single(padding),
+            _single(dilation),
+            False,
+            _single(0),
+            groups,
+            bias,
+            padding_mode,
+        )
 
     def forward(self, input):
-        return cplx.conv1d(input, self.weight, self.bias,
-                           self.stride[0], self.padding[0], self.dilation[0],
-                           self.groups, self.padding_mode)
+        return cplx.conv1d(
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.padding_mode,
+        )
 
 
 class CplxConv2d(CplxConvNd):
@@ -121,25 +156,44 @@ class CplxConv2d(CplxConvNd):
 
     See torch.nn.Conv2d for reference on the input dimensions.
     """
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 dilation=1,
-                 groups=1,
-                 bias=True,
-                 padding_mode="zeros"):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
         super().__init__(
-            in_channels, out_channels, _pair(kernel_size), _pair(stride),
-            _pair(padding), _pair(dilation), False, _pair(0), groups,
-            bias, padding_mode)
+            in_channels,
+            out_channels,
+            _pair(kernel_size),
+            _pair(stride),
+            _pair(padding),
+            _pair(dilation),
+            False,
+            _pair(0),
+            groups,
+            bias,
+            padding_mode,
+        )
 
     def forward(self, input):
-        return cplx.conv2d(input, self.weight, self.bias,
-                           self.stride, self.padding, self.dilation,
-                           self.groups, self.padding_mode)
+        return cplx.conv2d(
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.padding_mode,
+        )
 
 
 class CplxConv3d(CplxConvNd):
@@ -154,40 +208,83 @@ class CplxConv3d(CplxConvNd):
 
     See torch.nn.Conv2d for reference on the input dimensions.
     """
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 dilation=1,
-                 groups=1,
-                 bias=True,
-                 padding_mode="zeros"):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
         super().__init__(
-            in_channels, out_channels, _triple(kernel_size), _triple(stride),
-            _triple(padding), _triple(dilation), False, _triple(0), groups,
-            bias, padding_mode)
+            in_channels,
+            out_channels,
+            _triple(kernel_size),
+            _triple(stride),
+            _triple(padding),
+            _triple(dilation),
+            False,
+            _triple(0),
+            groups,
+            bias,
+            padding_mode,
+        )
 
     def forward(self, input):
-        return cplx.conv3d(input, self.weight, self.bias,
-                           self.stride, self.padding, self.dilation,
-                           self.groups, self.padding_mode)
+        return cplx.conv3d(
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+            self.padding_mode,
+        )
 
 
 class CplxConvTransposeNd(CplxConvNd):
     r"""An almost verbatim copy of `_ConvTransposeNd` from torch/nn/modules/conv.py"""
-    def __init__(self, in_channels, out_channels, kernel_size, stride,
-                 padding, dilation, transposed, output_padding, groups,
-                 bias, padding_mode):
-        if padding_mode not in ('zeros', 'circular'):
-            raise ValueError('Only "zeros" or "circular" padding mode are supported for {}'\
-                             .format(self.__class__.__name__))
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        transposed,
+        output_padding,
+        groups,
+        bias,
+        padding_mode,
+    ):
+        if padding_mode not in ("zeros", "circular"):
+            raise ValueError(
+                'Only "zeros" or "circular" padding mode are supported by `{}`'.format(
+                    self.__class__.__name__
+                )
+            )
 
         super().__init__(
-            in_channels, out_channels, _triple(kernel_size), _triple(stride),
-            _triple(padding), _triple(dilation), transposed, output_padding,
-            groups, bias, padding_mode)
+            in_channels,
+            out_channels,
+            _triple(kernel_size),
+            _triple(stride),
+            _triple(padding),
+            _triple(dilation),
+            transposed,
+            output_padding,
+            groups,
+            bias,
+            padding_mode,
+        )
 
     _output_padding = torch.nn.modules.conv._ConvTransposeNd._output_padding
 
@@ -197,32 +294,55 @@ class CplxConvTranspose1d(CplxConvTransposeNd):
 
     See torch.nn.ConvTranspose1d for reference on the input dimensions.
     """
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 dilation=1,
-                 output_padding=0,
-                 groups=1,
-                 bias=None,
-                 padding_mode="zeros"):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        output_padding=0,
+        groups=1,
+        bias=None,
+        padding_mode="zeros",
+    ):
         super().__init__(
-            in_channels, out_channels, _single(kernel_size), _single(stride),
-            _single(padding), _single(dilation), True, _single(output_padding),
-            groups, bias, padding_mode)
+            in_channels,
+            out_channels,
+            _single(kernel_size),
+            _single(stride),
+            _single(padding),
+            _single(dilation),
+            True,
+            _single(output_padding),
+            groups,
+            bias,
+            padding_mode,
+        )
 
     def forward(self, input, output_size=None):
-        if self.padding_mode not in ('zeros', 'circular'):
-            raise ValueError('Only `zeros` or `circular` padding mode are '
-                             'supported for CplxConvTranspose1d')
-        output_padding = self._output_padding(input, output_size, self.stride,
-                                              self.padding, self.kernel_size)
+        if self.padding_mode not in ("zeros", "circular"):
+            raise ValueError(
+                "Only `zeros` or `circular` padding mode are supported "
+                "by `CplxConvTranspose1d`"
+            )
+        output_padding = self._output_padding(
+            input, output_size, self.stride, self.padding, self.kernel_size
+        )
 
         return cplx.conv_transpose1d(
-            input, self.weight, self.bias, self.stride, self.padding,
-            output_padding, self.groups, self.dilation, self.padding_mode)
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            output_padding,
+            self.groups,
+            self.dilation,
+            self.padding_mode,
+        )
 
 
 class CplxConvTranspose2d(CplxConvTransposeNd):
@@ -230,32 +350,55 @@ class CplxConvTranspose2d(CplxConvTransposeNd):
 
     See torch.nn.ConvTranspose2d for reference on the input dimensions.
     """
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 dilation=1,
-                 output_padding=0,
-                 groups=1,
-                 bias=None,
-                 padding_mode="zeros"):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        output_padding=0,
+        groups=1,
+        bias=None,
+        padding_mode="zeros",
+    ):
         super().__init__(
-            in_channels, out_channels, _pair(kernel_size), _pair(stride),
-            _pair(padding), _pair(dilation), True, _pair(output_padding),
-            groups, bias, padding_mode)
+            in_channels,
+            out_channels,
+            _pair(kernel_size),
+            _pair(stride),
+            _pair(padding),
+            _pair(dilation),
+            True,
+            _pair(output_padding),
+            groups,
+            bias,
+            padding_mode,
+        )
 
     def forward(self, input, output_size=None):
-        if self.padding_mode not in ('zeros', 'circular'):
-            raise ValueError('Only `zeros` or `circular` padding mode are '
-                             'supported for CplxConvTranspose2d')
-        output_padding = self._output_padding(input, output_size, self.stride,
-                                              self.padding, self.kernel_size)
+        if self.padding_mode not in ("zeros", "circular"):
+            raise ValueError(
+                "Only `zeros` or `circular` padding mode are supported "
+                "by `CplxConvTranspose2d`"
+            )
+        output_padding = self._output_padding(
+            input, output_size, self.stride, self.padding, self.kernel_size
+        )
 
         return cplx.conv_transpose2d(
-            input, self.weight, self.bias, self.stride, self.padding,
-            output_padding, self.groups, self.dilation, self.padding_mode)
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            output_padding,
+            self.groups,
+            self.dilation,
+            self.padding_mode,
+        )
 
 
 class CplxConvTranspose3d(CplxConvTransposeNd):
@@ -263,29 +406,52 @@ class CplxConvTranspose3d(CplxConvTransposeNd):
 
     See torch.nn.ConvTranspose3d for reference on the input dimensions.
     """
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 dilation=1,
-                 output_padding=0,
-                 groups=1,
-                 bias=None,
-                 padding_mode="zeros"):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        output_padding=0,
+        groups=1,
+        bias=None,
+        padding_mode="zeros",
+    ):
         super().__init__(
-            in_channels, out_channels, _triple(kernel_size), _triple(stride),
-            _triple(padding), _triple(dilation), True, _triple(output_padding),
-            groups, bias, padding_mode)
+            in_channels,
+            out_channels,
+            _triple(kernel_size),
+            _triple(stride),
+            _triple(padding),
+            _triple(dilation),
+            True,
+            _triple(output_padding),
+            groups,
+            bias,
+            padding_mode,
+        )
 
     def forward(self, input, output_size=None):
-        if self.padding_mode not in ('zeros', 'circular'):
-            raise ValueError('Only `zeros` or `circular` padding mode are '
-                             'supported for CplxConvTranspose3d')
-        output_padding = self._output_padding(input, output_size, self.stride,
-                                              self.padding, self.kernel_size)
+        if self.padding_mode not in ("zeros", "circular"):
+            raise ValueError(
+                "Only `zeros` or `circular` padding mode are supported "
+                "by `CplxConvTranspose3d`"
+            )
+        output_padding = self._output_padding(
+            input, output_size, self.stride, self.padding, self.kernel_size
+        )
 
         return cplx.conv_transpose3d(
-            input, self.weight, self.bias, self.stride, self.padding,
-            output_padding, self.groups, self.dilation, self.padding_mode)
+            input,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            output_padding,
+            self.groups,
+            self.dilation,
+            self.padding_mode,
+        )
diff --git a/cplxmodule/nn/modules/extra.py b/cplxmodule/nn/modules/extra.py
index f39b67d..3438af2 100644
--- a/cplxmodule/nn/modules/extra.py
+++ b/cplxmodule/nn/modules/extra.py
@@ -10,6 +10,7 @@ class CplxDropout(torch.nn.Dropout2d, CplxToCplx):
 
     See torch.nn.Dropout1d for reference on the input dimensions and arguments.
     """
+
     def forward(self, input):
         *head, n_last = input.shape
 
diff --git a/cplxmodule/nn/modules/linear.py b/cplxmodule/nn/modules/linear.py
index a65326c..0c2f897 100644
--- a/cplxmodule/nn/modules/linear.py
+++ b/cplxmodule/nn/modules/linear.py
@@ -32,6 +32,7 @@ class CplxLinear(CplxToCplx):
                     + i (W_\mathrm{im} u + W_\mathrm{re} v)
         \,. $$
     """
+
     def __init__(self, in_features, out_features, bias=True):
         super().__init__()
         self.in_features = in_features
@@ -42,7 +43,7 @@ def __init__(self, in_features, out_features, bias=True):
         if bias:
             self.bias = CplxParameter(cplx.Cplx.empty(out_features))
         else:
-            self.register_parameter('bias', None)
+            self.register_parameter("bias", None)
 
         self.reset_parameters()
 
@@ -58,7 +59,7 @@ def forward(self, input):
         return cplx.linear(input, self.weight, self.bias)
 
     def extra_repr(self):
-        return 'in_features={}, out_features={}, bias={}'.format(
+        return "in_features={}, out_features={}, bias={}".format(
             self.in_features, self.out_features, self.bias is not None
         )
 
@@ -73,20 +74,23 @@ class CplxBilinear(CplxToCplx):
         \colon (u, v) \mapsto (u^\top A_j v)_{j=1}^{d_2}
         \,. $$
     """
-    def __init__(self, in1_features, in2_features, out_features, bias=True,
-                 conjugate=True):
+
+    def __init__(
+        self, in1_features, in2_features, out_features, bias=True, conjugate=True
+    ):
         super().__init__()
         self.in1_features = in1_features
         self.in2_features = in2_features
         self.out_features = out_features
 
-        self.weight = CplxParameter(cplx.Cplx.empty(
-            out_features, in1_features, in2_features))
+        self.weight = CplxParameter(
+            cplx.Cplx.empty(out_features, in1_features, in2_features)
+        )
 
         if bias:
             self.bias = CplxParameter(cplx.Cplx.empty(out_features))
         else:
-            self.register_parameter('bias', None)
+            self.register_parameter("bias", None)
 
         self.conjugate = conjugate
 
@@ -100,15 +104,17 @@ def reset_parameters(self):
             init.cplx_uniform_independent_(self.bias, -bound, bound)
 
     def forward(self, input1, input2):
-        return cplx.bilinear(input1, input2, self.weight,
-                             self.bias, self.conjugate)
+        return cplx.bilinear(input1, input2, self.weight, self.bias, self.conjugate)
 
     def extra_repr(self):
-        fmt = """in1_features={}, in2_features={}, out_features={}, """
-        fmt += """bias={}, conjugate={}"""
+        fmt = "in1_features={}, in2_features={}, out_features={}, bias={}, conjugate={}"
         return fmt.format(
-            self.in1_features, self.in2_features, self.out_features,
-            self.bias is not None, self.conjugate)
+            self.in1_features,
+            self.in2_features,
+            self.out_features,
+            self.bias is not None,
+            self.conjugate,
+        )
 
 
 class CplxPhaseShift(CplxToCplx):
@@ -126,6 +132,7 @@ class CplxPhaseShift(CplxToCplx):
     each complex feature in all channels by the same phase. Finally calling
     with CplxPhaseShift(1) shifts the inputs by the single common phase.
     """
+
     def __init__(self, *dim):
         super().__init__()
         self.phi = torch.nn.Parameter(torch.randn(*dim) * 0.02)
diff --git a/cplxmodule/nn/modules/pooling.py b/cplxmodule/nn/modules/pooling.py
index 09aef07..58dcb61 100644
--- a/cplxmodule/nn/modules/pooling.py
+++ b/cplxmodule/nn/modules/pooling.py
@@ -1,16 +1,19 @@
-import math
-import torch
-
-from torch.nn.modules.utils import _single, _pair, _triple
-
 from .base import CplxToCplx
 from ... import cplx
 
 
 class CplxMaxPoolNd(CplxToCplx):
     r"""An almost verbatim copy of `_MaxPoolNd` from torch/nn/modules/pooling.py"""
-    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
-                 return_indices=False, ceil_mode=False) -> None:
+
+    def __init__(
+        self,
+        kernel_size,
+        stride=None,
+        padding=0,
+        dilation=1,
+        return_indices=False,
+        ceil_mode=False,
+    ) -> None:
         super().__init__()
         self.kernel_size = kernel_size
         self.stride = stride if (stride is not None) else kernel_size
@@ -20,8 +23,10 @@ def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
         self.ceil_mode = ceil_mode
 
     def extra_repr(self) -> str:
-        return 'kernel_size={kernel_size}, stride={stride}, padding={padding}' \
-            ', dilation={dilation}, ceil_mode={ceil_mode}'.format(**self.__dict__)
+        return (
+            "kernel_size={kernel_size}, stride={stride}, padding={padding}"
+            ", dilation={dilation}, ceil_mode={ceil_mode}".format(**self.__dict__)
+        )
 
 
 class CplxMaxPool1d(CplxMaxPoolNd):
@@ -32,9 +37,16 @@ class CplxMaxPool1d(CplxMaxPoolNd):
     --------
     See docs for the parameters of `torch.nn.MaxPool1d`.
     """
+
     def forward(self, input):
-        return cplx.max_pool1d(input, self.kernel_size, self.stride,
-                               self.padding, self.dilation, self.ceil_mode)
+        return cplx.max_pool1d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.ceil_mode,
+        )
 
 
 class CplxMaxPool2d(CplxMaxPoolNd):
@@ -45,9 +57,16 @@ class CplxMaxPool2d(CplxMaxPoolNd):
     --------
     See docs for the parameters of `torch.nn.MaxPool2d`.
     """
+
     def forward(self, input):
-        return cplx.max_pool2d(input, self.kernel_size, self.stride,
-                               self.padding, self.dilation, self.ceil_mode)
+        return cplx.max_pool2d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.ceil_mode,
+        )
 
 
 class CplxMaxPool3d(CplxMaxPoolNd):
@@ -58,6 +77,13 @@ class CplxMaxPool3d(CplxMaxPoolNd):
     --------
     See docs for the parameters of `torch.nn.MaxPool3d`.
     """
+
     def forward(self, input):
-        return cplx.max_pool3d(input, self.kernel_size, self.stride,
-                               self.padding, self.dilation, self.ceil_mode)
+        return cplx.max_pool3d(
+            input,
+            self.kernel_size,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.ceil_mode,
+        )
diff --git a/cplxmodule/nn/relevance/ard.py b/cplxmodule/nn/relevance/ard.py
deleted file mode 100644
index 88c84eb..0000000
--- a/cplxmodule/nn/relevance/ard.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import warnings
-
-warnings.warn("Importing Automatic Relevance Determination layers from"
-              " `cplxmodule.nn.relevance.ard` has been deprecated since"
-              " version `2020.8` and will be removed in a later version."
-              " Please, import from `cplxmodule.nn.relevance.real.ard` or"
-              " `cplxmodule.nn.relevance.complex.ard`.",
-              DeprecationWarning)
-
-from .real import LinearARD, BilinearARD
-from .real import Conv1dARD, Conv2dARD, Conv3dARD
-
-from .complex import CplxLinearARD, CplxBilinearARD
-from .complex import CplxConv1dARD, CplxConv2dARD, CplxConv3dARD
diff --git a/cplxmodule/nn/relevance/base.py b/cplxmodule/nn/relevance/base.py
index 1ab8d02..e96980d 100644
--- a/cplxmodule/nn/relevance/base.py
+++ b/cplxmodule/nn/relevance/base.py
@@ -36,6 +36,7 @@ class BaseARD(torch.nn.Module):
         For $\tau=0.25$ and $\beta=0.66$ we have `threshold=2.96`.
 
     """
+
     @property
     def penalty(self):
         """Get the penalty induced by the variational approximation.
@@ -60,8 +61,7 @@ def penalty(self):
 
         # if the property in a derived class raises then it gets
         # silently undefined due to inheritance from torch.nn.Module.
-        raise NotImplementedError("Derived classes must compute "
-                                  "their own penalty.")
+        raise NotImplementedError("Derived classes must compute " "their own penalty.")
 
     def relevance(self, **kwargs):
         r"""Get the dropout mask based on the provided parameters.
@@ -80,8 +80,9 @@ def relevance(self, **kwargs):
             regularization. Soft masks would require cleaning up to eliminate
             the result of such co-adaptation (see `nn.masked.binarize_masks`).
         """
-        raise NotImplementedError("Derived classes must implement a float "
-                                  "mask of relevant coefficients.")
+        raise NotImplementedError(
+            "Derived classes must implement a float mask of relevant coefficients."
+        )
 
 
 def named_penalties(module, reduction="sum", prefix=""):
@@ -123,8 +124,9 @@ def named_penalties(module, reduction="sum", prefix=""):
     """
 
     if reduction is not None and reduction not in ("mean", "sum"):
-        raise ValueError(f"`reduction` must be either `None`,"
-                         f" `sum` or `mean`. Got {reduction}.")
+        raise ValueError(
+            f"`reduction` must be either `None`, `sum` or `mean`. Got {reduction}."
+        )
 
     # yields own penalty and penalties of all descendants
     for name, mod in module.named_modules(prefix=prefix):
@@ -158,7 +160,7 @@ def penalties(module, reduction="sum"):
     torch.Tensor
         A differentiable value of the penalty.
     """
-    for name, penalty in named_penalties(module, reduction=reduction):
+    for _, penalty in named_penalties(module, reduction=reduction):
         yield penalty
 
 
@@ -211,7 +213,4 @@ def compute_ard_masks(module, *, prefix="", **kwargs):
         return {}
 
     relevance = named_relevance(module, prefix=prefix, **kwargs)
-    return {
-        name + ("." if name else "") + "mask": mask
-        for name, mask in relevance
-    }
+    return {name + ("." if name else "") + "mask": mask for name, mask in relevance}
diff --git a/cplxmodule/nn/relevance/complex/ard.py b/cplxmodule/nn/relevance/complex/ard.py
index 4b69026..6656723 100644
--- a/cplxmodule/nn/relevance/complex/ard.py
+++ b/cplxmodule/nn/relevance/complex/ard.py
@@ -1,11 +1,11 @@
-import torch
+import torch  # noqa: F401
 import torch.nn.functional as F
 
 from .vd import CplxLinearVD, CplxBilinearVD
 from .vd import CplxConv1dVD, CplxConv2dVD, CplxConv3dVD
 
 
-class CplxARDMixin():
+class CplxARDMixin:
     @property
     def penalty(self):
         r"""Empricial Bayes penalty for complex layer with complex gaussian vi.
@@ -36,13 +36,14 @@ def penalty(self):
         """
 
         # `softplus` is $x \mapsto \log(1 + e^x)$
-        return F.softplus(- self.log_alpha)
+        return F.softplus(-self.log_alpha)
 
 
 class CplxLinearARD(CplxARDMixin, CplxLinearVD):
     """Complex-valued linear layer with automatic relevance
     determination.
     """
+
     pass
 
 
@@ -50,6 +51,7 @@ class CplxBilinearARD(CplxARDMixin, CplxBilinearVD):
     """Complex-valued bilinear layer with automatic relevance
     determination.
     """
+
     pass
 
 
@@ -57,6 +59,7 @@ class CplxConv1dARD(CplxARDMixin, CplxConv1dVD):
     """1D complex-valued convolution layer with automatic relevance
     determination.
     """
+
     pass
 
 
@@ -64,6 +67,7 @@ class CplxConv2dARD(CplxARDMixin, CplxConv2dVD):
     """2D complex-valued convolution layer with automatic relevance
     determination.
     """
+
     pass
 
 
@@ -71,4 +75,5 @@ class CplxConv3dARD(CplxARDMixin, CplxConv3dVD):
     """3D complex-valued convolution layer with automatic relevance
     determination.
     """
+
     pass
diff --git a/cplxmodule/nn/relevance/complex/base.py b/cplxmodule/nn/relevance/complex/base.py
index 756cb9e..49fb16d 100644
--- a/cplxmodule/nn/relevance/complex/base.py
+++ b/cplxmodule/nn/relevance/complex/base.py
@@ -20,6 +20,7 @@ class GaussianMixin:
         value of the weight. The higher the log-alpha the less relevant the
         parameter is.
     """
+
     def reset_variational_parameters(self):
         self.log_sigma2.data.uniform_(-10, -10)  # wtf?
 
@@ -46,8 +47,11 @@ def forward(self, input):
             return mu
 
         # \gamma = \sigma^2 (x \odot \bar{x})
-        s2 = F.linear(input.real * input.real + input.imag * input.imag,
-                      torch.exp(self.log_sigma2), None)
+        s2 = F.linear(
+            input.real * input.real + input.imag * input.imag,
+            torch.exp(self.log_sigma2),
+            None,
+        )
 
         return mu + cplx.randn_like(s2) * torch.sqrt(torch.clamp(s2, 1e-8))
 
@@ -55,10 +59,12 @@ def forward(self, input):
 class CplxBilinearGaussian(GaussianMixin, CplxBilinear):
     """Complex-valued bilinear layer with variational dropout."""
 
-    def __init__(self, in1_features, in2_features, out_features, bias=True,
-                 conjugate=True):
-        super().__init__(in1_features, in2_features, out_features,
-                         bias=bias, conjugate=conjugate)
+    def __init__(
+        self, in1_features, in2_features, out_features, bias=True, conjugate=True
+    ):
+        super().__init__(
+            in1_features, in2_features, out_features, bias=bias, conjugate=conjugate
+        )
 
         self.log_sigma2 = torch.nn.Parameter(torch.Tensor(*self.weight.shape))
         self.reset_variational_parameters()
@@ -68,24 +74,45 @@ def forward(self, input1, input2):
         if not self.training:
             return mu
 
-        s2 = F.bilinear(input1.real * input1.real + input1.imag * input1.imag,
-                        input2.real * input2.real + input2.imag * input2.imag,
-                        torch.exp(self.log_sigma2), None)
+        s2 = F.bilinear(
+            input1.real * input1.real + input1.imag * input1.imag,
+            input2.real * input2.real + input2.imag * input2.imag,
+            torch.exp(self.log_sigma2),
+            None,
+        )
 
         return mu + cplx.randn_like(s2) * torch.sqrt(torch.clamp(s2, 1e-8))
 
 
 class CplxConvNdGaussianMixin(GaussianMixin):
-    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
-                 padding=0, dilation=1, groups=1,
-                 bias=True, padding_mode='zeros'):
-        super().__init__(in_channels, out_channels, kernel_size, stride=stride,
-                         padding=padding, dilation=dilation, groups=groups,
-                         bias=bias, padding_mode=padding_mode)
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
 
         if self.padding_mode != "zeros":
-            raise ValueError(f"Only `zeros` padding mode is supported. "
-                             f"Got `{self.padding_mode}`.")
+            raise ValueError(
+                f"Only `zeros` padding mode is supported. Got `{self.padding_mode}`."
+            )
 
         self.log_sigma2 = torch.nn.Parameter(torch.Tensor(*self.weight.shape))
         self.reset_variational_parameters()
@@ -95,9 +122,15 @@ def _forward_impl(self, input, conv):
         if not self.training:
             return mu
 
-        s2 = conv(input.real * input.real + input.imag * input.imag,
-                  torch.exp(self.log_sigma2), None, self.stride,
-                  self.padding, self.dilation, self.groups)
+        s2 = conv(
+            input.real * input.real + input.imag * input.imag,
+            torch.exp(self.log_sigma2),
+            None,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
 
         return mu + cplx.randn_like(s2) * torch.sqrt(torch.clamp(s2, 1e-8))
 
diff --git a/cplxmodule/nn/relevance/complex/vd.py b/cplxmodule/nn/relevance/complex/vd.py
index b741ac8..2be452f 100644
--- a/cplxmodule/nn/relevance/complex/vd.py
+++ b/cplxmodule/nn/relevance/complex/vd.py
@@ -1,5 +1,3 @@
-import warnings
-
 import torch
 import torch.nn
 import scipy
@@ -28,6 +26,7 @@ class ExpiFunction(torch.autograd.Function):
     during the forward pass due to CPU-GPU device switch. Backward pass does
     not suffer from this issue and is computed on-device.
     """
+
     @staticmethod
     def forward(ctx, x):
         ctx.save_for_backward(x)
@@ -96,90 +95,35 @@ class CplxVDMixin:
     @property
     def penalty(self):
         """Exact complex KL divergence."""
-        n_log_alpha = - self.log_alpha
-        return euler_gamma + n_log_alpha - torch_expi(- torch.exp(n_log_alpha))
+        n_log_alpha = -self.log_alpha
+        return euler_gamma + n_log_alpha - torch_expi(-torch.exp(n_log_alpha))
 
 
 class CplxLinearVD(CplxVDMixin, RelevanceMixin, CplxLinearGaussian, BaseARD):
     """Complex-valued linear layer with variational dropout."""
+
     pass
 
 
 class CplxBilinearVD(CplxVDMixin, RelevanceMixin, CplxBilinearGaussian, BaseARD):
     """Complex-valued bilinear layer with variational dropout."""
+
     pass
 
 
 class CplxConv1dVD(CplxVDMixin, RelevanceMixin, CplxConv1dGaussian, BaseARD):
     """1D complex-valued convolution layer with variational dropout."""
+
     pass
 
 
 class CplxConv2dVD(CplxVDMixin, RelevanceMixin, CplxConv2dGaussian, BaseARD):
     """2D complex-valued convolution layer with variational dropout."""
+
     pass
 
 
 class CplxConv3dVD(CplxVDMixin, RelevanceMixin, CplxConv3dGaussian, BaseARD):
     """3D complex-valued convolution layer with variational dropout."""
-    pass
 
-
-class CplxLinearARD(object):
-    def __new__(cls, in_features, out_features, bias=True):
-        warnings.warn("Importing complex-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.complex` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.complex` submodule will export complex-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return CplxLinearVD(in_features, out_features, bias)
-
-
-class CplxBilinearARD(object):
-    def __new__(cls, in1_features, in2_features, out_features, bias=True,
-                conjugate=True):
-        warnings.warn("Importing complex-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.complex` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.complex` submodule will export complex-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return CplxBilinearVD(in1_features, in2_features, out_features, bias,
-                              conjugate)
-
-
-class CplxConv1dARD(object):
-    def __new__(cls, in_channels, out_channels, kernel_size, stride=1,
-                padding=0, dilation=1, groups=1,
-                bias=True, padding_mode='zeros'):
-        warnings.warn("Importing complex-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.complex` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.complex` submodule will export complex-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return CplxConv1dVD(in_channels, out_channels, kernel_size, stride,
-                            padding, dilation, groups, bias, padding_mode)
-
-
-class CplxConv2dARD(object):
-    def __new__(cls, in_channels, out_channels, kernel_size, stride=1,
-                padding=0, dilation=1, groups=1,
-                bias=True, padding_mode='zeros'):
-        warnings.warn("Importing complex-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.complex` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.complex` submodule will export complex-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return CplxConv2dVD(in_channels, out_channels, kernel_size, stride,
-                            padding, dilation, groups, bias, padding_mode)
+    pass
diff --git a/cplxmodule/nn/relevance/extensions/complex.py b/cplxmodule/nn/relevance/extensions/complex.py
index 6531c23..909eabd 100644
--- a/cplxmodule/nn/relevance/extensions/complex.py
+++ b/cplxmodule/nn/relevance/extensions/complex.py
@@ -15,7 +15,7 @@
 from ..complex import ExpiFunction, torch_expi
 
 
-class CplxVDScaleFreeMixin():
+class CplxVDScaleFreeMixin:
     @property
     def penalty(self):
         r"""The Kullback-Leibler divergence between the mean field approximate
@@ -40,36 +40,41 @@ def penalty(self):
         """
         log_abs_w = torch.log(abs(self.weight) + 1e-12)
         n_log_alpha = 2 * log_abs_w - self.log_sigma2
-        ei = torch_expi(- torch.exp(n_log_alpha))
+        ei = torch_expi(-torch.exp(n_log_alpha))
         return log_abs_w - self.log_sigma2 - 0.5 * ei
 
 
 class CplxLinearVDScaleFree(CplxVDScaleFreeMixin, CplxLinearVD):
     """Complex-valued linear layer with scale-free prior."""
+
     pass
 
 
 class CplxBilinearVDScaleFree(CplxVDScaleFreeMixin, CplxBilinearVD):
     """Complex-valued bilinear layer with scale-free prior."""
+
     pass
 
 
 class CplxConv1dVDScaleFree(CplxVDScaleFreeMixin, CplxConv1dVD):
     """1D complex-valued convolution layer with scale-free prior."""
+
     pass
 
 
 class CplxConv2dVDScaleFree(CplxVDScaleFreeMixin, CplxConv2dVD):
     """2D complex-valued convolution layer with scale-free prior."""
+
     pass
 
 
 class CplxConv3dVDScaleFree(CplxVDScaleFreeMixin, CplxConv3dVD):
     """3D complex-valued convolution layer with scale-free prior."""
+
     pass
 
 
-class CplxVDApproxMixin():
+class CplxVDApproxMixin:
     @property
     def penalty(self):
         r"""Softplus-sigmoid approximation of the complex KL divergence.
@@ -90,7 +95,7 @@ def penalty(self):
         See the accompanying notebook for the MC estimation of the k1-k3
         constants: `k1, k2, k3 = 0.57810091, 1.45926293, 1.36525956`
         """
-        n_log_alpha = - self.log_alpha
+        n_log_alpha = -self.log_alpha
         sigmoid = torch.sigmoid(1.36526 * n_log_alpha - 1.45926)
         return F.softplus(n_log_alpha) + 0.57810 * sigmoid
 
@@ -99,6 +104,7 @@ class CplxLinearVDApprox(CplxVDApproxMixin, CplxLinearVD):
     """Complex-valued linear layer with approximate
     var-dropout penalty.
     """
+
     pass
 
 
@@ -106,6 +112,7 @@ class CplxBilinearVDApprox(CplxVDApproxMixin, CplxBilinearVD):
     """Complex-valued bilinear layer with approximate
     var-dropout penalty.
     """
+
     pass
 
 
@@ -113,6 +120,7 @@ class CplxConv1dVDApprox(CplxVDApproxMixin, CplxConv1dVD):
     """1D complex-valued convolution layer with approximate
     var-dropout penalty.
     """
+
     pass
 
 
@@ -120,6 +128,7 @@ class CplxConv2dVDApprox(CplxVDApproxMixin, CplxConv2dVD):
     """2D complex-valued convolution layer with approximate
     var-dropout penalty.
     """
+
     pass
 
 
@@ -127,6 +136,7 @@ class CplxConv3dVDApprox(CplxVDApproxMixin, CplxConv3dVD):
     """3D complex-valued convolution layer with approximate
     var-dropout penalty.
     """
+
     pass
 
 
@@ -135,6 +145,7 @@ class BogusExpiFunction(ExpiFunction):
     but correct values on the backwards pass, provided there is no downstream
     dependence on its forward-pass output.
     """
+
     @staticmethod
     def forward(ctx, x):
         ctx.save_for_backward(x)
@@ -144,18 +155,19 @@ def forward(ctx, x):
 bogus_expi = BogusExpiFunction.apply
 
 
-class CplxVDBogusMixin():
+class CplxVDBogusMixin:
     @property
     def penalty(self):
         r"""KL-div with bogus forward output, but correct gradient."""
         log_alpha = self.log_alpha
-        return - log_alpha - bogus_expi(- torch.exp(- log_alpha))
+        return -log_alpha - bogus_expi(-torch.exp(-log_alpha))
 
 
 class CplxLinearVDBogus(CplxVDBogusMixin, CplxLinearVD):
     """Complex-valued linear layer with correct var dropout penalty
     gradient, but bogus penalty values.
     """
+
     pass
 
 
@@ -163,6 +175,7 @@ class CplxBilinearVDBogus(CplxVDBogusMixin, CplxBilinearVD):
     """Complex-valued bilinear layer with correct var dropout penalty
     gradient, but bogus penalty values.
     """
+
     pass
 
 
@@ -170,6 +183,7 @@ class CplxConv1dVDBogus(CplxVDBogusMixin, CplxConv1dVD):
     """1D complex-valued convolution layer with correct var dropout penalty
     gradient, but bogus penalty values.
     """
+
     pass
 
 
@@ -177,6 +191,7 @@ class CplxConv2dVDBogus(CplxVDBogusMixin, CplxConv2dVD):
     """2D complex-valued convolution layer with correct var dropout penalty
     gradient, but bogus penalty values.
     """
+
     pass
 
 
@@ -184,4 +199,5 @@ class CplxConv3dVDBogus(CplxVDBogusMixin, CplxConv3dVD):
     """3D complex-valued convolution layer with correct var dropout penalty
     gradient, but bogus penalty values.
     """
+
     pass
diff --git a/cplxmodule/nn/relevance/extensions/real/ell_zero.py b/cplxmodule/nn/relevance/extensions/real/ell_zero.py
index 5267e26..9e31eb3 100644
--- a/cplxmodule/nn/relevance/extensions/real/ell_zero.py
+++ b/cplxmodule/nn/relevance/extensions/real/ell_zero.py
@@ -1,5 +1,3 @@
-import warnings
-
 import math
 import torch
 import torch.nn
@@ -41,6 +39,7 @@ class LinearL0(torch.nn.Linear, BaseARD, SparsityStats):
            ICLR 2017
            https://arxiv.org/pdf/1611.00712.pdf
     """
+
     __sparsity_ignore__ = ("log_alpha",)
 
     beta, gamma, zeta = 0.66, -0.1, 1.1
@@ -84,7 +83,7 @@ def penalty(self):
         realtion $1 - \sigma(x) = \sigma(-x)$.
         """
         # compute P(z\neq 0) (minus, due to -ve log-alpha, c.f. eq. (12)).
-        shift = - self.beta * math.log(- self.gamma / self.zeta)
+        shift = -self.beta * math.log(-self.gamma / self.zeta)
         return torch.sigmoid(shift - self.log_alpha)
 
     def forward(self, input):
@@ -93,8 +92,13 @@ def forward(self, input):
         if self.training:
             if n == 1 or m == 1:
                 # this is a relatively "small" batch of unform rv.
-                u = torch.rand(*input.shape[:-1], n, m,
-                               dtype=input.dtype, device=input.device)
+                u = torch.rand(
+                    *input.shape[:-1],
+                    n,
+                    m,
+                    dtype=input.dtype,
+                    device=input.device,
+                )
             else:
                 # one unform sample for the whole batch! Very high gradient var.
                 u = torch.rand_like(self.log_alpha)
@@ -150,7 +154,7 @@ def gate(self, logit=None):
 
         else:
             # on inference in eq. (13) beta is fixed at 1.0, but not in their code
-            s = torch.sigmoid(- self.log_alpha)
+            s = torch.sigmoid(-self.log_alpha)
             # It seems that this beta is very important! But why?
 
         return torch.clamp((self.zeta - self.gamma) * s + self.gamma, 0, 1)
@@ -174,16 +178,3 @@ def relevance(self, *, hard, **kwargs):
     def sparsity(self, *, hard, **kwargs):
         n_relevant = float(self.relevance(hard=hard).sum().item())
         return [(id(self.weight), self.weight.numel() - n_relevant)]
-
-
-class LinearL0ARD(torch.nn.Linear, BaseARD, SparsityStats):
-    def __new__(self, in_features, out_features, bias=True, group=None):
-        warnings.warn("L0 layer learns probabilities of each parameter being"
-                      " nonzero and has little relation to Variational methods."
-                      " It was a serious oversight to name it as such. Thus"
-                      " starting with version `1.0` importing this layer under"
-                      " the incorrect name will be deprecated, and the name"
-                      " will be reverted to `LinearL0`.",
-                      FutureWarning)
-
-        return LinearL0(in_features, out_features, bias, group)
diff --git a/cplxmodule/nn/relevance/real/ard.py b/cplxmodule/nn/relevance/real/ard.py
index d95aa8c..63f3d84 100644
--- a/cplxmodule/nn/relevance/real/ard.py
+++ b/cplxmodule/nn/relevance/real/ard.py
@@ -1,11 +1,11 @@
-import torch
+import torch  # noqa: F401
 import torch.nn.functional as F
 
 from .vd import LinearVD, BilinearVD
 from .vd import Conv1dVD, Conv2dVD, Conv3dVD
 
 
-class RealARDMixin():
+class RealARDMixin:
     @property
     def penalty(self):
         r"""Penalty from arxiv:1811.00596.
@@ -36,29 +36,34 @@ def penalty(self):
         """
 
         # `softplus` is $x \mapsto \log(1 + e^x)$
-        return 0.5 * F.softplus(- self.log_alpha)
+        return 0.5 * F.softplus(-self.log_alpha)
 
 
 class LinearARD(RealARDMixin, LinearVD):
     """Linear layer with automatic relevance determination."""
+
     pass
 
 
 class Conv1dARD(RealARDMixin, Conv1dVD):
     """1D convolution layer with automatic relevance determination."""
+
     pass
 
 
 class Conv2dARD(RealARDMixin, Conv2dVD):
     """2D convolution layer with automatic relevance determination."""
+
     pass
 
 
 class Conv3dARD(RealARDMixin, Conv3dVD):
     """3D convolution layer with automatic relevance determination."""
+
     pass
 
 
 class BilinearARD(RealARDMixin, BilinearVD):
     """Bilinear layer with automatic relevance determination."""
+
     pass
diff --git a/cplxmodule/nn/relevance/real/base.py b/cplxmodule/nn/relevance/real/base.py
index 6595f7e..f918dda 100644
--- a/cplxmodule/nn/relevance/real/base.py
+++ b/cplxmodule/nn/relevance/real/base.py
@@ -15,6 +15,7 @@ class GaussianMixin:
         value of the weight. The higher the log-alpha the less relevant the
         parameter is.
     """
+
     def reset_variational_parameters(self):
         # initially everything is relevant
         self.log_sigma2.data.uniform_(-10, -10)
@@ -71,24 +72,43 @@ def forward(self, input1, input2):
         if not self.training:
             return mu
 
-        s2 = F.bilinear(input1 * input1, input2 * input2,
-                        torch.exp(self.log_sigma2), None)
+        s2 = F.bilinear(
+            input1 * input1, input2 * input2, torch.exp(self.log_sigma2), None
+        )
 
         # .normal reports a grad-fn, but weirdly does not pass grads!
         return mu + torch.randn_like(s2) * torch.sqrt(torch.clamp(s2, 1e-8))
 
 
 class ConvNdGaussianMixin(GaussianMixin):
-    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
-                 padding=0, dilation=1, groups=1,
-                 bias=True, padding_mode='zeros'):
-        super().__init__(in_channels, out_channels, kernel_size, stride=stride,
-                         padding=padding, dilation=dilation, groups=groups,
-                         bias=bias, padding_mode=padding_mode)
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        padding_mode="zeros",
+    ):
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
 
         if self.padding_mode != "zeros":
-            raise ValueError(f"Only `zeros` padding mode is supported. "
-                             f"Got `{self.padding_mode}`.")
+            raise ValueError(
+                f"Only `zeros` padding mode is supported. Got `{self.padding_mode}`."
+            )
 
         self.log_sigma2 = torch.nn.Parameter(torch.Tensor(*self.weight.shape))
         self.reset_variational_parameters()
@@ -131,8 +151,15 @@ def _forward_impl(self, input, conv):
         if not self.training:
             return mu
 
-        s2 = conv(input * input, torch.exp(self.log_sigma2), None,
-                  self.stride, self.padding, self.dilation, self.groups)
+        s2 = conv(
+            input * input,
+            torch.exp(self.log_sigma2),
+            None,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
         return mu + torch.randn_like(s2) * torch.sqrt(torch.clamp(s2, 1e-8))
 
 
diff --git a/cplxmodule/nn/relevance/real/vd.py b/cplxmodule/nn/relevance/real/vd.py
index 76765f0..895bdf8 100644
--- a/cplxmodule/nn/relevance/real/vd.py
+++ b/cplxmodule/nn/relevance/real/vd.py
@@ -1,5 +1,3 @@
-import warnings
-
 import torch
 import torch.nn
 
@@ -73,7 +71,7 @@ def penalty(self):
         accompanying notebook for the MC estimation of the constants:
         `k1, k2, k3 = 0.63576, 1.87320, 1.48695`
         """
-        n_log_alpha = - self.log_alpha
+        n_log_alpha = -self.log_alpha
         sigmoid = torch.sigmoid(1.48695 * n_log_alpha - 1.87320)
         return F.softplus(n_log_alpha) / 2 + 0.63576 * sigmoid
 
@@ -85,6 +83,7 @@ class LinearVD(RealVDMixin, RelevanceMixin, LinearGaussian, BaseARD):
     -------
     See `torch.nn.Linear` for reference on the dimensions and parameters.
     """
+
     pass
 
 
@@ -95,6 +94,7 @@ class BilinearVD(RealVDMixin, RelevanceMixin, BilinearGaussian, BaseARD):
     -------
     See `torch.nn.Bilinear` for reference on the dimensions and parameters.
     """
+
     pass
 
 
@@ -107,6 +107,7 @@ class Conv1dVD(RealVDMixin, RelevanceMixin, Conv1dGaussian, BaseARD):
     `cplxmodule.nn.relevance.ConvNdGaussianMixin` for details about the
     implementation of the reparameterization trick.
     """
+
     pass
 
 
@@ -119,6 +120,7 @@ class Conv2dVD(RealVDMixin, RelevanceMixin, Conv2dGaussian, BaseARD):
     `cplxmodule.nn.relevance.ConvNdGaussianMixin` for details about the
     implementation of the reparameterization trick.
     """
+
     pass
 
 
@@ -131,69 +133,5 @@ class Conv3dVD(RealVDMixin, RelevanceMixin, Conv3dGaussian, BaseARD):
     `cplxmodule.nn.relevance.ConvNdGaussianMixin` for details about the
     implementation of the reparameterization trick.
     """
-    pass
 
-
-class LinearARD(object):
-    def __new__(cls, in_features, out_features, bias=True):
-        # Due to incorrect naming real-valued Automatic Relevance Determination
-        # (ARD) layers from `.real` are and have always been Variational Dropout
-        # layers (VD). The only difference is in the prior: ARD uses Gaussian
-        # prior with adaptive precision, while VD uses log-uniform prior.
-        # While empirical evidence suggests that they both sprasify similary,
-        # and yield close compression levels, VD layers tended to have higher
-        # arithmetic complexity due to the pentaly term.
-        warnings.warn("Importing real-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.real` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.real` submodule will export real-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return LinearVD(in_features, out_features, bias)
-
-
-class Conv1dARD(object):
-    def __new__(cls, in_channels, out_channels, kernel_size, stride=1,
-                padding=0, dilation=1, groups=1,
-                bias=True, padding_mode='zeros'):
-        warnings.warn("Importing real-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.real` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.real` submodule will export real-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return Conv1dVD(in_channels, out_channels, kernel_size, stride,
-                        padding, dilation, groups, bias, padding_mode)
-
-
-class Conv2dARD(object):
-    def __new__(cls, in_channels, out_channels, kernel_size, stride=1,
-                padding=0, dilation=1, groups=1,
-                bias=True, padding_mode='zeros'):
-        warnings.warn("Importing real-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.real` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.real` submodule will export real-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return Conv2dVD(in_channels, out_channels, kernel_size, stride,
-                        padding, dilation, groups, bias, padding_mode)
-
-
-class BilinearARD(object):
-    def __new__(cls, in1_features, in2_features, out_features, bias=True):
-        warnings.warn("Importing real-valued Automatic Relevance Determination"
-                      " layers (ARD) from `cplxmodule.nn.relevance.real` has"
-                      " been deprecated due to misleading name. Starting with"
-                      " version `2021` the `.real` submodule will export real-"
-                      "valued Variational Dropout (VD) layers only. Please"
-                      " import ARD layers from `relevance.ard`.",
-                      FutureWarning)
-
-        return BilinearVD(in1_features, in2_features, out_features, bias)
+    pass
diff --git a/cplxmodule/nn/utils/sparsity.py b/cplxmodule/nn/utils/sparsity.py
index 314e6c0..478fae7 100644
--- a/cplxmodule/nn/utils/sparsity.py
+++ b/cplxmodule/nn/utils/sparsity.py
@@ -1,4 +1,4 @@
-import torch
+import torch  # noqa: F401
 import warnings
 
 
@@ -6,8 +6,9 @@ class SparsityStats(object):
     __sparsity_ignore__ = ()
 
     def sparsity(self, **kwargs):
-        raise NotImplementedError("Derived classes must implement "
-                                  "a method to estimate sparsity.")
+        raise NotImplementedError(
+            "Derived classes must implement a method to estimate sparsity."
+        )
 
 
 def named_sparsity(module, prefix="", **kwargs):
@@ -20,8 +21,10 @@ def named_sparsity(module, prefix="", **kwargs):
     incorrect estimation in case of parameter sharing.
     """
 
-    warnings.warn("Since v2020.06 module's buffers are also accounted "
-                  "by `named_sparsity`.", FutureWarning)
+    warnings.warn(
+        "Since v2020.06 module's buffers are also accounted by `named_sparsity`.",
+        FutureWarning,
+    )
 
     # gather the dropout statistics and service parameters to ignore
     n_dropout, p_service = {}, set()
@@ -38,11 +41,11 @@ def named_sparsity(module, prefix="", **kwargs):
 
     for name, par in module.named_parameters(prefix=prefix):
         if name not in p_service:
-            yield name, (n_dropout.get(id(par), 0.), par.numel())
+            yield name, (n_dropout.get(id(par), 0.0), par.numel())
 
     for name, buf in module.named_buffers(prefix=prefix):
         if name not in p_service:
-            yield name, (n_dropout.get(id(buf), 0.), buf.numel())
+            yield name, (n_dropout.get(id(buf), 0.0), buf.numel())
 
 
 def sparsity(module, **kwargs):
diff --git a/cplxmodule/utils/spectrum.py b/cplxmodule/utils/spectrum.py
index 325b3bb..19a5b51 100644
--- a/cplxmodule/utils/spectrum.py
+++ b/cplxmodule/utils/spectrum.py
@@ -4,7 +4,7 @@
 from .views import fix_dim, window_view
 
 
-def pwelch(x, dim, window, fs=1., scaling="density", n_overlap=None):
+def pwelch(x, dim, window, fs=1.0, scaling="density", n_overlap=None):
     r"""Estimate power spectral density using Welch's method.
 
     Arguments
@@ -46,46 +46,45 @@ def pwelch(x, dim, window, fs=1., scaling="density", n_overlap=None):
     scipy.signal.welch: the reference for this function.
     """
     if scaling not in ("density", "spectrum"):
-        raise ValueError(f"""Unrecognized `scaling` value {scaling}""")
+        raise ValueError(f"Unrecognized `scaling` value {scaling}")
 
-    if x.shape[-1] != 2:
-        raise TypeError("""The last dimension of the input must be 2:"""
-                        """x[..., 0] is real and x[..., 1] is imaginary.""")
+    assert x.is_complex()
+
+    # if x.shape[-1] != 2:
+    #     raise TypeError("The last dimension of the input must be 2:"
+    #                     "x[..., 0] is real and x[..., 1] is imaginary.")
 
     dim = fix_dim(dim, x.dim())
-    if not 0 <= dim < x.dim() - 1:
-        raise ValueError("""The last dimension of the input cannot contain """
-                         """the signal.""")
+    # if not 0 <= dim < x.dim() - 1:
+    #     raise ValueError("The last dimension of the input cannot contain "
+    #                      "the signal.")
 
     n_window = len(window)
     if n_overlap is None:
         n_overlap = n_window // 2
     assert n_window > n_overlap
 
-    # 1. make windowed view and dim-shuffle to the last but one dim
-    x_window = window_view(x, dim, n_window, n_window - n_overlap)
-    x_window = torch.transpose(x_window, dim + 1, -2)
-
-    # 2. Apply window and compute the 1d-fft with `arbitrary number
-    #   of leading batch dimensions`
-    xw = torch.mul(x_window, window.unsqueeze(-1))
-    fft = torch.fft(xw, signal_ndim=1, normalized=False)
+    # 1. make windowed view at the last dim
+    x_window = window_view(x, dim, n_window, n_window - n_overlap, at=-1)
 
-    # 3. undo the dim-shuffle on the fft result
-    fft = torch.transpose(fft, -2, dim + 1)
+    # 2. Apply window and compute batched 1d-fft over the last dim
+    xw = torch.mul(x_window, window)
+    fft = torch.fft.fft(xw, dim=-1, norm="backward")
 
-    # 4. compute the power spectrum with the proper scaling
+    # 3. compute the power spectrum with the proper scaling
     if scaling == "density":
         scale = fs * torch.sum(window**2)
+
     elif scaling == "spectrum":
-        scale = torch.sum(window)**2
+        scale = torch.sum(window) ** 2
 
     # used to have `/ x_window.shape[dim]`
-    Pxx = torch.sum(fft**2, dim=-1).mean(dim=dim) / scale
+    amplitude = abs(fft)
+    Pxx = amplitude.mul_(amplitude).mean(dim=dim) / scale
 
     # 5. get the frequencies
-    freq = np.fft.fftfreq(n_window, 1. / fs)
-    freq = torch.tensor(freq, dtype=x.dtype, device=x.device)
+    freq = np.fft.fftfreq(n_window, 1.0 / fs)
+    freq = torch.tensor(freq, dtype=x.real.dtype, device=x.device)
     return freq, Pxx
 
 
@@ -117,8 +116,9 @@ def fftshift(x, dim=-1):
     return torch.roll(x, x.shape[dim] // 2, dim)
 
 
-def bandwidth_power(x, fs, bands, dim=-2, n_overlap=None,
-                    nperseg=None, scaling="density"):
+def bandwidth_power(
+    x, fs, bands, dim=-2, n_overlap=None, nperseg=None, scaling="density"
+):
     r"""Compute the total power of a batch of signals in each band.
 
     Uses Welch's method (see `scipy.signal.welch`) with Hamming window
@@ -160,25 +160,39 @@ def bandwidth_power(x, fs, bands, dim=-2, n_overlap=None,
         distribution.
     """
     dim = fix_dim(dim, x.dim())
+    if not x.is_complex():
+        # convert the torch<1.8 "complex" format to a proper complex dtype
+        assert x.shape[-1] == 2  # the last dim has size 2 for re and im
+        x = x[..., 0] + 1j * x[..., 1]  # makes into ndim-1
+
     if nperseg is None:
         nperseg = x.shape[dim]
 
     # 1. Welch
-    window = torch.hamming_window(nperseg, periodic=False,
-                                  dtype=x.dtype, device=x.device)
-    ff, px = pwelch(x, dim, window, fs=fs, scaling=scaling,
-                    n_overlap=n_overlap)
+    window = torch.hamming_window(
+        nperseg, periodic=False, dtype=x.real.dtype, device=x.device
+    )
+    ff, px = pwelch(x, dim, window, fs=fs, scaling=scaling, n_overlap=n_overlap)
 
     ff, px = fftshift(ff), fftshift(px, dim=dim)
     if not bands:
         # this needs to return a tensor that has the expected shape
-        return ff, px, torch.empty(*px.shape[:dim], *px.shape[dim+1:], 0,
-                                   dtype=px.dtype, device=px.device)
+        return (
+            ff,
+            px,
+            torch.empty(
+                *px.shape[:dim],
+                *px.shape[dim + 1 :],
+                0,
+                dtype=px.dtype,
+                device=px.device,
+            ),
+        )
 
     # 2. Compute power within each band
-    channel, df = [], 1. / (nperseg * fs)
+    channel, df = [], 1.0 / (nperseg * fs)  # noqa: F841
     for lo, hi in bands:
-        index, = torch.nonzero(ff.gt(lo) & ff.lt(hi), as_tuple=True)
+        (index,) = torch.nonzero(ff.gt(lo) & ff.lt(hi), as_tuple=True)
         power = torch.index_select(px, dim, index)
         channel.append(power.sum(dim=dim))  # * df)
     # end for
@@ -187,8 +201,7 @@ def bandwidth_power(x, fs, bands, dim=-2, n_overlap=None,
     return ff, px, 10 * torch.log10(torch.stack(channel, dim=-1))
 
 
-def acpr_calc(signal, sample_rate, mcf, mcb, acf=None, acb=None,
-              nperseg=None, dim=-2):
+def acpr_calc(signal, sample_rate, mcf, mcb, acf=None, acb=None, nperseg=None, dim=-2):
     r"""Get the total power in the main and adjacent frequqency bands.
 
     Arguments
@@ -221,15 +234,15 @@ def acpr_calc(signal, sample_rate, mcf, mcb, acf=None, acb=None,
         acf, acb = [], []
 
     elif not isinstance(acf, (list, tuple)):
-        raise TypeError("Adjacent Channel Frequency offests must be a "
-                        "list or a tuple.")
+        raise TypeError(
+            "Adjacent Channel Frequency offests must be a " "list or a tuple."
+        )
 
     if isinstance(acb, (int, float)):
         acb = type(acf)([acb] * len(acf))
 
     elif not isinstance(acb, (list, tuple)):
-        raise TypeError("Adjacent Channel Bandwidth must be a list or "
-                        "a tuple.")
+        raise TypeError("Adjacent Channel Bandwidth must be a list or " "a tuple.")
 
     # form the bands
     bands = [(-0.5 * mcb + mcf, +0.5 * mcb + mcf)]
@@ -237,8 +250,14 @@ def acpr_calc(signal, sample_rate, mcf, mcb, acf=None, acb=None,
         bands.append((-0.5 * b + f, +0.5 * b + f))
 
     # compute the power in each band
-    ff, px, channel = bandwidth_power(signal, sample_rate, bands,
-                                      dim=dim, nperseg=nperseg,
-                                      n_overlap=0, scaling="spectrum")
+    ff, px, channel = bandwidth_power(
+        signal,
+        sample_rate,
+        bands,
+        dim=dim,
+        nperseg=nperseg,
+        n_overlap=0,
+        scaling="spectrum",
+    )
 
     return channel[..., :1], channel[..., 1:]
diff --git a/cplxmodule/utils/stats.py b/cplxmodule/utils/stats.py
deleted file mode 100644
index 9d53ff9..0000000
--- a/cplxmodule/utils/stats.py
+++ /dev/null
@@ -1,12 +0,0 @@
-import warnings
-
-warnings.warn("Importing sparsity helper function from this submodule"
-              " is deprecated, and `.stats` is to be removed in the"
-              " release version of `cplxmodule`. Please import from"
-              " `cplxmodule.nn.utils.sparsity`",
-              FutureWarning)
-
-from ..nn.utils.sparsity import SparsityStats
-
-from ..nn.utils.sparsity import named_sparsity
-from ..nn.utils.sparsity import sparsity
diff --git a/cplxmodule/utils/views.py b/cplxmodule/utils/views.py
index 25ee86f..e1ea0a6 100644
--- a/cplxmodule/utils/views.py
+++ b/cplxmodule/utils/views.py
@@ -6,7 +6,7 @@ def fix_dim(dim, n_dim):
     r"""For the given dimensionality make sure the `axis` is nonnegative."""
     axis = (n_dim + dim) if dim < 0 else dim
     if not 0 <= axis < n_dim:
-        raise ValueError(f"""Dimension {dim} is out of range for {n_dim}.""")
+        raise ValueError(f"Dimension {dim} is out of range for {n_dim}.")
 
     return axis
 
@@ -39,19 +39,22 @@ def complex_view(x, dim=-1, squeeze=True):
     # compute new shape and strides
     size, rem = shape[dim] // 2, shape[dim] % 2
     if rem != 0:
-        warnings.warn(f"Odd dimension size for the complex data unpacking: "
-                      f"taking the least size that fits.", RuntimeWarning)
+        warnings.warn(
+            "Odd dimension size for the complex data unpacking: "
+            "taking the least size that fits.",
+            RuntimeWarning,
+        )
 
     # new shape and stride structure
     if shape[dim] == 2 and squeeze:
         # if the complex dimension is exactly two, then just drop it
-        shape_view = shape[:dim] + shape[dim+1:]
-        strides_view = strides[:dim] + strides[dim+1:]
+        shape_view = shape[:dim] + shape[dim + 1 :]
+        strides_view = strides[:dim] + strides[dim + 1 :]
 
     else:
         # otherwise, half the size and double the stride
-        shape_view = shape[:dim] + [size] + shape[dim+1:]
-        strides_view = strides[:dim] + [2 * strides[dim]] + strides[dim+1:]
+        shape_view = shape[:dim] + [size] + shape[dim + 1 :]
+        strides_view = strides[:dim] + [2 * strides[dim]] + strides[dim + 1 :]
 
     # differentiable strided view into real and imaginary parts
     real = torch.as_strided(x, shape_view, strides_view, offset)
@@ -86,15 +89,17 @@ def window_view(x, dim, size, stride, at=None):
         general, contiguous in memory.
     """
     if size <= 0:
-        raise ValueError(f"""`size` must be a positive integer.""")
+        raise ValueError("`size` must be a positive integer.")
 
     if stride < 0:
-        raise ValueError(f"""`stride` must be a nonnegative integer.""")
+        raise ValueError("`stride` must be a nonnegative integer.")
 
     dim = fix_dim(dim, x.dim())
     if x.shape[dim] < size:
-        raise ValueError(f"""`x` at dim {dim} is too short ({x.shape[dim]}) """
-                         f"""for this window size ({size}).""")
+        raise ValueError(
+            f"`x` at dim {dim} is too short ({x.shape[dim]}) "
+            f"for this window size ({size})."
+        )
 
     if at is None:
         at = dim + 1
@@ -105,10 +110,10 @@ def window_view(x, dim, size, stride, at=None):
     strided_size = ((shape[dim] - size + 1) + stride - 1) // stride
 
     # new shape and stride structure
-    shape_view = shape[:dim] + [strided_size] + shape[dim+1:]
+    shape_view = shape[:dim] + [strided_size] + shape[dim + 1 :]
     shape_view.insert(at, size)
 
-    strides_view = strides[:dim] + [strides[dim] * stride] + strides[dim+1:]
+    strides_view = strides[:dim] + [strides[dim] * stride] + strides[dim + 1 :]
     strides_view.insert(at, strides[dim])
 
     # differentiable strided view
diff --git a/setup.py b/setup.py
index 371c6cf..bb66a1f 100644
--- a/setup.py
+++ b/setup.py
@@ -1,47 +1,59 @@
+import os
 from setuptools import setup
 
-setup(
-    name="cplxmodule",
-    version="2021.01.10",
-    description="Variational Dropout and Complex-valued Neural Networks in pytorch",
-    long_description="An extension for pytorch providing essential building blocks "
-                     "for Complex-valued Neural Networks and `torch.nn`-compatible "
-                     "real- and complex-valued Variational Dropout methods. "
-                     "Implements the tools necessary to train and sparsify "
-                     "both real and complex-valued models, and seamlessly "
-                     "integrate them into existing pipelines.",
-    url="https://github.com/ivannz/cplxmodule",
-    keywords=[
-        "pytorch",
-        "variational dropout",
-        "complex-valued networks",
-        "sparsification"
-    ],
-    license="MIT License",
-    author="Ivan Nazarov",
-    author_email="ivan.nazarov@skolkovotech.ru",
-    packages=[
-        "cplxmodule",
-        "cplxmodule.nn",
-        "cplxmodule.nn.modules",
-        "cplxmodule.nn.relevance",
-        "cplxmodule.nn.relevance.real",
-        "cplxmodule.nn.relevance.complex",
-        "cplxmodule.nn.relevance.extensions",
-        "cplxmodule.nn.relevance.extensions.real",
-        "cplxmodule.nn.masked",
-        "cplxmodule.nn.utils",
-        "cplxmodule.utils",
-    ],
-    install_requires=[
-        "torch>=1.7",
-        "numpy",
-        "scipy"
-    ],
-    tests_require=[
-        "tqdm",
-        "torchvision",
-        "matplotlib",
-        "pytest"
-    ]
-)
+
+if __name__ == "__main__":
+    # update the version number from the file at the root
+    version = open("VERSION", "r").read().strip()
+
+    cwd = os.path.dirname(os.path.abspath(__file__))
+    with open(os.path.join(cwd, "cplxmodule", "__version__.py"), "w") as f:
+        f.write(f'__version__ = "{version}"\n')
+
+    setup(
+        name="cplxmodule",
+        version=version,
+        description="Variational Dropout and Complex-valued Neural Networks in pytorch",
+        long_description="""An extension for pytorch providing essential building """
+        """blocks for Complex-valued Neural Networks and `torch.nn`-compatible """
+        """real- and complex-valued Variational Dropout methods. Implements the """
+        """tools necessary to train, sparsify and fine-tune both real- and """
+        """complex-valued models, and seamlessly integrate them into existing """
+        """pipelines.""",
+        url="https://github.com/ivannz/cplxmodule",
+        keywords=[
+            "pytorch",
+            "variational dropout",
+            "complex-valued networks",
+            "sparsification",
+        ],
+        license="MIT License",
+        author="Ivan Nazarov",
+        author_email="ivan.nazarov@skolkovotech.ru",
+        packages=[
+            "cplxmodule",
+            "cplxmodule.nn",
+            "cplxmodule.nn.modules",
+            "cplxmodule.nn.relevance",
+            "cplxmodule.nn.relevance.real",
+            "cplxmodule.nn.relevance.complex",
+            "cplxmodule.nn.relevance.extensions",
+            "cplxmodule.nn.relevance.extensions.real",
+            "cplxmodule.nn.masked",
+            "cplxmodule.nn.utils",
+            "cplxmodule.utils",
+        ],
+        python_requires=">=3.7",
+        install_requires=[
+            "torch>=1.8",
+            "numpy",
+            "scipy",
+        ],
+        tests_require=[
+            "tqdm",
+            "torchvision",
+            "matplotlib",
+            "pytest",
+            "scikit-learn",
+        ],
+    )


Test patch:
diff --git a/tests/draw_welch.py b/tests/draw_welch.py
deleted file mode 100644
index 92a05ab..0000000
--- a/tests/draw_welch.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-
-from cplxmodule.utils.spectrum import pwelch, fftshift
-from scipy.signal import welch
-
-
-from matplotlib.ticker import EngFormatter
-
-random_state = np.random.RandomState(479321083)
-
-# https://www.mathworks.com/help/signal/ref/pwelch.html#btulskp-6
-fs = 1000.
-tt = np.r_[:5 * fs - 1] / fs
-
-shape = 2, len(tt)
-
-epsilon = random_state.randn(*shape) + 1j * random_state.randn(*shape)
-np_x = np.cos(2 * np.pi * 100 * tt)[np.newaxis] + epsilon * 0.01
-
-tr_x = torch.tensor(np.stack([np_x.real, np_x.imag], axis=-1))
-tr_x.requires_grad = False
-
-tr_window = torch.hamming_window(500, periodic=False, dtype=tr_x.dtype)
-
-tr_ff, tr_px = pwelch(tr_x, 1, tr_window, fs=fs,
-                      scaling="density", n_overlap=300)
-np_ff, np_px = welch(np_x, fs=fs, axis=-1, window=tr_window.numpy(),
-                     nfft=None, nperseg=None, scaling="density",
-                     noverlap=300, detrend=False, return_onesided=False)
-
-fig = plt.figure(figsize=(14, 5))
-
-ax = fig.add_subplot(111, ylabel="Decibels (dB / Hz)", xlabel="Frequency")
-ax.xaxis.set_major_formatter(EngFormatter(unit="Hz"))
-
-ax.semilogy(np.fft.fftshift(np_ff), np.fft.fftshift(np_px[0]),
-            label="scipy", alpha=0.5)
-
-ax.semilogy(fftshift(tr_ff).numpy(), fftshift(tr_px[0]).numpy().T,
-            label="torch", alpha=0.5)
-
-ax.axvspan(90, 110, alpha=0.25, color="green")
-ax.axvline(100, lw=2, linestyle=":", color="black", alpha=0.5)
-
-ax.legend()
-
-plt.show()
diff --git a/tests/test_batchnorm.py b/tests/test_batchnorm.py
index d537b35..d6ef77b 100644
--- a/tests/test_batchnorm.py
+++ b/tests/test_batchnorm.py
@@ -9,7 +9,7 @@
 from cplxmodule import cplx
 from cplxmodule.nn import CplxBatchNorm1d, RealToCplx, CplxToReal
 from cplxmodule.nn.modules.batchnorm import cplx_batch_norm
-from cplxmodule.nn.modules.batchnorm import whiten2x2, whitendxd
+from cplxmodule.nn.modules.batchnorm import whiten2x2, whitendxd  # noqa: F401
 
 
 @pytest.fixture
@@ -48,7 +48,7 @@ def test_whitening(random_state, nugget=1e-8):
     res = np.einsum("ubfs, uvf->vbfs", c, R)
 
     m = res.mean(axis=axes, keepdims=True)
-    assert np.allclose(m, 0.)
+    assert np.allclose(m, 0.0)
 
     v = np.einsum("ubfs, vbfs->fuv", res - m, res - m) / n_samples
     assert np.allclose(v, np.eye(2)[np.newaxis], atol=1e-3)
@@ -64,12 +64,12 @@ def test_whitening(random_state, nugget=1e-8):
 def fit(X, y, model, n_steps=2000):
     model.train()
     feed = torch.utils.data.DataLoader(
-        torch.utils.data.TensorDataset(X, y),
-        batch_size=32, shuffle=True)
+        torch.utils.data.TensorDataset(X, y), batch_size=32, shuffle=True
+    )
 
     optim = torch.optim.Adam(model.parameters())
     with tqdm.tqdm(range(n_steps)) as bar:
-        for i in bar:
+        for _ in bar:
             for bx, by in feed:
                 optim.zero_grad()
 
@@ -83,14 +83,14 @@ def fit(X, y, model, n_steps=2000):
 def predict(X, model):
     model.eval()
     feed = torch.utils.data.DataLoader(
-        torch.utils.data.TensorDataset(X),
-        batch_size=256, shuffle=False)
+        torch.utils.data.TensorDataset(X), batch_size=256, shuffle=False
+    )
     with torch.no_grad():
         return torch.cat([*starmap(model, feed)], dim=0)
 
 
 def test_real_batchnorm(random_state):
-    x = random_state.randn(1000, 10) * 10. + 5.
+    x = random_state.randn(1000, 10) * 10.0 + 5.0
     x = torch.from_numpy(x).to(torch.float)
 
     model = torch.nn.BatchNorm1d(10, affine=True)
@@ -114,11 +114,11 @@ def test_cplx_batchnorm(random_state):
     out = cplx_batch_norm(z, None, None).numpy()
 
     re, im = out.real, out.imag
-    assert np.isclose(float(re.mean()), 0., atol=1e-1)
-    assert np.isclose(float(im.mean()), 0., atol=1e-1)
-    assert np.isclose(float((re*re).mean()), 1., atol=1e-1)
-    assert np.isclose(float((im*im).mean()), 1., atol=1e-1)
-    assert np.isclose(float((re*im).mean()), 0., atol=1e-1)
+    assert np.isclose(float(re.mean()), 0.0, atol=1e-1)
+    assert np.isclose(float(im.mean()), 0.0, atol=1e-1)
+    assert np.isclose(float((re * re).mean()), 1.0, atol=1e-1)
+    assert np.isclose(float((im * im).mean()), 1.0, atol=1e-1)
+    assert np.isclose(float((re * im).mean()), 0.0, atol=1e-1)
 
 
 def test_cplx_batchnorm_layer_noaffine(random_state):
@@ -137,23 +137,23 @@ def test_cplx_batchnorm_layer_noaffine(random_state):
 
     model.train()
     feed = torch.utils.data.DataLoader(
-        torch.utils.data.TensorDataset(x),
-        batch_size=32, shuffle=True)
+        torch.utils.data.TensorDataset(x), batch_size=32, shuffle=True
+    )
 
     # with torch.autograd.detect_anomaly():
-    for i in range(5):
-        for bx, in feed:
+    for _ in range(5):
+        for (bx,) in feed:
             F.mse_loss(model(bx), bx)
 
     model.eval()
 
     out = predict(x, model)
     re, im = out[..., 0::2], out[..., 1::2]
-    assert np.isclose(float(re.mean()), 0., atol=1e-1)
-    assert np.isclose(float(im.mean()), 0., atol=1e-1)
-    assert np.isclose(float((re*re).mean()), 1., atol=1e-1)
-    assert np.isclose(float((im*im).mean()), 1., atol=1e-1)
-    assert np.isclose(float((re*im).mean()), 0., atol=1e-1)
+    assert np.isclose(float(re.mean()), 0.0, atol=1e-1)
+    assert np.isclose(float(im.mean()), 0.0, atol=1e-1)
+    assert np.isclose(float((re * re).mean()), 1.0, atol=1e-1)
+    assert np.isclose(float((im * im).mean()), 1.0, atol=1e-1)
+    assert np.isclose(float((re * im).mean()), 0.0, atol=1e-1)
 
 
 def test_cplx_batchnorm_layer_affine(random_state):
@@ -176,11 +176,14 @@ def test_cplx_batchnorm_layer_affine(random_state):
     out = predict(x, model).reshape(*x.shape[:-1], -1, 2)
     with torch.no_grad():
         out = out - model[1].bias.reshape(1, 10, 1, 2)
-        res, _ = torch.solve(out.transpose(-1, -2), model[1].weight.permute(2, 0, 1))
+        res = torch.linalg.solve(
+            model[1].weight.permute(2, 0, 1),
+            out.transpose(-1, -2),
+        )
 
     re, im = res[..., 0, :], res[..., 1, :]
-    assert np.isclose(float(re.mean()), 0., atol=1e-1)
-    assert np.isclose(float(im.mean()), 0., atol=1e-1)
-    assert np.isclose(float((re*re).mean()), 1., atol=1e-1)
-    assert np.isclose(float((im*im).mean()), 1., atol=1e-1)
-    assert np.isclose(float((re*im).mean()), 0., atol=1e-1)
+    assert np.isclose(float(re.mean()), 0.0, atol=1e-1)
+    assert np.isclose(float(im.mean()), 0.0, atol=1e-1)
+    assert np.isclose(float((re * re).mean()), 1.0, atol=1e-1)
+    assert np.isclose(float((im * im).mean()), 1.0, atol=1e-1)
+    assert np.isclose(float((re * im).mean()), 0.0, atol=1e-1)
diff --git a/tests/test_cplx.py b/tests/test_cplx.py
index 61c6cfe..12e444f 100644
--- a/tests/test_cplx.py
+++ b/tests/test_cplx.py
@@ -10,15 +10,15 @@
 
 
 def cplx_allclose(input, other):
-    return torch.allclose(input.real, other.real) and \
-           torch.allclose(input.imag, other.imag)
+    return torch.allclose(input.real, other.real) and torch.allclose(
+        input.imag, other.imag
+    )
 
 
 def cplx_allclose_numpy(input, other):
     other = np.asarray(other)
-    return (
-        torch.allclose(input.real, torch.from_numpy(other.real))
-        and torch.allclose(input.imag, torch.from_numpy(other.imag))
+    return torch.allclose(input.real, torch.from_numpy(other.real)) and torch.allclose(
+        input.imag, torch.from_numpy(other.imag)
     )
 
 
@@ -148,16 +148,20 @@ def test_arithmetic_inplace(random_state):
     m = np.zeros_like(a)
 
     # test inplace __i*__
-    n += p ; m += a
+    n += p
+    m += a
     assert cplx_allclose_numpy(n, m)
 
-    n *= p ; m *= a
+    n *= p
+    m *= a
     assert cplx_allclose_numpy(n, m)
 
-    n -= p ; m -= a
+    n -= p
+    m -= a
     assert cplx_allclose_numpy(n, m)
 
-    n /= p ; m /= a
+    n /= p
+    m /= a
     assert cplx_allclose_numpy(n, m)
 
     with pytest.raises(RuntimeError, match=r"The expanded size of the tensor"):
@@ -306,29 +310,43 @@ def test_conv1d_transform(random_state):
     L = random_state.randn(14, 12, 7) + 1j * random_state.randn(14, 12, 7)
 
     p, U = map(cplx.Cplx.from_numpy, [a, L])
-    assert cplx_allclose(cplx.convnd_quick(F.conv1d, p, U),
-                         cplx.convnd_naive(F.conv1d, p, U))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv1d, p, U), cplx.convnd_naive(F.conv1d, p, U)
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv1d, p, U, stride=5),
-                         cplx.convnd_naive(F.conv1d, p, U, stride=5))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv1d, p, U, stride=5),
+        cplx.convnd_naive(F.conv1d, p, U, stride=5),
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv1d, p, U, padding=2),
-                         cplx.convnd_naive(F.conv1d, p, U, padding=2))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv1d, p, U, padding=2),
+        cplx.convnd_naive(F.conv1d, p, U, padding=2),
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv1d, p, U, dilation=3),
-                         cplx.convnd_naive(F.conv1d, p, U, dilation=3))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv1d, p, U, dilation=3),
+        cplx.convnd_naive(F.conv1d, p, U, dilation=3),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv1d, p, U),
-                         cplx.convnd_naive(F.conv1d, p, U))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv1d, p, U), cplx.convnd_naive(F.conv1d, p, U)
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv1d, p, U, stride=5),
-                         cplx.convnd_naive(F.conv1d, p, U, stride=5))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv1d, p, U, stride=5),
+        cplx.convnd_naive(F.conv1d, p, U, stride=5),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv1d, p, U, padding=2),
-                         cplx.convnd_naive(F.conv1d, p, U, padding=2))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv1d, p, U, padding=2),
+        cplx.convnd_naive(F.conv1d, p, U, padding=2),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv1d, p, U, dilation=3),
-                         cplx.convnd_naive(F.conv1d, p, U, dilation=3))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv1d, p, U, dilation=3),
+        cplx.convnd_naive(F.conv1d, p, U, dilation=3),
+    )
 
 
 def test_conv2d_transform(random_state):
@@ -372,29 +390,43 @@ def test_conv2d_transform(random_state):
     L = random_state.randn(14, 12, 7, 6) + 1j * random_state.randn(14, 12, 7, 6)
 
     p, U = map(cplx.Cplx.from_numpy, [a, L])
-    assert cplx_allclose(cplx.convnd_quick(F.conv2d, p, U),
-                         cplx.convnd_naive(F.conv2d, p, U))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv2d, p, U), cplx.convnd_naive(F.conv2d, p, U)
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv2d, p, U, stride=5),
-                         cplx.convnd_naive(F.conv2d, p, U, stride=5))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv2d, p, U, stride=5),
+        cplx.convnd_naive(F.conv2d, p, U, stride=5),
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv2d, p, U, padding=2),
-                         cplx.convnd_naive(F.conv2d, p, U, padding=2))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv2d, p, U, padding=2),
+        cplx.convnd_naive(F.conv2d, p, U, padding=2),
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv2d, p, U, dilation=3),
-                         cplx.convnd_naive(F.conv2d, p, U, dilation=3))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv2d, p, U, dilation=3),
+        cplx.convnd_naive(F.conv2d, p, U, dilation=3),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv2d, p, U),
-                         cplx.convnd_naive(F.conv2d, p, U))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv2d, p, U), cplx.convnd_naive(F.conv2d, p, U)
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv2d, p, U, stride=5),
-                         cplx.convnd_naive(F.conv2d, p, U, stride=5))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv2d, p, U, stride=5),
+        cplx.convnd_naive(F.conv2d, p, U, stride=5),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv2d, p, U, padding=2),
-                         cplx.convnd_naive(F.conv2d, p, U, padding=2))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv2d, p, U, padding=2),
+        cplx.convnd_naive(F.conv2d, p, U, padding=2),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv2d, p, U, dilation=3),
-                         cplx.convnd_naive(F.conv2d, p, U, dilation=3))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv2d, p, U, dilation=3),
+        cplx.convnd_naive(F.conv2d, p, U, dilation=3),
+    )
 
 
 def test_conv3d_transform(random_state):
@@ -419,7 +451,9 @@ def test_conv3d_transform(random_state):
 
     # check pure C: `correlate` uses conjugation
     #  https://docs.scipy.org/doc/numpy/reference/generated/numpy.correlate.html
-    x = random_state.randn(5, 12, 13, 21, 17) + 1j * random_state.randn(5, 12, 13, 21, 17)
+    x = random_state.randn(5, 12, 13, 21, 17) + 1j * random_state.randn(
+        5, 12, 13, 21, 17
+    )
     w = random_state.randn(1, 12, 7, 11, 5) + 1j * random_state.randn(1, 12, 7, 11, 5)
 
     nn = correlate(x, w.conj(), mode="valid")
@@ -434,33 +468,49 @@ def test_conv3d_transform(random_state):
     cc = cplx.convnd_3m(F.conv3d, cx, cw)
     assert cplx_allclose_numpy(cc, nn)
 
-    a = random_state.randn(5, 12, 14, 19, 27) + 1j * random_state.randn(5, 12, 14, 19, 27)
+    a = random_state.randn(5, 12, 14, 19, 27) + 1j * random_state.randn(
+        5, 12, 14, 19, 27
+    )
     L = random_state.randn(14, 12, 3, 4, 5) + 1j * random_state.randn(14, 12, 3, 4, 5)
 
     p, U = map(cplx.Cplx.from_numpy, [a, L])
-    assert cplx_allclose(cplx.convnd_quick(F.conv3d, p, U),
-                         cplx.convnd_naive(F.conv3d, p, U))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv3d, p, U), cplx.convnd_naive(F.conv3d, p, U)
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv3d, p, U, stride=5),
-                         cplx.convnd_naive(F.conv3d, p, U, stride=5))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv3d, p, U, stride=5),
+        cplx.convnd_naive(F.conv3d, p, U, stride=5),
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv3d, p, U, padding=2),
-                         cplx.convnd_naive(F.conv3d, p, U, padding=2))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv3d, p, U, padding=2),
+        cplx.convnd_naive(F.conv3d, p, U, padding=2),
+    )
 
-    assert cplx_allclose(cplx.convnd_quick(F.conv3d, p, U, dilation=3),
-                         cplx.convnd_naive(F.conv3d, p, U, dilation=3))
+    assert cplx_allclose(
+        cplx.convnd_quick(F.conv3d, p, U, dilation=3),
+        cplx.convnd_naive(F.conv3d, p, U, dilation=3),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv3d, p, U),
-                         cplx.convnd_naive(F.conv3d, p, U))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv3d, p, U), cplx.convnd_naive(F.conv3d, p, U)
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv3d, p, U, stride=5),
-                         cplx.convnd_naive(F.conv3d, p, U, stride=5))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv3d, p, U, stride=5),
+        cplx.convnd_naive(F.conv3d, p, U, stride=5),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv3d, p, U, padding=2),
-                         cplx.convnd_naive(F.conv3d, p, U, padding=2))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv3d, p, U, padding=2),
+        cplx.convnd_naive(F.conv3d, p, U, padding=2),
+    )
 
-    assert cplx_allclose(cplx.convnd_3m(F.conv3d, p, U, dilation=3),
-                         cplx.convnd_naive(F.conv3d, p, U, dilation=3))
+    assert cplx_allclose(
+        cplx.convnd_3m(F.conv3d, p, U, dilation=3),
+        cplx.convnd_naive(F.conv3d, p, U, dilation=3),
+    )
 
 
 def test_bilinear_transform(random_state):
@@ -474,46 +524,28 @@ def test_bilinear_transform(random_state):
     q = cplx.Cplx.from_numpy(b)
 
     base = np.einsum("bsi, bsj, fij ->bsf", a.conj(), z, L)
-    assert cplx_allclose_numpy(
-        cplx.bilinear(p, r, U, None, conjugate=True),
-        base)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_naive(p, r, U, None, conjugate=True),
-        base)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_cat(p, r, U, None, conjugate=True),
-        base)
+    assert cplx_allclose_numpy(cplx.bilinear(p, r, U, None, conjugate=True), base)
+    assert cplx_allclose_numpy(cplx.bilinear_naive(p, r, U, None, conjugate=True), base)
+    assert cplx_allclose_numpy(cplx.bilinear_cat(p, r, U, None, conjugate=True), base)
 
+    assert cplx_allclose_numpy(cplx.bilinear(p, r, U, q, conjugate=True), base + b)
     assert cplx_allclose_numpy(
-        cplx.bilinear(p, r, U, q, conjugate=True),
-        base + b)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_naive(p, r, U, q, conjugate=True),
-        base + b)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_cat(p, r, U, q, conjugate=True),
-        base + b)
+        cplx.bilinear_naive(p, r, U, q, conjugate=True), base + b
+    )
+    assert cplx_allclose_numpy(cplx.bilinear_cat(p, r, U, q, conjugate=True), base + b)
 
     base = np.einsum("bsi, bsj, fij ->bsf", a, z, L)
+    assert cplx_allclose_numpy(cplx.bilinear(p, r, U, None, conjugate=False), base)
     assert cplx_allclose_numpy(
-        cplx.bilinear(p, r, U, None, conjugate=False),
-        base)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_naive(p, r, U, None, conjugate=False),
-        base)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_cat(p, r, U, None, conjugate=False),
-        base)
+        cplx.bilinear_naive(p, r, U, None, conjugate=False), base
+    )
+    assert cplx_allclose_numpy(cplx.bilinear_cat(p, r, U, None, conjugate=False), base)
 
+    assert cplx_allclose_numpy(cplx.bilinear(p, r, U, q, conjugate=False), base + b)
     assert cplx_allclose_numpy(
-        cplx.bilinear(p, r, U, q, conjugate=False),
-        base + b)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_naive(p, r, U, q, conjugate=False),
-        base + b)
-    assert cplx_allclose_numpy(
-        cplx.bilinear_cat(p, r, U, q, conjugate=False),
-        base + b)
+        cplx.bilinear_naive(p, r, U, q, conjugate=False), base + b
+    )
+    assert cplx_allclose_numpy(cplx.bilinear_cat(p, r, U, q, conjugate=False), base + b)
 
 
 def test_type_conversion(random_state):
@@ -540,8 +572,9 @@ def test_type_conversion(random_state):
 
     # concatenated to cplx
     for dim in [0, 1, 2]:
-        stacked = torch.cat([torch.from_numpy(a.real),
-                             torch.from_numpy(a.imag)], dim=dim)
+        stacked = torch.cat(
+            [torch.from_numpy(a.real), torch.from_numpy(a.imag)], dim=dim
+        )
 
         q = cplx.from_concatenated_real(stacked, dim=dim)
         assert cplx_allclose_numpy(q, a)
@@ -556,10 +589,9 @@ def test_type_conversion(random_state):
     # cplx to interleaved
     for dim in [0, 1, 2]:
         q = cplx.from_interleaved_real(
-                cplx.to_interleaved_real(
-                    cplx.Cplx.from_numpy(a),
-                    flatten=True, dim=dim
-                ), dim=dim)
+            cplx.to_interleaved_real(cplx.Cplx.from_numpy(a), flatten=True, dim=dim),
+            dim=dim,
+        )
 
         assert cplx_allclose_numpy(q, a)
 
@@ -576,17 +608,20 @@ def test_enisum(random_state):
     equations = ["iij", "iji", "jii", "iii"]
     for eq in equations:
         assert cplx_allclose_numpy(cplx.einsum(eq, r), np.einsum(eq, c))
-        with pytest.raises(RuntimeError, match="dimension does not match"):
+        with pytest.raises(RuntimeError, match="but the sizes don't match"):
             cplx.einsum(eq, p)
 
     equations = [
-        "ijk, ikj", "ijk, ikj -> ij", "ijk, ikj -> k",
-        "ijk, lkj", "ijk, lkj -> li", "ijk, lkj -> lji",
+        "ijk, ikj",
+        "ijk, ikj -> ij",
+        "ijk, ikj -> k",
+        "ijk, lkj",
+        "ijk, lkj -> li",
+        "ijk, lkj -> lji",
         "ijk, lkp",
-        ]
+    ]
     for eq in equations:
-        assert cplx_allclose_numpy(cplx.einsum(eq, p, q),
-                                   np.einsum(eq, a, b))
+        assert cplx_allclose_numpy(cplx.einsum(eq, p, q), np.einsum(eq, a, b))
 
     with pytest.raises(RuntimeError, match="does not support more"):
         cplx.einsum("...", p, q, r)
@@ -596,20 +631,18 @@ def test_cat_stack(random_state):
     with pytest.raises(RuntimeError, match="a non-empty"):
         cplx.stack([], dim=0)
 
-    np_tensors = 10 * [
-        random_state.randn(5, 3, 7) + 1j * random_state.randn(5, 3, 7)
-    ]
+    np_tensors = 10 * [random_state.randn(5, 3, 7) + 1j * random_state.randn(5, 3, 7)]
     tr_tensors = [*map(cplx.Cplx.from_numpy, np_tensors)]
 
     for n in [0, 1, 2]:
         assert cplx_allclose_numpy(
-            cplx.cat(tr_tensors, dim=n),
-            np.concatenate(np_tensors, axis=n))
+            cplx.cat(tr_tensors, dim=n), np.concatenate(np_tensors, axis=n)
+        )
 
     for n in [0, 1, 2, 3]:
         assert cplx_allclose_numpy(
-            cplx.stack(tr_tensors, dim=n),
-            np.stack(np_tensors, axis=n))
+            cplx.stack(tr_tensors, dim=n), np.stack(np_tensors, axis=n)
+        )
 
     np_tensors = [
         random_state.randn(3, 7) + 1j * random_state.randn(3, 7),
@@ -633,8 +666,7 @@ def test_view(random_state):
 
     # Should be equivalent to reshape if in contiguous memory
     assert cplx_allclose(p.view(5, 64, 64), p.reshape(5, 64, 64))
-    assert cplx_allclose(p.view(5, 64, 64),
-                         cplx.Cplx.from_numpy(a.reshape(5, 64, 64)))
+    assert cplx_allclose(p.view(5, 64, 64), cplx.Cplx.from_numpy(a.reshape(5, 64, 64)))
 
     # Real and Imaginary parts should remain the same
     assert cplx_allclose_numpy(p.view(5, 64, 64), a.reshape(5, 64, 64))
@@ -651,9 +683,7 @@ def test_view(random_state):
         p.view(4, 64, 64)
 
 
-@pytest.mark.parametrize('shape', [
-    [10, 12, 31], [7, 3], [17], [2, 5, 6, 3, 7, 2, 4]
-])
+@pytest.mark.parametrize("shape", [[10, 12, 31], [7, 3], [17], [2, 5, 6, 3, 7, 2, 4]])
 def test_tensor_shape_dim_size(random_state, shape):
     # size, shape, and dim properties
     p = cplx.Cplx.empty(shape)
@@ -670,7 +700,7 @@ def test_tensor_shape_dim_size(random_state, shape):
 
 def test_scalar_shape_dim_size(random_state):
     # size, shape, and dim properties
-    p = cplx.Cplx(1+1j)
+    p = cplx.Cplx(1 + 1j)
     assert p.dim() == 0 and p.shape == p.size() == torch.Size([])
 
     with pytest.raises(IndexError, match="tensor has no dimensions"):
diff --git a/tests/test_cplxparameter.py b/tests/test_cplxparameter.py
index d52a5de..7d9aa28 100644
--- a/tests/test_cplxparameter.py
+++ b/tests/test_cplxparameter.py
@@ -10,10 +10,7 @@ def test_loading():
 
     # load complex parameter in full
     state_dict = {"real": torch.randn(*shape), "imag": torch.randn(*shape)}
-    par = CplxParameter(cplx.Cplx(
-        real=torch.ones(*shape),
-        imag=torch.zeros(*shape)
-    ))
+    par = CplxParameter(cplx.Cplx(real=torch.ones(*shape), imag=torch.zeros(*shape)))
     par.load_state_dict(state_dict, strict=True)
 
     assert torch.allclose(par.real, state_dict["real"])
@@ -29,10 +26,7 @@ def test_loading():
 
     # promote real tensor
     state_dict = {"": torch.randn(*shape)}
-    par = CplxParameter(cplx.Cplx(
-        real=torch.ones(*shape),
-        imag=torch.zeros(*shape)
-    ))
+    par = CplxParameter(cplx.Cplx(real=torch.ones(*shape), imag=torch.zeros(*shape)))
     par.load_state_dict(state_dict, strict=True)
 
     assert torch.allclose(par.real, state_dict[""])
@@ -42,10 +36,9 @@ def test_loading():
 def make_module(*shape):
     module = torch.nn.Module()
     module.mod = torch.nn.Module()
-    module.mod.par = CplxParameter(cplx.Cplx(
-        real=torch.ones(*shape),
-        imag=torch.zeros(*shape)
-    ))
+    module.mod.par = CplxParameter(
+        cplx.Cplx(real=torch.ones(*shape), imag=torch.zeros(*shape))
+    )
     return module
 
 
@@ -92,22 +85,32 @@ def test_malformed():
         par.load_state_dict({"imag": torch.randn(*shape)})
 
     with pytest.raises(RuntimeError, match=r"disallows redundant"):
-        par.load_state_dict({
-            "real": torch.randn(*shape),
-            "imag": torch.randn(*shape),
-            "bar": torch.randn(*shape),
-            "foo": torch.randn(*shape),
-        }, strict=True)
+        par.load_state_dict(
+            {
+                "real": torch.randn(*shape),
+                "imag": torch.randn(*shape),
+                "bar": torch.randn(*shape),
+                "foo": torch.randn(*shape),
+            },
+            strict=True,
+        )
 
     with pytest.raises(RuntimeError, match=r"size mismatch for"):
-        par.load_state_dict({
-            "real": torch.randn(1, 1), "imag": torch.randn(1, 1),
-        }, strict=True)
+        par.load_state_dict(
+            {
+                "real": torch.randn(1, 1),
+                "imag": torch.randn(1, 1),
+            },
+            strict=True,
+        )
 
     with pytest.raises(RuntimeError, match=r"size mismatch for"):
-        par.load_state_dict({
-            "": torch.randn(1, 1),
-        }, strict=True)
+        par.load_state_dict(
+            {
+                "": torch.randn(1, 1),
+            },
+            strict=True,
+        )
 
 
 def test_nested_malformed():
@@ -123,31 +126,43 @@ def test_nested_malformed():
 
     # redundant keys pertaining to the complex parameter are forbidden
     with pytest.raises(RuntimeError, match=r"disallows redundant"):
-        module.load_state_dict({
-            "mod.par.real": torch.randn(*shape),
-            "mod.par.imag": torch.randn(*shape),
-            "mod.par.bar": torch.randn(*shape),
-            "mod.par.foo": torch.randn(*shape),
-        }, strict=True)
+        module.load_state_dict(
+            {
+                "mod.par.real": torch.randn(*shape),
+                "mod.par.imag": torch.randn(*shape),
+                "mod.par.bar": torch.randn(*shape),
+                "mod.par.foo": torch.randn(*shape),
+            },
+            strict=True,
+        )
 
     # redundant keys unrelated to the complex parameter are ignored
-    module.load_state_dict({
-        "mod.par.real": torch.zeros(*shape),
-        "mod.par.imag": torch.ones(*shape),
-        "bar": torch.randn(*shape),
-        "foo": torch.randn(*shape),
-    }, strict=False)
+    module.load_state_dict(
+        {
+            "mod.par.real": torch.zeros(*shape),
+            "mod.par.imag": torch.ones(*shape),
+            "bar": torch.randn(*shape),
+            "foo": torch.randn(*shape),
+        },
+        strict=False,
+    )
 
     with pytest.raises(RuntimeError, match=r"size mismatch for"):
-        module.load_state_dict({
-            "mod.par.real": torch.randn(1, 1),
-            "mod.par.imag": torch.randn(1, 1),
-        }, strict=True)
+        module.load_state_dict(
+            {
+                "mod.par.real": torch.randn(1, 1),
+                "mod.par.imag": torch.randn(1, 1),
+            },
+            strict=True,
+        )
 
     with pytest.raises(RuntimeError, match=r"size mismatch for"):
-        module.load_state_dict({
-            "mod.par": torch.randn(1, 1),
-        }, strict=True)
+        module.load_state_dict(
+            {
+                "mod.par": torch.randn(1, 1),
+            },
+            strict=True,
+        )
 
     assert torch.allclose(module.mod.par.real, torch.zeros(*shape))
     assert torch.allclose(module.mod.par.imag, torch.ones(*shape))
diff --git a/tests/test_draw_welch.py b/tests/test_draw_welch.py
new file mode 100644
index 0000000..0aa8bc7
--- /dev/null
+++ b/tests/test_draw_welch.py
@@ -0,0 +1,65 @@
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+from cplxmodule.utils.spectrum import pwelch, fftshift
+from scipy.signal import welch
+
+from matplotlib.ticker import EngFormatter
+
+import pytest
+
+
+pytestmark = pytest.mark.skip(reason="this is not a test suite")
+
+
+if __name__ == "__main__":
+    random_state = np.random.RandomState(479321083)
+
+    # https://www.mathworks.com/help/signal/ref/pwelch.html#btulskp-6
+    fs = 1000.0
+    tt = np.r_[: 5 * fs - 1] / fs
+
+    shape = 2, len(tt)
+
+    epsilon = random_state.randn(*shape) + 1j * random_state.randn(*shape)
+    np_x = np.cos(2 * np.pi * 100 * tt)[np.newaxis] + epsilon * 0.01
+
+    tr_x = torch.from_numpy(np_x)
+    tr_x.requires_grad = False
+
+    tr_window = torch.hamming_window(500, periodic=False, dtype=tr_x.real.dtype)
+
+    tr_ff, tr_px = pwelch(tr_x, 1, tr_window, fs=fs, scaling="density", n_overlap=300)
+    np_ff, np_px = welch(
+        np_x,
+        fs=fs,
+        axis=-1,
+        window=tr_window.numpy(),
+        nfft=None,
+        nperseg=None,
+        scaling="density",
+        noverlap=300,
+        detrend=False,
+        return_onesided=False,
+    )
+
+    fig = plt.figure(figsize=(14, 5))
+
+    ax = fig.add_subplot(111, ylabel="Decibels (dB / Hz)", xlabel="Frequency")
+    ax.xaxis.set_major_formatter(EngFormatter(unit="Hz"))
+
+    ax.semilogy(
+        np.fft.fftshift(np_ff), np.fft.fftshift(np_px[0]), label="scipy", alpha=0.5
+    )
+
+    ax.semilogy(
+        fftshift(tr_ff).numpy(), fftshift(tr_px[0]).numpy().T, label="torch", alpha=0.5
+    )
+
+    ax.axvspan(90, 110, alpha=0.25, color="green")
+    ax.axvline(100, lw=2, linestyle=":", color="black", alpha=0.5)
+
+    ax.legend()
+
+    plt.show()
diff --git a/tests/test_init.py b/tests/test_init.py
index 112c4c6..4f036b0 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -9,21 +9,23 @@
 
 def cplx_allclose_numpy(input, other):
     other = np.asarray(other)
-    return (
-        torch.allclose(input.real, torch.from_numpy(other.real))
-        and torch.allclose(input.imag, torch.from_numpy(other.imag))
+    return torch.allclose(input.real, torch.from_numpy(other.real)) and torch.allclose(
+        input.imag, torch.from_numpy(other.imag)
     )
 
 
-@pytest.mark.parametrize('initializer', [
-    init.cplx_kaiming_normal_,
-    init.cplx_xavier_normal_,
-    init.cplx_kaiming_uniform_,
-    init.cplx_xavier_uniform_,
-    init.cplx_trabelsi_standard_,
-    init.cplx_trabelsi_independent_,
-    init.cplx_uniform_independent_,
-])
+@pytest.mark.parametrize(
+    "initializer",
+    [
+        init.cplx_kaiming_normal_,
+        init.cplx_xavier_normal_,
+        init.cplx_kaiming_uniform_,
+        init.cplx_xavier_uniform_,
+        init.cplx_trabelsi_standard_,
+        init.cplx_trabelsi_independent_,
+        init.cplx_uniform_independent_,
+    ],
+)
 def test_initializer(initializer):
     initializer(Cplx.empty(500, 1250))
     initializer(Cplx.empty(1250, 500))
diff --git a/tests/test_masked.py b/tests/test_masked.py
index d3e2124..35d81ec 100644
--- a/tests/test_masked.py
+++ b/tests/test_masked.py
@@ -39,12 +39,24 @@ def test_setattr_interface():
     assert torch.allclose(lin.weight_masked, mask * lin.weight)
 
     # input masking
-    lin.mask = mask = torch.randint(2, size=(1, shape[0],)).float()
+    lin.mask = mask = torch.randint(
+        2,
+        size=(
+            1,
+            shape[0],
+        ),
+    ).float()
     assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask))
     assert torch.allclose(lin.weight_masked, mask * lin.weight)
 
     # unstructured masking
-    lin.mask = mask = torch.randint(2, size=(shape[1], shape[0],)).float()
+    lin.mask = mask = torch.randint(
+        2,
+        size=(
+            shape[1],
+            shape[0],
+        ),
+    ).float()
     assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask))
     assert torch.allclose(lin.weight_masked, mask * lin.weight)
 
@@ -86,13 +98,25 @@ def test_method_interface():
     assert torch.allclose(lin.weight_masked, mask * lin.weight)
 
     # input masking
-    mask = torch.randint(2, size=(1, shape[0],)).float()
+    mask = torch.randint(
+        2,
+        size=(
+            1,
+            shape[0],
+        ),
+    ).float()
     lin.mask_(mask)
     assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask))
     assert torch.allclose(lin.weight_masked, mask * lin.weight)
 
     # unstructured masking
-    mask = torch.randint(2, size=(shape[1], shape[0],)).float()
+    mask = torch.randint(
+        2,
+        size=(
+            shape[1],
+            shape[0],
+        ),
+    ).float()
     lin.mask_(mask)
     assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask))
     assert torch.allclose(lin.weight_masked, mask * lin.weight)
@@ -130,7 +154,7 @@ def test_state_dict_loading():
     masked.mask = None
     origin.load_state_dict(masked.state_dict(), strict=True)
 
-    masked.mask = torch.tensor(1.)
+    masked.mask = torch.tensor(1.0)
     with pytest.raises(RuntimeError, match=r"Unexpected key\(s\)"):
         origin.load_state_dict(masked.state_dict(), strict=True)
 
@@ -142,12 +166,24 @@ def test_state_dict_loading():
     masked.load_state_dict({**origin.state_dict(), "mask": None})
 
     # unsetting a mask with state_dict having `None` should work, too
-    masked.mask = mask = torch.randint(2, size=(shape[1], shape[0],)).float()
+    masked.mask = mask = torch.randint(
+        2,
+        size=(
+            shape[1],
+            shape[0],
+        ),
+    ).float()
     masked.load_state_dict({**origin.state_dict(), "mask": None})
     assert masked.mask is None
 
     # see if the existing mask is unaffected when mask is missing
-    masked.mask = mask = torch.randint(2, size=(shape[1], shape[0],)).float()
+    masked.mask = mask = torch.randint(
+        2,
+        size=(
+            shape[1],
+            shape[0],
+        ),
+    ).float()
     masked.load_state_dict(origin.state_dict(), strict=False)
     assert torch.allclose(masked.mask, mask)
 
diff --git a/tests/test_mnist.py b/tests/test_mnist.py
index 535a6ef..afb390e 100644
--- a/tests/test_mnist.py
+++ b/tests/test_mnist.py
@@ -14,16 +14,28 @@
 
 from cplxmodule.nn.relevance import compute_ard_masks
 from cplxmodule.nn.masked import binarize_masks, deploy_masks
-from cplxmodule.nn.masked import named_masks
 from cplxmodule.nn.utils.sparsity import named_sparsity
 
+import pytest
 
-def model_fit(model, feed, optim, n_steps=100, threshold=1.0,
-              klw=1e-3, reduction="mean", verbose=True):
+
+pytestmark = pytest.mark.skip(reason="this is not a test suite")
+
+
+def model_fit(
+    model,
+    feed,
+    optim,
+    n_steps=100,
+    threshold=1.0,
+    klw=1e-3,
+    reduction="mean",
+    verbose=True,
+):
     losses = []
     with tqdm.tqdm(range(n_steps)) as bar:
         model.train()
-        for i in bar:
+        for _ in bar:
             for data, target in feed:
                 optim.zero_grad()
 
@@ -37,8 +49,7 @@ def model_fit(model, feed, optim, n_steps=100, threshold=1.0,
 
                 losses.append(float(loss))
                 if verbose:
-                    f_sparsity = sparsity(model, hard=True,
-                                          threshold=threshold)
+                    f_sparsity = sparsity(model, hard=True, threshold=threshold)
                 else:
                     f_sparsity = float("nan")
 
@@ -108,6 +119,7 @@ class FeedWrapper(object):
     **kwargs : keyword arguments
         The keyword arguments to be passed to `torch.Tensor.to()`.
     """
+
     def __init__(self, feed, **kwargs):
         assert isinstance(feed, torch.utils.data.DataLoader)
         self.feed, self.kwargs = feed, kwargs
@@ -121,12 +133,12 @@ def __iter__(self):
 
         else:
             for batch in iter(self.feed):
-                yield tuple(b.to(**self.kwargs)
-                            for b in batch)
+                yield tuple(b.to(**self.kwargs) for b in batch)
 
 
 class Model(torch.nn.Module):
     """A convolutional net."""
+
     def __init__(self, conv2d=Conv2d, linear=Linear):
         super().__init__()
 
@@ -142,34 +154,48 @@ def forward(self, x):
         return F.log_softmax(self.fc2(x), dim=1)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     threshold = 3.0
-    device_ = torch.device("cuda:3")
+    device_ = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.1307,), (0.3081,))
-    ])
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+    )
 
     # reverse the roles of the MNIST train-test split
-    mnist_test = datasets.MNIST('./data', transform=transform,
-                                train=True, download=True)
-    mnist_train = datasets.MNIST('./data', transform=transform,
-                                 train=False)
+    mnist_test = datasets.MNIST(
+        "./data",
+        transform=transform,
+        train=True,
+        download=True,
+    )
+    mnist_train = datasets.MNIST(
+        "./data",
+        transform=transform,
+        train=False,
+    )
 
     # define wrapped data feeds
     feeds = {
         "train": FeedWrapper(
             torch.utils.data.DataLoader(
-                mnist_train, batch_size=32, shuffle=True),
-            device=device_),
+                mnist_train,
+                batch_size=32,
+                shuffle=True,
+            ),
+            device=device_,
+        ),
         "test": FeedWrapper(
             torch.utils.data.DataLoader(
-                mnist_test, batch_size=256, shuffle=False),
-            device=device_),
+                mnist_test,
+                batch_size=256,
+                shuffle=False,
+            ),
+            device=device_,
+        ),
     }
 
-    # Models and training setttings
+    # Models and training settings
     models = {
         "none": None,
         "dense": Model(Conv2d, Linear),
@@ -177,6 +203,8 @@ def forward(self, x):
         "masked": Model(Conv2dMasked, LinearMasked),
     }
 
+    # prunig schedule settings (n_epochs: int, kl_div_coeff: float)
+    # XXX higher `kl_div_coeff` means stronger pruning, at the cost of performance
     phases = {
         "dense": (40, 0.0),
         "ard": (40, 1e-2),
@@ -194,8 +222,7 @@ def forward(self, x):
         if models[src] is not None:
             # compute the dropout masks and normalize them
             state_dict = models[src].state_dict()
-            masks = compute_ard_masks(models[src], hard=False,
-                                      threshold=threshold)
+            masks = compute_ard_masks(models[src], hard=False, threshold=threshold)
 
             state_dict, masks = binarize_masks(state_dict, masks)
 
@@ -208,8 +235,14 @@ def forward(self, x):
         model.to(device_)
         optim = torch.optim.Adam(model.parameters())
         model, losses[dst] = model_fit(
-            model, feeds["train"], optim, n_steps=n_steps,
-            threshold=threshold, klw=klw, reduction="mean")
+            model,
+            feeds["train"],
+            optim,
+            n_steps=n_steps,
+            threshold=threshold,
+            klw=klw,
+            reduction="mean",
+        )
 
     # run tests
     for key, model in models.items():
diff --git a/tests/test_modules.py b/tests/test_modules.py
index a694f68..25aa7f6 100644
--- a/tests/test_modules.py
+++ b/tests/test_modules.py
@@ -11,8 +11,9 @@
 
 
 def cplx_allclose(input, other):
-    return torch.allclose(input.real, other.real) and \
-           torch.allclose(input.imag, other.imag)
+    return torch.allclose(input.real, other.real) and torch.allclose(
+        input.imag, other.imag
+    )
 
 
 def _cplx_emulate_module(Module):
@@ -28,7 +29,7 @@ def forward(self, input):
             im = self.imag(input.real) + self.real(input.imag)
 
             # undo the double application of bias
-            if getattr(self.real, 'bias', None) is not None:
+            if getattr(self.real, "bias", None) is not None:
                 broadcast = []
                 if isinstance(self.real, torch.nn.modules.conv._ConvNd):
                     broadcast = len(self.real.kernel_size) * [1]
@@ -41,14 +42,16 @@ def state_dict(self):
             state_dict = super().state_dict()
             # weight
             out = {
-                'weight.real': state_dict.pop('real.weight'),
-                'weight.imag': state_dict.pop('imag.weight'),
+                "weight.real": state_dict.pop("real.weight"),
+                "weight.imag": state_dict.pop("imag.weight"),
             }
-            if getattr(self.real, 'bias', None) is not None:
-                out.update({
-                    'bias.real': state_dict.pop('real.bias'),
-                    'bias.imag': state_dict.pop('imag.bias'),
-                })
+            if getattr(self.real, "bias", None) is not None:
+                out.update(
+                    {
+                        "bias.real": state_dict.pop("real.bias"),
+                        "bias.imag": state_dict.pop("imag.bias"),
+                    }
+                )
 
             return {**out, **state_dict}
 
@@ -61,6 +64,7 @@ def state_dict(self):
 
 class _EmulatedCplxModuleMeta(type):
     """Meta class for promoting torch.nn Modules to complex ones."""
+
     @lru_cache(maxsize=None)
     def __getitem__(self, Base):
         if isinstance(Base, type) and issubclass(Base, torch.nn.Module):
@@ -108,18 +112,17 @@ def test_linear():
         assert cplx_allclose(duplicated(z), implemented(z))
 
 
-def do_test_conv(Layer, CplxLayer, in_channels, out_channels, kernel_size,
-                 **kwargs):
+def do_test_conv(Layer, CplxLayer, in_channels, out_channels, kernel_size, **kwargs):
     emulated = EmulatedCplxModule[Layer](
-        in_channels, out_channels, kernel_size, **kwargs).double()
-    implemented = CplxLayer(
-        in_channels, out_channels, kernel_size, **kwargs).double()
+        in_channels, out_channels, kernel_size, **kwargs
+    ).double()
+    implemented = CplxLayer(in_channels, out_channels, kernel_size, **kwargs).double()
 
     # ensure identical parameters
     implemented.load_state_dict(emulated.state_dict())
 
     # use float64 for testing
-    shape = 4, in_channels, *[16]*len(emulated.real.kernel_size)
+    shape = 4, in_channels, *[16] * len(emulated.real.kernel_size)
     with torch.no_grad():
         z = cplx.randn(*shape, dtype=torch.double)
         assert cplx_allclose(implemented(z), emulated(z))
@@ -134,11 +137,14 @@ def do_test_conv(Layer, CplxLayer, in_channels, out_channels, kernel_size,
         assert cplx_allclose(duplicated(z), implemented(z))
 
 
-@pytest.mark.parametrize('case', [
-    (torch.nn.Conv1d, nn.CplxConv1d, (3,)),
-    (torch.nn.Conv2d, nn.CplxConv2d, (3, 4)),
-    (torch.nn.Conv3d, nn.CplxConv3d, (3, 4, 5)),
-])
+@pytest.mark.parametrize(
+    "case",
+    [
+        (torch.nn.Conv1d, nn.CplxConv1d, (3,)),
+        (torch.nn.Conv2d, nn.CplxConv2d, (3, 4)),
+        (torch.nn.Conv3d, nn.CplxConv3d, (3, 4, 5)),
+    ],
+)
 def test_conv(case):
     Layer, CplxLayer, kernel_size = case
 
@@ -164,18 +170,22 @@ def test_conv(case):
 
     # test padding=2 w. circular
     with pytest.xfail(reason="always xfail"):
-        do_test_conv(*common, kernel_size, bias=True, padding=2,
-                     padding_mode="circular")
+        do_test_conv(
+            *common, kernel_size, bias=True, padding=2, padding_mode="circular"
+        )
 
     # test groups=2
     do_test_conv(*common, kernel_size, bias=True, groups=2)
 
 
-@pytest.mark.parametrize('case', [
-    (torch.nn.ConvTranspose1d, nn.CplxConvTranspose1d, (3,)),
-    (torch.nn.ConvTranspose2d, nn.CplxConvTranspose2d, (3, 4)),
-    (torch.nn.ConvTranspose3d, nn.CplxConvTranspose3d, (3, 4, 5)),
-])
+@pytest.mark.parametrize(
+    "case",
+    [
+        (torch.nn.ConvTranspose1d, nn.CplxConvTranspose1d, (3,)),
+        (torch.nn.ConvTranspose2d, nn.CplxConvTranspose2d, (3, 4)),
+        (torch.nn.ConvTranspose3d, nn.CplxConvTranspose3d, (3, 4, 5)),
+    ],
+)
 def test_conv_transpose(case):
     Layer, CplxLayer, kernel_size = case
 
@@ -206,16 +216,21 @@ def test_conv_transpose(case):
     do_test_conv(*common, kernel_size, bias=True, stride=2, output_padding=1)
 
 
-@pytest.mark.parametrize('pair', [
-    (nn.modules.casting.InterleavedRealToCplx,
-     nn.modules.casting.CplxToInterleavedReal),
-    (nn.modules.casting.ConcatenatedRealToCplx,
-     nn.modules.casting.CplxToConcatenatedReal),
-    (nn.modules.casting.TensorToCplx,
-     nn.modules.casting.CplxToTensor),
-    (nn.modules.casting.AsTypeCplx,
-     nn.modules.linear.CplxReal),
-])
+@pytest.mark.parametrize(
+    "pair",
+    [
+        (
+            nn.modules.casting.InterleavedRealToCplx,
+            nn.modules.casting.CplxToInterleavedReal,
+        ),
+        (
+            nn.modules.casting.ConcatenatedRealToCplx,
+            nn.modules.casting.CplxToConcatenatedReal,
+        ),
+        (nn.modules.casting.TensorToCplx, nn.modules.casting.CplxToTensor),
+        (nn.modules.casting.AsTypeCplx, nn.modules.linear.CplxReal),
+    ],
+)
 def test_casting(pair):
     to_cplx, to_real = pair
     tensor = torch.randn(32, 31, 1024, 2)
diff --git a/tests/test_onnx.py b/tests/test_onnx.py
index 3914cff..b81a439 100644
--- a/tests/test_onnx.py
+++ b/tests/test_onnx.py
@@ -17,16 +17,22 @@
 from cplxmodule.nn.modules import casting
 
 
-def onnx_export_to(filename, module, input, *,
-                   training=False, opset_version=12):
+def onnx_export_to(filename, module, input, *, training=False, opset_version=11):
     # jit compile and export into ONNX
     torch.onnx.export(
-        module, (input,), filename, verbose=False,
-        export_params=True, training=training,
-        opset_version=opset_version, do_constant_folding=True,
-        input_names=['input'], output_names=['output'],
+        module,
+        (input,),
+        filename,
+        verbose=False,
+        export_params=True,
+        training=training,
+        opset_version=opset_version,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["output"],
         # all axes are static
-        dynamic_axes={'input': [], 'output': []})
+        dynamic_axes={"input": [], "output": []},
+    )
 
 
 def do_onnx_export_test(module, input, *, training=False):
@@ -44,13 +50,12 @@ def do_onnx_inference_test(module, input, *, training=False):
 
         # run inference through the exported model
         input = torch.randn_like(input)
-        output, = InferenceSession(file.name).run(
-            ['output'], {'input': input.numpy()}
+        (output,) = InferenceSession(file.name).run(
+            ["output"], {"input": input.numpy()}
         )
 
         assert torch.allclose(
-            torch.from_numpy(output),
-            module(input), rtol=5e-5, atol=1e-6
+            torch.from_numpy(output), module(input), rtol=1e-4, atol=1e-5
         )
 
 
@@ -59,14 +64,10 @@ def wrap_cplxtoreal(*modules):
 
 
 def wrap_cplxtocplx(*modules):
-    return torch.nn.Sequential(casting.TensorToCplx(),
-                               *modules,
-                               casting.CplxToTensor())
+    return torch.nn.Sequential(casting.TensorToCplx(), *modules, casting.CplxToTensor())
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_linear_onnx_export(training):
     module, input = torch.nn.Linear(16, 32), torch.randn(2, 3, 5, 16)
     do_onnx_export_test(module, input, training=training)
@@ -83,9 +84,7 @@ def test_linear_onnx_export(training):
     do_onnx_export_test(module, input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_conv1d_onnx_export(training):
     module, input = torch.nn.Conv1d(16, 32, 5), torch.randn(2, 16, 25)
     do_onnx_export_test(module, input, training=training)
@@ -102,9 +101,7 @@ def test_conv1d_onnx_export(training):
     do_onnx_export_test(module, input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_conv2d_onnx_export(training):
     module, input = torch.nn.Conv2d(16, 32, 5), torch.randn(2, 16, 25, 25)
     do_onnx_export_test(module, input, training=training)
@@ -121,9 +118,7 @@ def test_conv2d_onnx_export(training):
     do_onnx_export_test(module, input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_conv3d_onnx_export(training):
     module, input = torch.nn.Conv3d(16, 32, 5), torch.randn(2, 16, 25, 25, 25)
     do_onnx_export_test(module, input, training=training)
@@ -140,9 +135,7 @@ def test_conv3d_onnx_export(training):
     do_onnx_export_test(module, input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_batchnorm_onnx_export(training):
     input = torch.randn(2, 32, 256)
     # with pytest.skip():
@@ -153,9 +146,7 @@ def test_batchnorm_onnx_export(training):
     do_onnx_export_test(module, input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_linear_onnx_export(training):
     module, input = nn.CplxLinear(16, 32), torch.randn(2, 3, 5, 16, 2)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
@@ -172,15 +163,12 @@ def test_cplx_linear_onnx_export(training):
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_conv1d_onnx_export(training):
     module, input = nn.CplxConv1d(16, 32, 5), torch.randn(3, 16, 25, 2)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
-    do_onnx_inference_test(wrap_cplxtocplx(module), input,
-                           training=training)
+    do_onnx_inference_test(wrap_cplxtocplx(module), input, training=training)
 
     module = relevance.CplxConv1dVD(16, 32, 5)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
@@ -193,15 +181,12 @@ def test_cplx_conv1d_onnx_export(training):
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_conv2d_onnx_export(training):
     module, input = nn.CplxConv2d(16, 32, 5), torch.randn(3, 16, 25, 25, 2)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
-    do_onnx_inference_test(wrap_cplxtocplx(module), input,
-                           training=training)
+    do_onnx_inference_test(wrap_cplxtocplx(module), input, training=training)
 
     module = relevance.CplxConv2dVD(16, 32, 5)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
@@ -214,15 +199,12 @@ def test_cplx_conv2d_onnx_export(training):
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_conv3d_onnx_export(training):
     module, input = nn.CplxConv3d(16, 32, 5), torch.randn(3, 16, 25, 25, 25, 2)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
-    do_onnx_inference_test(wrap_cplxtocplx(module), input,
-                           training=training)
+    do_onnx_inference_test(wrap_cplxtocplx(module), input, training=training)
 
     module = relevance.CplxConv3dVD(16, 32, 5)
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
@@ -235,49 +217,62 @@ def test_cplx_conv3d_onnx_export(training):
     do_onnx_export_test(wrap_cplxtocplx(module), input, training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_batchnorm_onnx_export(training):
     module = nn.CplxBatchNorm1d(16, track_running_stats=False)
-    do_onnx_export_test(wrap_cplxtocplx(module).float(),
-                        torch.randn(2, 16, 256, 2).float(), training=training)
+    do_onnx_export_test(
+        wrap_cplxtocplx(module).float(),
+        torch.randn(2, 16, 256, 2).float(),
+        training=training,
+    )
 
     module = nn.CplxBatchNorm1d(32, track_running_stats=True)
-    do_onnx_export_test(wrap_cplxtocplx(module).float(),
-                        torch.randn(2, 32, 256, 2).float(), training=training)
+    do_onnx_export_test(
+        wrap_cplxtocplx(module).float(),
+        torch.randn(2, 32, 256, 2).float(),
+        training=training,
+    )
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_activations_onnx_export(training):
     module = nn.CplxModReLU()
-    do_onnx_export_test(wrap_cplxtocplx(module).float(),
-                        torch.randn(2, 3, 5, 16, 2).float(), training=training)
+    do_onnx_export_test(
+        wrap_cplxtocplx(module).float(),
+        torch.randn(2, 3, 5, 16, 2).float(),
+        training=training,
+    )
 
     module = nn.CplxAdaptiveModReLU()
-    do_onnx_export_test(wrap_cplxtocplx(module).float(),
-                        torch.randn(2, 3, 5, 16, 2).float(), training=training)
+    do_onnx_export_test(
+        wrap_cplxtocplx(module).float(),
+        torch.randn(2, 3, 5, 16, 2).float(),
+        training=training,
+    )
 
     module = nn.CplxModulus()
-    do_onnx_export_test(wrap_cplxtoreal(module).float(),
-                        torch.randn(2, 3, 5, 16, 2).float(), training=training)
+    do_onnx_export_test(
+        wrap_cplxtoreal(module).float(),
+        torch.randn(2, 3, 5, 16, 2).float(),
+        training=training,
+    )
 
     with pytest.raises(RuntimeError, match="the operator atan2 to ONNX"):
         module = nn.CplxAngle()
-        do_onnx_export_test(wrap_cplxtoreal(module).float(),
-                            torch.randn(2, 3, 5, 16, 2).float(),
-                            training=training)
+        do_onnx_export_test(
+            wrap_cplxtoreal(module).float(),
+            torch.randn(2, 3, 5, 16, 2).float(),
+            training=training,
+        )
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_concatenated_casting_float_onnx_export(training):
-    module = torch.nn.Sequential(casting.ConcatenatedRealToCplx(),
-                                 nn.CplxIdentity(),
-                                 casting.CplxToConcatenatedReal())
+    module = torch.nn.Sequential(
+        casting.ConcatenatedRealToCplx(),
+        nn.CplxIdentity(),
+        casting.CplxToConcatenatedReal(),
+    )
     input = torch.randn(2, 16, 256)
 
     do_onnx_export_test(module.float(), input.float(), training=training)
@@ -285,43 +280,37 @@ def test_cplx_concatenated_casting_float_onnx_export(training):
     do_onnx_inference_test(module.float(), input.float(), training=training)
 
     with pytest.xfail(reason="double is not implemented in ONNX"):
-        do_onnx_export_test(module.double(), input.double(),
-                            training=training)
+        do_onnx_export_test(module.double(), input.double(), training=training)
 
     with pytest.xfail(reason="double is not implemented in ONNX"):
-        do_onnx_inference_test(module.double(), input.double(),
-                               training=training)
+        do_onnx_inference_test(module.double(), input.double(), training=training)
 
 
 @pytest.mark.xfail(reason="torch.as_strided is not supported by ONNX")
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_interleaved_casting_onnx_export(training):
-    module = torch.nn.Sequential(casting.InterleavedRealToCplx(),
-                                 nn.CplxIdentity(),
-                                 casting.CplxToInterleavedReal())
+    module = torch.nn.Sequential(
+        casting.InterleavedRealToCplx(),
+        nn.CplxIdentity(),
+        casting.CplxToInterleavedReal(),
+    )
     input = torch.randn(2, 16, 256)
 
     do_onnx_export_test(module.float(), input.float(), training=training)
     do_onnx_inference_test(module.float(), input.float(), training=training)
 
     with pytest.xfail(reason="double is not implemented in ONNX"):
-        do_onnx_export_test(module.double(), input.double(),
-                            training=training)
+        do_onnx_export_test(module.double(), input.double(), training=training)
 
     with pytest.xfail(reason="double is not implemented in ONNX"):
-        do_onnx_inference_test(module.double(), input.double(),
-                               training=training)
+        do_onnx_inference_test(module.double(), input.double(), training=training)
 
 
-@pytest.mark.parametrize('training', [
-    True, False, None
-])
+@pytest.mark.parametrize("training", [True, False, None])
 def test_cplx_astype_casting_float_onnx_export(training):
-    module = torch.nn.Sequential(casting.AsTypeCplx(),
-                                 nn.CplxIdentity(),
-                                 nn.modules.linear.CplxReal())
+    module = torch.nn.Sequential(
+        casting.AsTypeCplx(), nn.CplxIdentity(), nn.modules.linear.CplxReal()
+    )
     input = torch.randn(2, 16, 256)
 
     do_onnx_export_test(module.float(), input.float(), training=training)
@@ -330,9 +319,7 @@ def test_cplx_astype_casting_float_onnx_export(training):
         do_onnx_inference_test(module.float(), input.float(), training=training)
 
     with pytest.xfail(reason="double is not implemented in ONNX"):
-        do_onnx_export_test(module.double(), input.double(),
-                            training=training)
+        do_onnx_export_test(module.double(), input.double(), training=training)
 
     with pytest.xfail(reason="double is not implemented in ONNX"):
-        do_onnx_inference_test(module.double(), input.double(),
-                               training=training)
+        do_onnx_inference_test(module.double(), input.double(), training=training)
diff --git a/tests/test_relevance.py b/tests/test_relevance.py
index d46f6b8..e4e1ab5 100644
--- a/tests/test_relevance.py
+++ b/tests/test_relevance.py
@@ -7,7 +7,7 @@
 
 import torch.nn.functional as F
 
-from cplxmodule import Cplx
+from cplxmodule import Cplx  # noqa: F401
 
 from torch.nn import Linear
 from cplxmodule.nn import CplxLinear
@@ -49,8 +49,9 @@ def test_torch_expi(random_state):
     assert torch.autograd.gradcheck(torch_expi, trx.requires_grad_(True))
 
 
-def model_train(X, y, model, n_steps=20000, threshold=1.0,
-                reduction="mean", klw=1e-3, verbose=True):
+def model_train(
+    X, y, model, n_steps=20000, threshold=1.0, reduction="mean", klw=1e-3, verbose=True
+):
     import tqdm
 
     model.train()
@@ -58,7 +59,7 @@ def model_train(X, y, model, n_steps=20000, threshold=1.0,
 
     losses = []
     with tqdm.tqdm(range(n_steps)) as bar:
-        for i in bar:
+        for _ in bar:
             optim.zero_grad()
 
             y_pred = model(X)
@@ -73,8 +74,7 @@ def model_train(X, y, model, n_steps=20000, threshold=1.0,
 
             losses.append(float(loss))
             if verbose:
-                f_sparsity = sparsity(model, hard=True,
-                                      threshold=threshold)
+                f_sparsity = sparsity(model, hard=True, threshold=threshold)
             else:
                 f_sparsity = float("nan")
 
@@ -101,28 +101,48 @@ def example(kind="cplx"):
     def construct_real(linear):
         from collections import OrderedDict
 
-        return torch.nn.Sequential(OrderedDict([
-            ("body", torch.nn.Sequential(OrderedDict([
-                # ("linear", linear(n_features, n_features, bias=True)),
-                # ("relu", torch.nn.LeakyReLU()),
-            ]))),
-            ("final", linear(n_features, n_output, bias=False)),
-        ]))
+        body = torch.nn.Sequential(
+            OrderedDict(
+                [
+                    # ("linear", linear(n_features, n_features, bias=True)),
+                    # ("relu", torch.nn.LeakyReLU()),
+                ]
+            )
+        )
+
+        return torch.nn.Sequential(
+            OrderedDict(
+                [
+                    ("body", body),
+                    ("final", linear(n_features, n_output, bias=False)),
+                ]
+            )
+        )
 
     def construct_cplx(linear):
         from collections import OrderedDict
         from cplxmodule.nn import RealToCplx, CplxToReal
-        from cplxmodule.nn import CplxAdaptiveModReLU
-
-        return torch.nn.Sequential(OrderedDict([
-            ("cplx", RealToCplx()),
-            ("body", torch.nn.Sequential(OrderedDict([
-                # ("linear", linear(n_features // 2, n_features // 2, bias=True)),
-                # ("relu", CplxAdaptiveModReLU(n_features // 2)),
-            ]))),
-            ("final", linear(n_features // 2, n_output // 2, bias=False)),
-            ("real", CplxToReal()),
-        ]))
+        from cplxmodule.nn import CplxAdaptiveModReLU  # noqa: F401
+
+        body = torch.nn.Sequential(
+            OrderedDict(
+                [
+                    # ("linear", linear(n_features // 2, n_features // 2, bias=True)),
+                    # ("relu", CplxAdaptiveModReLU(n_features // 2)),
+                ]
+            )
+        )
+
+        return torch.nn.Sequential(
+            OrderedDict(
+                [
+                    ("cplx", RealToCplx()),
+                    ("body", body),
+                    ("final", linear(n_features // 2, n_output // 2, bias=False)),
+                    ("real", CplxToReal()),
+                ]
+            )
+        )
 
     device_ = torch.device("cpu")
     if kind == "cplx":
@@ -132,7 +152,7 @@ def construct_cplx(linear):
         phases = {
             "CplxLinear": (1000, 0.0),
             "CplxLinearARD": (4000, 1e-1),
-            "CplxLinearMasked": (500, 0.0)
+            "CplxLinearMasked": (500, 0.0),
         }
 
     elif kind == "real-ard":
@@ -142,7 +162,7 @@ def construct_cplx(linear):
         phases = {
             "Linear": (1000, 0.0),
             "LinearARD": (4000, 1e-1),
-            "LinearMasked": (500, 0.0)
+            "LinearMasked": (500, 0.0),
         }
 
     elif kind == "real-l0":
@@ -152,7 +172,7 @@ def construct_cplx(linear):
         phases = {
             "Linear": (1000, 0.0),
             "LinearL0": (4000, 2e-2),
-            "LinearMasked": (500, 0.0)
+            "LinearMasked": (500, 0.0),
         }
 
     elif kind == "real-lasso":
@@ -162,7 +182,7 @@ def construct_cplx(linear):
         phases = {
             "Linear": (1000, 0.0),
             "LinearLASSO": (4000, 1e-1),
-            "LinearMasked": (500, 0.0)
+            "LinearMasked": (500, 0.0),
         }
 
     if kind == "real-lasso":
@@ -177,7 +197,7 @@ def construct_cplx(linear):
 
     # a simple dataset
     X = torch.randn(10100, n_features)
-    y = - X[:, :n_output].clone()
+    y = -X[:, :n_output].clone()
     X, y = X.to(device_), y.to(device_)
 
     train_X, train_y = X[:100], y[:100]
@@ -185,9 +205,7 @@ def construct_cplx(linear):
 
     # construct models
     models = {"none": None}
-    models.update({
-        l.__name__: construct(l) for l in layers
-    })
+    models.update({ell.__name__: construct(ell) for ell in layers})
 
     # train a sequence of models
     names, losses = list(models.keys()), {}
@@ -200,8 +218,7 @@ def construct_cplx(linear):
         if models[src] is not None:
             # compute the dropout masks and normalize them
             state_dict = models[src].state_dict()
-            masks = compute_ard_masks(models[src], hard=False,
-                                      threshold=threshold)
+            masks = compute_ard_masks(models[src], hard=False, threshold=threshold)
 
             state_dict, masks = binarize_masks(state_dict, masks)
 
@@ -213,9 +230,15 @@ def construct_cplx(linear):
 
         model.to(device_)
 
-        model, losses[dst] = model_train(train_X, train_y, model,
-                                         n_steps=n_steps, threshold=threshold,
-                                         klw=klw, reduction=reduction)
+        model, losses[dst] = model_train(
+            train_X,
+            train_y,
+            model,
+            n_steps=n_steps,
+            threshold=threshold,
+            klw=klw,
+            reduction=reduction,
+        )
     # end for
 
     # get scores on test
@@ -263,7 +286,7 @@ def forward(self, input):
         phases = {
             "CplxBilinear": (1000, 0.0),
             "CplxBilinearARD": (10000, 1e-1),
-            "CplxBilinearMasked": (500, 0.0)
+            "CplxBilinearMasked": (500, 0.0),
         }
 
     elif kind == "real":
@@ -271,7 +294,7 @@ def forward(self, input):
         phases = {
             "Bilinear": (1000, 0.0),
             "BilinearARD": (10000, 1e-1),
-            "BilinearMasked": (500, 0.0)
+            "BilinearMasked": (500, 0.0),
         }
         construct = BilinearTest
 
@@ -286,10 +309,10 @@ def forward(self, input):
     out = X[:, :n_output]
     if "cplx" in kind:
         z = from_real(out, copy=False)
-        y = - to_real(z.conj * z, flatten=False).mean(dim=-2)
+        y = -to_real(z.conj * z, flatten=False).mean(dim=-2)
 
     else:
-        y = - (out * out).mean(dim=-1, keepdim=True)
+        y = -(out * out).mean(dim=-1, keepdim=True)
 
     X, y = X.to(device_), y.to(device_)
 
@@ -298,9 +321,7 @@ def forward(self, input):
 
     # construct models
     models = {"none": None}
-    models.update({
-        l.__name__: construct(l) for l in layers
-    })
+    models.update({ell.__name__: construct(ell) for ell in layers})
 
     # train a sequence of models
     names, losses = list(models.keys()), {}
@@ -313,8 +334,7 @@ def forward(self, input):
         if models[src] is not None:
             # compute the dropout masks and normalize them
             state_dict = models[src].state_dict()
-            masks = compute_ard_masks(models[src], hard=False,
-                                      threshold=threshold)
+            masks = compute_ard_masks(models[src], hard=False, threshold=threshold)
 
             state_dict, masks = binarize_masks(state_dict, masks)
 
@@ -326,9 +346,15 @@ def forward(self, input):
 
         model.to(device_)
 
-        model, losses[dst] = model_train(train_X, train_y, model,
-                                         n_steps=n_steps, threshold=threshold,
-                                         klw=klw, reduction=reduction)
+        model, losses[dst] = model_train(
+            train_X,
+            train_y,
+            model,
+            n_steps=n_steps,
+            threshold=threshold,
+            klw=klw,
+            reduction=reduction,
+        )
     # end for
 
     # get scores on test
@@ -343,7 +369,7 @@ def forward(self, input):
         print([*named_sparsity(model, hard=True, threshold=threshold)])
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     example("real-ard")
     example("real-l0")
     example("real-lasso")
diff --git a/tests/test_spectrum.py b/tests/test_spectrum.py
index 5ec8604..69181ac 100644
--- a/tests/test_spectrum.py
+++ b/tests/test_spectrum.py
@@ -12,14 +12,12 @@ def test_complex_fft(random_state):
     shape, axis = (2, 3, 256, 2, 2), 2
 
     np_x = random_state.randn(*shape) + 1j * random_state.randn(*shape)
-    tr_x = torch.tensor(np.stack([np_x.real, np_x.imag], axis=-1))
+    tr_x = torch.from_numpy(np_x)
 
     np_fft = np.fft.fft(np_x, axis=axis)
-    tr_fft = torch.fft(tr_x.transpose(axis, -2),
-                       signal_ndim=1).transpose(axis, -2)
+    tr_fft = torch.fft.fft(tr_x, dim=axis)
 
-    assert torch.allclose(tr_fft[..., 0], torch.from_numpy(np_fft.real))
-    assert torch.allclose(tr_fft[..., 1], torch.from_numpy(np_fft.imag))
+    assert torch.allclose(tr_fft, torch.from_numpy(np_fft))
 
 
 def test_hamming_window(random_state=None):
@@ -28,14 +26,12 @@ def test_hamming_window(random_state=None):
     n_window = 1024
 
     np_window = hamming(n_window, False).astype(np.float64)
-    tr_window = torch.hamming_window(n_window, periodic=True,
-                                     dtype=torch.float64)
+    tr_window = torch.hamming_window(n_window, periodic=True, dtype=torch.float64)
 
     assert torch.allclose(tr_window, torch.from_numpy(np_window))
 
     np_window = hamming(n_window, True).astype(np.float64)
-    tr_window = torch.hamming_window(n_window, periodic=False,
-                                     dtype=torch.float64)
+    tr_window = torch.hamming_window(n_window, periodic=False, dtype=torch.float64)
 
     assert torch.allclose(tr_window, torch.from_numpy(np_window))
 
@@ -45,33 +41,50 @@ def test_pwelch(random_state):
     from scipy.signal import welch
 
     # https://www.mathworks.com/help/signal/ref/pwelch.html#btulskp-6
-    fs = 1000.
-    tt = np.r_[:5 * fs - 1] / fs
+    fs = 1000.0
+    tt = np.r_[: 5 * fs - 1] / fs
 
     shape = 2, len(tt)
 
     epsilon = random_state.randn(*shape) + 1j * random_state.randn(*shape)
     np_x = np.cos(2 * np.pi * 100 * tt)[np.newaxis] + epsilon * 0.01
 
-    tr_x = torch.tensor(np.stack([np_x.real, np_x.imag], axis=-1))
+    # tr_x = torch.tensor(np.stack([np_x.real, np_x.imag], axis=-1))
+    tr_x = torch.from_numpy(np_x)
     tr_x.requires_grad = False
 
-    tr_window = torch.hamming_window(500, periodic=False, dtype=tr_x.dtype)
-
-    tr_ff, tr_px = pwelch(tr_x, 1, tr_window, fs=fs,
-                          scaling="density", n_overlap=300)
-    np_ff, np_px = welch(np_x, fs=fs, axis=-1, window=tr_window.numpy(),
-                         nfft=None, nperseg=None, scaling="density",
-                         noverlap=300, detrend=False, return_onesided=False)
+    tr_window = torch.hamming_window(500, periodic=False, dtype=tr_x.real.dtype)
+
+    tr_ff, tr_px = pwelch(tr_x, 1, tr_window, fs=fs, scaling="density", n_overlap=300)
+    np_ff, np_px = welch(
+        np_x,
+        fs=fs,
+        axis=-1,
+        window=tr_window.numpy(),
+        nfft=None,
+        nperseg=None,
+        scaling="density",
+        noverlap=300,
+        detrend=False,
+        return_onesided=False,
+    )
 
     assert torch.allclose(tr_px, torch.from_numpy(np_px))
     assert torch.allclose(tr_ff, torch.from_numpy(np_ff))
 
-    tr_ff, tr_px = pwelch(tr_x, 1, tr_window, fs=fs,
-                          scaling="spectrum", n_overlap=499)
-    np_ff, np_px = welch(np_x, fs=fs, axis=-1, window=tr_window.numpy(),
-                         nfft=None, nperseg=None, scaling="spectrum",
-                         noverlap=499, detrend=False, return_onesided=False)
+    tr_ff, tr_px = pwelch(tr_x, 1, tr_window, fs=fs, scaling="spectrum", n_overlap=499)
+    np_ff, np_px = welch(
+        np_x,
+        fs=fs,
+        axis=-1,
+        window=tr_window.numpy(),
+        nfft=None,
+        nperseg=None,
+        scaling="spectrum",
+        noverlap=499,
+        detrend=False,
+        return_onesided=False,
+    )
 
     assert torch.allclose(tr_px, torch.from_numpy(np_px))
     assert torch.allclose(tr_ff, torch.from_numpy(np_ff))
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 1ca5d2a..3f385fb 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -21,13 +21,14 @@ def test_window_view(random_state):
 
     tr_x_view = window_view(tr_x, dim, size, stride)
     for i in range(tr_x_view.shape[dim]):
-        slice_ = np.r_[i * stride:i * stride + size]
+        slice_ = np.r_[i * stride : i * stride + size]
         a = tr_x_view.index_select(dim, torch.tensor(i)).squeeze(dim)
         b = tr_x.index_select(dim, torch.tensor(slice_))
         assert_allclose(a, b)
 
-    assert_allclose(window_view(tr_x, dim, size, stride, at=-1),
-                    tr_x.unfold(dim, size, stride))
+    assert_allclose(
+        window_view(tr_x, dim, size, stride, at=-1), tr_x.unfold(dim, size, stride)
+    )
 
 
 def test_complex_view(random_state):