tf.keras.distribution.LayoutMap

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

tf.keras.distribution.LayoutMap(
    device_mesh=None
)

LayoutMap uses a string as key and a TensorLayout 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 TensorLayout, and then retrieve the corresponding TensorLayout instance.

In the normal case, the key to query is usually the variable.path, which is the idenifier of the variable.

As shortcut, tuple or list of axis names are also allowed when inserting as value, and will be converted to TensorLayout.

layout_map = LayoutMap(device_mesh=None)
layout_map['dense.*kernel'] = (None, 'model')         # layout_2d
layout_map['dense.*bias'] = ('model',)                # layout_1d
layout_map['conv2d.*kernel'] = TensorLayout((None, None, None, 'model'))
layout_map['conv2d.*bias'] = TensorLayout(('model',))  # layout_1d

layout_1 = layout_map['dense_1.kernel']             # layout_1 == layout_2d
layout_2 = layout_map['dense_1.bias']               # layout_2 == layout_1d
layout_3 = layout_map['dense_2.kernel']             # layout_3 == layout_2d
layout_4 = layout_map['dense_2.bias']               # layout_4 == layout_1d
layout_5 = layout_map['my_model/conv2d_123/kernel'] # layout_5 == layout_4d
layout_6 = layout_map['my_model/conv2d_123/bias']   # layout_6 == layout_1d
layout_7 = layout_map['my_model/conv3d_1/kernel']   # layout_7 == None
layout_8 = layout_map['my_model/conv3d_1/bias']     # layout_8 == None

device_mesh An optional DeviceMesh that can be used to populate the TensorLayout.device_mesh if TensorLayout.device_mesh is not set.

device_mesh

Methods

clear

clear()

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

get

get(
    key, default=None
)

Retrieves 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. When there are multiple matches for the regex, an ValueError will be raised. Returns None if there isn't any match found.

Args
key String key to query a layout.

Returns
Corresponding layout based on the 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(
    other, /, **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
)

Retrieves 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. When there are multiple matches for the regex, an ValueError will be raised. Returns None if there isn't any match found.

Args
key String key to query a layout.

Returns
Corresponding layout based on the query.

__iter__

View source

__iter__()

__len__

View source

__len__()