INNER CODE UNIT · Python

test_einops

HenryNdubuaku/nanodl · nanodl/__init__.py:328

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)

    except ImportError as e:
        print(e)
        sys.exit(1)

View source record →