tf.distribute.experimental.partitioners.FixedShardsPartitioner