Keras 3 + JAX 使得通过数据并行和模型并行（变量分割）在任意多的设备上高效分布大型模型成为可能。

它不需要修改您的模型代码或训练代码——您需要做的只是单独指定一个“LayoutMap”，它告诉框架如何分割或复制模型中的每个变量。