Hi.
I am trying to configure the LayoutMap for using DTensors with a tf.keras.applications model. Since most of the models share components like layers.Conv2D, layers.Dense I think there has to be a generic LayoutMap that could serve most of the situations. This is more so because LayoutMap treats the keys as regex expressions.
I have investigated this a bit in this Colab Notebook. I am still unable to figure out a proper LayoutMap that works. Any help is appreciated.