Why not Jax?
- JAX Windows support is lacking
- CuPy is much closer to CUDA than JAX, so you can get better performance
- CuPy is generally more mature than JAX (fewer bugs)
- CuPy is more flexible thanks to cp.RawKernel
- (For those familiar with NumPy) CuPy is closer to NumPy than jax.numpy
But CuPy does not support automatic gradient computation, so if you do deep learning, use JAX instead. Or PyTorch, if you do not trust Google to maintain a project for a prolonged period of time https://killedbygoogle.com/
Just prepare the input on NumPy or CuPy, and then you can just feed it to NumPy APIs. NumPy functions will handle itself if the input is NumPy ndarray, or dispatch the execution to CuPy if the input is CuPy ndarray.
CuPy is an (almost) drop-in replacement for NumPy, so the following works surprisingly often:
if use_cpu:
import numpy as np
else:
import cupy as npThis is the problem with these kind of methods. It works, until it doesn't in an unknown way.
DL is major use case; is CuPy planning on adding auto gradient comp?
Jax also provides back-propagation wherever possible, so you can optimize.
Cool to see that it's still kicking!