tf.keras.distribution.DeviceMesh

A cluster of computation devices for distributed computation.

tf.keras.distribution.DeviceMesh(
    shape, axis_names, devices=None
)

This API is aligned with jax.sharding.Mesh and tf.dtensor.Mesh, which represents the computation devices in the global context.

Args
`shape`	tuple of list of integers. The shape of the overall `DeviceMesh`, e.g. `(8,)` for a data parallel only distribution, or `(4, 2)` for a model+data parallel distribution.
`axis_names`	List of string. The logical name of the each axis for the `DeviceMesh`. The length of the `axis_names` should match to the rank of the `shape`. The `axis_names` will be used to match/create the `TensorLayout` when distribute the data and variables.
`devices`	Optional list of devices. Defaults to all the available devices locally from `keras.distribution.list_devices()`.

Attributes

axis_names

devices

shape

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates. Some content is licensed under the numpy license.

Last updated 2024-06-07 UTC.