How to load in and export as CSV file with JAX

[stromal]

https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.load.html

GOAL:

1. To load in a CSV file with JAX to a variable
2. To export a variable via JAX to a CSV file

• It is a 6 GB file with 3 columns *(Object, int, float)

CODE

CODE 1

import jax
import jax.numpy as jnp
d = jnp.load('my_data.csv')  

ERROR 1

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In [12], line 2
      1 # jax load in my_data.csv to array
----> 2 djax = jnp.load('my_data.csv')  
      3 djax

File ~/anaconda3/envs/dask_etl/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:307, in load(*args, **kwargs)
    303 @_wraps(np.load, update_doc=False)
    304 def load(*args: Any, **kwargs: Any) -> Array:
    305   # The main purpose of this wrapper is to recover bfloat16 data types.
    306   # Note: this will only work for files created via np.save(), not np.savez().
--> 307   out = np.load(*args, **kwargs)
    308   if isinstance(out, np.ndarray):
    309     # numpy does not recognize bfloat16, so arrays are serialized as void16
    310     if out.dtype == 'V2':

File ~/anaconda3/envs/dask_etl/lib/python3.10/site-packages/numpy/lib/npyio.py:438, in load(file, mmap_mode, allow_pickle, fix_imports, encoding, max_header_size)
    435 else:
    436     # Try a pickle
    437     if not allow_pickle:
--> 438         raise ValueError("Cannot load file containing pickled data "
    439                          "when allow_pickle=False")
    440     try:
    441         return pickle.load(fid, **pickle_kwargs)

ValueError: Cannot load file containing pickled data when allow_pickle=False

CODE 2

import jax
import jax.numpy as jnp
d = jnp.load('my_data.csv',allow_pickle=True)  

ERROR 2

---------------------------------------------------------------------------
UnpicklingError                           Traceback (most recent call last)
File ~/anaconda3/envs/dask_etl/lib/python3.10/site-packages/numpy/lib/npyio.py:441, in load(file, mmap_mode, allow_pickle, fix_imports, encoding, max_header_size)
    440 try:
--> 441     return pickle.load(fid, **pickle_kwargs)
    442 except Exception as e:

UnpicklingError: could not find MARK

The above exception was the direct cause of the following exception:

UnpicklingError                           Traceback (most recent call last)
Cell In [11], line 2
      1 # jax load in my_data.csv to array
----> 2 djax = jnp.load('my_data.csv',allow_pickle=True)  
      3 djax

File ~/anaconda3/envs/dask_etl/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:307, in load(*args, **kwargs)
    303 @_wraps(np.load, update_doc=False)
    304 def load(*args: Any, **kwargs: Any) -> Array:
    305   # The main purpose of this wrapper is to recover bfloat16 data types.
    306   # Note: this will only work for files created via np.save(), not np.savez().
--> 307   out = np.load(*args, **kwargs)
    308   if isinstance(out, np.ndarray):
    309     # numpy does not recognize bfloat16, so arrays are serialized as void16
    310     if out.dtype == 'V2':

File ~/anaconda3/envs/dask_etl/lib/python3.10/site-packages/numpy/lib/npyio.py:443, in load(file, mmap_mode, allow_pickle, fix_imports, encoding, max_header_size)
    441     return pickle.load(fid, **pickle_kwargs)
    442 except Exception as e:
--> 443     raise pickle.UnpicklingError(
    444         f"Failed to interpret file {file!r} as a pickle") from e

UnpicklingError: Failed to interpret file 'my_data.csv' as a pickle

[jakevdp]

jnp.load, like its counterpart np.load, is designed to work with npy and .npz files, not csv files.

Any CSV data loading solution will happen on the host side, and so the best approach in JAX will likely be to use your preferred Python CSV loader and then convert to a JAX array; for example you could use pandas:

import pandas as pd
arr = jnp.asarray(pd.read_csv('my_data.csv').values)

[stromal]

@jakevdp
Well my main issue is that I doesn't want to use pandas because my data file is too large.
That is why I want to use JAX to get a quick GPU accelerated load in.

[jakevdp]

JAX does not offer any functionality for GPU-accelerated CSV parsing, and I'm not sure what other options there are. You may be able to use some other tool built for this, and then convert the resulting buffer to JAX via dlpack.

[stromal]

Modin is another option but that has also some issues (the load in is fin, but after)