[nnx] add data parallel toy example#4354

copybara-service[bot] merged 1 commit intomaingoogle/flax:mainfrom

nnx-data-parallel-examplegoogle/flax:nnx-data-parallel-exampleCopy head branch name to clipboard