flaxlm
flaxlm copied to clipboard
use shard_map for fsdp