Update XLA and fix infeed on the GPU

elixir-nx / nx

Multi-dimensional arrays (tensors) and numerical definitions for Elixir

2.62k stars 190 forks source link

Update XLA and fix infeed on the GPU #1487

Closed jonatanklosko closed 4 months ago

jonatanklosko commented 4 months ago

This updates XLA to the latest revision (xla v0.7.0 already released).

When testing on the GPU, I noticed that infeed is broken. So far we had different logic for CPU and GPU, and I managed to unify it into a single code path. See the comment in exla_client.cc for details.