tf.keras.dtensor.experimental.LayoutMap

A dict-like object that maps string to Layout instances.

tf.keras.dtensor.experimental.LayoutMap(
    mesh=None
)

LayoutMap uses a string as key and a Layout as value. There is a behavior difference between a normal Python dict and this class. The string key will be treated as a regex when retrieving the value. See the docstring of get for more details.

See below for a usage example. You can define the naming schema of the Layout, and then retrieve the corresponding Layout instance.

To use the LayoutMap with a Model, please see the docstring of tf.keras.dtensor.experimental.layout_map_scope.

map = LayoutMap(mesh=None)
map['.*dense.*kernel'] = layout_2d
map['.*dense.*bias'] = layout_1d
map['.*conv2d.*kernel'] = layout_4d
map['.*conv2d.*bias'] = layout_1d

layout_1 = map['dense_1.kernel']    #   layout_1 == layout_2d
layout_2 = map['dense_1.bias']      #   layout_2 == layout_1d
layout_3 = map['dense_2.kernel']    #   layout_3 == layout_2d
layout_4 = map['dense_2.bias']      #   layout_4 == layout_1d
layout_5 = map['my_model/conv2d_123/kernel']    #   layout_5 == layout_4d
layout_6 = map['my_model/conv2d_123/bias']      #   layout_6 == layout_1d

Args
mesh	An optional Mesh that can be used to create all replicated layout as default when there isn't a layout found based on the input string query.

Methods

clear

clear()

D.clear() -> None. Remove all items from D.

get

get(
    key, default=None
)

Retrieve the corresponding layout by the string key.

When there isn't an exact match, all the existing keys in the layout map will be treated as a regex and map against the input key again. The first match will be returned, based on the key insertion order. Return None if there isn't any match found.

Args
key	the string key as the query for the layout.

Returns
Corresponding layout based on the query.

get_default_mesh

View source

get_default_mesh()

Return the default Mesh set at instance creation.

The Mesh can be used to create default replicated Layout when there isn't a match of the input string query.

items

items()

D.items() -> a set-like object providing a view on D's items

keys

keys()

D.keys() -> a set-like object providing a view on D's keys

pop

pop(
    key, default=__marker
)

D.pop(k[,d]) -> v, remove specified key and return the corresponding value. If key is not found, d is returned if given, otherwise KeyError is raised.

popitem

popitem()

D.popitem() -> (k, v), remove and return some (key, value) pair as a 2-tuple; but raise KeyError if D is empty.

setdefault

setdefault(
    key, default=None
)

D.setdefault(k[,d]) -> D.get(k,d), also set D[k]=d if k not in D

update

update(
    *args, **kwds
)

D.update([E, ]**F) -> None. Update D from mapping/iterable E and F. If E present and has a .keys() method, does: for k in E: D[k] = E[k] If E present and lacks .keys() method, does: for (k, v) in E: D[k] = v In either case, this is followed by: for k, v in F.items(): D[k] = v

values

values()

D.values() -> an object providing a view on D's values

__contains__

__contains__(
    key
)

__eq__

__eq__(
    other
)

Return self==value.

__getitem__

View source

__getitem__(
    key
)

Retrieve the corresponding layout by the string key.

When there isn't an exact match, all the existing keys in the layout map will be treated as a regex and map against the input key again. The first match will be returned, based on the key insertion order. Return None if there isn't any match found.

Args
key	the string key as the query for the layout.

Returns
Corresponding layout based on the query.

__iter__

View source

__iter__()

__len__

View source

__len__()