tf.keras.distribution.LayoutMap

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

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

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

get

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

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

keys

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

pop

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

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

setdefault

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

update

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

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

__contains__

__eq__

Return self==value.

__getitem__

View source

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

__len__

View source