INNER CODE UNIT · Python
test_optax
HenryNdubuaku/nanodl · nanodl/__init__.py:324
def test_optax(optax):
optimizer = optax.sgd(learning_rate=0.1)
def test_einops(einops):
arr = einops.rearrange([1, 2, 3], "a b c -> b a c")
def main():
try:
flax = check_library_installed("flax")
jax = check_library_installed("jax")
optax = check_library_installed("optax")
einops = check_library_installed("einops")
test_flax(flax)
test_jax(jax)
test_optax(optax)