JAX – NumPy on the CPU, GPU, and TPU, with great automatic differentiation

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by cl3misch on Friday September 29, 2023 @11:27AM (#37705645)

Yes, but is it really the same? Afaik the `unroll=n` parameter translates `n` iterations into a vanilla `for` loop which is then unrolled into sequential statements (in contrast to a JAX `fori` loop). There still is no loop on the accelerator, strictly speaking?

[ Parent | on HN ]

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by marmaduke on Friday September 29, 2023 @01:11PM (#37707214)

I think this is up to XLA to handle not Jax. The whole selling point in TF of the tf.function decorator (which uses XLA underneath as well) is that it fuses arithmetic to lower launch count.

[ Parent | on HN ]

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by marmaduke on Friday September 29, 2023 @03:53PM (#37709309)

there's a paper about it that I just found, enjoy https://arxiv.org/pdf/2301.13062

[ Parent | on HN ]