google / google/flax

Show sharding information in Module.tabulate

Open
#3,325 0 comments 0 reactions 1 assignee Claimed by @cgarciae View on GitHub
Priority: P2 - no schedule
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Show params and input/output sharding information in `Module.tabulate`.

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.