def get_library_fns(dtype):
    """Given a dtype e.g. jnp.float32 or torch.complex64, returns the appropriate
    namespace for standard array functionality (either torch_fns or jax_fns)."""
    try:
        from jax import numpy as jnp

        if check_valid_dtype(jnp.array(0.1), dtype=dtype, alloc_fn=lambda x, y: x.astype(y)):
            from cola.backends import jax_fns as fns

            return fns

    except ImportError:
        pass

    try:
        import torch

        if check_valid_dtype(torch.tensor(0.1), dtype=dtype, alloc_fn=lambda x, y: x.to(y)):
            from cola.backends import torch_fns as fns

            return fns

    except ImportError:
        pass

    if check_valid_dtype(np.array(0.1), dtype=dtype, alloc_fn=lambda x, y: x.astype(y)):
        from cola.backends import np_fns as fns

        return fns

    raise ImportError("No supported array library found")
Given a dtype e.g. jnp.float32 or torch.complex64, returns the appropriate
    namespace for standard array functionality (either torch_fns or jax_fns).
cola/backends/backends.py