Improve test coverage

[josevalim]

[seanmor5]

Overall more test coverage is required for every module.

Axon

• [x] Test layer creation for all layers
• [x] Test Axon.init
  • [x] inside defn
  • [x] outside defn
• [x] Test Axon.predict
  • [x] inside defn
  • [x] outside defn
• [x] Test inspection
• [x] Test compile

Axon.Updates

Integration is handled in Optimizers. It may be useful to reference how optax tests their transforms.

• [ ] Test initialization of each update
• [ ] Test update fn of each update

Axon.Optimizers

Reference Optax

• [ ] Test initialization of each optimizer
• [ ] Test update fn of each optimizer

Axon.Training

These should probably be exercised as integration tests. The XOR example is a good example as an integration test, and we can start from a deterministic point.

[seanmor5]

Axon's test coverage has increased greatly here, and despite some shortcomings of specific APIs (Axon.Loop), Axon is tested well against PT in other libraries, so I am closing this :)