Code

Add MNIST example with SPMD for JAX#2346

Code

Open

sandipanpanda wants to merge 5 commits intokubeflow:masterkubeflow/training-operator:masterfrom sandipanpanda:jax-examplesandipanpanda/training-operator:jax-exampleCopy head branch name to clipboard

+453-7

Commits

Commits on Dec 25, 2024

• Add MNIST example with SPMD for JAX

  Show description for 619aa20
  sandipanpanda

  committed
  619aa20

  Copy full SHA for 619aa20

• Update CONTRIBUTING.md

  Show description for ca78118
  sandipanpanda

  committed
  ca78118

  Copy full SHA for ca78118

• Add JAXJob output

  Show description for 7f523e4
  sandipanpanda

  committed
  7f523e4

  Copy full SHA for 7f523e4

• Update JAXJob CI images

  Show description for 1ed9b7d
  sandipanpanda

  committed
  1ed9b7d

  Copy full SHA for 1ed9b7d

Commits on Dec 28, 2024

• Adjust jaxjob spmd example batch size

  Show description for 904dbd6
  sandipanpanda

  committed
  904dbd6

  Copy full SHA for 904dbd6