Adding tutorial for data loaders on gpu with jax

[selamw1]

Adding Introduction to Data Loaders on GPU with JAX.

[selamw1]

Headings adjusted and new overview paragraph added into the introduction.

[trallard]

@selamw1 there seem to be a couple of merge conflicts here would you be able to fix them, please?

[selamw1]

File conflict resolved.

[jakevdp]

I see the new introductory paragraph – thanks for adding that!

I'm still somewhat concerned here about duplication between the three "data loader" tutorials: one for CPU, one for GPU, and one for TPU. Putting myself in the shoes of a reader trying to become familiar with JAX, it's hard to know without reading each in detail what the key differences are, or why I should care.

Maybe we could work on de-duplicating the three tutorials, so that they build on each other sequentially rather than repeating the same concepts three times with slight differences?

[selamw1]

I see the new introductory paragraph – thanks for adding that!

I'm still somewhat concerned here about duplication between the three "data loader" tutorials: one for CPU, one for GPU, and one for TPU. Putting myself in the shoes of a reader trying to become familiar with JAX, it's hard to know without reading each in detail what the key differences are, or why I should care.

Maybe we could work on de-duplicating the three tutorials, so that they build on each other sequentially rather than repeating the same concepts three times with slight differences?

The notebooks have been adjusted to highlight the key differences between each other, and their respective references have been added to make them standalone tutorials.

[jakevdp]

Thanks - we'll have to rebase this PR to reflect the fact that docs has moved to docs/source.

[selamw1]

Rebase has been done to move files from docs to docs/source.

[jakevdp]

Thanks - it looks like we ended up with two copies of each file, one in docs/ and one in docs/source/. Can you remove the ones in docs/? Thanks!

[jakevdp]

Stray git conflict marker; we need to remove this.

[selamw1]

done.

[selamw1]

Thanks - it looks like we ended up with two copies of each file, one in docs/ and one in docs/source/. Can you remove the ones in docs/? Thanks!

Yes, its done.

[jakevdp]

Thanks!