News for hackers. Stuff that matters.

A comment on JAX – NumPy on the CPU, GPU, and TPU, with great automatic differentiation. Parent comment

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by krasin on Friday September 29, 2023 @12:54AM (#37699423)
An alternative is to write most of the program in Python + JAX + implement a few custom XLA ops in CUDA / Triton. That way, the program is very readable and can interoperate with the larger ecosystem, while still being fast to run.
[ Parent | on HN ]

There are no comments on this comment yet.