4 ms·

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.