topology
topology
¶
Hardware information and topology representation for distributed JAX setups. Defines a Topology class that encapsulates device and process information, as well as JAX Mesh configuration.
Topology
¶
Bases: BaseModel
new(chip: Chip, shard: ShardingPolicy | None = None) -> Topology
classmethod
¶
Create a Topology instance based on the current JAX device configuration.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
chip
|
Chip
|
The chip type being used. |
required |
shard
|
ShardingPolicy | None
|
Parallelism policy. The batch axis uses all remaining devices. |
None
|