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.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.