fix: randomize all layer types

[aski02]

References #18.

However, the test for transformer-specific layers could be cleaned up once we support language models for explainers as well.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.06%. Comparing base (055695f) to head (b887c52).
Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #252   +/-   ##
=======================================
  Coverage   92.06%   92.06%           
=======================================
  Files          55       55           
  Lines        2331     2331           
=======================================
  Hits         2146     2146           
  Misses        185      185           

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

[dilyabareeva]

Hi @aski02. Thanks for your work. I have two major comments:

1. The tests are only checking if parameters have been changed after randomization. Could we add tests that do inference on a batch of random data and see if the model output NaNs? That was the main reason we limited randomization to Linear layers previously.
2. As far as I remember, the current approach didn't work for BatchNorm layers. Could we add tests for them as well?