def METHOD_NAME(mesh):
    layout_map = layout_map_lib.LayoutMap(mesh=mesh)
    layout_4d = dtensor.Layout.replicated(mesh, rank=4)
    layout_2d = dtensor.Layout.replicated(mesh, rank=2)
    layout_1d = dtensor.Layout.replicated(mesh, rank=1)
    layout_map["conv2d.*kernel"] = layout_4d
    layout_map["conv2d.*bias"] = layout_1d
    layout_map["dense.*kernel"] = layout_2d
    layout_map["dense.*bias"] = layout_1d
    return layout_map