import jax.numpy as jnp


def save(filename, **kwargs) -> None:
    """Save JAX arrays to a compressed npz file.

    Args:
        filename: Path to the .npz file to save.
        **kwargs: Arrays to save, with names as keyword arguments.
    """
    jnp.savez(filename, **kwargs)


def load(filename) -> dict:
    """Load JAX arrays from an npz file.

    Args:
        filename: Path to the .npz file to load.

    Returns:
        Dictionary mapping array names to JAX arrays.
    """
    data = jnp.load(filename, allow_pickle=True)  # type: ignore[union-attr]
    return {key: data[key] for key in data.files}


if __name__ == "__main__":
    # Simple test
    arr_jax = jnp.array([1.0, 2.0, 3.0])
    dict_jax = {"a": arr_jax, "b": jnp.array([[1, 2], [3, 4]])}
    save("test_data.npz", my_array=arr_jax, my_dict=dict_jax)

    loaded_data = load("test_data.npz")
    print("Loaded data:", loaded_data)
