#

shmap

Here are 2 public repositories matching this topic...

fattorib / transformer_shmap

Tensor Parallelism with JAX + Shard Map

transformers gpt tpu jax tensor-parallelism pjit shmap

Updated Sep 29, 2023
Python

azzeddineCH / flash-nanoGPT

Jax/Flax re-write of @karpathy 🐐 NanoGPT using some of the common Jax libraries/features (shmap, pallas, jmp, optax, orbax)

transformers flax jax gpt-2 shmap

Updated Jun 7, 2024
Python

Improve this page

Add a description, image, and links to the shmap topic page so that developers can more easily learn about it.

Curate this topic

Add this topic to your repo

To associate your repository with the shmap topic, visit your repo's landing page and select "manage topics."