jax
86698fba - Simplify pyval result handler.

Commit
6 years ago
Simplify pyval result handler. The previous change to the pyval result handler attempted to prefetch values to the host. However, this didn't actually work because decompose() doesn't propagate prefetched values. However, there's a much simpler thing we can do here: just call to_py() on the entire value. The XLA library knows how to return (nested) tuple values and we only need convert them to JaxTuples. This achieves everything we wanted from prefetching (namely that we don't synchronously fetch the tuple leaves one by one).
Author
Parents
Loading