sharding_util

Utilities for working with sharded arrays and parameters in Penzai.

Classes

ConstrainSharding	A layer that constrains the sharding of a tree of arrays.
ConstrainShardingByName	A layer that constrains the sharding of a tree of NamedArrays by name.

Functions

name_to_name_device_put(tree, mesh[, ...])	Shards a tree of pz.nx.NamedArray objects based on their axis names.
name_to_name_sharding(tree, mesh[, ...])	Shards a tree of pz.nx.NamedArray objects based on their axis names.
sharded_init(initializer, *init_args, mesh)	Initializes a model, with constants and variables sharded based on a mesh.