google / google/flax

SPMD for initializing model using nnx.jit

Open
#4,129 2 comments 0 reactions 0 assignees View on GitHub
Priority: P1 - soon
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

I would like to ask an example for initializing and updating a model using nnx.jit with SPMD.
Is there any relevant example?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.