View source on GitHub |
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
Args | |
---|---|
device_mesh
|
An optional DeviceMesh that can be used to populate the
TensorLayout.device_mesh if TensorLayout.device_mesh is not set.
|
Attributes | |
---|---|
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__
__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__
__iter__()
__len__
__len__()