Skip to content

Jax example test

TestJaxExample #

Bases: LearningAlgorithmTests[JaxExample]

Tests for the Jax example algorithm.

This simply reuses all the tests in the base test suite, specifying that the datamodule passed to the JaxExample should be for image classification and the network should be a flax.linen.Module.