4 ms·

If you ever need to reap the benefits of CUDA & GPU computations without getting into the details, check out JAX by our corporate overlords™ (https://github.com

If you ever need to reap the benefits of CUDA & GPU computations without getting into the details, check out JAX by our corporate overlords™ (https://github.com/google/jax https://github.com/google/jax), it has a NumPy like syntax and super fast to get started

p1esk 5y ago
Why would you suggest JAX? CuPy seems like an obvious choice here (simpler and a lot more mature). Jax is only needed if you want automatic differentiation.
- hamilyon2 5y ago
  TIL CuPy exists and is stable and mature. This is why i read this forum, every now and then there is serendipitous connection, new angle or discovery.
- dekhn 5y ago
  I just want pyXLA, a tool to directly construct XLA programs.