chore: import upstream snapshot with attribution

tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc

@@ -0,0 +1,175 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+StatusOr<Layout> ComputeResultLayout(mlir::Operation* op,
+                                     const Layout& input_layout) {
+  if (!mlir::isa<mlir::TF::ArgMaxOp>(op))
+    return absl::UnimplementedError(absl::StrCat(
+        "SPMD expansion for op type: ", OpName(op), " not yet implemented."));
+
+  auto argmax_op = llvm::cast<mlir::TF::ArgMaxOp>(op);
+  const auto input_rank = ValueRank(argmax_op.getInput());
+  TF_ASSIGN_OR_RETURN(int64_t axis,
+                      ExtractConstIntFromValue(argmax_op.getDimension()));
+
+  if (axis < 0) axis += input_rank;
+
+  std::vector<std::string> sharding_specs;
+  for (int i = 0; i < input_rank; ++i) {
+    if (i != axis) sharding_specs.push_back(input_layout.sharding_spec(i));
+  }
+  return Layout::GetLayout(sharding_specs, input_layout.mesh());
+}
+}  // namespace
+
+StatusOr<mlir::Operation*> ArgMaxSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto argmax_op = llvm::cast<mlir::TF::ArgMaxOp>(op);
+  TF_ASSIGN_OR_RETURN(int64_t axis,
+                      ExtractConstIntFromValue(argmax_op.getDimension()));
+  TF_ASSIGN_OR_RETURN(auto input_layout,
+                      ExtractLayoutFromOperand(argmax_op.getInput()));
+  TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(argmax_op));
+  if (!input_layout || !output_layout)
+    return absl::InvalidArgumentError(
+        absl::StrCat(OpName(op), " is missing layouts during SPMD Expansion."));
+
+  mlir::Value input = argmax_op.getInput();
+  const auto input_rank = ValueRank(input);
+
+  TF_ASSIGN_OR_RETURN(auto input_shape, GetShapeOfValue(input));
+
+  if (input_rank == -1)
+    return absl::UnimplementedError("missing rank for input.");
+  if (axis < 0) axis += input_rank;
+
+  mlir::OpBuilder builder(op);
+  {
+    std::vector<std::string> tgt_input_sharding_specs;
+    tgt_input_sharding_specs.reserve(input_shape.size());
+    Mesh mesh = input_layout->mesh();
+
+    for (int i = 0; i < input_shape.size(); ++i) {
+      // const auto dim_name
+      if (i == axis) {
+        // Set replicated for `axis` dim.
+        tgt_input_sharding_specs.push_back(Layout::kUnshardedDim);
+      } else {
+        // Keep the rest dimension.
+        tgt_input_sharding_specs.push_back(input_layout->sharding_spec(i));
+      }
+    }
+
+    if (!Layout::IsUnshardedDimension(input_layout->sharding_spec(axis))) {
+      TF_ASSIGN_OR_RETURN(Layout layout,
+                          Layout::GetLayout(input_layout->type(),
+                                            tgt_input_sharding_specs, mesh));
+      TF_ASSIGN_OR_RETURN(input,
+                          EmitAllGather(builder, input, *input_layout, layout));
+    }
+  }
+
+  auto new_argmax = mlir::TF::ArgMaxOp::create(builder, argmax_op.getLoc(),
+                                               argmax_op.getResult().getType(),
+                                               input, argmax_op.getDimension());
+  op->getResult(0).replaceAllUsesWith(new_argmax.getOutput());
+  op->erase();
+
+  return InferSPMDExpandedLocalShape(new_argmax);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ArgMaxSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If the input layout is missing, don't return an output layout.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout& input_layout = input_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(auto result_layout,
+                      ComputeResultLayout(op, input_layout));
+  if (result_layout.rank() != input_layout.rank() - 1)
+    return absl::FailedPreconditionError(absl::StrCat(
+        OpName(op), " derived output layout rank is ", result_layout.rank(),
+        " not ", input_layout.rank() - 1, " as expected."));
+
+  return llvm::DenseMap<int, Layout>({{0, result_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ArgMaxSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If no output layout, then do not infer any operand layouts.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  auto argmax_op = llvm::cast<mlir::TF::ArgMaxOp>(op);
+  TF_ASSIGN_OR_RETURN(int64_t axis,
+                      ExtractConstIntFromValue(argmax_op.getDimension()));
+  auto input = argmax_op.getInput();
+  const auto input_rank = ValueRank(input);
+
+  // Handle the case of negative axis.
+  if (axis < 0) axis += input_rank;
+
+  const Layout& output_layout = output_layouts.lookup(0);
+
+  TF_ASSIGN_OR_RETURN(auto input_shape, GetShapeOfValue(input));
+
+  std::vector<std::string> layout_sharding;
+
+  int output_dim = 0;
+  for (int i = 0; i < input_shape.size(); ++i) {
+    if (i == axis) {
+      layout_sharding.emplace_back(Layout::kUnshardedDim);
+    } else {
+      layout_sharding.emplace_back(output_layout.sharding_spec(output_dim));
+      output_dim += 1;
+    }
+  }
+
+  // Add Layout for first input attribute, while the second one is axis as a
+  // scalar, we don't need to set its layout.
+  TF_ASSIGN_OR_RETURN(const Layout result_layout,
+                      Layout::GetLayout(layout_sharding, output_layout.mesh()));
+  return llvm::DenseMap<int, Layout>({{0, result_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ARGMAX_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ARGMAX_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class ArgMaxSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ARGMAX_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/bias_add_spmd_expander.cc

@@ -0,0 +1,173 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/bias_add_spmd_expander.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+int get_c_dimension_idx(const Layout& layout, llvm::StringRef data_format) {
+  // If format is "N...C", the bias is added to the last dimension.
+  int c_dim_idx = layout.sharding_spec_strs().size() - 1;
+  if (data_format.starts_with("NC")) {
+    // If format is "NC...", the bias is added to the 'C' dimension.
+    c_dim_idx = layout.sharding_spec_strs().size() - 3;
+  }
+  return c_dim_idx;
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> BiasAddExpander::ExpandOp(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  mlir::TF::BiasAddOp bias_add_op = llvm::cast<mlir::TF::BiasAddOp>(op);
+  const llvm::StringRef data_format = bias_add_op.getDataFormat();
+  const int c_dim_idx = get_c_dimension_idx(output_layout, data_format);
+
+  // Bias add op has 2 inputs: value and bias.
+  assert(op->getOpOperands().size() == 2);
+  mlir::OpOperand& input = op->getOpOperand(0);
+  TF_ASSIGN_OR_RETURN(Layout input_layout,
+                      ExtractRequiredLayoutFromOperand(input.get()));
+
+  mlir::OpOperand& bias = op->getOpOperand(1);
+
+  TF_ASSIGN_OR_RETURN(const Layout bias_layout,
+                      ExtractRequiredLayoutFromOperand(bias.get()));
+
+  // Check if output is sharded more, change input layout to match output
+  // layout.
+  int64_t num_input_shards = input_layout.num_shards_for_dim(c_dim_idx);
+  int64_t num_output_shards = output_layout.num_shards_for_dim(c_dim_idx);
+
+  if (num_input_shards < num_output_shards) {
+    mlir::Value output;
+    std::vector<std::string> input_new_specs =
+        output_layout.sharding_spec_strs();
+    TF_ASSIGN_OR_RETURN(
+        const Layout new_input_layout,
+        Layout::GetLayout(input_new_specs, input_layout.mesh()));
+    TF_ASSIGN_OR_RETURN(
+        output, EmitRelayout(input.get(), input_layout, new_input_layout));
+    input.set(output);
+    input_layout = new_input_layout;
+  }
+
+  // Map bias layout sharding to match sharding for 'c' dimension of input, if
+  // not same already.
+  if (bias_layout.sharding_spec(0) != input_layout.sharding_spec(c_dim_idx)) {
+    mlir::Value output;
+
+    std::vector<std::string> bias_new_specs = {
+        input_layout.sharding_spec_strs()[c_dim_idx]};
+    TF_ASSIGN_OR_RETURN(const Layout new_bias_layout,
+                        Layout::GetLayout(bias_new_specs, bias_layout.mesh()));
+    TF_ASSIGN_OR_RETURN(output,
+                        EmitRelayout(bias.get(), bias_layout, new_bias_layout));
+    bias.set(output);
+  }
+
+  // Perform SPMD operation locally
+  mlir::Operation* new_local_op = InferSPMDExpandedLocalShape(op);
+
+  // Convert result layout to output layout.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  TF_ASSIGN_OR_RETURN(mlir::Value relayout_output,
+                      EmitRelayout(new_local_op->getOpResult(0), input_layout,
+                                   output_layout, &newly_created_ops));
+  op->getResult(0).replaceAllUsesExcept(relayout_output, newly_created_ops);
+  return relayout_output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> BiasAddExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If we do not have an input layout then do not infer an output layout.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  Layout input_layout = input_layouts.lookup(0);
+  mlir::TF::BiasAddOp bias_add_op = llvm::cast<mlir::TF::BiasAddOp>(op);
+  llvm::StringRef data_format = bias_add_op.getDataFormat();
+  int c_dim_idx = get_c_dimension_idx(input_layout, data_format);
+
+  std::vector<std::string> new_output_layout_specs =
+      input_layout.sharding_spec_strs();
+  if (Layout::IsUnshardedDimension(new_output_layout_specs[c_dim_idx]) &&
+      input_layouts.find(1) != input_layouts.end()) {
+    // Shard c_dim using bias sharding as long as the sharding spec is not
+    // already used in input for some other dimension.
+    Layout bias_layout = input_layouts.lookup(1);
+    std::string bias_sharding = bias_layout.sharding_spec(0);
+    if (std::find(new_output_layout_specs.begin(),
+                  new_output_layout_specs.end(),
+                  bias_sharding) == new_output_layout_specs.end()) {
+      new_output_layout_specs[c_dim_idx] = bias_layout.sharding_spec(0);
+    }
+  }
+  TF_ASSIGN_OR_RETURN(
+      Layout new_output_layout,
+      Layout::GetLayout(new_output_layout_specs, input_layout.mesh()));
+
+  return llvm::DenseMap<int, Layout>({{0, new_output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> BiasAddExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  llvm::DenseMap<int, Layout> input_layouts;
+
+  // If output layout is given, match input_layout and bias layout to match
+  // it.
+  Layout output_layout = output_layouts.lookup(0);
+
+  // Bias layout should match 'C' dimension of input layout.
+  mlir::TF::BiasAddOp bias_add_op = llvm::cast<mlir::TF::BiasAddOp>(op);
+  llvm::StringRef data_format = bias_add_op.getDataFormat();
+  const int c_dim_idx = get_c_dimension_idx(output_layout, data_format);
+
+  std::vector<std::string> bias_new_specs = {
+      output_layout.sharding_spec_strs()[c_dim_idx]};
+  TF_ASSIGN_OR_RETURN(Layout new_bias_layout,
+                      Layout::GetLayout(bias_new_specs, output_layout.mesh()));
+
+  return llvm::DenseMap<int, Layout>(
+      {{0, output_layout}, {1, new_bias_layout}});
+}
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/bias_add_spmd_expander.h

@@ -0,0 +1,48 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BIAS_ADD_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BIAS_ADD_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Expander for `BiasAdd` Op.
+// `BiasAdd` uses non-standard broadcasting rules if 'NC...' data_format is
+// used, thus we can't just reuse ElementwiseSPMDExpander as it is.
+class BiasAddExpander : public SPMDExpanderBase {
+ protected:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BIAS_ADD_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.cc

@@ -0,0 +1,220 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> BroadcastToSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  auto broadcast_op = llvm::cast<mlir::TF::BroadcastToOp>(op);
+  TF_ASSIGN_OR_RETURN(
+      const Layout shape_layout,
+      ExtractRequiredLayoutFromOperand(broadcast_op.getShape()));
+  if (!shape_layout.IsFullyReplicated()) {
+    return absl::InvalidArgumentError(
+        "Error during BroadcastOp SPMD Expansion. Shape input of broadcast op "
+        "must be fully replicated.");
+  }
+
+  TF_ASSIGN_OR_RETURN(
+      const Layout input_layout,
+      ExtractRequiredLayoutFromOperand(broadcast_op.getInput()));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(broadcast_op));
+
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> input_global_size,
+      GetGlobalShapeOfValueFromDTensorLayout(broadcast_op.getInput()));
+
+  llvm::SmallVector<int64_t, 4> broadcast_to_shape;
+  TF_RETURN_IF_ERROR(ExtractConstVectorFromValue(
+      GetForwardedDTensorLayoutInput(broadcast_op.getShape()),
+      &broadcast_to_shape));
+
+  // Input to BroadcastTo op requires all to all if non-broadcasted-dimensions
+  // are not same.
+  const int broadcasted_dimensions = output_layout.rank() - input_layout.rank();
+  bool requires_all_to_all = false;
+  const auto output_num_shards = output_layout.num_shards();
+  for (int i = 0; i < input_layout.rank(); ++i) {
+    const int output_dim_index = i + broadcasted_dimensions;
+    const std::string& output_layout_dim =
+        output_layout.sharding_spec(output_dim_index);
+    if (input_global_size[i] > 1 &&
+        input_layout.sharding_spec(i) != output_layout_dim) {
+      requires_all_to_all = true;
+    }
+    if (output_layout_dim != Layout::kUnshardedDim) {
+      broadcast_to_shape[output_dim_index] /=
+          output_num_shards[output_dim_index];
+    }
+  }
+
+  // Insert all-to-all operations just before Broadcast op to ensure all inputs
+  // in correct local values.
+  mlir::OpBuilder builder(op);
+  mlir::Value input_data = broadcast_op.getInput();
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  const Layout all_to_all_input_layout =
+      Layout::ReplicatedOnMesh(mesh, input_layout.rank());
+
+  if (requires_all_to_all) {
+    TF_ASSIGN_OR_RETURN(auto input_data,
+                        EmitAllGather(builder, input_data, input_layout,
+                                      all_to_all_input_layout));
+    op->setOperand(0, input_data);
+  } else {
+    // When all-to-all is not needed, output of BroadcastTo operation may be
+    // sharded. In that case, we must ensure that `shape` input of BroadcastTo
+    // op has correct local sharded shape.
+    // Note that we include the sharding on the first
+    for (int i = 0; i < broadcasted_dimensions; ++i)
+      if (output_layout.sharding_spec(i) != Layout::kUnshardedDim)
+        broadcast_to_shape[i] /= output_num_shards[i];
+    mlir::Value new_broadcast_to_shape =
+        Int64Const(builder, op->getLoc(), broadcast_to_shape);
+    op->setOperand(1, new_broadcast_to_shape);
+  }
+
+  op = InferSPMDExpandedLocalShape(op);
+  if (!requires_all_to_all) return op;
+
+  // If we all-to-all'ed, we may need to split after the local BroadcastTo op
+  // has been created in graph.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  builder.setInsertionPointAfter(op);
+  TF_ASSIGN_OR_RETURN(
+      auto final_output,
+      EmitAllScatter(builder, op->getOpResult(0),
+                     all_to_all_input_layout.LeftPad(output_layout.rank()),
+                     output_layout, &newly_created_ops));
+  op->getOpResult(0).replaceAllUsesExcept(final_output, newly_created_ops);
+  return final_output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+BroadcastToSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If we do not have an input layout then do not infer an output layout.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  auto broadcast_op = llvm::cast<mlir::TF::BroadcastToOp>(op);
+  TF_ASSIGN_OR_RETURN(
+      const auto broadcasted_output_shape,
+      GetShapeOfValue(broadcast_op.getOutput(), /*fail_on_dynamic=*/true));
+  TF_ASSIGN_OR_RETURN(
+      const auto input_shape,
+      GetShapeOfValue(broadcast_op.getInput(), /*fail_on_dynamic=*/true));
+
+  // Broadcasting works from trailing dimensions and dimensions are broadcasted
+  // in forward direction.
+  const int output_shape_rank = broadcasted_output_shape.size();
+  const int input_shape_rank = input_shape.size();
+  const int broadcasted_dimensions = output_shape_rank - input_shape_rank;
+
+  if (broadcasted_dimensions < 0)
+    return absl::FailedPreconditionError(
+        "Broadcasted dimension was less than 0.");
+
+  Layout input_layout = input_layouts.lookup(0);
+
+  std::vector<std::string> layout_sharding;
+  for (int i = 0; i < output_shape_rank; ++i) {
+    if (i < broadcasted_dimensions) {
+      layout_sharding.push_back(Layout::kUnshardedDim);
+    } else {
+      layout_sharding.push_back(
+          input_layout.sharding_spec(i - broadcasted_dimensions));
+    }
+  }
+  TF_ASSIGN_OR_RETURN(Layout inferred_output_layout,
+                      Layout::GetLayout(layout_sharding, mesh));
+  return llvm::DenseMap<int, Layout>({{0, inferred_output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+BroadcastToSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // If output layout is not set, then we can only infer the `shape` input
+  // which should always be replicated.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>(
+        {{1, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)))}});
+
+  auto output_layout = output_layouts.lookup(0);
+
+  auto broadcast_op = llvm::cast<mlir::TF::BroadcastToOp>(op);
+  TF_ASSIGN_OR_RETURN(
+      const auto broadcasted_output_shape,
+      GetShapeOfValue(broadcast_op.getOutput(), /*fail_on_dynamic=*/true));
+  TF_ASSIGN_OR_RETURN(
+      const auto input_shape,
+      GetShapeOfValue(broadcast_op.getInput(), /*fail_on_dynamic=*/true));
+
+  // Broadcasting works from trailing dimensions and dimensions are broadcasted
+  // in forward direction.
+  const int output_shape_rank = broadcasted_output_shape.size();
+  const int input_shape_rank = input_shape.size();
+  const int broadcasted_dimensions = output_shape_rank - input_shape_rank;
+
+  std::vector<std::string> sharding_specs;
+  for (int i = 0; i < input_shape_rank; ++i) {
+    if (input_shape[i] == 1) {
+      sharding_specs.push_back(Layout::kUnshardedDim);
+    } else {
+      sharding_specs.push_back(
+          output_layout.sharding_spec(i + broadcasted_dimensions));
+    }
+  }
+  TF_ASSIGN_OR_RETURN(Layout inferred_operand_layout,
+                      Layout::GetLayout(sharding_specs, mesh));
+  // `shape` input of BroadcastTo is always set as replicated.
+  return llvm::DenseMap<int, Layout>(
+      {{0, inferred_operand_layout},
+       {1, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)))}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BROADCAST_TO_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BROADCAST_TO_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class BroadcastToSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_BROADCAST_TO_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/concat_spmd_expander.cc

@@ -0,0 +1,169 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/concat_spmd_expander.h"
+
+#include <cstdint>
+#include <optional>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+absl::Status VerifyConcatLayout(mlir::Value concat_dim_operand,
+                                const Layout& concat_layout) {
+  TF_ASSIGN_OR_RETURN(int64_t concat_dim_value,
+                      ExtractConstIntFromValue(concat_dim_operand));
+  for (const auto& shard_and_dimension :
+       llvm::enumerate(concat_layout.num_shards())) {
+    if (shard_and_dimension.index() == concat_dim_value &&
+        shard_and_dimension.value() > 1) {
+      return absl::InvalidArgumentError(
+          "Concat op SPMD with concat dimension in sharded dimension is "
+          "not supported.");
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+StatusOr<Layout> ReduceForConcatOutputLayout(mlir::Value concat_dim_operand,
+                                             const Layout& layout) {
+  TF_ASSIGN_OR_RETURN(int64_t concat_dim_value,
+                      ExtractConstIntFromValue(concat_dim_operand));
+  // Set concatenated dimension to replicated.
+  return layout.GetLayoutWithReducedDims({concat_dim_value},
+                                         /*keep_dims=*/true);
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> ConcatSPMDExpander::ExpandOp(mlir::Operation* op) {
+  if (!llvm::isa<mlir::TF::ConcatOp, mlir::TF::ConcatV2Op>(op))
+    return absl::InvalidArgumentError(
+        "Requested SPMD Expansion for op that is not Concat or ConcatV2.");
+
+  // Extract the concat dim. ConcatOp and ConcatV2Op define the dim at
+  // different position.
+  bool is_concat_v1 = llvm::isa<mlir::TF::ConcatOp>(op);
+  const int concat_dim_operand_idx =
+      is_concat_v1 ? 0 : op->getNumOperands() - 1;
+  mlir::Value concat_dim = op->getOperand(concat_dim_operand_idx);
+
+  // Ensure that Concat op is not sharded on concat-dimension.
+  TF_ASSIGN_OR_RETURN(auto concat_output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  TF_RETURN_IF_ERROR(VerifyConcatLayout(concat_dim, concat_output_layout));
+
+  // Relayout all inputs to match output layout before concating.
+  // This will ensure that a local concat is the correct thing to do.
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    // Skip relayout for the concat dim operand.
+    if (i == concat_dim_operand_idx) continue;
+
+    TF_ASSIGN_OR_RETURN(Layout layout,
+                        ExtractRequiredLayoutFromOperand(op->getOperand(i)));
+    TF_ASSIGN_OR_RETURN(
+        mlir::Value new_input,
+        EmitRelayout(op->getOperand(i), layout, concat_output_layout));
+
+    op->setOperand(i, new_input);
+  }
+
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ConcatSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // Verify operand layouts to ensure that there are no conflicting concat
+  // operand layouts.
+  const bool is_concat_v1 = llvm::isa<mlir::TF::ConcatOp>(op);
+  auto begin_idx = is_concat_v1 ? 1 : 0;
+  auto end_idx =
+      is_concat_v1 ? op->getNumOperands() - 1 : op->getNumOperands() - 2;
+
+  llvm::DenseMap<int, Layout> concat_operands_layouts;
+  for (auto idx = begin_idx; idx <= end_idx; ++idx) {
+    if (input_layouts.find(idx) != input_layouts.end())
+      concat_operands_layouts[idx] = input_layouts.lookup(idx);
+  }
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> concat_operand_layout,
+                      GetMergedOperandLayout(concat_operands_layouts, op));
+  // Concat/ConcatV2 has different operand index for concat dim. Retrieve the
+  // correct concat dim value.
+  const int concat_dim_operand_idx =
+      is_concat_v1 ? 0 : op->getNumOperands() - 1;
+  mlir::Value concat_dim = op->getOperand(concat_dim_operand_idx);
+
+  // If consistent operand layout exists, propagate the operand layout to output
+  // layout.
+  if (concat_operand_layout) {
+    TF_ASSIGN_OR_RETURN(
+        const Layout reduced_concat_layout,
+        ReduceForConcatOutputLayout(concat_dim, *concat_operand_layout));
+    return llvm::DenseMap<int, Layout>({{0, reduced_concat_layout}});
+  } else {
+    return llvm::DenseMap<int, Layout>();
+  }
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ConcatSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If output layout does not exist then do not infer any operand layouts.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  Layout output_layout = output_layouts.lookup(0);
+  const bool is_concat_v1 = llvm::isa<mlir::TF::ConcatOp>(op);
+  // Concat/ConcatV2 has different operand index for concat dim. Retrieve the
+  // correct concat dim value.
+  const int concat_dim_operand_idx =
+      is_concat_v1 ? 0 : op->getNumOperands() - 1;
+  mlir::Value concat_dim = op->getOperand(concat_dim_operand_idx);
+
+  // If suggested output layout exists, verify that concatenated dimension is
+  // replicated to ensure no cross device communication is needed, then
+  // propagate the output layout to all concat tensor operands' layout.
+  // Set concatenated dimension to replicated.
+  TF_ASSIGN_OR_RETURN(const Layout inferred_input_layout,
+                      ReduceForConcatOutputLayout(concat_dim, output_layout));
+
+  llvm::DenseMap<int, Layout> operand_layouts;
+
+  for (auto index = 0; index < op->getNumOperands(); ++index) {
+    if (index == concat_dim_operand_idx) continue;
+    operand_layouts[index] = inferred_input_layout;
+  }
+  return operand_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/concat_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONCAT_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONCAT_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class ConcatSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONCAT_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.cc

@@ -0,0 +1,123 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.h"
+
+#include <cassert>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> WhileRegionSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  assert(op->getNumOperands() == op->getNumResults());
+  // Set the type for the results of the WhileRegion explicitly.
+  //
+  // Normally we would use InferSPMDExpandedLocalShape for this, but that
+  // function requires the op to either have a type inference interface
+  // (which WhileRegion does not) or a TensorFlow ShapeFn (WhileRegion is not
+  // a TensorFlow op). During the standard MLIR shape inference pass this op
+  // is handled by a special case in InferShapeForSingleOperation.
+  for (int i = 0; i < op->getNumOperands(); ++i)
+    op->getResult(i).setType(op->getOperand(i).getType());
+
+  auto while_op = llvm::cast<mlir::TF::WhileRegionOp>(op);
+  for (const auto& data :
+       llvm::enumerate(llvm::zip(while_op.getCond().front().getArguments(),
+                                 while_op.getBody().front().getArguments()))) {
+    const int index = data.index();
+    mlir::BlockArgument cond_arg = std::get<0>(data.value());
+    mlir::BlockArgument body_arg = std::get<1>(data.value());
+    cond_arg.setType(while_op.getOperand(index).getType());
+    body_arg.setType(while_op.getOperand(index).getType());
+  }
+
+  return op;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+WhileRegionSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  return absl::UnimplementedError(
+      "WhileRegion does not support compute layouts. This should not be "
+      "called.");
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+WhileRegionSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  return absl::UnimplementedError(
+      "WhileRegion does not support compute layouts. This should not be "
+      "called.");
+}
+
+StatusOr<mlir::Operation*> IfRegionSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto if_op = llvm::cast<mlir::TF::IfRegionOp>(op);
+  for (mlir::Value result : if_op->getResults()) {
+    auto result_layout_op = llvm::dyn_cast_or_null<mlir::TF::DTensorLayout>(
+        *result.getUsers().begin());
+    if (!result_layout_op)
+      return absl::InvalidArgumentError(
+          "Missing layout of If op result during SPMD expansion.");
+
+    const Layout layout = result_layout_op.getLayout();
+    if (!layout.IsFullyReplicated()) {
+      const auto global_shape = result_layout_op.getGlobalShape();
+      if (!global_shape)
+        return absl::InvalidArgumentError(
+            "Shape of If op must be statically known for SPMD expansion.");
+
+      result.setType(mlir::RankedTensorType::get(
+          layout.LocalShapeFromGlobalShape(*global_shape),
+          mlir::cast<mlir::TensorType>(result.getType()).getElementType()));
+    }
+  }
+  return op;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IfRegionSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // No-op for forward propagation.
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IfRegionSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Layout propagation for TF::IfRegion op is no-op. Actual layout
+  // propagation logic depends on layout propgation of ops inside the
+  // then/else regions of the IfRegion op.
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)))}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.h

@@ -0,0 +1,63 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONTROL_FLOW_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONTROL_FLOW_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class WhileRegionSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class IfRegionSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONTROL_FLOW_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc

@@ -0,0 +1,815 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/conv_spmd_expander.h"
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/dtensor_dialect/ir/dtensor_attributes.h"
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+template <typename ConvOp>
+absl::Status VerifyConvLayout(const Layout& input_layout,
+                              const Layout& filter_layout, ConvOp conv_op) {
+  if (!filter_layout.IsFullyReplicated())
+    return absl::InvalidArgumentError(
+        "Filter for convolution must have fully replicated layout.");
+
+  // Data format "NCHW" or "NCDHW".
+  int channel_dim = 1;
+  if (conv_op.getDataFormat() == "NHWC")
+    channel_dim = 3;
+  else if (conv_op.getDataFormat() == "NDHWC")
+    channel_dim = 4;
+
+  if (input_layout.sharding_spec(channel_dim) != Layout::kUnshardedDim)
+    return absl::InvalidArgumentError(
+        "Conv input's channel dimension must be replicated.");
+
+  if (input_layout.IsBatchParallel() || input_layout.IsFullyReplicated())
+    // No further checks needed for replicated case.
+    return absl::OkStatus();
+
+  if (conv_op.getPadding() == "EXPLICIT")
+    return absl::InvalidArgumentError(
+        "Explicit padding not supported for convolution with spatial "
+        "partitions.");
+
+  const int num_non_default_dilations =
+      llvm::count_if(conv_op.getDilations(), [](mlir::Attribute dilation) {
+        return mlir::cast<mlir::IntegerAttr>(dilation).getInt() != 1;
+      });
+  if (num_non_default_dilations > 0)
+    return absl::InvalidArgumentError(
+        "Only dilation rate 1 is supported for convolution with spatial "
+        "partitions.");
+
+  // TODO(b/208700444): support convolution with strides greater than 1.
+  const int num_non_default_strides =
+      llvm::count_if(conv_op.getStrides(), [](mlir::Attribute stride) {
+        return mlir::cast<mlir::IntegerAttr>(stride).getInt() != 1;
+      });
+  if (num_non_default_strides > 0)
+    return absl::InvalidArgumentError(
+        "Only stride 1 is supported for convolution with spatial partitions.");
+
+  mlir::Value input = conv_op.getInput();
+  auto input_type = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
+  if (!input_type || !input_type.hasStaticShape())
+    return absl::InvalidArgumentError(
+        "Input must have static shapes for convolution with spatial "
+        "partitions.");
+
+  mlir::Value filter = conv_op.getFilter();
+  auto filter_type = mlir::dyn_cast<mlir::RankedTensorType>(filter.getType());
+  if (!filter_type || !filter_type.hasStaticShape())
+    return absl::InvalidArgumentError(
+        "Filter must have static shapes for convolution with spatial "
+        "partitions.");
+
+  llvm::ArrayRef<int64_t> filter_shape = filter_type.getShape();
+  for (auto it = filter_shape.begin(); it != filter_shape.end() - 2; ++it) {
+    if (*it % 2 != 1)
+      return absl::InvalidArgumentError(
+          "Filter dimensions must be odd numbers for convolution with "
+          "spatial partitions.");
+  }
+
+  return absl::OkStatus();
+}
+
+mlir::Value PadInputOnUnshardedDim(mlir::OpBuilder& builder,
+                                   mlir::Location location,
+                                   mlir::Value input_tensor, int curr_input_dim,
+                                   int64_t curr_filter_dim_size) {
+  auto input_tensor_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(input_tensor.getType());
+  auto input_tensor_shape = input_tensor_type.getShape();
+
+  const size_t paddings_flat_length = input_tensor_type.getRank() * 2;
+  llvm::SmallVector<int64_t, 4> paddings_flat_vec(paddings_flat_length, 0);
+  int64_t padding_size = curr_filter_dim_size - 1;
+  paddings_flat_vec[2 * curr_input_dim] = padding_size / 2;
+  paddings_flat_vec[2 * curr_input_dim + 1] = padding_size / 2;
+
+  llvm::SmallVector<int64_t, 4> paddings_shape(input_tensor_shape.begin(),
+                                               input_tensor_shape.end());
+  paddings_shape[curr_input_dim] += padding_size;
+
+  mlir::Value paddings_flat = Int64Const(builder, location, paddings_flat_vec);
+  mlir::RankedTensorType paddings_type = mlir::RankedTensorType::get(
+      paddings_shape, input_tensor_type.getElementType());
+  mlir::Value paddings = mlir::TF::ReshapeOp::create(
+      builder, location, paddings_flat,
+      Int64Const(builder, location, {input_tensor_type.getRank(), 2}));
+  return mlir::TF::PadOp::create(builder, location, paddings_type, input_tensor,
+                                 paddings);
+}
+
+template <typename ConvOp>
+StatusOr<mlir::Operation*> HandleConv(ConvOp conv_op) {
+  mlir::OpBuilder builder(conv_op);
+  TF_ASSIGN_OR_RETURN(const Layout input_layout,
+                      ExtractRequiredLayoutFromOperand(conv_op.getInput()));
+  TF_ASSIGN_OR_RETURN(const Layout filter_layout,
+                      ExtractRequiredLayoutFromOperand(conv_op.getFilter()));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(conv_op));
+
+  TF_RETURN_IF_ERROR(VerifyConvLayout(input_layout, filter_layout, conv_op));
+
+  if (input_layout.IsBatchParallel() || input_layout.IsFullyReplicated())
+    // No special handling needed for replicated case.
+    return InferSPMDExpandedLocalShape(conv_op);
+
+  mlir::tf_device::ClusterOp cluster =
+      conv_op->template getParentOfType<mlir::tf_device::ClusterOp>();
+  TF_ASSIGN_OR_RETURN(mlir::Value mesh_coordinates,
+                      GetMeshCoordinatesFromCluster(cluster));
+  const Mesh& mesh = input_layout.mesh();
+  mlir::Location location = conv_op->getLoc();
+
+  const std::vector<std::string> input_sharding_spec =
+      input_layout.sharding_spec_strs();
+  const std::vector<std::string> output_sharding_spec =
+      output_layout.sharding_spec_strs();
+  llvm::StringRef format = conv_op.getDataFormat();
+  llvm::StringRef padding = conv_op.getPadding();
+
+  const auto input_num_shards = input_layout.num_shards();
+  const auto output_num_shards = output_layout.num_shards();
+
+  auto filter_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(conv_op.getFilter().getType());
+  auto filter_shape = filter_type.getShape();
+
+  int begin_input_dim = -1, end_input_dim = -1;
+  if (format == "NCHW") {
+    begin_input_dim = 2;
+    end_input_dim = 3;
+  } else if (format == "NHWC") {
+    begin_input_dim = 1;
+    end_input_dim = 2;
+  } else if (format == "NCDHW") {
+    begin_input_dim = 2;
+    end_input_dim = 4;
+  } else if (format == "NDHWC") {
+    begin_input_dim = 1;
+    end_input_dim = 3;
+  }
+
+  // For non-batch, non-channel dimension sharding, conduct halo exchange.
+  for (int curr_input_dim = begin_input_dim; curr_input_dim <= end_input_dim;
+       ++curr_input_dim) {
+    int curr_filter_dim = curr_input_dim - begin_input_dim;
+
+    auto input_type =
+        mlir::dyn_cast<mlir::RankedTensorType>(conv_op.getInput().getType());
+    auto input_shape = input_type.getShape();
+
+    if (input_sharding_spec[curr_input_dim] == Layout::kUnshardedDim) {
+      if (padding == "SAME") {
+        // Since we always emit a Conv op with "VALID" padding, we need to
+        // manually pad the input tensor.
+        conv_op->setOperand(
+            0, PadInputOnUnshardedDim(builder, location, conv_op.getInput(),
+                                      curr_input_dim,
+                                      filter_shape[curr_filter_dim]));
+      }
+      // No halo exchange is needed for unsharded dims.
+      continue;
+    }
+
+    TF_ASSIGN_OR_RETURN(const int mesh_dim_index,
+                        mesh.idx_for_dim(input_sharding_spec[curr_input_dim]));
+    TF_ASSIGN_OR_RETURN(mlir::Value scalar_mesh_coordinate,
+                        SelectScalarValueFromArray(builder, mesh_dim_index,
+                                                   location, mesh_coordinates));
+
+    int halo_size;
+    if (padding == "SAME") {
+      halo_size = std::floor(filter_shape[curr_filter_dim] / 2);
+    } else if (padding == "VALID") {
+      int input_local_size = input_shape[curr_input_dim];
+      int input_size = input_local_size * input_num_shards[curr_input_dim];
+      int output_size = input_size - (filter_shape[curr_filter_dim] - 1);
+      int output_local_size = output_size / output_num_shards[curr_input_dim];
+      halo_size = output_local_size + (filter_shape[curr_filter_dim] - 1) -
+                  input_local_size;
+    } else {
+      return absl::UnimplementedError(
+          absl::StrCat("Spatially partitioned convolution with padding \"",
+                       padding.str(), "\" is not supported."));
+    }
+
+    if (halo_size == 0)
+      // No exchange is needed for empty halos.
+      continue;
+
+    builder.setInsertionPoint(conv_op);
+    TF_ASSIGN_OR_RETURN(
+        mlir::Value halo_exchanged_input,
+        EmitHaloExchange(builder, halo_size,
+                         input_sharding_spec[curr_input_dim], input_layout,
+                         mesh_coordinates, cluster, location,
+                         conv_op.getInput()));
+
+    if (padding == "SAME") {
+      conv_op->setOperand(0, halo_exchanged_input);
+    } else if (padding == "VALID") {
+      // Slice the halo exchanged tensor to the desired size based on the index
+      // of the shard on the current dimension.
+
+      llvm::SmallVector<int32_t, 4> halo_sizes(input_layout.rank(), 0);
+      halo_sizes[curr_input_dim] = halo_size;
+      mlir::Value halo_sizes_const = IntConst(builder, location, halo_sizes);
+
+      llvm::SmallVector<int32_t, 4> halo_increments(input_layout.rank(), 0);
+      halo_increments[curr_input_dim] =
+          halo_size / (input_num_shards[curr_input_dim] - 1);
+      mlir::Value halo_increments_const =
+          IntConst(builder, location, halo_increments);
+
+      mlir::Value offset = mlir::TF::MulOp::create(
+          builder, location, halo_increments_const.getType(),
+          scalar_mesh_coordinate, halo_increments_const);
+      mlir::Value slice_begin =
+          mlir::TF::SubOp::create(builder, location, halo_sizes_const, offset);
+
+      llvm::SmallVector<int64_t, 4> slice_size(input_shape.begin(),
+                                               input_shape.end());
+      slice_size[curr_input_dim] += halo_size;
+      mlir::Value slice_size_const = Int64Const(builder, location, slice_size);
+      // slice_size_const and slize_begin_int64 has to be same type.
+      mlir::Value slice_begin_int64 = mlir::TF::CastOp::create(
+          builder, location,
+          mlir::RankedTensorType::get({input_layout.rank()},
+                                      builder.getI64Type()),
+          slice_begin);
+
+      mlir::RankedTensorType sliced_input_type =
+          mlir::RankedTensorType::get(slice_size, input_type.getElementType());
+      mlir::Value sliced_input = mlir::TF::SliceOp::create(
+          builder, location, sliced_input_type, /*input=*/halo_exchanged_input,
+          /*begin=*/slice_begin_int64, /*size=*/slice_size_const);
+      conv_op->setOperand(0, sliced_input);
+    }
+
+    // Spatially partitioned convolution always uses VALID padding after halo
+    // exchange.
+    conv_op.setPaddingAttr(builder.getStringAttr("VALID"));
+  }
+
+  return InferSPMDExpandedLocalShape(conv_op);
+}
+
+template <typename ConvBackpropInputOp>
+StatusOr<mlir::Operation*> HandleConvBackpropInput(
+    const Layout& output_layout, ConvBackpropInputOp conv_op) {
+  TF_ASSIGN_OR_RETURN(std::vector<Layout> input_layouts,
+                      ExtractRequiredLayoutFromOperands(conv_op));
+
+  const Layout& input_shape_layout = input_layouts[0];
+  const Layout& filter_layout = input_layouts[1];
+  const Layout& grad_layout = input_layouts[2];
+
+  // We only support batch sharding for these. In this case the output and input
+  // gradient must both be batch sharded. The filter input must be replicated.
+  if (!(output_layout.IsBatchParallel() || output_layout.IsFullyReplicated()) ||
+      !(grad_layout.IsBatchParallel() || grad_layout.IsFullyReplicated())) {
+    return errors::InvalidArgument("{0} only supports batch parallel layouts.",
+                                   conv_op->getName().getStringRef().str());
+  }
+  if (!filter_layout.IsFullyReplicated()) {
+    return errors::InvalidArgument("{0} only supports replicated filters.",
+                                   conv_op->getName().getStringRef().str());
+  }
+  if (!input_shape_layout.IsFullyReplicated()) {
+    return errors::InvalidArgument(
+        "Layout of the input shape (parameter 0) of {0} must be replicated.",
+        conv_op->getName().getStringRef().str());
+  }
+
+  llvm::SmallVector<int64_t, 4> global_shape;
+  absl::Status extract_status =
+      ExtractConstVectorFromValue(conv_op.getInputSizes(), &global_shape);
+
+  // If the input is dynamic size, we expect the output is all so dynamic size
+  // since they should roughly be the same shape. Don't support this for right
+  // now. The easy way to support this is to all gather the gradient input and
+  // compute this as a large local convolution and then slice to the output
+  // layout.
+  if (!extract_status.ok()) {
+    return errors::InvalidArgument("{0} requires static shape for input size.",
+                                   conv_op->getName().getStringRef().str());
+  }
+
+  // Compute the 'true' input/output layout of the operation. E.g. batch sharded
+  // vs non-batch sharded. If at least one of the the input gradient or output
+  // gradient is batch sharded, use that dimension.
+  std::string batch_sharding_dimension = grad_layout.sharding_spec(0);
+  if (batch_sharding_dimension == Layout::kUnshardedDim) {
+    batch_sharding_dimension = output_layout.sharding_spec(0);
+  } else if ((output_layout.sharding_spec(0) != Layout::kUnshardedDim) &&
+             (batch_sharding_dimension != output_layout.sharding_spec(0))) {
+    return errors::InvalidArgument(
+        "Input and output layout to {2} have incompatible sharding dimensions: "
+        "\"{0}\" and \"{1}\".",
+        grad_layout.sharding_spec(0), output_layout.sharding_spec(0),
+        conv_op->getName().getStringRef().str());
+  }
+
+  const Layout desired_input_gradient_layout =
+      Layout::BatchShardedLike(grad_layout, batch_sharding_dimension);
+  const Layout desired_output_gradient_layout =
+      Layout::BatchShardedLike(output_layout, batch_sharding_dimension);
+
+  const Layout desired_input_layout = Layout::BatchShardedOnMesh(
+      grad_layout.mesh(), global_shape.size(), batch_sharding_dimension);
+  const std::vector<int64_t> local_shape =
+      desired_input_layout.LocalShapeFromGlobalShape(global_shape);
+
+  mlir::OpBuilder builder(conv_op.getOperation());
+  mlir::Value new_const = IntConst(
+      builder, conv_op->getLoc(),
+      llvm::SmallVector<int32_t, 4>(local_shape.begin(), local_shape.end()));
+  conv_op.getInputSizesMutable().assign(new_const);
+
+  TF_ASSIGN_OR_RETURN(mlir::Value local_input_gradient,
+                      EmitRelayout(conv_op.getOutBackprop(), grad_layout,
+                                   desired_input_gradient_layout));
+  conv_op.getOutBackpropMutable().assign(local_input_gradient);
+
+  InferSPMDExpandedLocalShape(conv_op);
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value local_output_gradient,
+      EmitRelayout(conv_op.getOutput(), desired_output_gradient_layout,
+                   output_layout, &newly_created_ops));
+  conv_op.getOutput().replaceAllUsesExcept(local_output_gradient,
+                                           newly_created_ops);
+  return local_output_gradient.getDefiningOp();
+}
+
+// This expands backprop ops which take tensor inputs into those which take
+// sizes. We first convert the (currently local) input shape to global and use
+// the const rather than input. The shape will be converted back to local in
+// HandleConvBackpropInput, but this is the correct behavior as
+// HandleConvBackpropInput will decided how it wants to expand the op.
+template <typename To, typename From>
+StatusOr<mlir::Operation*> HandleConvBackpropInputTensor(
+    const Layout& output_layout, From conv_op) {
+  TF_ASSIGN_OR_RETURN(llvm::SmallVector<int64_t> local_shape,
+                      GetTFShapeFromType(conv_op.getInput().getType()));
+
+  TF_ASSIGN_OR_RETURN(const Layout input_layout,
+                      ExtractRequiredLayoutFromOperand(conv_op.getInput()));
+
+  const std::vector<int64_t> global_shape =
+      input_layout.GlobalShapeFromLocalShape(local_shape);
+
+  mlir::OpBuilder builder(conv_op);
+  mlir::Value global_input_shape = IntConst(
+      builder, conv_op->getLoc(),
+      llvm::SmallVector<int32_t, 4>(global_shape.begin(), global_shape.end()));
+
+  // Insert a replicated layout along this edge, so that we can call
+  // HandleConvBackpropInput which expects there to be a layout here.
+  mlir::TF::ShapeAttr global_input_shape_shape = mlir::TF::ShapeAttr::get(
+      builder.getContext(),
+      mlir::cast<mlir::TensorType>(global_input_shape.getType()));
+  mlir::TF::DTensorLayout global_input_shape_with_layout =
+      mlir::TF::DTensorLayout::create(
+          builder, conv_op->getLoc(), global_input_shape,
+          mlir::dtensor::LayoutAttr::get(
+              builder.getContext(),
+              Layout::ReplicatedOnMesh(input_layout.mesh(), 1)),
+          global_input_shape_shape);
+
+  To new_conv = To::create(
+      builder, conv_op->getLoc(), conv_op->getResultTypes(),
+      mlir::ValueRange({global_input_shape_with_layout, conv_op.getFilter(),
+                        conv_op.getOutBackprop()}),
+      conv_op->getAttrs());
+
+  conv_op.getOutput().replaceAllUsesWith(new_conv.getOutput());
+  conv_op.erase();
+
+  return HandleConvBackpropInput(output_layout, new_conv);
+}
+
+template <typename ConvBackpropFilterOp>
+StatusOr<mlir::Operation*> HandleConvBackpropFilter(
+    const Layout& output_layout, ConvBackpropFilterOp conv_op) {
+  TF_ASSIGN_OR_RETURN(std::vector<Layout> input_layouts,
+                      ExtractRequiredLayoutFromOperands(conv_op));
+
+  const Layout& input_layout = input_layouts[0];
+  const Layout& filter_shape_layout = input_layouts[1];
+  const Layout& grad_layout = input_layouts[2];
+
+  // We only support batch sharding for these. In this case the input
+  // activations and input gradient should both be batch sharded and
+  // the output (the filter gradient) should be replicated.
+  if (!(output_layout.IsBatchParallel() || output_layout.IsFullyReplicated()) ||
+      !(grad_layout.IsBatchParallel() || grad_layout.IsFullyReplicated())) {
+    return errors::InvalidArgument("{0} only supports batch parallel layouts.",
+                                   conv_op->getName().getStringRef().str());
+  }
+  if (!output_layout.IsFullyReplicated()) {
+    return errors::InvalidArgument("{0} only supports replicated filters.",
+                                   conv_op->getName().getStringRef().str());
+  }
+  if (!filter_shape_layout.IsFullyReplicated()) {
+    return errors::InvalidArgument(
+        "Filter shape input (parameter 1) for {0} must have replicated layout.",
+        conv_op->getName().getStringRef().str());
+  }
+
+  // Compute the 'true' input layouts of the operation. E.g. batch sharded
+  // vs non-batch sharded. Basically we get the batch sharding dimension from
+  // one of the inputs and check that the other is potentially sharded on the
+  // same dimension.
+  // TODO(b/262417847): if batch_sharding_dimension is Layout::kUnsharded, then
+  // we should consider sharding the input here. It may be faster to spread
+  // the convolution out and then all reduce after vs running it all locally.
+  std::string batch_sharding_dimension = input_layout.sharding_spec(0);
+  if (batch_sharding_dimension == Layout::kUnshardedDim) {
+    batch_sharding_dimension = grad_layout.sharding_spec(0);
+  } else if ((grad_layout.sharding_spec(0) != Layout::kUnshardedDim) &&
+             (batch_sharding_dimension != grad_layout.sharding_spec(0))) {
+    return errors::InvalidArgument(
+        "Input and gradient layouts for {2} have incompatible batch sharding "
+        "dimensions: \"{0}\" and \"{1}\".",
+        input_layouts[0].sharding_spec(0), input_layouts[0].sharding_spec(0),
+        conv_op->getName().getStringRef().str());
+  }
+
+  const Layout desired_input_activation_layout =
+      Layout::BatchShardedLike(input_layout, batch_sharding_dimension);
+  const Layout desired_input_gradient_layout =
+      Layout::BatchShardedLike(grad_layout, batch_sharding_dimension);
+
+  TF_ASSIGN_OR_RETURN(mlir::Value local_input_activation,
+                      EmitRelayout(conv_op.getInput(), input_layout,
+                                   desired_input_activation_layout));
+  conv_op.getInputMutable().assign(local_input_activation);
+
+  TF_ASSIGN_OR_RETURN(mlir::Value local_input_gradient,
+                      EmitRelayout(conv_op.getOutBackprop(), grad_layout,
+                                   desired_input_gradient_layout));
+  conv_op.getOutBackpropMutable().assign(local_input_gradient);
+
+  InferSPMDExpandedLocalShape(conv_op);
+
+  // Output shall be replicated. If we were batch sharded, we need to
+  // all-reduce the partial results.
+
+  if (batch_sharding_dimension != Layout::kUnshardedDim) {
+    mlir::OpBuilder builder(conv_op.getOperation());
+    builder.setInsertionPointAfter(conv_op);
+    return DT_CTX(EmitAllReduce(builder, output_layout,
+                                {batch_sharding_dimension}, conv_op,
+                                kReduceOpAdd));
+  }
+
+  return conv_op.getOperation();
+}
+
+// This expands backprop ops which take tensor inputs into those which take
+// sizes. We check that the filter input shape is global and then make that a
+// const, and replace the op with the version taking shapes.
+template <typename To, typename From>
+StatusOr<mlir::Operation*> HandleConvBackpropFilterTensor(
+    const Layout& output_layout, From conv_op) {
+  TF_ASSIGN_OR_RETURN(const Layout filter_layout,
+                      ExtractRequiredLayoutFromOperand(conv_op.getFilter()));
+
+  if (!filter_layout.IsFullyReplicated()) {
+    return absl::InvalidArgumentError(
+        "Convolution backpropation ops only support replicated filters.");
+  }
+
+  TF_ASSIGN_OR_RETURN(llvm::SmallVector<int64_t> global_filter_shape,
+                      GetTFShapeFromType(conv_op.getFilter().getType()));
+
+  mlir::OpBuilder builder(conv_op);
+  mlir::Value global_filter_shape_const =
+      IntConst(builder, conv_op->getLoc(),
+               llvm::SmallVector<int32_t, 4>(global_filter_shape.begin(),
+                                             global_filter_shape.end()));
+
+  // Insert a replicated layout along this edge, so that we can call
+  // HandleConvBackpropInput which expects there to be a layout here.
+  mlir::TF::ShapeAttr global_filter_shape_shape = mlir::TF::ShapeAttr::get(
+      builder.getContext(),
+      mlir::cast<mlir::TensorType>(global_filter_shape_const.getType()));
+  mlir::TF::DTensorLayout global_filter_shape_with_layout =
+      mlir::TF::DTensorLayout::create(
+          builder, conv_op->getLoc(), global_filter_shape_const,
+          mlir::dtensor::LayoutAttr::get(
+              builder.getContext(),
+              Layout::ReplicatedOnMesh(filter_layout.mesh(), 1)),
+          global_filter_shape_shape);
+
+  To new_conv = To::create(
+      builder, conv_op->getLoc(), conv_op->getResultTypes(),
+      mlir::ValueRange({conv_op.getInput(), global_filter_shape_with_layout,
+                        conv_op.getOutBackprop()}),
+      conv_op->getAttrs());
+
+  conv_op.getOutput().replaceAllUsesWith(new_conv.getOutput());
+  conv_op.erase();
+
+  return HandleConvBackpropFilter(output_layout, new_conv);
+}
+
+StatusOr<mlir::Operation*> HandleMaxPoolGradOp(
+    const Layout& output_layout, mlir::TF::MaxPoolGradOp max_pool_grad_op) {
+  // MaxPoolGrad has 3 inputs: Original Input to MaxPool, Output of MaxPool and
+  // Gradients.
+  assert(max_pool_grad_op->getOpOperands().size() == 3);
+
+  // Relayout gradient input to match layout of output of maxpool.
+  mlir::OpOperand& max_pool_output = max_pool_grad_op->getOpOperand(1);
+  TF_ASSIGN_OR_RETURN(Layout max_pool_output_layout,
+                      ExtractRequiredLayoutFromOperand(max_pool_output.get()));
+
+  mlir::OpOperand& grad_input = max_pool_grad_op->getOpOperand(2);
+  TF_ASSIGN_OR_RETURN(Layout grad_input_layout,
+                      ExtractRequiredLayoutFromOperand(grad_input.get()));
+  TF_ASSIGN_OR_RETURN(mlir::Value new_grad_input,
+                      EmitRelayout(grad_input.get(), grad_input_layout,
+                                   max_pool_output_layout));
+  grad_input.set(new_grad_input);
+
+  return InferSPMDExpandedLocalShape(max_pool_grad_op);
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> ConvSPMDExpander::ExpandOp(mlir::Operation* op) {
+  // The first argument to Conv2DBackpropInputOp is the shape of the input we
+  // are generating. Since this is almost always the output of a call to
+  // `shape`, we lose the ability to infer the original input layout. (c.f if
+  // Conv2DBackpropInput accepted the input _tensor_ instead of the shape).
+  // Since in eager execution, we cannot look ahead at consumer operations, we
+  // instead attach the original input layout as a secondary attribute on the
+  // output of the shape operation, and use this to infer the desired layout for
+  // this op.
+
+  TF_ASSIGN_OR_RETURN(const auto output_layout, ExtractSingleLayoutFromOp(op));
+
+  // Forward prop ops.
+  if (llvm::isa<mlir::TF::Conv2DOp>(op))
+    return HandleConv<>(llvm::cast<mlir::TF::Conv2DOp>(op));
+  if (llvm::isa<mlir::TF::Conv3DOp>(op))
+    return HandleConv<>(llvm::cast<mlir::TF::Conv3DOp>(op));
+
+  // Backward prop input ops.
+  if (llvm::isa<mlir::TF::Conv2DBackpropInputOp>(op))
+    return HandleConvBackpropInput<>(
+        *output_layout, llvm::cast<mlir::TF::Conv2DBackpropInputOp>(op));
+  if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv2DBackpropInputV2Op>(op))
+    return HandleConvBackpropInputTensor<mlir::TF::Conv2DBackpropInputOp>(
+        *output_layout, conv_op);
+  if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv3DBackpropInputOp>(op))
+    return HandleConvBackpropInputTensor<mlir::TF::Conv3DBackpropInputV2Op>(
+        *output_layout, conv_op);
+  if (llvm::isa<mlir::TF::Conv3DBackpropInputV2Op>(op))
+    return HandleConvBackpropInput<>(
+        *output_layout, llvm::cast<mlir::TF::Conv3DBackpropInputV2Op>(op));
+
+  // Backward prop filter ops.
+  if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp>(op))
+    return HandleConvBackpropFilter<>(
+        *output_layout, llvm::cast<mlir::TF::Conv2DBackpropFilterOp>(op));
+  if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv2DBackpropFilterV2Op>(op))
+    return HandleConvBackpropFilterTensor<mlir::TF::Conv2DBackpropFilterOp>(
+        *output_layout, conv_op);
+  if (auto conv_op = llvm::dyn_cast<mlir::TF::Conv3DBackpropFilterOp>(op))
+    return HandleConvBackpropFilterTensor<mlir::TF::Conv3DBackpropFilterV2Op>(
+        *output_layout, conv_op);
+  if (llvm::isa<mlir::TF::Conv3DBackpropFilterV2Op>(op))
+    return HandleConvBackpropFilter<>(
+        *output_layout, llvm::cast<mlir::TF::Conv3DBackpropFilterV2Op>(op));
+
+  // For all other ops, only batch sharded or fully replicated sharding is
+  // supported for now.
+  if (!output_layout->IsFullyReplicated() && !output_layout->IsBatchParallel())
+    return absl::UnimplementedError(
+        llvm::formatv(
+            "Only replicated or batch parallel layout is supported in "
+            "expansion of {0}, but got output layout: {1}",
+            op->getName().getStringRef().str(), output_layout->ToString())
+            .str());
+
+  if (auto max_pool_grad = mlir::dyn_cast<mlir::TF::MaxPoolGradOp>(op))
+    return HandleMaxPoolGradOp(*output_layout, max_pool_grad);
+
+  // Local expansion only for all other ops.
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ConvSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
+
+  if (llvm::isa<mlir::TF::Conv2DOp, mlir::TF::Conv3DOp, mlir::TF::MaxPoolOp,
+                mlir::TF::MaxPoolGradOp>(op)) {
+    // Conv2d/Conv3d and MaxPool ops are grouped together as they all try to
+    // propagate layout from input image (operand 0).
+
+    // If requested 'input' layout exist, try to request same layout for output.
+    if (input_layouts.find(0) != input_layouts.end()) {
+      output_layouts[0] = input_layouts.lookup(0);
+    } else {
+      // For MaxPoolGrad, request same layout as 'orig_output' or 'grad'
+      // whatever is present.
+      if (llvm::isa<mlir::TF::MaxPoolGradOp>(op)) {
+        if (input_layouts.find(1) != input_layouts.end())
+          output_layouts[0] = input_layouts.lookup(1);
+        else if (input_layouts.find(2) != input_layouts.end())
+          output_layouts[0] = input_layouts.lookup(2);
+      }
+    }
+  } else if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
+                       mlir::TF::Conv2DBackpropInputV2Op,
+                       mlir::TF::Conv3DBackpropInputOp,
+                       mlir::TF::Conv3DBackpropInputV2Op>(op)) {
+    if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
+                  mlir::TF::Conv3DBackpropInputV2Op>(op)) {
+      if (input_layouts.find(2) != input_layouts.end()) {
+        // The propagate the gradient layout to the new gradient, e.g. respect
+        // the spatial partitioning of the input gradient.
+        output_layouts[0] = input_layouts.lookup(2);
+      }
+    } else {
+      if (input_layouts.find(0) != input_layouts.end()) {
+        // The propagate the gradient layout to the new gradient, e.g. respect
+        // the spatial partitioning of the input.
+        output_layouts[0] = input_layouts.lookup(0);
+      }
+    }
+  } else if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
+                       mlir::TF::Conv2DBackpropFilterV2Op,
+                       mlir::TF::Conv3DBackpropFilterOp,
+                       mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
+    if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
+                  mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
+      // For the ops which take filter shape as input, just return a replicated
+      // output shape.
+      output_layouts[0] =
+          Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOpResult(0)));
+    } else if (input_layouts.find(1) != input_layouts.end()) {
+      // For the ops taking a real filter, just copy the filter layout.
+      output_layouts[0] = input_layouts.lookup(1);
+    }
+  } else {
+    return absl::InvalidArgumentError(
+        llvm::formatv(
+            "Layout propagation for unrecognized convolution op {0} not "
+            "supported.",
+            OpName(op))
+            .str());
+  }
+
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ConvSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands());
+
+  if (llvm::isa<mlir::TF::Conv2DOp, mlir::TF::Conv3DOp, mlir::TF::MaxPoolOp,
+                mlir::TF::MaxPoolGradOp>(op)) {
+    // If suggested output layout exists, try to request input image to have the
+    // same layout so that all computation would be local.
+    if (output_layouts.find(0) != output_layouts.end()) {
+      const Layout output_layout = output_layouts.lookup(0);
+
+      input_layouts[0] = output_layout;
+
+      // Request replicated for filter input if Conv2D/Conv3D.
+      if (llvm::isa<mlir::TF::Conv2DOp, mlir::TF::Conv3DOp>(op)) {
+        input_layouts[1] =
+            Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
+      }
+      if (llvm::isa<mlir::TF::MaxPoolGradOp>(op)) {
+        input_layouts[1] = output_layout;  // 'orig_output'
+        input_layouts[2] = output_layout;  // 'grad'
+      }
+    }
+  } else if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
+                       mlir::TF::Conv2DBackpropInputV2Op,
+                       mlir::TF::Conv3DBackpropInputOp,
+                       mlir::TF::Conv3DBackpropInputV2Op>(op)) {
+    // Generally mark the filter as replicated.
+    input_layouts[1] =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
+    if (llvm::isa<mlir::TF::Conv2DBackpropInputOp,
+                  mlir::TF::Conv3DBackpropInputV2Op>(op)) {
+      // This input is a shape.
+      input_layouts[0] =
+          Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)));
+    }
+
+    if (output_layouts.find(0) != output_layouts.end()) {
+      Layout output_layout = output_layouts.lookup(0);
+      // Ask for the grad to have the same layout as the output. The reasoning
+      // here is that the if the output is spatially partitioned, we expect
+      // that the grad is spatially partitioned as well.
+      input_layouts[2] = output_layout;
+      if (llvm::isa<mlir::TF::Conv2DBackpropInputV2Op,
+                    mlir::TF::Conv3DBackpropInputOp>(op)) {
+        input_layouts[0] = output_layout;
+      }
+    }
+  } else if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
+                       mlir::TF::Conv2DBackpropFilterV2Op,
+                       mlir::TF::Conv3DBackpropFilterOp,
+                       mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
+    // Note for Filter op, we generally expect that the output layout would
+    // match the variable layout for the filter which is generally replicated.
+    // The gradient layout most likely needs to agree with the input layout,
+    // e.g. both spatially partitioned or not. This is somewhat similar to
+    // MatMul, for now just set both to replicated.
+
+    input_layouts[0] =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)));
+    input_layouts[2] =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(2)));
+
+    if (llvm::isa<mlir::TF::Conv2DBackpropFilterOp,
+                  mlir::TF::Conv3DBackpropFilterV2Op>(op)) {
+      // For ops taking filter shape as input, just use a replicated input
+      // layout.
+      input_layouts[1] =
+          Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
+    } else if (output_layouts.find(0) != output_layouts.end()) {
+      // For ops taking filer directly as input copy the output layout to the
+      // filter layout.
+      input_layouts[1] = output_layouts.lookup(0);
+    }
+  } else {
+    return absl::InvalidArgumentError(
+        llvm::formatv(
+            "Layout propagation for unrecognized convolution op {0} not "
+            "supported.",
+            OpName(op))
+            .str());
+  }
+
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/conv_spmd_expander.h

@@ -0,0 +1,51 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONV_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONV_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Implement Layout propagation and SPMD expansion for Convolution ops.
+//
+// The extended class will be registered in spmd_expander.cc for Conv2D/3D and
+// Conv2D/3D Backprop ops to enable proper DTensor behavior of them. This
+// implementation is internal and specific to DTensor while upstream(python)
+// users won't need to use this class directly in any fashion.
+class ConvSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CONV_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc

@@ -0,0 +1,139 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
+
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+// Extract `axis` tensor from Cumsum op and return it's positive value, since
+// it can be a negative index.
+StatusOr<int64_t> GetAxisDimension(mlir::Operation* op) {
+  auto cumsum = llvm::dyn_cast<mlir::TF::CumsumOp>(op);
+  if (cumsum == nullptr) {
+    return absl::InternalError(
+        absl::StrCat("Expected Cumsum op but got : ", OpName(op)));
+  }
+  TF_ASSIGN_OR_RETURN(int64_t axis_dim,
+                      ExtractConstIntFromValue(cumsum.getAxis()));
+  int64_t tensor_rank = ValueRank(cumsum.getX());
+  // Axis can be in range [-tensor_rank, tensor_rank), so we add tensor_rank
+  // to wrap it around.
+  if (axis_dim >= -tensor_rank && axis_dim < 0) {
+    axis_dim += tensor_rank;
+  } else if (axis_dim < -tensor_rank || axis_dim >= tensor_rank) {
+    return absl::InvalidArgumentError(
+        "Invalid axis; expected a value in [-tensor_rank, tensor_rank)");
+  }
+  return axis_dim;
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> CumsumSPMDExpander::ExpandOp(mlir::Operation* op) {
+  StatusOr<int64_t> axis_dim = GetAxisDimension(op);
+  if (!axis_dim.ok()) return axis_dim.status();
+
+  TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
+  assert(output_layout);
+
+  // Our intermediate computation layout is the output layout with
+  // the axis dimension replicated. So set both the operand and output layout
+  // to this intermediate layout.
+  TF_ASSIGN_OR_RETURN(Layout intermediate_layout,
+                      output_layout->GetLayoutWithReducedDims(
+                          {axis_dim.value()}, /*keep_dims=*/true));
+
+  // Relayout operand to intermediate layout.
+  mlir::OpBuilder builder(op);
+  const auto operand = op->getOperand(0);
+  TF_ASSIGN_OR_RETURN(auto operand_layout, ExtractLayoutFromOperand(operand));
+  if (!operand_layout)
+    return absl::InvalidArgumentError(
+        "input layout of Cumsum op must be known before SPMD "
+        "expansion.");
+
+  TF_ASSIGN_OR_RETURN(
+      const auto new_operand,
+      EmitRelayout(operand, operand_layout.value(), intermediate_layout));
+  op->setOperand(0, new_operand);
+
+  op = InferSPMDExpandedLocalShape(op);
+
+  // Relayout output to intermediate layout.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  builder.setInsertionPointAfter(op);
+  TF_ASSIGN_OR_RETURN(auto final_output,
+                      EmitRelayout(op->getOpResult(0), intermediate_layout,
+                                   output_layout.value(), &newly_created_ops));
+  op->getOpResult(0).replaceAllUsesExcept(final_output, newly_created_ops);
+  return final_output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> CumsumSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  TF_ASSIGN_OR_RETURN(int64_t axis_dim, GetAxisDimension(op));
+
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  auto input_layout = input_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(
+      Layout input_layout_reduced_dims,
+      input_layout.GetLayoutWithReducedDims({axis_dim},
+                                            /*keep_dims=*/true));
+  return llvm::DenseMap<int, Layout>({{0, input_layout_reduced_dims}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> CumsumSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  TF_ASSIGN_OR_RETURN(int64_t axis_dim, GetAxisDimension(op));
+
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+  auto output_layout = output_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(
+      Layout output_layout_reduced_dims,
+      output_layout.GetLayoutWithReducedDims({axis_dim},
+                                             /*keep_dims=*/true));
+  return llvm::DenseMap<int, Layout>({{0, output_layout_reduced_dims}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CUMSUM_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CUMSUM_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class CumsumSPMDExpander : public SPMDExpanderBase {
+ private:
+  // SPMD expand the op for local computation.
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_CUMSUM_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.cc

@@ -0,0 +1,365 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+// Check all tensors are batch parallel
+bool AllBatchParallel(const std::vector<Layout>& layouts,
+                      const llvm::DenseMap<int, int>& batchable_indices) {
+  for (int i = 0; i < layouts.size(); ++i) {
+    if (!layouts[i].IsBatchParallel(batchable_indices.lookup(i))) return false;
+  }
+  return true;
+}
+// Check all Layouts have the same batch rank
+bool SameBatchRank(const std::vector<Layout>& layouts,
+                   const llvm::DenseMap<int, int>& batchable_indices) {
+  absl::flat_hash_set<int> batch_ranks;
+  // Add all batch ranks of layouts
+  for (auto const& idx_and_non_batch_rank : batchable_indices) {
+    auto const& idx = idx_and_non_batch_rank.first;
+    auto const& non_batch_rank = idx_and_non_batch_rank.second;
+    batch_ranks.insert(layouts[idx].rank() - non_batch_rank);
+  }
+  return batch_ranks.size() <= 1;
+}
+
+// Check if any layout from a set of indices is not a nullopt
+bool AnyLayoutExist(const llvm::DenseMap<int, Layout>& layouts,
+                    const llvm::DenseMap<int, int>& indices) {
+  for (auto const& idx_and_unused : indices) {
+    auto const& idx = idx_and_unused.first;
+    if (layouts.find(idx) != layouts.end()) return true;
+  }
+  return false;
+}
+// Given layouts to merge and a map of {indices of the batchable layouts, rank
+// of non batch dimensions} merges those batchable layouts to produce one single
+// layout. Assumes all batch ranks are the same for all batchable layouts, which
+// is enforced before this
+//
+// Merged together so that we have the layout which is sharded in a tensor dim
+// if and only if all layouts are sharded in the same sharding_spec.
+StatusOr<Layout> MergeBatchLayouts(
+    const llvm::DenseMap<int, Layout>& layouts,
+    const llvm::DenseMap<int, int>& batchable_args, const Mesh& mesh) {
+  // Get the batch rank
+  int layout_idx = -1;
+  for (auto const& idx_and_unused : batchable_args) {
+    auto const& idx = idx_and_unused.first;
+    if (layouts.find(idx) != layouts.end()) layout_idx = idx;
+  }
+
+  int batch_rank =
+      layouts.lookup(layout_idx).rank() - batchable_args.lookup(layout_idx);
+  // Initialize with replicated
+  std::vector<std::string> merged_specs(batch_rank, Layout::kUnshardedDim);
+
+  // Merge layouts. If any dimension don't agree on sharding dim, then replicate
+  for (int i = 0; i < batch_rank; ++i) {
+    absl::flat_hash_set<std::string> spec_set;
+    for (auto const& arg_idx_and_unused : batchable_args) {
+      auto const& arg_idx = arg_idx_and_unused.first;
+      if (layouts.find(arg_idx) == layouts.end()) continue;
+      const std::string spec = layouts.lookup(arg_idx).sharding_spec(i);
+      if (spec != Layout::kUnshardedDim) {
+        spec_set.insert(spec);
+      }
+    }
+    if (spec_set.size() == 1) {
+      merged_specs[i] = *spec_set.begin();
+    } else {
+      merged_specs[i] = Layout::kUnshardedDim;
+    }
+  }
+
+  // Deduplicate same usage of mesh dims. [x,x] -> [unsharded, unsharded]
+  absl::flat_hash_map<std::string, int> counter;
+  for (const std::string& spec : merged_specs) counter[spec] += 1;
+  for (std::string& spec : merged_specs) {
+    if (counter[spec] > 1) {
+      spec = Layout::kUnshardedDim;
+    }
+  }
+  return Layout::GetLayout(merged_specs, mesh);
+}
+
+// Choose an intermediate layout to relayout. Picks the most frequently
+// sharded mesh dimension for every batch dimension, then deduplicates (n-1)
+// of all repeated mesh dimensions, leaving the rightmost duplicate sharded
+//
+// Note that this assumes the number of batch dims for every batchable
+// tensor is the same and is enforced before this
+//
+// Examples:
+// Given layouts: [x,y],[x,z],[y,z], produces [x,z]
+// Deduplication: [x,x] -> will become [*, x]
+StatusOr<Layout> IntermediateBatchLayout(
+    const std::vector<Layout>& operand_layouts,
+    const llvm::DenseMap<int, int>& batchable_operands,
+    const std::vector<Layout>& output_layouts,
+    const llvm::DenseMap<int, int>& batchable_outputs, const Mesh& mesh) {
+  if (batchable_operands.empty()) {
+    return absl::UnimplementedError(
+        llvm::formatv("There must be at least one batchable operand").str());
+  }
+  int first_batcharg_index = batchable_outputs.begin()->first;
+  int batch_rank = operand_layouts[first_batcharg_index].rank() -
+                   batchable_operands.find(first_batcharg_index)->second;
+
+  std::vector<std::string> batch_specs(batch_rank, Layout::kUnshardedDim);
+
+  // For each batch dimension, finds the most commonly used mesh dimension
+  // and sets that to batch_specs[i].
+  for (int i = 0; i < batch_rank; ++i) {
+    std::string mesh_dim = Layout::kUnshardedDim;
+    int max_count = 0;
+    absl::flat_hash_map<std::string, int> counter;
+    // add operand counts
+    for (auto const& idx_and_unused : batchable_operands) {
+      auto const& idx = idx_and_unused.first;
+      std::string spec = operand_layouts[idx].sharding_spec(i);
+      if (spec != Layout::kUnshardedDim) counter[spec]++;
+      if (counter[spec] > max_count) {
+        max_count = counter[spec];
+        mesh_dim = spec;
+      }
+    }
+    // add output counts
+    for (auto const& idx_and_unused : batchable_outputs) {
+      auto const& idx = idx_and_unused.first;
+      std::string spec = output_layouts[idx].sharding_spec(i);
+      if (spec != Layout::kUnshardedDim) counter[spec]++;
+      if (counter[spec] > max_count) {
+        max_count = counter[spec];
+        mesh_dim = spec;
+      }
+    }
+    batch_specs[i] = mesh_dim;
+  }
+  // deduplicate
+  absl::flat_hash_map<std::string, int> counter;
+  for (const std::string& spec : batch_specs) counter[spec] += 1;
+  for (std::string& spec : batch_specs) {
+    if (counter[spec] > 1) {
+      counter[spec]--;
+      spec = Layout::kUnshardedDim;
+    }
+  }
+  return Layout::GetLayout(batch_specs, mesh);
+}
+}  // namespace
+
+// Relayout all operands that have batch dimensions to batch sharded
+// The outputs will get the correct inferred shape from the operands
+StatusOr<mlir::Operation*> DataparallelSPMDExpander::RelayoutOperandsAndOutputs(
+    mlir::Operation* op, const std::vector<Layout>& operand_layouts,
+    const std::vector<Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  TF_ASSIGN_OR_RETURN(
+      const Layout intermediate_batch_layout,
+      IntermediateBatchLayout(operand_layouts, batchable_operands_,
+                              output_layouts, batchable_outputs_, mesh));
+  // Relayout batchable operands
+  for (auto i = 0; i < operand_layouts.size(); ++i) {
+    // Relayout operands that have a batch dimension to intermediate layout
+    if (batchable_operands_.find(i) != batchable_operands_.end()) {
+      int replicated_rank =
+          ValueRank(op->getOperand(i)) - intermediate_batch_layout.rank();
+      TF_ASSIGN_OR_RETURN(
+          auto new_layout,
+          ConcatenateLayouts(intermediate_batch_layout,
+                             Layout::ReplicatedOnMesh(mesh, replicated_rank)));
+      TF_ASSIGN_OR_RETURN(
+          const auto new_operand,
+          EmitRelayout(op->getOperand(i), operand_layouts[i], new_layout));
+      op->setOperand(i, new_operand);
+    }
+  }
+  // Expand to local shape
+  op = InferSPMDExpandedLocalShape(op);
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  llvm::SmallVector<mlir::Value, 4> generated_outputs;
+  llvm::SmallVector<mlir::Type, 4> generated_types;
+
+  // Track the op that comes last after splitting.
+  mlir::Operation* last_op_after_splitting = op;
+
+  // Relayout batchable outputs
+  for (auto i = 0; i < output_layouts.size(); ++i) {
+    // Relayout to batch shard if tensor has batch dim
+    if (batchable_outputs_.find(i) != batchable_outputs_.end()) {
+      int replicated_rank =
+          ValueRank(op->getResult(i)) - intermediate_batch_layout.rank();
+      TF_ASSIGN_OR_RETURN(
+          auto new_layout,
+          ConcatenateLayouts(intermediate_batch_layout,
+                             Layout::ReplicatedOnMesh(mesh, replicated_rank)));
+      TF_ASSIGN_OR_RETURN(auto new_output,
+                          EmitRelayout(op->getOpResult(i), new_layout,
+                                       output_layouts[i], &newly_created_ops));
+      generated_outputs.emplace_back(new_output);
+      generated_types.emplace_back(new_output.getType());
+      if (last_op_after_splitting->isBeforeInBlock(
+              new_output.getDefiningOp())) {
+        last_op_after_splitting = new_output.getDefiningOp();
+      }
+    } else {
+      generated_outputs.push_back(op->getResult(i));
+      generated_types.push_back(op->getResult(i).getType());
+    }
+  }
+  mlir::OpBuilder builder(op);
+  builder.setInsertionPointAfter(last_op_after_splitting);
+
+  // Tie all outputs together with identity_n
+  auto identity_op = mlir::TF::IdentityNOp::create(
+      builder, op->getLoc(), generated_types, generated_outputs);
+  newly_created_ops.insert(identity_op);
+  for (int i = 0; i < output_layouts.size(); ++i) {
+    op->getOpResult(i).replaceAllUsesExcept(identity_op.getResult(i),
+                                            newly_created_ops);
+  }
+
+  return identity_op.getOperation();
+}
+
+StatusOr<mlir::Operation*> DataparallelSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const auto output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(const auto operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+  // Check all input and output are batch parallel
+  if (!AllBatchParallel(operand_layouts, batchable_operands_) ||
+      !AllBatchParallel(output_layouts, batchable_outputs_)) {
+    return absl::UnimplementedError(
+        llvm::formatv("All operands and outputs must be batch parallel.")
+            .str());
+  }
+  // Check that the rank of batch dimensions are same for all batchable tensors
+  if (!SameBatchRank(operand_layouts, batchable_operands_) ||
+      !SameBatchRank(output_layouts, batchable_outputs_)) {
+    return absl::UnimplementedError(
+        llvm::formatv("All operands and outputs with batch dimensions must "
+                      "have same batch dimension rank")
+            .str());
+  }
+  if (AllReplicated(output_layouts) && AllReplicated(operand_layouts))
+    return InferSPMDExpandedLocalShape(op);
+  return RelayoutOperandsAndOutputs(op, operand_layouts, output_layouts);
+}
+
+// Take all layouts of batchable operands, and merge them to produce a single
+// layout for all batchable outputs.
+StatusOr<llvm::DenseMap<int, Layout>>
+DataparallelSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> output_layouts;
+  // Compute output layouts
+  if (AnyLayoutExist(input_layouts, batchable_operands_)) {
+    TF_ASSIGN_OR_RETURN(
+        const Layout& batch_output_layout,
+        MergeBatchLayouts(input_layouts, batchable_operands_, mesh));
+    for (const auto& output_and_index : llvm::enumerate(op->getOpResults())) {
+      const int output_index = output_and_index.index();
+      auto output = output_and_index.value();
+      int rank = ValueRank(output);
+      if (batchable_outputs_.find(output_index) != batchable_outputs_.end()) {
+        int replicated_rank = batchable_outputs_[output_index];
+        TF_ASSIGN_OR_RETURN(auto new_layout,
+                            ConcatenateLayouts(batch_output_layout,
+                                               Layout::ReplicatedOnMesh(
+                                                   mesh, replicated_rank)));
+        output_layouts[output_index] = new_layout;
+      } else {
+        output_layouts[output_index] = Layout::ReplicatedOnMesh(mesh, rank);
+      }
+    }
+  }
+  return output_layouts;
+}
+
+// Take all layouts of batchable outputs, and merge them to produce a single
+// layout for all batchable operands.
+StatusOr<llvm::DenseMap<int, Layout>>
+DataparallelSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts;
+  // Compute input layouts in the following way: For operand indices that
+  // have a batch dimension, batch shard it the same way as the output layouts.
+  // Otherwise, replicate.
+  if (AnyLayoutExist(output_layouts, batchable_outputs_)) {
+    TF_ASSIGN_OR_RETURN(
+        const Layout& batch_operand_layout,
+        MergeBatchLayouts(output_layouts, batchable_outputs_, mesh));
+
+    for (const auto& operand_and_index : llvm::enumerate(op->getOperands())) {
+      const int operand_index = operand_and_index.index();
+      auto operand = operand_and_index.value();
+      int rank = ValueRank(operand);
+      if (batchable_operands_.find(operand_index) !=
+          batchable_operands_.end()) {
+        int replicated_rank = batchable_operands_[operand_index];
+        TF_ASSIGN_OR_RETURN(auto new_layout,
+                            ConcatenateLayouts(batch_operand_layout,
+                                               Layout::ReplicatedOnMesh(
+                                                   mesh, replicated_rank)));
+        input_layouts[operand_index] = new_layout;
+      } else {
+        input_layouts[operand_index] = Layout::ReplicatedOnMesh(mesh, rank);
+      }
+    }
+  }
+  return input_layouts;
+}
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.h

@@ -0,0 +1,70 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DATAPARALLEL_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DATAPARALLEL_SPMD_EXPANDER_H_
+#include <utility>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// General SPMD Expander for data parallel ops.
+
+// We define data parallel ops as ops that have tensors possibly with a batch
+// dimension. Assumes batch dimensions start from the left. Tensors may
+// may have multiple batch dimensions, including zero
+class DataparallelSPMDExpander : public SPMDExpanderBase {
+ protected:
+  // These maps contain {arg_index, non_batch_rank}
+  // Example is for TF:FFT2D, the batchable_operands and batchable_outputs has
+  // {0, 2} because the first argument is batchable and the last 2 dimensions
+  // are the non-batch dimensions
+  llvm::DenseMap<int, int> batchable_operands_;
+  llvm::DenseMap<int, int> batchable_outputs_;
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ public:
+  explicit DataparallelSPMDExpander(llvm::DenseMap<int, int> batchable_operands,
+                                    llvm::DenseMap<int, int> batchable_outputs)
+      : batchable_operands_(std::move(batchable_operands)),
+        batchable_outputs_(std::move(batchable_outputs)) {}
+
+ private:
+  // Relayouts all operands and outputs with a batch dimensions to a batch
+  // sharded layout. This should only be called when there is at least one
+  // batch sharded operand or batch sharded output
+  StatusOr<mlir::Operation*> RelayoutOperandsAndOutputs(
+      mlir::Operation* op, const std::vector<Layout>& operand_layouts,
+      const std::vector<Layout>& output_layouts);
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DATAPARALLEL_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/disable_copy_on_read_spmd_expander.cc

@@ -0,0 +1,48 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/disable_copy_on_read_spmd_expander.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> DisableCopyOnReadSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+DisableCopyOnReadSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // DisableCopyOnRead has no outputs;
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+DisableCopyOnReadSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& operand_layouts,
+    const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Prefer the layout from operand zero.
+  return operand_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/disable_copy_on_read_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DISABLE_COPY_ON_READ_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DISABLE_COPY_ON_READ_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class DisableCopyOnReadSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op, const llvm::DenseMap<int, Layout>& operand_layouts,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DISABLE_COPY_ON_READ_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/dtensor_op_spmd_expander.cc

@@ -0,0 +1,342 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/dtensor_op_spmd_expander.h"
+
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/optional.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/dtensor_send_recv.h"
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+// FIXME(feyu): This function should take layouts as arguments. It doesn't need
+// the ops.
+// Validates send/recv layout and mesh configurations. Among other things, this
+// checks for below constraints.
+// 1. Src/target layouts have non empty mesh.
+// 2. Src/target layouts have the same host.
+// 3. Src/target layouts are from different mesh.
+// 4. One of scr/target layout is from host mesh cluster.
+// 5. CPU host cluster mesh has 1 device.
+absl::Status ValidateSendRecvLayoutConfiguration(
+    mlir::TF::DTensorSend dtensor_send, mlir::TF::DTensorRecv dtensor_recv) {
+  // If either one of the send/recv ops has already been lowered, then send/recv
+  // configuration has already been verified.
+  if (!dtensor_send || !dtensor_recv) return absl::OkStatus();
+
+  TF_ASSIGN_OR_RETURN(const absl::optional<Layout> send_layout_or_null,
+                      ExtractLayoutFromOperand(dtensor_send.getInput()));
+
+  if (!send_layout_or_null.has_value())
+    return absl::InvalidArgumentError(
+        "Input to DTensorSend must have specified layout.");
+
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(dtensor_recv));
+
+  const Layout& send_layout = send_layout_or_null.value();
+  const Layout& recv_layout = output_layout;
+  const Mesh& recv_mesh = dtensor_recv.getMesh();
+
+  const Mesh& send_mesh = send_layout.mesh();
+
+  // If any one of send/recv mesh are empty, return error.
+  if (send_mesh.IsEmpty() || recv_mesh.IsEmpty())
+    return absl::InvalidArgumentError(
+        "Found empty mesh when sending/receiving tensor across clusters.");
+
+  // If send host not found in list of receiving hosts, return error.
+  std::vector<std::string> send_hosts = send_layout.ReducedMesh().hosts();
+  std::vector<std::string> recv_hosts = recv_layout.ReducedMesh().hosts();
+  if (send_hosts != recv_hosts)
+    return absl::InvalidArgumentError("Send and receive hosts don't match");
+
+  // Check shards in sending host match those in the receiving host.
+  const auto send_host_shard_map = send_layout.HostShardMap();
+  const auto recv_host_shard_map = recv_layout.HostShardMap();
+  for (const std::string& host : send_hosts) {
+    const ShardVector& shards_in_send_host =
+        send_host_shard_map.find(host)->second;
+    ShardVector shards_in_recv_host = recv_host_shard_map.find(host)->second;
+    if (shards_in_send_host != shards_in_recv_host)
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Send and receive host shard vectors don't match. Send shard_vector:",
+          shards_in_send_host.ToString(),
+          " / Recv host spec : ", shards_in_recv_host.ToString()));
+  }
+
+  // Send/Recv mesh must be different.
+  if (recv_mesh == send_mesh)
+    return absl::InvalidArgumentError(
+        "Found CopyToMesh op sending tensor to same mesh. Only use "
+        "CopyToMesh to transfer data across different mesh cluster. For "
+        "changing layout within the same mesh, use tf.Relayout op.");
+
+  // Either one of send/recv pair must be to/from CPU mesh.
+  // For example, TPU mesh -> GPU mesh or TPU mesh -> another TPU mesh
+  // is disallowed.
+  if (!send_mesh.is_cpu_mesh() && !recv_mesh.is_cpu_mesh())
+    return absl::InvalidArgumentError(
+        "tf.CopyToMesh op must be used to send data from/to host mesh.");
+
+  return absl::OkStatus();
+}
+
+template <typename RelayoutOp>
+StatusOr<mlir::Operation*> ExpandRelayoutOp(RelayoutOp relayout,
+                                            Layout target_layout,
+                                            Layout input_layout,
+                                            Layout output_layout) {
+  bool match_present = false;
+  for (const std::string& sharding_spec : target_layout.sharding_spec_strs())
+    if (sharding_spec == Layout::kMatch) match_present = true;
+
+  if (!match_present && output_layout != target_layout)
+    return absl::InternalError(
+        "output layout of Relayout op after layout propagation does not match "
+        "layout specified by Relayout op.");
+  auto value_or_status =
+      EmitRelayout(relayout.getInput(), input_layout, output_layout);
+  if (!value_or_status.ok())
+    return absl::InvalidArgumentError(
+        llvm::formatv("Unsupported layout received for {0} op. Trying "
+                      "to set tensor "
+                      "to layout : {1}. Found error {2}",
+                      relayout->getName().getStringRef(),
+                      target_layout.ToString(),
+                      value_or_status.status().message())
+            .str());
+  mlir::Value output = value_or_status.value();
+  relayout.getOutput().replaceAllUsesWith(output);
+  relayout.erase();
+  return output.getDefiningOp();
+}
+
+// Takes relayout which may have kMatch dimensions and uses it to mask input.
+// Here source_layout
+StatusOr<Layout> MergeLayouts(
+    const absl::flat_hash_set<std::string>& used_mesh_dimensions,
+    const Layout& mask_layout, const Layout& target_layout) {
+  std::vector<std::string> sharding_specs(mask_layout.sharding_spec_strs());
+  for (int i = 0; i < target_layout.rank(); ++i) {
+    if (sharding_specs[i] == Layout::kMatch &&
+        !used_mesh_dimensions.contains(target_layout.sharding_spec(i)))
+      sharding_specs[i] = target_layout.sharding_spec(i);
+  }
+  return Layout::GetLayout(sharding_specs, target_layout.mesh());
+}
+
+// Computes the layout of Relayout's (or RelayoutLike's) input or output, based
+// on the layout from the corresponding output or input (as `incoming_layout`).
+// Note that this implies that we compute the same layout for the
+// operand and output.
+// `mask_layout` is set to the user-supplied layout attribute on the op.
+StatusOr<llvm::DenseMap<int, Layout>> ComputeRelayoutLayout(
+    const Layout& mask_layout, std::optional<const Layout> incoming_layout) {
+  absl::flat_hash_set<std::string> used_dimensions;
+  bool match_present = false;
+  for (const std::string& sharding_spec : mask_layout.sharding_spec_strs()) {
+    if (sharding_spec == Layout::kMatch)
+      match_present = true;
+    else if (Layout::IsShardedDimension(sharding_spec))
+      used_dimensions.insert(sharding_spec);
+  }
+  if (!match_present) {
+    return llvm::DenseMap<int, Layout>({{0, mask_layout}});
+  }
+
+  if (incoming_layout) {
+    TF_ASSIGN_OR_RETURN(
+        Layout new_layout,
+        MergeLayouts(used_dimensions, mask_layout, *incoming_layout));
+    return llvm::DenseMap<int, Layout>({{0, new_layout}});
+  }
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> RelayoutSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto relayout = mlir::cast<mlir::TF::RelayoutOp>(op);
+  TF_ASSIGN_OR_RETURN(const Layout target_layout,
+                      Layout::FromString(relayout.getLayout().str()));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(const Layout input_layout,
+                      ExtractRequiredLayoutFromOperand(relayout.getInput()));
+
+  return ExpandRelayoutOp<mlir::TF::RelayoutOp>(relayout, target_layout,
+                                                input_layout, output_layout);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+RelayoutSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  auto relayout = llvm::cast<mlir::TF::RelayoutOp>(op);
+  TF_ASSIGN_OR_RETURN(const Layout mask_layout,
+                      Layout::FromString(relayout.getLayout().str()));
+  std::optional<const Layout> incoming_layout;
+  if (input_layouts.find(0) != input_layouts.end())
+    incoming_layout.emplace(input_layouts.lookup(0));
+
+  return ComputeRelayoutLayout(mask_layout, incoming_layout);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+RelayoutSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  auto relayout = llvm::cast<mlir::TF::RelayoutOp>(op);
+  TF_ASSIGN_OR_RETURN(const Layout mask_layout,
+                      Layout::FromString(relayout.getLayout().str()));
+  std::optional<const Layout> incoming_layout;
+  if (output_layouts.find(0) != output_layouts.end())
+    incoming_layout.emplace(output_layouts.lookup(0));
+
+  return ComputeRelayoutLayout(mask_layout, incoming_layout);
+}
+
+StatusOr<mlir::Operation*> RelayoutLikeSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  auto relayout_grad = mlir::cast<mlir::TF::RelayoutLikeOp>(op);
+  TF_ASSIGN_OR_RETURN(
+      const Layout target_layout,
+      ExtractRequiredLayoutFromOperand(relayout_grad.getLayoutInput()));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(
+      const Layout input_layout,
+      ExtractRequiredLayoutFromOperand(relayout_grad.getInput()));
+
+  return ExpandRelayoutOp<mlir::TF::RelayoutLikeOp>(
+      relayout_grad, target_layout, input_layout, output_layout);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+RelayoutLikeSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // RelayoutLike's output has the same layout as the corresponding Relayout's
+  // input operand.
+  if (input_layouts.find(1) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+  return llvm::DenseMap<int, Layout>({{0, input_layouts.lookup(1)}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+RelayoutLikeSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout output_layout = output_layouts.lookup(0);
+  return llvm::DenseMap<int, Layout>({
+      // Return replicated layout for the input operand since we do not want to
+      // enforce any particular layout on it.
+      {0, Layout::ReplicatedLike(output_layout)},
+      // Set layout for the forward pass's input operand to match the output of
+      // the RelayoutLike op.
+      {1, output_layout},
+  });
+}
+
+StatusOr<mlir::Operation*> DTensorSendSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  mlir::ModuleOp module = op->getParentOfType<mlir::ModuleOp>();
+  auto dtensor_send = llvm::cast<mlir::TF::DTensorSend>(op);
+
+  TF_ASSIGN_OR_RETURN(mlir::Operation * recv_op,
+                      GetCorrespondingDTensorSendRecvOp<mlir::TF::DTensorSend>(
+                          module, dtensor_send));
+  auto dtensor_recv = llvm::dyn_cast<mlir::TF::DTensorRecv>(recv_op);
+
+  TF_RETURN_IF_ERROR(
+      ValidateSendRecvLayoutConfiguration(dtensor_send, dtensor_recv));
+
+  return LowerDTensorSend(op, recv_op);
+}
+
+// DTensorSend op respects input layout from input operations and does not
+// set any preferred inputs layouts. During SPMD expansion, however, tensor
+// values are changed to replicated layout before transferring data across mesh
+// cluster.
+StatusOr<llvm::DenseMap<int, Layout>>
+DTensorSendSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+DTensorSendSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<mlir::Operation*> DTensorRecvSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  mlir::ModuleOp module = op->getParentOfType<mlir::ModuleOp>();
+  auto dtensor_recv = llvm::cast<mlir::TF::DTensorRecv>(op);
+
+  TF_ASSIGN_OR_RETURN(mlir::Operation * send_op,
+                      GetCorrespondingDTensorSendRecvOp<mlir::TF::DTensorRecv>(
+                          module, dtensor_recv));
+  auto dtensor_send = llvm::dyn_cast<mlir::TF::DTensorSend>(send_op);
+
+  TF_RETURN_IF_ERROR(
+      ValidateSendRecvLayoutConfiguration(dtensor_send, dtensor_recv));
+
+  return LowerDTensorRecv(send_op, op);
+}
+
+// DTensorRecv always returns tensors with fully replicated layout.
+StatusOr<llvm::DenseMap<int, Layout>>
+DTensorRecvSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  mlir::TF::DTensorRecv dtensor_recv =
+      mlir::dyn_cast<mlir::TF::DTensorRecv>(op);
+  if (!dtensor_recv) {
+    return absl::InvalidArgumentError(
+        llvm::formatv("Expecting DTensorRecvOp but got {0}", OpName(op)).str());
+  }
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+DTensorRecvSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/dtensor_op_spmd_expander.h

@@ -0,0 +1,98 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DTENSOR_OP_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DTENSOR_OP_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Converts layout of input tensor to target layout inserting split or reduction
+// ops if necessary.
+class RelayoutSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// Converts layout of gradient tensor to the layout of the original Relayout's
+// input tensor, using the same expansion logic as RelayoutOp.
+class RelayoutLikeSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// Lowers DTensorSend op to backend specific TF send/ xla send operation.
+// Following is the semantics for DTensorSend/Recv.
+// a) Both replicated/sharded DTensors can be sent/received.
+// b) When sharded DTensor is sent to another mesh, the DTensor is first
+//    all-to-all'ed to replicated tensor and sent to target mesh.
+// c) Send/Recv mesh must be from or to CPU mesh. That is, TPU<->TPU or
+//    GPU<->GTU is not supported.
+// d) Cross host send/recv is not supported. That is, sending tensor from
+//    TPU device of TPUWorker 0 to host of TPUWorker 1 is unsupported.
+class DTensorSendSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// Lowers DTensorRecv op to backend specific TF recv/ xla recv operation.
+class DTensorRecvSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_DTENSOR_OP_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc

@@ -0,0 +1,571 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.h"
+
+#include <cassert>
+#include <cstddef>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Einsum, like reductions, is implemented as a local operation followed by
+// an all-reduce over dimensions that have been reduced.
+StatusOr<mlir::Operation*> EinsumSPMDExpander::ExpandOp(mlir::Operation* op) {
+  std::vector<Layout> input_layouts(op->getNumOperands());
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    TF_ASSIGN_OR_RETURN(auto layout,
+                        ExtractLayoutFromOperand(op->getOperand(i)));
+    if (!layout)
+      return absl::InvalidArgumentError(
+          absl::StrCat("missing layout for input ", i));
+    input_layouts[i] = layout.value();
+  }
+  TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
+  if (!output_layout)
+    return absl::InvalidArgumentError("is missing output layout.");
+
+  std::vector<mlir::Value> new_inputs;
+  Layout layout_after_einsum;
+  absl::flat_hash_set<std::string> reduce_dims;
+  TF_RETURN_IF_ERROR(MaybeRelayoutInputs(input_layouts, op,
+                                         output_layout.value(), reduce_dims,
+                                         layout_after_einsum, new_inputs));
+
+  mlir::OpBuilder builder(op);
+  mlir::IRMapping mapping;
+  for (int i = 0; i < op->getNumOperands(); ++i)
+    mapping.map(op->getOperand(i), new_inputs[i]);
+  mlir::Operation* new_op = builder.clone(*op, mapping);
+  // Note that the output shape of new_op is cloned from op, so we need to
+  // update to the local shape.
+  new_op = InferSPMDExpandedLocalShape(new_op);
+
+  if (!reduce_dims.empty()) {
+    TF_ASSIGN_OR_RETURN(
+        new_op, EmitAllReduce(builder, layout_after_einsum, reduce_dims, new_op,
+                              kReduceOpAdd));
+  }
+
+  TF_ASSIGN_OR_RETURN(auto final_output,
+                      EmitRelayout(new_op->getOpResult(0), layout_after_einsum,
+                                   output_layout.value()));
+
+  op->getOpResult(0).replaceAllUsesWith(final_output);
+  op->erase();
+
+  return final_output.getDefiningOp();
+}
+
+// TODO(power) -- we use a simplified equation parser here. consider
+// refactoring einsum_spmd_expander and reusing the TF parser.
+//
+// Given the input equation, this has 3 outputs:
+// reduced_dims: The set of mesh dimesions that we need to all reduce over.
+// input_mappings: for each equation input, the map from the equation labels
+//   to the tensor dimension of that label.
+// output_mapping: as above, but for the equation output.
+absl::Status ExtractEquationRelations(
+    absl::string_view equation, absl::flat_hash_set<char>& reduced_dims,
+    std::vector<absl::flat_hash_map<char, std::vector<int>>>& input_mappings,
+    absl::flat_hash_map<char, std::vector<int>>& output_mapping) {
+  std::pair<std::string, std::string> parts = absl::StrSplit(equation, "->");
+  absl::flat_hash_set<char> non_reduced_dims;
+
+  // Mark kept dimensions from the output.
+  for (const auto& char_and_index : llvm::enumerate(parts.second)) {
+    // TODO(b/172691887): Support Broadcasting for einsum.
+    if (char_and_index.value() == '.')
+      return absl::UnimplementedError(absl::StrCat(
+          "Broadcasting is unimplemented for einsum. Received equation ",
+          equation));
+    non_reduced_dims.insert(char_and_index.value());
+
+    // Construct the output mapping, note that output is not allowed to have
+    // duplicate labels. This would mean that the datatype of output_mapping
+    // should really be absl::flat_hash_map<char, int>, but having the same
+    // type as the input_mapping keeps GetSpecsFromLabelsAndMap simpler.
+    if (output_mapping.contains(char_and_index.value()))
+      return errors::InvalidArgument("received label ", char_and_index.value(),
+                                     " multiple times in the "
+                                     "output of einsum equation ",
+                                     equation);
+
+    output_mapping[char_and_index.value()].emplace_back(char_and_index.index());
+  }
+
+  std::vector<std::string> inputs = absl::StrSplit(parts.first, ',');
+  // Note that the TF einsum op only supports at most 2 inputs. This is slightly
+  // confusing as the tf.einsum interface actually supports > 2 inputs.
+  if (inputs.size() > 2)
+    return absl::InvalidArgumentError(
+        absl::StrCat("einsum only supports at most 2 inputs received equation ",
+                     equation, " which has ", inputs.size(), " inputs"));
+
+  input_mappings.resize(inputs.size());
+
+  // Compute the input mappings and keep track of labels which are reduced.
+  for (int i = 0; i < inputs.size(); ++i) {
+    for (const auto& char_and_index : llvm::enumerate(inputs[i])) {
+      input_mappings[i][char_and_index.value()].emplace_back(
+          char_and_index.index());
+      if (!non_reduced_dims.contains(char_and_index.value()))
+        reduced_dims.insert(char_and_index.value());
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+// For a set of layouts and mappings from labels to offsets in the layouts,
+// return a mappings of labels to ShardingSpecs.
+// If the label appears multiples with different mesh dimensions in the
+// sharding specs we raise an error if replicate_incompatible_dimensions is
+// false. Otherwise we treat the dimension as if it were unsharded.
+// Labels with unsharded dimensions are not recorded in the output.
+StatusOr<absl::flat_hash_map<char, std::string>> GetLabelToShardingSpec(
+    bool replicate_incompatible_dimensions, const std::vector<Layout>& layouts,
+    const std::vector<absl::flat_hash_map<char, std::vector<int>>>& mappings) {
+  absl::flat_hash_map<char, std::string> label_to_sharding_spec;
+  absl::flat_hash_set<char> incompatible_labels;
+
+  // For each mapping, identify the mesh dimension and whether it has been
+  // reduced away.
+  for (int index = 0; index < layouts.size(); ++index) {
+    for (const auto& mapping : mappings[index]) {
+      for (int offset : mapping.second) {
+        if (offset >= layouts[index].rank())
+          return absl::InvalidArgumentError(
+              llvm::formatv(
+                  "specified einsum equation for operand {0} tried to "
+                  "read layout at offset {1}, but layout is {2} with rank "
+                  "{3}",
+                  index, offset, layouts[index].ToString(),
+                  layouts[index].rank())
+                  .str());
+
+        const std::string& sharding_spec = layouts[index].sharding_spec(offset);
+
+        if (label_to_sharding_spec.contains(mapping.first)) {
+          if (Layout::IsShardedDimension(sharding_spec) &&
+              label_to_sharding_spec[mapping.first] != sharding_spec) {
+            if (!replicate_incompatible_dimensions)
+              return absl::InvalidArgumentError(
+                  llvm::formatv(
+                      "incompatible mesh dimensions in equation, label '{0}' "
+                      "is mapped to mesh dimension '{1}' and '{2}'",
+                      mapping.first, sharding_spec,
+                      label_to_sharding_spec[mapping.first])
+                      .str());
+            else
+              incompatible_labels.insert(mapping.first);
+          }
+        } else if (Layout::IsShardedDimension(sharding_spec)) {
+          label_to_sharding_spec[mapping.first] = sharding_spec;
+        }
+      }
+    }
+  }
+
+  // For labels that had incompatible dimensions, treat them as replicated.
+  // We would need to insert some all to all in the SPMD expansion for these.
+  for (char label : incompatible_labels) label_to_sharding_spec.erase(label);
+
+  return label_to_sharding_spec;
+}
+
+// The layout we generated may be invalid as the same dimension may be used
+// multiple times. E.g. ab,bc->ac (i.e. matmul) with a and c sharded over the
+// same dim. In this case we mark all such dimensions as replicated.
+StatusOr<Layout> VerifyOrFixLayout(
+    std::pair<std::vector<std::string>, absl::flat_hash_map<std::string, int>>
+        pair,
+    const Mesh& mesh) {
+  std::vector<std::string> sharding_specs = pair.first;
+  absl::flat_hash_map<std::string, int> dimension_use_count = pair.second;
+  for (int i = 0; i < sharding_specs.size(); ++i)
+    if (Layout::IsShardedDimension(sharding_specs[i]) &&
+        dimension_use_count[sharding_specs[i]] > 1)
+      sharding_specs[i] = Layout::kUnshardedDim;
+  return Layout::GetLayout(sharding_specs, mesh);
+}
+
+// Construct a layout on a given mesh from the label to tensor dimension map
+// and the label to mesh_dimension map.
+std::pair<std::vector<std::string>, absl::flat_hash_map<std::string, int>>
+GetSpecsFromLabelsAndMap(
+    const absl::flat_hash_map<char, std::vector<int>>& label_to_index,
+    const absl::flat_hash_map<char, std::string>& label_to_sharding_spec) {
+  int layout_rank = 0;
+  for (const auto& label_and_indices : label_to_index)
+    layout_rank += label_and_indices.second.size();
+
+  std::vector<std::string> sharding_specs(layout_rank);
+  absl::flat_hash_map<std::string, int> dimension_use_count;
+  for (const auto& label_and_indices : label_to_index) {
+    const auto& loc = label_to_sharding_spec.find(label_and_indices.first);
+    if (loc != label_to_sharding_spec.end()) {
+      const std::string& sharding_spec = loc->second;
+      for (int index : label_and_indices.second)
+        sharding_specs[index] = sharding_spec;
+      dimension_use_count[sharding_spec] += label_and_indices.second.size();
+    } else {
+      for (int index : label_and_indices.second)
+        sharding_specs[index] = Layout::kUnshardedDim;
+    }
+  }
+  return std::make_pair(sharding_specs, dimension_use_count);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> EinsumSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  if (input_layouts.empty()) return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Need the mapping of input and output labels from the equation.
+  auto einsum_op = mlir::cast<mlir::TF::EinsumOp>(op);
+  size_t num_inputs = einsum_op.getNumOperands();
+  std::string equation = einsum_op.getEquation().str();
+  absl::flat_hash_set<char> reduced_dim_labels;
+  std::vector<absl::flat_hash_map<char, std::vector<int>>> input_mappings;
+  absl::flat_hash_map<char, std::vector<int>> output_mapping;
+
+  TF_RETURN_IF_ERROR(ExtractEquationRelations(equation, reduced_dim_labels,
+                                              input_mappings, output_mapping));
+  if (input_mappings.size() != num_inputs)
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Einsum equation ", equation, " has ", input_mappings.size(),
+        " inputs but this op has ", num_inputs, " inputs."));
+
+  // GetLabelToShardingSpec requires two inputs if the einsum equation needs
+  // two inputs. We may only have one layout, so make other replicated. This
+  // will have the same effect as only using the defined layout and using
+  // replicated for all the missing dimensions.
+  std::vector<Layout> layouts;
+  for (int k = 0; k < num_inputs; ++k) {
+    if (input_layouts.find(k) != input_layouts.end()) {
+      layouts.emplace_back(input_layouts.lookup(k));
+    } else {
+      int rank = ValueRank(op->getOperand(k));
+      if (rank < 0)
+        return absl::InvalidArgumentError(
+            absl::StrCat("No rank for input ", k));
+      // This case can only happen when there are two inputs. Input 1 - k
+      // is the other input. In this case of the if, input k is missing, so
+      // this means that input 1 - k must be there.
+      layouts.emplace_back(Layout::ReplicatedOnMesh(mesh, rank));
+    }
+  }
+
+  // For each input, identify the mesh dimension
+  TF_ASSIGN_OR_RETURN(
+      auto input_label_to_sharding_spec,
+      GetLabelToShardingSpec(
+          /*replicate_incompatible_dimensions=*/true, layouts, input_mappings));
+  // Compute output layout based on retained mesh dimensions
+  TF_ASSIGN_OR_RETURN(
+      const auto& output_layout,
+      VerifyOrFixLayout(GetSpecsFromLabelsAndMap(output_mapping,
+                                                 input_label_to_sharding_spec),
+                        mesh));
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> EinsumSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout output_layout = output_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Need the mapping of input and output labels from the equation.
+  auto einsum_op = mlir::cast<mlir::TF::EinsumOp>(op);
+  size_t num_inputs = einsum_op.getNumOperands();
+  std::string equation = einsum_op.getEquation().str();
+  absl::flat_hash_set<char> reduced_dim_labels;
+  std::vector<absl::flat_hash_map<char, std::vector<int>>> input_mappings;
+  absl::flat_hash_map<char, std::vector<int>> output_mapping;
+
+  TF_RETURN_IF_ERROR(ExtractEquationRelations(equation, reduced_dim_labels,
+                                              input_mappings, output_mapping));
+  if (input_mappings.size() != num_inputs)
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Einsum equation ", equation, " has ", input_mappings.size(),
+        " inputs but this op has ", num_inputs, " inputs."));
+
+  // Using the output mapping, construct an equation label to mesh dimension
+  // mapping.
+  TF_ASSIGN_OR_RETURN(auto output_label_to_sharding_spec,
+                      GetLabelToShardingSpec(
+                          /*replicate_incompatible_dimensions=*/false,
+                          {output_layout}, {output_mapping}));
+
+  // Defines a set of labels that could be set to Any. The conditions for an
+  // operand label to be set to any are that 1) is not present in the output
+  // and 2) is not repeated in any operand.
+  absl::flat_hash_set<char> labels_for_any;
+  for (const auto& operand_mapping : input_mappings)
+    for (const auto& label_to_indices : operand_mapping)
+      labels_for_any.insert(label_to_indices.first);
+
+  // Filter repeated labels.
+  for (const auto& operand_mapping : input_mappings)
+    for (const auto& label_to_indices : operand_mapping)
+      if (label_to_indices.second.size() > 1)
+        labels_for_any.erase(label_to_indices.first);
+
+  // Filter labels in output.
+  for (const auto& label_to_indices : output_mapping)
+    labels_for_any.erase(label_to_indices.first);
+
+  llvm::DenseMap<int, Layout> input_layouts(num_inputs);
+
+  // Derive operand sharding specs from output's sharding specs.
+  for (size_t i = 0; i < num_inputs; ++i) {
+    absl::flat_hash_map<char, std::vector<int>> labels_to_indices =
+        input_mappings[i];
+    std::pair<std::vector<std::string>, absl::flat_hash_map<std::string, int>>
+        sharding_specs_and_dim_count = GetSpecsFromLabelsAndMap(
+            labels_to_indices, output_label_to_sharding_spec);
+
+    std::vector<std::string> sharding_specs =
+        sharding_specs_and_dim_count.first;
+    absl::flat_hash_map<std::string, int> dim_count =
+        sharding_specs_and_dim_count.second;
+
+    // Flip "unsharded" specs to "any" if they are present in the set.
+    for (const auto& label_to_indices : labels_to_indices) {
+      char label = label_to_indices.first;
+      if (labels_for_any.contains(label)) {
+        int index = label_to_indices.second[0];
+        sharding_specs[index] = Layout::kAny;
+      }
+    }
+    TF_ASSIGN_OR_RETURN(
+        const auto& layout,
+        VerifyOrFixLayout(std::make_pair(sharding_specs, dim_count), mesh));
+    input_layouts[i] = layout;
+  }
+
+  return input_layouts;
+}
+
+// A few things we don't support, but could:
+// * "xx->" or "xx->x": Trace like operation where at least one input dimension
+//   for x is sharded. If both are sharded, we can compute the einsum on the
+//   diagonal machines in the mesh and 0s on the off diagonals and then all
+//   the much smaller matrix.
+absl::Status EinsumSPMDExpander::MaybeRelayoutInputs(
+    const std::vector<Layout>& input_layouts, mlir::Operation* op,
+    const Layout& output_layout, absl::flat_hash_set<std::string>& reduce_dims,
+    Layout& einsum_layout, std::vector<mlir::Value>& new_inputs) {
+  if (!mlir::isa<mlir::TF::EinsumOp>(op))
+    return absl::InvalidArgumentError(
+        "called einsum spmd expander but op is not Einsum.");
+
+  mlir::TF::EinsumOp einsum = mlir::cast<mlir::TF::EinsumOp>(op);
+  std::vector<absl::flat_hash_map<char, std::vector<int>>> input_mappings;
+  absl::flat_hash_map<char, std::vector<int>> output_mapping;
+  absl::flat_hash_set<char> contracting_labels;
+  absl::flat_hash_set<char> all_labels;
+  TF_RETURN_IF_ERROR(ExtractEquationRelations(einsum.getEquation().str(),
+                                              contracting_labels,
+                                              input_mappings, output_mapping));
+
+  for (const auto& input_mapping : input_mappings)
+    for (const auto& char_and_positions : input_mapping)
+      all_labels.emplace(char_and_positions.first);
+
+  // We will update this array throughout this function with the following rules
+  // 1. The sharding of a label which is not in the map is unknown.
+  // 2. Once the sharding of label becomes known and is unsharded, we
+  //    won't change that.
+  TF_ASSIGN_OR_RETURN(auto input_label_to_sharding_spec,
+                      GetLabelToShardingSpec(
+                          /*replicate_incompatible_dimensions=*/false,
+                          input_layouts, input_mappings));
+
+  TF_ASSIGN_OR_RETURN(const auto output_label_to_sharding_spec,
+                      GetLabelToShardingSpec(
+                          /*replicate_incompatible_dimensions=*/false,
+                          {output_layout}, {output_mapping}));
+
+  for (const char label : all_labels) {
+    if (input_label_to_sharding_spec.contains(label) &&
+        output_label_to_sharding_spec.contains(label) &&
+        input_label_to_sharding_spec[label] !=
+            output_label_to_sharding_spec.find(label)->second)
+      return errors::InvalidArgument(
+          "for label ", label, " input and output layouts are sharded on ",
+          " non-equal dimensions ", input_label_to_sharding_spec[label],
+          " and ", output_label_to_sharding_spec.find(label)->second,
+          "respectively");
+  }
+
+  // First priority is to ensure that labels which occur at least twice on one
+  // side never get sharded, as we cannot deal with that. This corresponds to
+  // taking a trace on that input, which will require us to be unsharded.
+  for (const auto& input_mapping : input_mappings)
+    for (const auto& char_and_positions : input_mapping)
+      if (char_and_positions.second.size() > 1)
+        input_label_to_sharding_spec[char_and_positions.first] =
+            Layout::kUnshardedDim;
+
+  absl::flat_hash_map<std::string, absl::flat_hash_set<char>>
+      sharding_dim_to_non_contracting_labels;
+  absl::flat_hash_map<std::string, absl::flat_hash_set<char>>
+      sharding_dim_to_contracting_labels;
+  for (const auto& label_and_spec : input_label_to_sharding_spec) {
+    if (Layout::IsShardedDimension(label_and_spec.second)) {
+      if (contracting_labels.contains(label_and_spec.first))
+        sharding_dim_to_contracting_labels[label_and_spec.second].insert(
+            label_and_spec.first);
+      else
+        sharding_dim_to_non_contracting_labels[label_and_spec.second].insert(
+            label_and_spec.first);
+    }
+  }
+
+  // If a non-contracting dimension is sharded in the output and non-sharded
+  // in the input and no other label is sharded on that dimension, then shard
+  // it.
+  // This handles the *,x . x,* -> *,y case and also if batch dimensions are
+  // sharded on the output but not the input.
+  for (const char label : all_labels) {
+    // Note that only sharded labels are in output_label_to_sharding_spec, so
+    // there is no need to check that the spec is sharded.
+    if (!contracting_labels.contains(label) &&
+        output_label_to_sharding_spec.contains(label) &&
+        !input_label_to_sharding_spec.contains(label)) {
+      const std::string& string_spec =
+          output_label_to_sharding_spec.find(label)->second;
+      if (!sharding_dim_to_non_contracting_labels.contains(string_spec) &&
+          !sharding_dim_to_contracting_labels.contains(string_spec)) {
+        input_label_to_sharding_spec[label] = string_spec;
+        sharding_dim_to_non_contracting_labels[string_spec].insert(label);
+      }
+    }
+  }
+
+  // Handle the case when two non-contracting dimensions are have the same
+  // sharding spec.
+  // Note that the case of three non-contracting dimensions having the same
+  // sharding spec is impossible. Since there are at most two inputs, at least
+  // one input would have two dimensions with the same sharing spec.
+  // This handles the y,x . x,y -> *,y case.
+  absl::flat_hash_set<std::string> dims_with_multiple_labels;
+  for (const auto& spec_and_labels : sharding_dim_to_non_contracting_labels) {
+    if (spec_and_labels.second.size() > 1) {
+      assert(spec_and_labels.second.size() == 2);
+      dims_with_multiple_labels.insert(spec_and_labels.first);
+    }
+  }
+  for (const auto& dim : dims_with_multiple_labels) {
+    // TODO(bfontain): Update this to pick default label to keep based on shape.
+    char label_to_keep = 0xFF;
+    // Note that all these conditions evaluated in the loop below are mutually
+    // as exclusive as no two labels in the output have the same sharding
+    // spec.
+    // If the no label is found we choose the lexicographically least label to
+    // keep this stable with respect to ordering.
+    for (const char label : sharding_dim_to_non_contracting_labels[dim]) {
+      if (output_label_to_sharding_spec.contains(label) &&
+          output_label_to_sharding_spec.find(label)->second == dim) {
+        label_to_keep = label;
+        break;
+      } else if (label < label_to_keep) {
+        label_to_keep = label;
+      }
+    }
+    for (const char label : sharding_dim_to_non_contracting_labels[dim])
+      if (label != label_to_keep)
+        input_label_to_sharding_spec[label] = Layout::kUnshardedDim;
+    sharding_dim_to_non_contracting_labels[dim].clear();
+    sharding_dim_to_non_contracting_labels[dim].insert(label_to_keep);
+  }
+
+  // Handle the case where a non-contracting and contracting dim have the same
+  // sharding spec. For now we always unshard the contracting axis. Note that
+  // this is safe.
+  // This handles the case x,y . *,y -> x,y
+  for (const auto& spec_and_labels : sharding_dim_to_contracting_labels) {
+    if (!spec_and_labels.second.empty() &&
+        !sharding_dim_to_non_contracting_labels[spec_and_labels.first]
+             .empty()) {
+      assert(spec_and_labels.second.size() == 1);
+      assert(sharding_dim_to_non_contracting_labels[spec_and_labels.first]
+                 .size() == 1);
+      input_label_to_sharding_spec[*spec_and_labels.second.begin()] =
+          Layout::kUnshardedDim;
+    }
+  }
+
+  // Relayout the inputs
+  mlir::OpBuilder builder(op);
+  new_inputs.resize(input_mappings.size());
+  for (int i = 0; i < input_mappings.size(); ++i) {
+    TF_ASSIGN_OR_RETURN(
+        const Layout new_input_layout,
+        VerifyOrFixLayout(GetSpecsFromLabelsAndMap(
+                              input_mappings[i], input_label_to_sharding_spec),
+                          output_layout.mesh()));
+
+    TF_ASSIGN_OR_RETURN(
+        new_inputs[i],
+        EmitRelayout(op->getOperand(i), input_layouts[i], new_input_layout));
+  }
+
+  TF_ASSIGN_OR_RETURN(
+      einsum_layout,
+      VerifyOrFixLayout(GetSpecsFromLabelsAndMap(output_mapping,
+                                                 input_label_to_sharding_spec),
+                        output_layout.mesh()));
+
+  for (const auto& contracting : contracting_labels)
+    reduce_dims.emplace(input_label_to_sharding_spec[contracting]);
+
+  return absl::OkStatus();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.h

@@ -0,0 +1,67 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EINSUM_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EINSUM_SPMD_EXPANDER_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class EinsumSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ private:
+  // This function will take the inputs, possibly slice or add AllConcat along
+  // various mesh dimensions before the einsum operation takes place. This also
+  // returns:
+  // * The mesh dimensions, if any, that the output of the einsum should be
+  //   summed along.
+  // * The resulting output layout of the einsum operation, so we can insert an
+  //   AllConcat/split to make the output have the desired layout.
+  // * The new inputs to fed into the einsum.
+  absl::Status MaybeRelayoutInputs(
+      const std::vector<Layout>& input_layouts, mlir::Operation* op,
+      const Layout& output_layout,
+      absl::flat_hash_set<std::string>& reduce_dims, Layout& einsum_layout,
+      std::vector<mlir::Value>& new_inputs);
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EINSUM_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc

@@ -0,0 +1,184 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.h"
+
+#include <cassert>
+#include <cstdint>
+#include <optional>
+
+#include "absl/container/flat_hash_set.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/UseDefLists.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+StatusOr<llvm::SmallVector<int64_t, 4>> GetShape(mlir::Value value) {
+  TF_ASSIGN_OR_RETURN(const auto shape, GetShapeOfValue(value));
+  return llvm::SmallVector<int64_t, 4>{shape.begin(), shape.end()};
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> ElementwiseSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
+  assert(output_layout);
+
+  mlir::OpBuilder builder(op);
+
+  for (auto& operand : op->getOpOperands()) {
+    // Verify that all output dimensions (including the dimensions added by
+    // broadcasting) is more sharded then the correspdonding layout
+    // configuration of the same dimension of every operands.
+    TF_ASSIGN_OR_RETURN(auto operand_layout,
+                        ExtractLayoutFromOperand(operand.get()));
+    if (!operand_layout)
+      return absl::InvalidArgumentError(
+          "input layout of elementwise op must be known before SPMD "
+          "expansion.");
+
+    // For scalar operands, splitting is not needed.
+    if (operand_layout->rank() == 0) continue;
+
+    llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+    // Note that due to broacasting, inputs and output tensors are aligned to
+    // the right. Therefore, we truncate the output layout.
+    const int rank_offset = output_layout->rank() - operand_layout->rank();
+
+    // Get the desired layout for this operand.
+    //
+    // - Truncate: It should be the same layout as the output, truncated to
+    //   take into account broadcasting on the rank and removing sharding in
+    //   dimensions where the operand has size 1 (where we have broadcasting on
+    //   the dimension size).
+    //
+    // - Relayout: If the output and operand have different sharding spec, we
+    //   adjust the operands. For example, if operand is 'z,*' and output is
+    //   '*.y', relayout operand to conform output. This means the SPMD safer
+    //   and easeier. In future, we might do certain optimization to save FLops.
+    //   For example, if all operands are 'x,y' and output is '*,*', relayouting
+    //   output could be the choice (saving communications).
+    auto truncated_layout = output_layout->Truncate(rank_offset, /*end=*/true);
+    mlir::Value output;
+    TF_ASSIGN_OR_RETURN(const auto& shape, ExtractGlobalInputShape(operand));
+    absl::flat_hash_set<int> size_one_dims;
+    for (int i = 0; i < shape.size(); ++i)
+      if (shape[i] == 1) size_one_dims.emplace(i);
+    TF_ASSIGN_OR_RETURN(truncated_layout,
+                        truncated_layout.GetLayoutWithReducedDims(
+                            size_one_dims, /*keep_dims=*/true));
+    TF_ASSIGN_OR_RETURN(
+        output, EmitRelayout(operand.get(), *operand_layout, truncated_layout));
+    operand.set(output);
+  }
+
+  // If result is a resource, the shape of the result should be adjusted to
+  // local value of the resource, based on the layout for output.
+  // This logic is similar to VarHandle op SPMD expansion.
+  //
+  // Resource output is only likely to be for identity op. However, keeping
+  // the checkgeneric here.
+  auto op_result = op->getOpResult(0);
+  if (IsResourceType(op_result)) {
+    TF_RETURN_IF_ERROR(InferSPMDExpandedLocalShapeForResourceOutput(
+        &op_result, output_layout.value(), builder.getContext()));
+  }
+
+  // For element-wise op SPMD expansion, given that operand layouts are
+  // compatible to op's layout, op can simply be executed without any changes.
+  return InferSPMDExpandedLocalShape(op);
+}
+
+// Computes output layouts of elementwise operation using broadcast logic for
+// operands.
+StatusOr<llvm::DenseMap<int, Layout>>
+ElementwiseSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> merged_operand_layout,
+                      GetMergedOperandLayout(input_layouts, op));
+
+  if (merged_operand_layout) {
+    const int output_rank = ValueRank(op->getOpResult(0));
+    if (output_rank == -1)
+      return absl::InvalidArgumentError("Output has unknown rank");
+
+    // We assume that all elementwise operations output a single tensor.
+    return llvm::DenseMap<int, Layout>(
+        {{0, merged_operand_layout->LeftPad(output_rank)}});
+  }
+  return llvm::DenseMap<int, Layout>();
+}
+
+// Computes input layouts of elementwise operation using broadcast logic for
+// operands.
+StatusOr<llvm::DenseMap<int, Layout>>
+ElementwiseSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Do not infer any operand layout if no output layout is present.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  Layout output_layout = output_layouts.lookup(0);
+  llvm::DenseMap<int, Layout> input_layouts;
+
+  for (const auto& operand_and_index : llvm::enumerate(op->getOperands())) {
+    const int operand_index = operand_and_index.index();
+    auto operand = operand_and_index.value();
+
+    TF_ASSIGN_OR_RETURN(auto operand_shape, GetShape(operand));
+    Layout output_layout_truncated = output_layout.Truncate(
+        output_layout.sharding_spec_strs().size() - operand_shape.size(),
+        /*end=*/true);
+    auto inferred_operand_layout_strs =
+        output_layout_truncated.sharding_spec_strs();
+
+    if (inferred_operand_layout_strs.size() != operand_shape.size())
+      return absl::FailedPreconditionError(
+          "Mismatch of operand shape size and layout size.");
+    for (const auto& dim_shape_and_index : llvm::enumerate(operand_shape)) {
+      const int dim_index = dim_shape_and_index.index();
+      const int dim_shape = dim_shape_and_index.value();
+      if (dim_shape <= 1) {
+        inferred_operand_layout_strs[dim_index] = Layout::kUnshardedDim;
+      }
+    }
+    TF_ASSIGN_OR_RETURN(
+        auto inferred_operand_layout,
+        Layout::GetLayout(inferred_operand_layout_strs, output_layout.mesh()));
+    input_layouts[operand_index] = inferred_operand_layout;
+  }
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.h

@@ -0,0 +1,46 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ELEMENTWISE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ELEMENTWISE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Base case for elementwise ops.
+// This includes some of the Unary Ops like Neg, Abs, etc.
+class ElementwiseSPMDExpander : public SPMDExpanderBase {
+ protected:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ELEMENTWISE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/expanddims_spmd_expander.cc

@@ -0,0 +1,145 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/expanddims_spmd_expander.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/types/optional.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> ExpandDimsExpander::ExpandOp(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const absl::optional<Layout> output_layout,
+                      ExtractSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(const absl::optional<Layout> operand_layout,
+                      ExtractLayoutFromOperand(op->getOperand(0)));
+  mlir::TF::ExpandDimsOp expand_dims_op =
+      mlir::cast<mlir::TF::ExpandDimsOp>(op);
+
+  InferSPMDExpandedLocalShape(op);
+
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> global_output_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(expand_dims_op.getOutput()));
+
+  // Compute current output layout (just input layout with unsharded on the
+  // new dim);
+  TF_ASSIGN_OR_RETURN(int64_t dim,
+                      ExtractConstIntFromValue(expand_dims_op.getDim()));
+
+  if (dim < 0) dim += global_output_shape.size();
+  std::vector<std::string> sharding_specs(global_output_shape.size());
+  for (int i = 0; i < global_output_shape.size(); ++i) {
+    if (i < dim)
+      sharding_specs[i] = operand_layout->sharding_spec(i);
+    else if (i == dim)
+      sharding_specs[i] = Layout::kUnshardedDim;
+    else
+      sharding_specs[i] = operand_layout->sharding_spec(i - 1);
+  }
+  TF_ASSIGN_OR_RETURN(const Layout current_output_layout,
+                      Layout::GetLayout(sharding_specs, output_layout->mesh()));
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value output_value,
+      EmitRelayout(expand_dims_op.getOutput(), current_output_layout,
+                   *output_layout, &newly_created_ops));
+
+  expand_dims_op.getOutput().replaceAllUsesExcept(output_value,
+                                                  newly_created_ops);
+
+  return output_value.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ExpandDimsExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  auto expand_dims_op = mlir::cast<mlir::TF::ExpandDimsOp>(op);
+
+  TF_ASSIGN_OR_RETURN(int64_t dim,
+                      ExtractConstIntFromValue(expand_dims_op.getDim()));
+
+  // Do not infer any output layout if no operand layout is present.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  auto input_layout = input_layouts.lookup(0);
+
+  if (dim < 0) dim += input_layout.rank() + 1;
+  std::vector<std::string> layout_sharding;
+
+  for (int i = 0; i <= input_layout.rank(); ++i) {
+    if (i == dim) layout_sharding.push_back(Layout::kUnshardedDim);
+    if (i < input_layout.rank())
+      layout_sharding.push_back(input_layout.sharding_spec(i));
+  }
+  TF_ASSIGN_OR_RETURN(auto inferred_output_layout,
+                      Layout::GetLayout(layout_sharding, mesh));
+  return llvm::DenseMap<int, Layout>({{0, inferred_output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ExpandDimsExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+  auto output_layout = output_layouts.lookup(0);
+
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  auto expand_dims_op = mlir::cast<mlir::TF::ExpandDimsOp>(op);
+
+  TF_ASSIGN_OR_RETURN(int64_t dim,
+                      ExtractConstIntFromValue(expand_dims_op.getDim()));
+
+  if (dim < 0) dim += output_layout.rank();
+
+  std::vector<std::string> layout_sharding;
+
+  for (int i = 0; i < output_layout.rank(); ++i) {
+    if (i == dim) continue;
+    layout_sharding.push_back(output_layout.sharding_spec(i));
+  }
+
+  TF_ASSIGN_OR_RETURN(auto inferred_input_layout,
+                      Layout::GetLayout(layout_sharding, mesh));
+  auto input_axis_rank = ValueRank(expand_dims_op->getOperand(1));
+
+  return llvm::DenseMap<int, Layout>(
+      {{0, inferred_input_layout},
+       {1, Layout::ReplicatedOnMesh(mesh, /*rank=*/input_axis_rank)}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/expanddims_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EXPANDDIMS_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EXPANDDIMS_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class ExpandDimsExpander : public SPMDExpanderBase {
+ public:
+  // Runs SPMD expansion for tf.ExpandDims. If needed, inserts a relayout.
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_EXPANDDIMS_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/fft_spmd_expander.cc

@@ -0,0 +1,504 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/fft_spmd_expander.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+// Return the last unsharded axis. Return -1 for fully replicated dtensor
+int LastUnshardedAxis(const std::vector<std::string>& sharding_specs) {
+  for (int i = sharding_specs.size() - 1; i >= 0; --i)
+    if (sharding_specs[i] == Layout::kUnshardedDim) return i;
+  return -1;
+}
+
+// Return false for non-distributed *FFTN.
+bool IsDistributedFFTN(int num_transform_axes, const Layout& layout) {
+  std::vector<std::string> sharding_specs = layout.sharding_spec_strs();
+  for (int i = sharding_specs.size() - num_transform_axes;
+       i < sharding_specs.size(); ++i)
+    if (sharding_specs[i] != Layout::kUnshardedDim) {
+      return true;
+    }
+  return false;
+}
+
+bool IsComplexFFT(mlir::Value input) {
+  auto data_type =
+      mlir::dyn_cast<mlir::TensorType>(input.getType()).getElementType();
+  return mlir::isa<mlir::ComplexType>(data_type);
+}
+
+absl::Status IsProperFFTLength(
+    mlir::Operation* op, const llvm::SmallVector<int64_t, 4>& fft_length_vec) {
+  TF_ASSIGN_OR_RETURN(auto input_layout,
+                      ExtractRequiredLayoutFromOperand(op->getOperand(0)));
+  const Mesh& mesh = input_layout.mesh();
+  int axes = fft_length_vec.size();
+  // RFFT in DTensor requires axes except -1 to have the same shape as input.
+  llvm::ArrayRef<int64_t> input_shape =
+      mlir::dyn_cast<mlir::TensorType>(op->getOperand(0).getType()).getShape();
+  std::vector<int64_t> input_shape_vec = input_shape.vec();
+  for (int i = 0; i < axes - 1; ++i)
+    if (fft_length_vec[i] != input_shape_vec[input_shape_vec.size() - axes + i])
+      return absl::InvalidArgumentError(
+          "DTensor RFFTOps are not suitable for current 'fft_length'.");
+
+  // fft_length[-1] should be divisible by the corresponding device number.
+  int num_of_devices_last_dim = mesh.dim_sizes()[input_shape_vec.size() - 1];
+  if (axes > 1 && fft_length_vec[axes - 1] % num_of_devices_last_dim == 1)
+    return absl::InvalidArgumentError(
+        "The values with current 'fft_length' are not shardable.");
+  return absl::OkStatus();
+}
+
+StatusOr<llvm::SmallVector<int64_t, 4>> ExtractFFTLengthFromOp(
+    mlir::Operation* op) {
+  mlir::Value fft_length = op->getOperand(1);
+  llvm::SmallVector<int64_t, 4> fft_length_vec;
+  TF_RETURN_IF_ERROR(ExtractConstVectorFromValue(fft_length, &fft_length_vec));
+  TF_RETURN_IF_ERROR(IsProperFFTLength(op, fft_length_vec));
+  return fft_length_vec;
+}
+
+// Forward flow for FFT and backward flow for iFFT
+// sharding_specs has at least one element
+void PropagateFFTLayout(std::vector<std::string>& sharding_specs, int axes) {
+  int last_unsharded_axis = LastUnshardedAxis(sharding_specs);
+  if (last_unsharded_axis == -1)
+    sharding_specs[sharding_specs.size() - 1] = Layout::kUnshardedDim;
+  else if (last_unsharded_axis != sharding_specs.size() - 1)
+    std::iter_swap(sharding_specs.end() - 1,
+                   sharding_specs.begin() + last_unsharded_axis);
+  std::string last_sharding_spec = sharding_specs.back();
+  sharding_specs.pop_back();
+  sharding_specs.insert(sharding_specs.end() - axes + 1, last_sharding_spec);
+}
+
+// Backward flow for FFT and forward flow for iFFT
+// sharding_specs has at least one element
+void PropagateIFFTLayout(std::vector<std::string>& sharding_specs, int axes) {
+  int last_unsharded_axis = LastUnshardedAxis(sharding_specs);
+  if (last_unsharded_axis == -1)
+    sharding_specs[sharding_specs.size() - axes] = Layout::kUnshardedDim;
+  else if (last_unsharded_axis != sharding_specs.size() - axes)
+    std::iter_swap(sharding_specs.end() - axes,
+                   sharding_specs.begin() + last_unsharded_axis);
+  std::string unsharded_axis = sharding_specs[sharding_specs.size() - axes];
+  sharding_specs.erase(sharding_specs.end() - axes);
+  sharding_specs.push_back(unsharded_axis);
+}
+
+StatusOr<mlir::Value> EmitTransposeRelayout(mlir::OpBuilder& builder,
+                                            mlir::Location location,
+                                            mlir::Value input,
+                                            const Layout& init_layout,
+                                            const Mesh& mesh,
+                                            std::pair<int, int>& perm_axes) {
+  std::vector<int64_t> perm_for_transpose;
+  int input_rank = ValueRank(input);
+  perm_for_transpose.reserve(input_rank);
+  for (int ax = 0; ax < input_rank; ++ax) {
+    perm_for_transpose.push_back(ax);
+  }
+  std::iter_swap(perm_for_transpose.begin() + perm_axes.first,
+                 perm_for_transpose.begin() + perm_axes.second);
+  mlir::Operation* transpose_op =
+      EmitTransposeOp(builder, location, input, perm_for_transpose);
+  mlir::Value transposed_input = transpose_op->getResult(0);
+
+  std::vector<std::string> sharding_specs = init_layout.sharding_spec_strs();
+  std::iter_swap(sharding_specs.begin() + perm_axes.first,
+                 sharding_specs.begin() + perm_axes.second);
+  TF_ASSIGN_OR_RETURN(
+      auto transposed_input_layout,
+      Layout::GetLayout(init_layout.type(), sharding_specs, mesh));
+  TF_ASSIGN_OR_RETURN(
+      transposed_input,
+      EmitRelayout(transposed_input, transposed_input_layout, init_layout));
+  return transposed_input;
+}
+
+absl::Status NormalizeAxes(std::vector<int>& transform_axes, int input_rank) {
+  std::sort(transform_axes.begin(), transform_axes.end());
+  for (int i = 0; i < transform_axes.size(); ++i) {
+    if (transform_axes[i] >= input_rank) {
+      return absl::InvalidArgumentError("Axes to perform FFTN on are invalid.");
+    } else if (transform_axes[i] < 0) {
+      transform_axes[i] += input_rank;
+    }
+  }
+  if (transform_axes.empty()) {
+    transform_axes.reserve(input_rank);
+    for (int i = 0; i < input_rank; ++i) transform_axes.push_back(i);
+  }
+  return absl::OkStatus();
+}
+
+// Make the last axis of the layout unsharded by swapping with another
+// axis or forcing the last axis to be unsharded if fully sharded.
+StatusOr<Layout> UnshardLastAxis(int input_rank, const Layout& layout,
+                                 const Mesh& mesh) {
+  if (input_rank < 1)
+    return absl::InvalidArgumentError("input_rank must be >= 1");
+  std::vector<std::string> input_sharding_specs = layout.sharding_spec_strs();
+  int last_unsharded_axis = LastUnshardedAxis(input_sharding_specs);
+  if (last_unsharded_axis == -1)
+    input_sharding_specs[input_rank - 1] = Layout::kUnshardedDim;
+  else if (last_unsharded_axis != input_rank - 1)
+    std::iter_swap(input_sharding_specs.end() - 1,
+                   input_sharding_specs.begin() + last_unsharded_axis);
+  return Layout::GetLayout(layout.type(), input_sharding_specs, mesh);
+}
+
+// Lowering FFTN/RFFTN operation for the first N-1 axes to N-1 1-d FFTOps.
+StatusOr<mlir::Operation*> ExpandFFTNImpl(
+    mlir::Operation* xfft_op, mlir::Operation* fft_op,
+    std::vector<int>& transform_axes, const int input_rank,
+    mlir::OpBuilder& builder, mlir::Location location,
+    const Layout& input_layout, const Mesh& mesh) {
+  SetSingleLayoutOnOp(xfft_op, input_layout);
+  mlir::Value output = xfft_op->getResult(0);
+
+  if (transform_axes.empty()) {
+    fft_op->getResult(0).replaceAllUsesWith(xfft_op->getResult(0));
+    fft_op->erase();
+    return InferSPMDExpandedLocalShape(xfft_op);
+  }
+
+  std::vector<int64_t> perm;
+  perm.reserve(input_rank);
+  for (int i = 0; i < input_rank - 1; ++i) {
+    perm.push_back(i);
+  }
+  perm.insert(perm.end() - transform_axes.size(), input_rank - 1);
+
+  mlir::TF::FFTOp new_fft_op;
+  Layout intermediate_layout = input_layout;
+  int ax;
+  while (!transform_axes.empty()) {
+    ax = transform_axes.back();
+    transform_axes.pop_back();
+
+    std::pair<int, int> perm_axes = {ax, input_rank - 1};
+    TF_ASSIGN_OR_RETURN(
+        mlir::Value transposed_input,
+        EmitTransposeRelayout(builder, location, output, intermediate_layout,
+                              mesh, perm_axes));
+
+    new_fft_op = mlir::TF::FFTOp::create(
+        builder, location, transposed_input.getType(), transposed_input);
+    SetSingleLayoutOnOp(new_fft_op, intermediate_layout);
+    output = new_fft_op.getOutput();
+  }
+
+  mlir::Operation* transpose_op =
+      EmitTransposeOp(builder, location, output, perm);
+  mlir::Value transposed_output = transpose_op->getResult(0);
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  builder.setInsertionPointAfter(new_fft_op);
+  TF_ASSIGN_OR_RETURN(auto final_output,
+                      EmitRelayout(transposed_output, intermediate_layout,
+                                   intermediate_layout, &newly_created_ops));
+  fft_op->getOpResult(0).replaceAllUsesExcept(final_output, newly_created_ops);
+  fft_op->erase();
+  return InferSPMDExpandedLocalShape(final_output.getDefiningOp());
+}
+
+StatusOr<mlir::Operation*> ExpandFFTN(mlir::Operation* fft_op,
+                                      std::vector<int>& transform_axes) {
+  mlir::OpBuilder builder(fft_op);
+  mlir::Value input = fft_op->getOperand(0);
+  TF_ASSIGN_OR_RETURN(auto input_layout,
+                      ExtractRequiredLayoutFromOperand(input));
+  const int input_rank = ValueRank(input);
+  const Mesh& mesh = input_layout.mesh();
+  mlir::Location location = fft_op->getLoc();
+
+  TF_RETURN_IF_ERROR(NormalizeAxes(transform_axes, input_rank));
+  int num_transform_axes = transform_axes.size();
+
+  if (!IsDistributedFFTN(num_transform_axes, input_layout))
+    return InferSPMDExpandedLocalShape(fft_op);
+
+  // FIXME(b/292286720): Since the last axis must be one of the transform_axes
+  // in current 1/2/3d transform ops, we don't need to find the last
+  // transofrm_axes and can just use -1. Need to be fixed by adding transpose op
+  // prior to this or unshard the last transform_axes.
+  TF_ASSIGN_OR_RETURN(Layout intermediate_layout,
+                      UnshardLastAxis(input_rank, input_layout, mesh));
+  TF_ASSIGN_OR_RETURN(mlir::Value intermediate,
+                      EmitRelayout(input, input_layout, intermediate_layout));
+
+  if (IsComplexFFT(input)) {
+    // FFT for the last axis.
+    mlir::TF::FFTOp fft_output_op = mlir::TF::FFTOp::create(
+        builder, location, intermediate.getType(), intermediate);
+    transform_axes.pop_back();
+    return ExpandFFTNImpl(fft_output_op, fft_op, transform_axes, input_rank,
+                          builder, location, intermediate_layout, mesh);
+  } else {
+    TF_ASSIGN_OR_RETURN(auto fft_length_vec, ExtractFFTLengthFromOp(fft_op));
+    mlir::Value fft_length = IntConst(
+        builder, location, (int32_t)fft_length_vec[num_transform_axes - 1]);
+    llvm::ArrayRef<int64_t> rfft_shape =
+        mlir::dyn_cast<mlir::TensorType>(intermediate.getType()).getShape();
+    std::vector<int64_t> rfft_shape_vec = rfft_shape.vec();
+    int num_of_devices_last_dim = mesh.dim_sizes()[input_rank - 1];
+    rfft_shape_vec[input_rank - 1] =
+        fft_length_vec[num_transform_axes - 1] / 2 + 1;
+    if (fft_length_vec.size() > 1 &&
+        rfft_shape_vec[input_rank - 1] % num_of_devices_last_dim != 0)
+      return absl::InvalidArgumentError(
+          "No suitable algorithm in DTensor found for current 'fft_length'.");
+
+    mlir::Type output_type = mlir::RankedTensorType::get(
+        rfft_shape_vec,
+        mlir::dyn_cast<mlir::TensorType>(fft_op->getResult(0).getType())
+            .getElementType());
+    // Real FFT for the last axis.
+    mlir::TF::RFFTOp rfft_output_op = mlir::TF::RFFTOp::create(
+        builder, location, output_type, intermediate, fft_length);
+    transform_axes.pop_back();
+    return ExpandFFTNImpl(rfft_output_op, fft_op, transform_axes, input_rank,
+                          builder, location, intermediate_layout, mesh);
+  }
+}
+
+StatusOr<mlir::Operation*> ExpandIFFTN(mlir::Operation* ifft_op,
+                                       std::vector<int>& transform_axes) {
+  mlir::OpBuilder builder(ifft_op);
+  mlir::Value input = ifft_op->getOperand(0);
+  TF_ASSIGN_OR_RETURN(auto input_layout,
+                      ExtractRequiredLayoutFromOperand(input));
+  const auto input_rank = ValueRank(input);
+  const Mesh& mesh = input_layout.mesh();
+  mlir::Location location = ifft_op->getLoc();
+  std::vector<std::string> input_sharding_specs =
+      input_layout.sharding_spec_strs();
+
+  TF_RETURN_IF_ERROR(NormalizeAxes(transform_axes, input_rank));
+  int num_transform_axes = transform_axes.size();
+
+  if (!IsDistributedFFTN(num_transform_axes, input_layout))
+    return InferSPMDExpandedLocalShape(ifft_op);
+
+  input_sharding_specs.push_back(
+      input_sharding_specs[input_rank - num_transform_axes]);
+  input_sharding_specs.erase(input_sharding_specs.begin() + input_rank -
+                             num_transform_axes);
+  TF_ASSIGN_OR_RETURN(
+      input_layout,
+      Layout::GetLayout(input_layout.type(), input_sharding_specs, mesh));
+  TF_ASSIGN_OR_RETURN(Layout intermediate_layout,
+                      UnshardLastAxis(input_rank, input_layout, mesh));
+
+  std::vector<int64_t> perm;
+  perm.reserve(input_rank);
+  for (int i = 0; i < input_rank; ++i)
+    if (i != input_rank - num_transform_axes) {
+      perm.push_back(i);
+    }
+  perm.push_back(input_rank - num_transform_axes);
+  mlir::Operation* transpose_op =
+      EmitTransposeOp(builder, location, input, perm);
+  mlir::Value transposed_output = transpose_op->getResult(0);
+  TF_ASSIGN_OR_RETURN(
+      transposed_output,
+      EmitRelayout(transposed_output, input_layout, intermediate_layout));
+
+  mlir::TF::IFFTOp fft_new_op;
+  int ax;  // current axis
+  while (transform_axes.size() > 1) {
+    ax = transform_axes[1] - 1;
+    transform_axes.erase(transform_axes.begin() + 1);
+    fft_new_op = mlir::TF::IFFTOp::create(
+        builder, location, transposed_output.getType(), transposed_output);
+    SetSingleLayoutOnOp(fft_new_op, intermediate_layout);
+    transposed_output = fft_new_op.getOutput();
+    // Swap and relayout
+    std::pair<int, int> perm_axes = {ax, input_rank - 1};
+    TF_ASSIGN_OR_RETURN(
+        transposed_output,
+        EmitTransposeRelayout(builder, location, transposed_output,
+                              intermediate_layout, mesh, perm_axes));
+  }
+
+  if (IsComplexFFT(ifft_op->getResult(0))) {
+    // IFFT for the last axis.
+    mlir::TF::IFFTOp ifft_output_op = mlir::TF::IFFTOp::create(
+        builder, location, transposed_output.getType(), transposed_output);
+    SetSingleLayoutOnOp(ifft_output_op, intermediate_layout);
+    builder.setInsertionPointAfter(ifft_output_op);
+
+    ifft_op->getResult(0).replaceAllUsesWith(ifft_output_op);
+    ifft_op->erase();
+    return InferSPMDExpandedLocalShape(ifft_output_op);
+  } else {
+    TF_ASSIGN_OR_RETURN(auto complex_fft_length_vec,
+                        ExtractFFTLengthFromOp(ifft_op));
+    mlir::Value ifft_length =
+        IntConst(builder, location,
+                 (int32_t)complex_fft_length_vec[num_transform_axes - 1]);
+    // IRFFT for the last axis.
+    mlir::TF::IRFFTOp irfft_output_op = mlir::TF::IRFFTOp::create(
+        builder, location, ifft_op->getResult(0).getType(), transposed_output,
+        ifft_length);
+    SetSingleLayoutOnOp(irfft_output_op, intermediate_layout);
+    builder.setInsertionPointAfter(irfft_output_op);
+    ifft_op->getResult(0).replaceAllUsesWith(irfft_output_op.getOutput());
+    ifft_op->erase();
+    return InferSPMDExpandedLocalShape(irfft_output_op);
+  }
+}
+}  // namespace
+
+StatusOr<mlir::Operation*> FFTSPMDExpander::ExpandOp(mlir::Operation* op) {
+  std::vector<int> last_axis{-1};
+  std::vector<int> last_2_axes{-2, -1};
+  std::vector<int> last_3_axes{-3, -2, -1};
+  return llvm::TypeSwitch<mlir::Operation*, StatusOr<mlir::Operation*>>(op)
+      // Forward prop ops.
+      .Case<mlir::TF::FFTOp, mlir::TF::RFFTOp>(
+          [&](auto op) { return ExpandFFTN(op, last_axis); })
+      .Case<mlir::TF::FFT2DOp, mlir::TF::RFFT2DOp>(
+          [&](auto op) { return ExpandFFTN(op, last_2_axes); })
+      .Case<mlir::TF::FFT3DOp, mlir::TF::RFFT3DOp>(
+          [&](auto op) { return ExpandFFTN(op, last_3_axes); })
+
+      // Backward prop ops.
+      .Case<mlir::TF::IFFTOp, mlir::TF::IRFFTOp>(
+          [&](auto op) { return ExpandIFFTN(op, last_axis); })
+      .Case<mlir::TF::IFFT2DOp, mlir::TF::IRFFT2DOp>(
+          [&](auto op) { return ExpandIFFTN(op, last_2_axes); })
+      .Case<mlir::TF::IFFT3DOp, mlir::TF::IRFFT3DOp>(
+          [&](auto op) { return ExpandIFFTN(op, last_3_axes); })
+      .Default([&](auto op) { return InferSPMDExpandedLocalShape(op); });
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> FFTSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  auto iter = input_layouts.find(0);
+  if (iter == input_layouts.end()) return llvm::DenseMap<int, Layout>();
+
+  const Layout& input_layout = iter->second;
+  std::vector<std::string> sharding_specs = input_layout.sharding_spec_strs();
+  if (sharding_specs.empty())
+    return absl::FailedPreconditionError(
+        absl::StrCat(OpName(op), " has no sharding specs."));
+  llvm::TypeSwitch<mlir::Operation*>(op)
+      .Case<mlir::TF::FFTOp, mlir::TF::RFFTOp>(
+          [&sharding_specs](auto op) { PropagateFFTLayout(sharding_specs, 1); })
+      .Case<mlir::TF::FFT2DOp, mlir::TF::RFFT2DOp>(
+          [&sharding_specs](auto op) { PropagateFFTLayout(sharding_specs, 2); })
+      .Case<mlir::TF::FFT3DOp, mlir::TF::RFFT3DOp>(
+          [&sharding_specs](auto op) { PropagateFFTLayout(sharding_specs, 3); })
+
+      .Case<mlir::TF::IFFTOp, mlir::TF::IRFFTOp>([&sharding_specs](auto op) {
+        PropagateIFFTLayout(sharding_specs, 1);
+      })
+      .Case<mlir::TF::IFFT2DOp, mlir::TF::IRFFT2DOp>(
+          [&sharding_specs](auto op) {
+            PropagateIFFTLayout(sharding_specs, 2);
+          })
+      .Case<mlir::TF::IFFT3DOp, mlir::TF::IRFFT3DOp>(
+          [&sharding_specs](auto op) {
+            PropagateIFFTLayout(sharding_specs, 3);
+          });
+
+  TF_ASSIGN_OR_RETURN(auto result_layout,
+                      Layout::GetLayout(input_layout.type(), sharding_specs,
+                                        input_layout.mesh()));
+  if (result_layout.rank() != input_layout.rank())
+    return absl::FailedPreconditionError(absl::StrCat(
+        OpName(op), " derived output layout rank is ", result_layout.rank(),
+        " not ", input_layout.rank(), " as expected."));
+
+  return llvm::DenseMap<int, Layout>({{0, result_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> FFTSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  auto iter = output_layouts.find(0);
+  if (iter == output_layouts.end()) return llvm::DenseMap<int, Layout>();
+
+  const Layout& output_layout = iter->second;
+  std::vector<std::string> sharding_specs = output_layout.sharding_spec_strs();
+  if (sharding_specs.empty())
+    return absl::FailedPreconditionError(
+        absl::StrCat(OpName(op), " has no sharding specs."));
+
+  llvm::TypeSwitch<mlir::Operation*>(op)
+      .Case<mlir::TF::FFTOp, mlir::TF::RFFTOp>([&sharding_specs](auto op) {
+        PropagateIFFTLayout(sharding_specs, 1);
+      })
+      .Case<mlir::TF::FFT2DOp, mlir::TF::RFFT2DOp>([&sharding_specs](auto op) {
+        PropagateIFFTLayout(sharding_specs, 2);
+      })
+      .Case<mlir::TF::FFT3DOp, mlir::TF::RFFT3DOp>([&sharding_specs](auto op) {
+        PropagateIFFTLayout(sharding_specs, 3);
+      })
+
+      .Case<mlir::TF::IFFTOp, mlir::TF::IRFFTOp>(
+          [&sharding_specs](auto op) { PropagateFFTLayout(sharding_specs, 1); })
+      .Case<mlir::TF::IFFT2DOp, mlir::TF::IRFFT2DOp>(
+          [&sharding_specs](auto op) { PropagateFFTLayout(sharding_specs, 2); })
+      .Case<mlir::TF::IFFT3DOp, mlir::TF::IRFFT3DOp>(
+          [&sharding_specs](auto op) {
+            PropagateFFTLayout(sharding_specs, 3);
+          });
+
+  TF_ASSIGN_OR_RETURN(auto result_layout,
+                      Layout::GetLayout(output_layout.type(), sharding_specs,
+                                        output_layout.mesh()));
+  if (result_layout.rank() != output_layout.rank())
+    return absl::FailedPreconditionError(absl::StrCat(
+        OpName(op), " derived output layout rank is ", result_layout.rank(),
+        " not ", output_layout.rank(), " as expected."));
+
+  return llvm::DenseMap<int, Layout>({{0, result_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/fft_spmd_expander.h

@@ -0,0 +1,47 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_FFT_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_FFT_SPMD_EXPANDER_H_
+
+#include <string>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Implement Layout propagation and SPMD expansion for FFT ops.
+class FFTSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_FFT_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/fill_spmd_expander.cc

@@ -0,0 +1,123 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/fill_spmd_expander.h"
+
+#include <cstdint>
+#include <optional>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "xla/mlir_hlo/utils/convert_op_folder.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/dtensor_location.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> FillSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto original_fill = mlir::cast<mlir::TF::FillOp>(op);
+  TF_ASSIGN_OR_RETURN(auto dims_layout,
+                      ExtractLayoutFromOperand(original_fill.getDims()));
+  if (!dims_layout.has_value()) {
+    return absl::InvalidArgumentError(
+        "Failed during SPMD expansion of tf.FillOp. Layout of dimension "
+        "input must be known.");
+  }
+
+  if (!dims_layout->IsFullyReplicated()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Expected the layout for fill's `dims` argument to be fully "
+        "replicated. Got ",
+        dims_layout->ToString()));
+  }
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> output_layout,
+                      ExtractSingleLayoutFromOp(op));
+  if (!output_layout.has_value())
+    return absl::InternalError(
+        "FillOp doesn't have a layout after layout propagation");
+  if (output_layout->IsFullyReplicated()) {
+    // For fully replicated layouts the local shape on each device is the same
+    // as the global shape.
+    return InferSPMDExpandedLocalShape(op);
+  }
+
+  // For sharded outputs, the `dims` just needs to be translated from the
+  // global to the local shape.
+  mlir::OpBuilder builder(op->getBlock(), ++mlir::Block::iterator(op));
+
+  // Create a tensor from the sharding spec, with the dtype of the original
+  // attribute.
+  auto shard_values = output_layout->num_shards();
+  auto int_type = mlir::RankedTensorType::get(
+      static_cast<int64_t>(shard_values.size()), builder.getIntegerType(32));
+  auto int_attr = mlir::DenseIntElementsAttr::get(int_type, shard_values);
+  auto target_type_attr = mlir::hlo::convertElementsAttr(
+      int_attr, mlir::cast<mlir::TensorType>(original_fill.getDims().getType())
+                    .getElementType());
+
+  auto location = DT_LOC(op);
+  auto num_shards =
+      mlir::TF::ConstOp::create(builder, location, target_type_attr);
+  // Divide the global shape by the sharding spec.
+  auto div = mlir::TF::DivOp::create(builder, location, original_fill.getDims(),
+                                     num_shards.getResult());
+  // Copy over static shape information if available
+  auto global_output_type =
+      mlir::cast<mlir::TensorType>(original_fill.getResult().getType());
+  TF_ASSIGN_OR_RETURN(
+      auto local_type,
+      LocalTypeFromGlobalType(output_layout.value(), global_output_type));
+
+  auto new_fill = mlir::TF::FillOp::create(
+      builder, location, local_type, div.getResult(), original_fill.getValue());
+  original_fill.getResult().replaceAllUsesWith(new_fill.getOutput());
+  original_fill.erase();
+  return InferSPMDExpandedLocalShape(new_fill);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> FillSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  // Always set replicated layout for output.
+  return llvm::DenseMap<int, Layout>(
+      {{0,
+        Layout::ReplicatedOnMesh(
+            mesh, ValueRank(llvm::cast<mlir::TF::FillOp>(op).getOutput()))}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> FillSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  // Always set replicated layout for dims / value operand of Fill op.
+  return llvm::DenseMap<int, Layout>({{0, Layout::ReplicatedOnMesh(mesh, 1)},
+                                      {1, Layout::ReplicatedOnMesh(mesh, 0)}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/fill_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_FILL_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_FILL_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class FillSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_FILL_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/gather_spmd_expander.cc

@@ -0,0 +1,449 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h"
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/types/optional.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> GatherV2SPMDExpander::ExpandOp(mlir::Operation* op) {
+  return ExpandOpHelper<mlir::TF::GatherV2Op>(op);
+}
+
+StatusOr<int64_t> GatherV2SPMDExpander::GetAxis(mlir::Operation* op) {
+  return ExtractConstIntFromValue(
+      llvm::cast<mlir::TF::GatherV2Op>(op).getAxis());
+}
+
+StatusOr<uint64_t> GatherV2SPMDExpander::GetBatchDim(mlir::Operation* op) {
+  return llvm::cast<mlir::TF::GatherV2Op>(op).getBatchDims();
+}
+
+StatusOr<mlir::Operation*> ResourceGatherSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  return ExpandOpHelper<mlir::TF::ResourceGatherOp>(op);
+}
+
+StatusOr<int64_t> ResourceGatherSPMDExpander::GetAxis(mlir::Operation* op) {
+  return 0;
+}
+
+StatusOr<uint64_t> ResourceGatherSPMDExpander::GetBatchDim(
+    mlir::Operation* op) {
+  return llvm::cast<mlir::TF::ResourceGatherOp>(op).getBatchDims();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+GatherCommonSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(int64_t axis, GetAxis(op));
+  TF_ASSIGN_OR_RETURN(uint64_t batch_dims, GetBatchDim(op));
+
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  std::optional<Layout> params_layout;
+  if (input_layouts.find(0) != input_layouts.end()) {
+    params_layout.emplace(input_layouts.lookup(0));
+  }
+  std::optional<Layout> indices_layout;
+  if (input_layouts.find(1) != input_layouts.end()) {
+    indices_layout.emplace(input_layouts.lookup(1));
+  }
+
+  const int params_rank = ValueRank(op->getOperand(0));
+  const int indices_rank = ValueRank(op->getOperand(1));
+  if (params_rank == -1)
+    return absl::InvalidArgumentError("missing rank for params input");
+  if (indices_rank == -1)
+    return absl::InvalidArgumentError("missing rank for indices input");
+
+  // Handle the case of negative axis.
+  if (axis < 0) axis += params_rank;
+  if (batch_dims < 0) batch_dims += indices_rank;
+  if (!params_layout && !indices_layout) {
+    return llvm::DenseMap<int, Layout>();
+  }
+  std::vector<std::string> output_layout_specs;
+
+  // Get a list of mesh dims that params uses, other than the dim for axis.
+  llvm::DenseSet<llvm::StringRef> params_mesh_dims;
+  if (params_layout) {
+    for (int i = 0; i < params_rank; ++i)
+      if (i != axis &&
+          !Layout::IsUnshardedDimension(params_layout->sharding_spec(i)))
+        params_mesh_dims.insert(params_layout->sharding_spec(i));
+  }
+
+  auto add_mesh_dim_if = [&](const absl::optional<Layout>& input_layout,
+                             int64_t dim, bool indices = false) {
+    // Only add the mesh dimension to the output_layout if 1) the input layout
+    // exists and 2) when the input is indices and the params dims don't
+    // contain the mesh dim we are adding (to avoid two different tensor dims
+    // being sharded over the same mesh dim).
+    // Note that params->sharding_spec(axis) is specifically excluded from the
+    // params_mesh_dims during construction above. This means that if we are
+    // processing the indices layout and it contains
+    // params->sharding_spec(axis) we will still add that sharding_spec to the
+    // output layout.
+    if (input_layout && (!indices || !params_mesh_dims.contains(
+                                         input_layout->sharding_spec(dim))))
+      output_layout_specs.push_back(input_layout->sharding_spec(dim));
+    else
+      output_layout_specs.push_back(Layout::kUnshardedDim);
+  };
+
+  for (int i = 0; i < axis; ++i) add_mesh_dim_if(params_layout, i);
+  for (int i = batch_dims; i < indices_rank; ++i)
+    add_mesh_dim_if(indices_layout, i, /*indices=*/true);
+  for (int i = axis + 1; i < params_rank; ++i)
+    add_mesh_dim_if(params_layout, i);
+
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      Layout::GetLayout(output_layout_specs, mesh));
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+GatherCommonSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(int64_t axis, GetAxis(op));
+  TF_ASSIGN_OR_RETURN(uint64_t batch_dims, GetBatchDim(op));
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands());
+  // Axis is a constant so replicate it. ResourceGatherOp does not have axis.
+  if (llvm::isa<mlir::TF::GatherV2Op>(op)) {
+    input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/0);
+  }
+
+  auto it = output_layouts.find(0);
+  if (it == output_layouts.end()) {
+    return input_layouts;
+  }
+
+  // This will always exist since there is only one output.
+  const Layout output_layout = it->getSecond();
+
+  const int params_rank = ValueRank(op->getOperand(0));
+  const int indices_rank = ValueRank(op->getOperand(1));
+  if (params_rank == -1)
+    return absl::InvalidArgumentError("missing rank for params input");
+  if (indices_rank == -1)
+    return absl::InvalidArgumentError("missing rank for indices input");
+
+  // Handle the case of negative axis.
+  if (axis < 0) axis += params_rank;
+  if (batch_dims < 0) batch_dims += indices_rank;
+  std::vector<std::string> params_layout_specs;
+  std::vector<std::string> indices_layout_specs;
+  params_layout_specs.reserve(params_rank);
+  indices_layout_specs.reserve(indices_rank);
+
+  // Extract the params layout. We will request that axis is replicated as
+  // that gives the least issues with spmd expansion.
+  // E.g. If we had axis = 1 and parmas layout [p0 p1 p2 p3]
+  // input layout [i0 i1] then the output layout would have been
+  // [p0 i0 i1 p2 p3] so to go backwards, we extract the ranges [0, axis)
+  // and [axis + indices.rank(), output.rank()) for params (with a replicated
+  // dim for the missing dimension inbetween). Indices layout is based on the
+  // range [axis, axis+indices)/
+  for (int i = 0; i < axis; ++i)
+    params_layout_specs.push_back(output_layout.sharding_spec(i));
+  params_layout_specs.push_back(Layout::kUnshardedDim);
+  for (int i = axis + indices_rank - batch_dims; i < output_layout.rank(); ++i)
+    params_layout_specs.push_back(output_layout.sharding_spec(i));
+
+  // Extract the indices layout.
+  for (int i = 0; i < batch_dims; ++i)
+    indices_layout_specs.push_back(output_layout.sharding_spec(i));
+  for (int i = axis; i < axis + indices_rank - batch_dims; ++i)
+    indices_layout_specs.push_back(output_layout.sharding_spec(i));
+
+  TF_ASSIGN_OR_RETURN(const Layout params_layout,
+                      Layout::GetLayout(params_layout_specs, mesh));
+  TF_ASSIGN_OR_RETURN(const Layout indices_layout,
+                      Layout::GetLayout(indices_layout_specs, mesh));
+  input_layouts[0] = params_layout;
+  input_layouts[1] = indices_layout;
+
+  return input_layouts;
+}
+
+namespace {
+
+StatusOr<Layout> GatherNdGetOutputLayoutFromInput(
+    const absl::optional<Layout>& params_layout, int params_rank,
+    const absl::optional<Layout>& indices_layout, int indices_rank,
+    int index_dimensions, const Mesh& mesh) {
+  // The layout of the output should be the layout of the first rank-1
+  // dimensions of the indices plus the layout of the last params.rank -
+  // indices.dims[-1] dimensions of params. If one of the two is missing we
+  // replace them with replicated.
+  // If sharding dimension is used by both params and indices, the params
+  // layout will be respected as generally params is larger than indices.
+  std::vector<std::string> output_specs(params_rank - index_dimensions +
+                                        indices_rank - 1);
+  absl::flat_hash_set<std::string> used_dimensions;
+  const int params_offset = -index_dimensions + indices_rank - 1;
+
+  for (int i = index_dimensions; i < params_rank; ++i) {
+    if (params_layout &&
+        Layout::IsShardedDimension(params_layout->sharding_spec(i))) {
+      const auto& params_spec = params_layout->sharding_spec(i);
+      output_specs[i + params_offset] = params_spec;
+      used_dimensions.emplace(params_spec);
+    } else {
+      output_specs[i + params_offset] = Layout::kUnshardedDim;
+    }
+  }
+  for (int i = 0; i < indices_rank - 1; ++i) {
+    if (indices_layout &&
+        Layout::IsShardedDimension(indices_layout->sharding_spec(i)) &&
+        !used_dimensions.contains(indices_layout->sharding_spec(i)))
+      output_specs[i] = indices_layout->sharding_spec(i);
+    else
+      output_specs[i] = Layout::kUnshardedDim;
+  }
+  return Layout::GetLayout(output_specs, mesh);
+}
+
+absl::Status GatherNdGetInputLayoutFromOutput(
+    const Layout& output_layout, Layout* params_layout, int params_rank,
+    Layout* indices_layout, int indices_rank, int index_dimensions,
+    const Mesh& mesh) {
+  // We copy the first indices_rank - 1 dimensions of the output layout to
+  // indices_layout (with the last dimensions replicated) and the remaining
+  // dimensions to params_layout (with the first index_dimensions dimensions
+  // replicated).
+  std::vector<std::string> params_specs(params_rank);
+  std::vector<std::string> indices_specs(indices_rank);
+
+  for (int i = 0; i < index_dimensions; ++i)
+    params_specs[i] = Layout::kUnshardedDim;
+
+  const int params_offset = -index_dimensions + indices_rank - 1;
+  for (int i = index_dimensions; i < params_rank; ++i)
+    params_specs[i] = output_layout.sharding_spec(i + params_offset);
+
+  for (int i = 0; i < indices_rank - 1; ++i)
+    indices_specs[i] = output_layout.sharding_spec(i);
+
+  indices_specs[indices_rank - 1] = Layout::kUnshardedDim;
+
+  TF_ASSIGN_OR_RETURN(*params_layout, Layout::GetLayout(params_specs, mesh));
+  TF_ASSIGN_OR_RETURN(*indices_layout, Layout::GetLayout(indices_specs, mesh));
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> GatherNdSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto gather_op = llvm::cast<mlir::TF::GatherNdOp>(op);
+  TF_ASSIGN_OR_RETURN(Layout params_layout,
+                      ExtractRequiredLayoutFromOperand(gather_op.getParams()));
+  TF_ASSIGN_OR_RETURN(Layout indices_layout,
+                      ExtractRequiredLayoutFromOperand(gather_op.getIndices()));
+  TF_ASSIGN_OR_RETURN(Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(gather_op));
+
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> indices_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(gather_op.getIndices()));
+  const int index_dimensions = indices_shape.back();
+
+  const auto params_rank = ValueRank(gather_op.getParams());
+  const auto indices_rank = ValueRank(gather_op.getIndices());
+  if (params_rank == -1)
+    return absl::InvalidArgumentError("missing rank for params input");
+  if (indices_rank == -1)
+    return absl::InvalidArgumentError("missing rank for indices input");
+
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Note that there may be conflicts between the two input layouts.
+  // To resolve these we:
+  // 1) Compute output layout from the current input layouts. This ignores any
+  //    sharding on the 'lookup' axis in params and the last axis in indices.
+  //    If params and indices are sharded on same dimension, the indices will
+  //    be unsharded in that dimension.
+  // 2) Merge any new sharding from the final output layout into this layout.
+  //    E.g. if the two inputs are unsharded, but the output is, this will be
+  //    more efficient as the local computations will be smaller since we will
+  //    more finely shard the inputs.
+  // 3) Finally, we use this new output layout to compute what input layouts we
+  //    need to get the given output.
+
+  // Step 1)
+  TF_ASSIGN_OR_RETURN(const Layout pre_output_layout,
+                      GatherNdGetOutputLayoutFromInput(
+                          params_layout, params_rank, indices_layout,
+                          indices_rank, index_dimensions, mesh));
+
+  // Step 2)
+  llvm::DenseSet<llvm::StringRef> used_dimensions;
+  for (const auto& spec : pre_output_layout.sharding_spec_strs())
+    if (Layout::IsShardedDimension(spec)) used_dimensions.insert(spec);
+
+  std::vector<std::string> sharding_specs(output_layout.rank());
+  for (int i = 0; i < sharding_specs.size(); ++i) {
+    if (Layout::IsShardedDimension(pre_output_layout.sharding_spec(i)))
+      sharding_specs[i] = pre_output_layout.sharding_spec(i);
+    // Merge in sharded dimensions from the output which aren't already used
+    // by the pre_output_layout.
+    else if (Layout::IsShardedDimension(output_layout.sharding_spec(i)) &&
+             !used_dimensions.contains(output_layout.sharding_spec(i)))
+      sharding_specs[i] = output_layout.sharding_spec(i);
+    else
+      sharding_specs[i] = Layout::kUnshardedDim;
+  }
+
+  // Step 3)
+  TF_ASSIGN_OR_RETURN(
+      const Layout merged_output_layout,
+      Layout::GetLayout(sharding_specs, pre_output_layout.mesh()));
+
+  Layout new_params_layout;
+  Layout new_indices_layout;
+  TF_RETURN_IF_ERROR(GatherNdGetInputLayoutFromOutput(
+      merged_output_layout, &new_params_layout, params_rank,
+      &new_indices_layout, indices_rank, index_dimensions, mesh));
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value params,
+      EmitRelayout(gather_op.getParams(), params_layout, new_params_layout));
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value indices,
+      EmitRelayout(gather_op.getIndices(), indices_layout, new_indices_layout));
+
+  mlir::OpBuilder builder(op);
+
+  // mlir::TF::GatherNdOp does not have a builder that does shape inference
+  // so we have to provided the output type. This output type is incorrect and
+  // we manually run shape inference after.
+  mlir::TF::GatherNdOp new_gather = builder.create<mlir::TF::GatherNdOp>(
+      op->getLoc(), gather_op.getOutput().getType(), params, indices);
+  InferSPMDExpandedLocalShape(new_gather);
+
+  TF_ASSIGN_OR_RETURN(mlir::Value output,
+                      EmitRelayout(new_gather.getOutput(), merged_output_layout,
+                                   output_layout));
+
+  gather_op.getOutput().replaceAllUsesWith(output);
+  gather_op.erase();
+  return output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+GatherNdSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  auto gather_op = llvm::cast<mlir::TF::GatherNdOp>(op);
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  std::optional<Layout> params_layout;
+  if (input_layouts.find(0) != input_layouts.end())
+    params_layout.emplace(input_layouts.lookup(0));
+  std::optional<Layout> indices_layout;
+  if (input_layouts.find(1) != input_layouts.end())
+    indices_layout.emplace(input_layouts.lookup(1));
+
+  TF_ASSIGN_OR_RETURN(llvm::ArrayRef<int64_t> indices_shape,
+                      ExtractGlobalInputShape(op->getOpOperand(1)));
+  const int index_dimensions = indices_shape.back();
+  if (index_dimensions < 0)
+    return absl::UnimplementedError(
+        "dynamic last dimension for index is not supported");
+
+  const int params_rank = ValueRank(gather_op.getParams());
+  const int indices_rank = ValueRank(gather_op.getIndices());
+  if (params_rank == -1)
+    return absl::InvalidArgumentError("missing rank for params input");
+  if (indices_rank == -1)
+    return absl::InvalidArgumentError("missing rank for indices input");
+
+  if (params_layout || indices_layout) {
+    TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                        GatherNdGetOutputLayoutFromInput(
+                            params_layout, params_rank, indices_layout,
+                            indices_rank, index_dimensions, mesh));
+    return llvm::DenseMap<int, Layout>({{0, output_layout}});
+  }
+
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+GatherNdSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If there is no output layout present then do not infer any operand layouts.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  auto gather_op = llvm::cast<mlir::TF::GatherNdOp>(op);
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  TF_ASSIGN_OR_RETURN(llvm::ArrayRef<int64_t> indices_shape,
+                      ExtractGlobalInputShape(op->getOpOperand(1)));
+  const int index_dimensions = indices_shape.back();
+  if (index_dimensions < 0)
+    return absl::UnimplementedError(
+        "dynamic last dimension for index is not supported");
+
+  const int params_rank = ValueRank(gather_op.getParams());
+  const int indices_rank = ValueRank(gather_op.getIndices());
+  if (params_rank == -1)
+    return absl::InvalidArgumentError("missing rank for params input");
+  if (indices_rank == -1)
+    return absl::InvalidArgumentError("missing rank for indices input");
+
+  const Layout output_layout = output_layouts.lookup(0);
+  Layout params_layout, indices_layout;
+  TF_RETURN_IF_ERROR(GatherNdGetInputLayoutFromOutput(
+      output_layout, &params_layout, params_rank, &indices_layout, indices_rank,
+      index_dimensions, mesh));
+  return llvm::DenseMap<int, Layout>({
+      {0, params_layout},
+      {1, indices_layout},
+  });
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h

@@ -0,0 +1,247 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_GATHER_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_GATHER_SPMD_EXPANDER_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class GatherCommonSPMDExpander : public SPMDExpanderBase {
+ public:
+  template <typename OpType>
+  StatusOr<mlir::Operation*> ExpandOpHelper(mlir::Operation* op) {
+    TF_ASSIGN_OR_RETURN(int64_t axis, GetAxis(op));
+    TF_ASSIGN_OR_RETURN(uint64_t batch_dims, GetBatchDim(op));
+
+    mlir::Value params = op->getOperand(0);
+    mlir::Value indices = op->getOperand(1);
+
+    TF_ASSIGN_OR_RETURN(std::vector<Layout> operand_layouts,
+                        ExtractRequiredLayoutFromOperands(op));
+    TF_ASSIGN_OR_RETURN(const Layout& output_layout,
+                        ExtractRequiredSingleLayoutFromOp(op));
+
+    const Layout& params_layout = operand_layouts[0];
+    const Layout& indices_layout = operand_layouts[1];
+
+    const int params_rank = ValueRank(params);
+    const int indices_rank = ValueRank(indices);
+
+    if (params_rank == -1)
+      return absl::InvalidArgumentError("Missing rank for params input.");
+    if (indices_rank == -1)
+      return absl::InvalidArgumentError("Missing rank for indices input.");
+
+    // Handle the case of negative axis.
+    if (axis < 0) axis += params_rank;
+    if (batch_dims < 0) batch_dims += indices_rank;
+    mlir::OpBuilder builder(op);
+
+    // Step 1: If the params are sharded on dim axis, an unconditional
+    // all-concat is generated. Alternatively, we could do: all-concating
+    // indices, followed by tf.Gather + slicing with correct masks.
+    //
+    // Currently we only support the case that the output layout matching the
+    // params layout for all non-axis dim. Other cases needs either a slicing or
+    // all-concat, which can be added later.
+    {
+      Mesh mesh = params_layout.mesh();
+      std::vector<std::string> tgt_params_sharding_specs;
+      tgt_params_sharding_specs.reserve(params_rank);
+      // check the first half
+      for (int i = 0; i < axis; ++i) {
+        const std::string& dim_name = params_layout.sharding_spec(i);
+        if (dim_name != output_layout.sharding_spec(i)) {
+          return absl::InvalidArgumentError(
+              llvm::formatv(
+                  "input and output layout do not agree on non-axis dim {0}. "
+                  "\n  params: {1}\n  output: {2}, axis: {3}",
+                  i, params_layout.ToString(), output_layout.ToString(), axis)
+                  .str());
+        }
+        tgt_params_sharding_specs.push_back(dim_name);
+      }
+      // Set replicated for `axis` dim.
+      tgt_params_sharding_specs.push_back(Layout::kUnshardedDim);
+      // Check the second half
+      for (int i = axis + 1; i < params_rank; ++i) {
+        const std::string& dim_name = params_layout.sharding_spec(i);
+        // To align the param dim with output, we can think we insert
+        // indices_rank
+        // - batch_dims dims from indices and remove one from param (axis), so
+        // the shifting is indices_rank - batch_dims - 1.
+        if (dim_name !=
+            output_layout.sharding_spec(i + indices_rank - batch_dims - 1)) {
+          return absl::InvalidArgumentError(
+              llvm::formatv(
+                  "input and output layout do not agree on non-axis dim {0}. "
+                  "\n  params: {1}\n  output: {2}, axis: {3}",
+                  i, params_layout.ToString(), output_layout.ToString(), axis)
+                  .str());
+        }
+        tgt_params_sharding_specs.push_back(dim_name);
+      }
+
+      if (!Layout::IsUnshardedDimension(params_layout.sharding_spec(axis))) {
+        if (llvm::isa<mlir::TF::ResourceGatherOp>(op)) {
+          return absl::InvalidArgumentError(absl::StrCat(
+              "DTensor does not support sharded 0th dimension for the resource "
+              "tensor for ResourceGatherOp. Please unshard dimension ",
+              axis));
+        }
+        TF_ASSIGN_OR_RETURN(Layout tgt_params_layout,
+                            Layout::GetLayout(params_layout.type(),
+                                              tgt_params_sharding_specs, mesh));
+        TF_ASSIGN_OR_RETURN(
+            params,
+            EmitAllGather(builder, params, params_layout, tgt_params_layout));
+      }
+    }
+
+    // Step 2: Check the output layout. If it requires all-relayouting indices.
+    // Do it.
+    //
+    // Indices shape is not big typically. Relayouting is expected to be cheap.
+    {
+      bool indices_relayout_needed = false;
+      Mesh mesh = output_layout.mesh();
+      std::vector<std::string> tgt_indices_sharding_specs;
+      tgt_indices_sharding_specs.reserve(indices_rank);
+      for (int i = 0; i < indices_rank; ++i) {
+        int index_in_output;
+        int index_in_indices;
+        if (i < batch_dims) {
+          // For dim within batch_dims, indices dim is aligning at the same
+          // index as output.
+          index_in_output = i;
+          index_in_indices = i;
+        } else {
+          // For dim after batch_dims, we can remove batch_dims from outputs and
+          // indices first, i.e., (i - batch_dims), add axis back, i.e., axis -
+          // batch_dims, and then put batch_dims back, so the target position in
+          // output is
+          //
+          //   i - batch_dims + axis - batch_dims + batch_dims
+          //
+          // which is as follows:
+          index_in_output = i + axis - batch_dims;
+          index_in_indices = i;
+        }
+        tgt_indices_sharding_specs.push_back(
+            output_layout.sharding_spec(index_in_output));
+
+        if (output_layout.sharding_spec(index_in_output) !=
+            indices_layout.sharding_spec(index_in_indices)) {
+          indices_relayout_needed = true;
+        }
+      }
+
+      if (indices_relayout_needed) {
+        TF_ASSIGN_OR_RETURN(
+            Layout tgt_indices_layout,
+            Layout::GetLayout(indices_layout.type(), tgt_indices_sharding_specs,
+                              mesh));
+        TF_ASSIGN_OR_RETURN(
+            indices, EmitRelayout(indices, indices_layout, tgt_indices_layout));
+      }
+    }
+
+    auto new_operands = llvm::to_vector<4>(op->getOperands());
+    new_operands[0] = params;
+    new_operands[1] = indices;
+
+    mlir::Operation* new_gather =
+        OpType::create(builder, op->getLoc(), op->getResultTypes(),
+                       mlir::ValueRange(new_operands), op->getAttrs())
+            .getOperation();
+
+    op->getResult(0).replaceAllUsesWith(new_gather->getResult(0));
+    op->erase();
+
+    return InferSPMDExpandedLocalShape(new_gather);
+  }
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+  virtual StatusOr<int64_t> GetAxis(mlir::Operation* op) = 0;
+  virtual StatusOr<uint64_t> GetBatchDim(mlir::Operation* op) = 0;
+};
+
+class GatherV2SPMDExpander : public GatherCommonSPMDExpander {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<int64_t> GetAxis(mlir::Operation* op) override;
+
+  StatusOr<uint64_t> GetBatchDim(mlir::Operation* op) override;
+};
+
+class ResourceGatherSPMDExpander : public GatherCommonSPMDExpander {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<int64_t> GetAxis(mlir::Operation* op) override;
+
+  StatusOr<uint64_t> GetBatchDim(mlir::Operation* op) override;
+};
+
+class GatherNdSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_GATHER_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.cc

@@ -0,0 +1,98 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> IdentityNSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto layouts, ExtractLayoutFromOp(op));
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  llvm::SmallVector<mlir::Value, 4> generated_outputs;
+  llvm::SmallVector<mlir::Type, 4> generated_types;
+
+  mlir::OpBuilder builder(op);
+  // Track the op that comes last after splitting.
+  mlir::Operation* last_op_after_splitting = op;
+  for (int i = 0; i < layouts.size(); ++i) {
+    auto output_layout = layouts[i];
+    if (!output_layout)
+      return absl::InvalidArgumentError(absl::StrCat(
+          "layout of (", i,
+          "-th output of IdentityNOp must be known before SPMD expansion."));
+
+    TF_ASSIGN_OR_RETURN(auto operand_layout,
+                        ExtractLayoutFromOperand(op->getOperand(i)));
+    if (!operand_layout)
+      return absl::InvalidArgumentError(absl::StrCat(
+          "layout of (", i,
+          "-th input of IdentityNOp must be known before SPMD expansion."));
+
+    TF_ASSIGN_OR_RETURN(const mlir::Value output,
+                        EmitRelayout(op->getOperand(i), *operand_layout,
+                                     *output_layout, &newly_created_ops));
+    generated_outputs.emplace_back(output);
+    generated_types.emplace_back(output.getType());
+    // InsertionPoint has to come after all newly created Ops.
+    if (last_op_after_splitting->isBeforeInBlock(output.getDefiningOp())) {
+      last_op_after_splitting = output.getDefiningOp();
+    }
+  }
+
+  builder.setInsertionPointAfter(last_op_after_splitting);
+  auto identity_op = mlir::TF::IdentityNOp::create(
+      builder, op->getLoc(), generated_types, generated_outputs);
+
+  for (int i = 0; i < layouts.size(); ++i)
+    op->getOpResult(i).replaceAllUsesExcept(identity_op.getResult(i),
+                                            newly_created_ops);
+
+  op->erase();
+
+  return InferSPMDExpandedLocalShape(identity_op.getOperation());
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IdentityNSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  return input_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IdentityNSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  return output_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IDENTITY_N_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IDENTITY_N_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// SPMD Expander for IdentityNOp.
+class IdentityNSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IDENTITY_N_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/in_top_k_spmd_expander.cc

@@ -0,0 +1,200 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/in_top_k_spmd_expander.h"
+
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+// Creates a new layout for the "predictions" input based on the provided
+// layout, ensuring that the 2nd dimension is replicated.
+StatusOr<Layout> GetSuggestedPredictionsLayout(const Layout& layout) {
+  // predictions is a rank-2 tensor (batch_size x num_classes)
+  std::vector<std::string> layout_specs(2);
+  layout_specs[0] = layout.sharding_spec(0);
+  layout_specs[1] = Layout::kUnshardedDim;
+
+  return Layout::GetLayout(layout_specs, layout.mesh());
+}
+
+// Creates a new layout based on the input "layout" but matching the batch dim
+// of "other_layout".
+StatusOr<Layout> MatchBatchDim(const Layout& layout,
+                               const Layout& other_layout) {
+  std::vector<std::string> layout_specs(layout.rank());
+  layout_specs[0] = other_layout.sharding_spec(0);
+  for (int i = 1; i < layout.rank(); ++i) {
+    layout_specs[i] = layout.sharding_spec(i);
+  }
+
+  return Layout::GetLayout(layout_specs, layout.mesh());
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> InTopKSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto in_top_k_op = mlir::cast<mlir::TF::InTopKV2Op>(op);
+
+  mlir::Value predictions = in_top_k_op.getPredictions();
+  TF_ASSIGN_OR_RETURN(const Layout predictions_layout,
+                      ExtractRequiredLayoutFromOperand(predictions));
+  mlir::Value targets = in_top_k_op.getTargets();
+  TF_ASSIGN_OR_RETURN(const Layout targets_layout,
+                      ExtractRequiredLayoutFromOperand(targets));
+
+  TF_ASSIGN_OR_RETURN(const Layout precision_layout,
+                      ExtractRequiredSingleLayoutFromOp(in_top_k_op));
+
+  bool relayout_predictions = false;
+  bool relayout_targets = false;
+  Layout new_predictions_layout = predictions_layout;
+  Layout new_targets_layout = targets_layout;
+
+  // Last dim (classes) of "predictions" is required to be replicated.
+  if (!predictions_layout.IsLastDimReplicated()) {
+    TF_ASSIGN_OR_RETURN(new_predictions_layout,
+                        GetSuggestedPredictionsLayout(new_predictions_layout));
+    relayout_predictions = true;
+  }
+
+  // If the batch dims of "targets" and "predictions" don't match, relayout the
+  // less-sharded input to match the sharding of the other.
+  const std::string& predictions_batch_dim =
+      new_predictions_layout.sharding_spec(0);
+  const std::string& targets_batch_dim = new_targets_layout.sharding_spec(0);
+  if (predictions_batch_dim != targets_batch_dim) {
+    if (Layout::IsShardedDimension(targets_batch_dim) &&
+        Layout::IsShardedDimension(predictions_batch_dim)) {
+      // Since "targets" is the smaller tensor, relayout it to match
+      // "predictions".
+      TF_ASSIGN_OR_RETURN(
+          new_targets_layout,
+          MatchBatchDim(new_targets_layout, new_predictions_layout));
+      relayout_targets = true;
+    } else if (Layout::IsShardedDimension(targets_batch_dim)) {
+      TF_ASSIGN_OR_RETURN(
+          new_predictions_layout,
+          MatchBatchDim(new_predictions_layout, new_targets_layout));
+      relayout_predictions = true;
+    } else if (Layout::IsShardedDimension(predictions_batch_dim)) {
+      TF_ASSIGN_OR_RETURN(
+          new_targets_layout,
+          MatchBatchDim(new_targets_layout, new_predictions_layout));
+      relayout_targets = true;
+    }
+  }
+
+  mlir::OpBuilder builder(op);
+  mlir::IRMapping mapping;
+  // Apply any input relayouts.
+  if (relayout_predictions) {
+    TF_ASSIGN_OR_RETURN(
+        mlir::Value new_predictions,
+        EmitRelayout(predictions, predictions_layout, new_predictions_layout));
+    mapping.map(op->getOperand(0), new_predictions);
+  }
+  if (relayout_targets) {
+    TF_ASSIGN_OR_RETURN(
+        mlir::Value new_targets,
+        EmitRelayout(targets, targets_layout, new_targets_layout));
+    mapping.map(op->getOperand(1), new_targets);
+  }
+
+  mlir::Operation* new_op = builder.clone(*op, mapping);
+  new_op = InferSPMDExpandedLocalShape(new_op);
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_precision,
+                      EmitRelayout(new_op->getOpResult(0), new_targets_layout,
+                                   precision_layout));
+  op->getResult(0).replaceAllUsesWith(new_precision);
+  op->erase();
+
+  return new_precision.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> InTopKSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // output is a rank-1 tensor (batch_size,)
+  std::vector<std::string> layout_specs(1);
+  layout_specs[0] = Layout::kUnshardedDim;
+
+  // Set the output (precision) layout based on the inputs predictions and
+  // targets. Prefer whichever is more sharded on the batch dimension.
+  // Since "targets" is the smaller tensor, match its batch dimension first,
+  // and override it with the batch dimension of "predictions" if that is
+  // (possibly differently) sharded.
+  if (input_layouts.find(1) != input_layouts.end()) {
+    const Layout& targets_layout = input_layouts.lookup(1);
+    const std::string& targets_batch_dim = targets_layout.sharding_spec(0);
+    if (Layout::IsShardedDimension(targets_batch_dim)) {
+      layout_specs[0] = targets_batch_dim;
+    }
+  }
+  if (input_layouts.find(0) != input_layouts.end()) {
+    const Layout& predictions_layout = input_layouts.lookup(0);
+    const std::string& predictions_batch_dim =
+        predictions_layout.sharding_spec(0);
+    if (Layout::IsShardedDimension(predictions_batch_dim)) {
+      layout_specs[0] = predictions_batch_dim;
+    }
+  }
+
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      Layout::GetLayout(layout_specs, mesh));
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> InTopKSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If no output layout is present then do not infer any operand layouts.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout& output_layout = output_layouts.lookup(0);
+
+  // Set the layouts of predictions and targets based on the output. Last
+  // dimension of predictions needs to be replicated.
+  TF_ASSIGN_OR_RETURN(const Layout predictions_layout,
+                      GetSuggestedPredictionsLayout(output_layout));
+  const Layout targets_layout = output_layout;
+  const Layout k_layout =
+      Layout::ReplicatedOnMesh(output_layout.mesh(), /*rank=*/0);
+
+  return llvm::DenseMap<int, Layout>(
+      {{0, predictions_layout}, {1, targets_layout}, {2, k_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/in_top_k_spmd_expander.h

@@ -0,0 +1,46 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IN_TOP_K_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IN_TOP_K_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Implements SPMD Expansion for the TF::InTopKV2 op. See base class
+// SPMDExpanderBase for function documentation.
+class InTopKSPMDExpander final : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IN_TOP_K_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.cc

@@ -0,0 +1,175 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.h"
+
+#include <algorithm>
+#include <vector>
+
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/device_utils.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+template <typename T>
+StatusOr<mlir::Operation*> Expand(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const std::vector<Layout> output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(const std::vector<Layout> operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+  if (!AllReplicated(output_layouts) || !AllReplicated(operand_layouts)) {
+    return absl::UnimplementedError(
+        llvm::formatv("Expecting {0} to have input and output layouts to be "
+                      "fully replicated but was not. ",
+                      OpName(op))
+            .str());
+  }
+
+  // Build an if op that only runs the op on device 0. Every other device
+  // will run a no-op.
+  mlir::ModuleOp module = op->getParentOfType<mlir::ModuleOp>();
+  mlir::SymbolTable symbol_table(module);
+  mlir::Location location = op->getLoc();
+  mlir::OpBuilder builder(op);
+
+  auto func_type =
+      mlir::FunctionType::get(builder.getContext(), op->getOperandTypes(),
+                              llvm::ArrayRef<mlir::Type>{});
+  // Build then_func that is the branch of device_id != 0, which only contains a
+  // single NoOp.
+  mlir::func::FuncOp then_func = mlir::func::FuncOp::create(
+      location,
+      llvm::formatv("{0}_then_func_{1}", OpName(op), OpHash(op)).str(),
+      func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+  // Set function visibility to private to indicate that it is only used in
+  // this module.
+  then_func.setVisibility(mlir::SymbolTable::Visibility::Private);
+  mlir::Block* then_fn_block = then_func.addEntryBlock();
+  mlir::OpBuilder then_fn_builder =
+      mlir::OpBuilder::atBlockBegin(then_fn_block);
+  mlir::TF::NoOp::create(then_fn_builder, location);
+  mlir::func::ReturnOp::create(then_fn_builder, location);
+
+  // Build else_func that is the branch of device_id == 0.
+  // The else func is just the original op.
+  mlir::func::FuncOp else_func = mlir::func::FuncOp::create(
+      location,
+      llvm::formatv("{0}_else_func_{1}", OpName(op), OpHash(op)).str(),
+      func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+  // Set function visibility to private to indicate that it is only used in
+  // this module.
+  else_func.setVisibility(mlir::SymbolTable::Visibility::Private);
+
+  mlir::Block* else_fn_block = else_func.addEntryBlock();
+  mlir::OpBuilder else_fn_builder =
+      mlir::OpBuilder::atBlockBegin(else_fn_block);
+
+  T::create(else_fn_builder, location, op->getResultTypes(),
+            else_fn_block->getArguments());
+  mlir::func::ReturnOp::create(else_fn_builder, location);
+
+  symbol_table.insert(then_func);
+  symbol_table.insert(else_func);
+
+  TF_ASSIGN_OR_RETURN(mlir::Value device_id, DeviceId(op));
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value zero_scalar,
+      CreateZeroScalarConst(
+          builder, location,
+          mlir::cast<mlir::TensorType>(device_id.getType()).getElementType()));
+
+  mlir::TF::NotEqualOp not_equal = mlir::TF::NotEqualOp::create(
+      builder, location, device_id, zero_scalar,
+      /*incompatible_shape_error=*/builder.getBoolAttr(false));
+
+  mlir::Operation* if_op = mlir::TF::IfOp::create(
+      builder, location, then_func.getFunctionType().getResults(),
+      /*cond=*/not_equal.getResult(),
+      /*input=*/op->getOperands(),
+      /*then_branch=*/then_func.getSymName(),
+      /*else_branch=*/else_func.getSymName(), /*is_stateless=*/false);
+
+  op->replaceAllUsesWith(if_op);
+  op->erase();
+  return if_op;
+}
+
+}  // namespace
+StatusOr<mlir::Operation*> IOOpSPMDExpander::ExpandOp(mlir::Operation* op) {
+  if (llvm::isa<mlir::TF::WriteSummaryOp>(op)) {
+    return Expand<mlir::TF::WriteSummaryOp>(op);
+  } else if (llvm::isa<mlir::TF::FlushSummaryWriterOp>(op)) {
+    return Expand<mlir::TF::FlushSummaryWriterOp>(op);
+  }
+  return absl::UnimplementedError(
+      llvm::formatv("SPMD for op : {0} is not implemented ", OpName(op)).str());
+}
+
+// Always return a set of replicated layouts for now. If there is a case where
+// a dtensor user is writing a large tensor that is sharded, then, we can
+// support that in the future.
+StatusOr<llvm::DenseMap<int, Layout>> IOOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    output_layouts[i] =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getResult(i)));
+  }
+  return output_layouts;
+}
+
+// Always return a set of replicated layouts. IO ops usually either have
+// no output or a scalar output, in which case it is replicated.
+StatusOr<llvm::DenseMap<int, Layout>> IOOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands());
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    int rank = std::max(0, ValueRank(op->getOperand(i)));
+    input_layouts[i] = Layout::ReplicatedOnMesh(mesh, rank);
+  }
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IO_OP_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IO_OP_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// SPMD Expander for IO ops, such as tf.WriteSummary.
+class IOOpSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_IO_OP_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.cc

@@ -0,0 +1,151 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/constants.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/dtensor_location.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> IteratorGetNextSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  mlir::TF::IteratorGetNextOp original_op =
+      mlir::cast<mlir::TF::IteratorGetNextOp>(op);
+  mlir::OpBuilder builder(op);
+
+  TF_ASSIGN_OR_RETURN(std::vector<Layout> output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+
+  llvm::SmallVector<mlir::Type, 4> local_types(original_op->getNumResults());
+
+  for (int i = 0; i < original_op->getNumResults(); ++i) {
+    mlir::TensorType global_output_type =
+        mlir::cast<mlir::TensorType>(original_op.getResult(i).getType());
+    std::vector<int64_t> local_shape =
+        output_layouts[i].LocalShapeFromGlobalShape(
+            global_output_type.getShape());
+    local_types[i] = mlir::RankedTensorType::get(
+        local_shape, global_output_type.getElementType());
+  }
+
+  auto new_op = mlir::TF::IteratorGetNextOp::create(
+      builder, DT_LOC(op->getLoc()), local_types, original_op->getOperand(0));
+
+  for (int i = 0; i < original_op->getNumResults(); ++i) {
+    original_op.getResult(i).replaceAllUsesWith(new_op.getResult(i));
+  }
+  original_op.erase();
+  return InferSPMDExpandedLocalShape(new_op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IteratorGetNextSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // Extract the output element layouts from the `tf._element_layouts` attribute
+  // of the iterator resource tensor.
+  TF_ASSIGN_OR_RETURN(const auto layouts,
+                      ExtractElementLayoutsFromOperand(op->getOpOperand(0)));
+
+  llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    output_layouts[i] = layouts[i];
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IteratorGetNextSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Iterator resource tensors are always 0-dimensional.
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, /*rank=*/0)}});
+}
+
+StatusOr<mlir::Operation*> IteratorGetNextAsOptionalSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  // Extract the output element layouts from the `tf._element_layouts` attribute
+  // of the iterator resource tensor.
+  TF_ASSIGN_OR_RETURN(const auto output_layouts,
+                      ExtractElementLayoutsFromOperand(op->getOpOperand(0)));
+
+  auto array_attr = op->getAttrOfType<mlir::ArrayAttr>(kIteratorOutputShapes);
+  if (!array_attr)
+    return absl::InvalidArgumentError(
+        llvm::formatv("Could not find `{0}` attribute of op: {1}",
+                      kIteratorOutputShapes, op->getName())
+            .str());
+
+  llvm::SmallVector<mlir::Attribute, 4> output_shape_attrs(array_attr.size());
+  for (int i = 0; i < array_attr.size(); ++i) {
+    std::vector<int64_t> local_shape =
+        output_layouts[i].LocalShapeFromGlobalShape(
+            mlir::cast<mlir::TF::ShapeAttr>(array_attr[i]).getShape());
+    output_shape_attrs[i] = mlir::cast<mlir::Attribute>(
+        mlir::TF::ShapeAttr::get(op->getContext(), {local_shape}));
+  }
+
+  // Update the `output_shapes` attribute on the op to match the local shape
+  // based on the iterator element layouts.
+  op->setAttr(kIteratorOutputShapes,
+              mlir::ArrayAttr::get(op->getContext(), output_shape_attrs));
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IteratorGetNextAsOptionalSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Variant tensors are always 0-dimensional.
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, /*rank=*/0)}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+IteratorGetNextAsOptionalSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Iterator resource tensors are always 0-dimensional.
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, /*rank=*/0)}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.h

@@ -0,0 +1,57 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ITERATOR_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ITERATOR_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class IteratorGetNextSPMDExpander final : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class IteratorGetNextAsOptionalSPMDExpander final : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_ITERATOR_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.cc

@@ -0,0 +1,476 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+void GetTransposeSettings(mlir::Operation* op, bool* left_transposed,
+                          bool* right_transposed) {
+  if (mlir::isa<mlir::TF::BatchMatMulV2Op>(op)) {
+    mlir::TF::BatchMatMulV2Op mm = mlir::cast<mlir::TF::BatchMatMulV2Op>(op);
+    // Adjoint is just conjugate transpose.
+    *left_transposed = mm.getAdjX();
+    *right_transposed = mm.getAdjY();
+  } else if (mlir::isa<mlir::TF::MatMulOp>(op)) {
+    mlir::TF::MatMulOp mm = mlir::cast<mlir::TF::MatMulOp>(op);
+    *left_transposed = mm.getTransposeA();
+    *right_transposed = mm.getTransposeB();
+  }
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> MatMulSPMDExpander::ExpandOp(mlir::Operation* op) {
+  bool left_transposed;
+  bool right_transposed;
+  TF_ASSIGN_OR_RETURN(const Layout left_layout,
+                      ExtractRequiredLayoutFromOperand(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const Layout right_layout,
+                      ExtractRequiredLayoutFromOperand(op->getOperand(1)));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+
+  GetTransposeSettings(op, &left_transposed, &right_transposed);
+
+  std::string reduce_dim;
+  Layout layout_after_matmul;
+  mlir::Value left, right;
+
+  TF_RETURN_IF_ERROR(MaybeRelayoutInputs(
+      op, left_layout, left_transposed, right_layout, right_transposed,
+      output_layout, reduce_dim, layout_after_matmul, left, right));
+
+  mlir::OpBuilder builder(op);
+
+  mlir::IRMapping mapping;
+  mapping.map(op->getOperand(0), left);
+  mapping.map(op->getOperand(1), right);
+  mlir::Operation* new_op = builder.clone(*op, mapping);
+  // Note that the output shape of new_op is cloned from op, so we need to
+  // update to the local shape.
+  new_op = InferSPMDExpandedLocalShape(new_op);
+
+  if (Layout::IsShardedDimension(reduce_dim)) {
+    TF_ASSIGN_OR_RETURN(
+        new_op, EmitAllReduce(builder, layout_after_matmul, {reduce_dim},
+                              new_op, kReduceOpAdd));
+  }
+
+  TF_ASSIGN_OR_RETURN(
+      auto final_output,
+      EmitRelayout(new_op->getOpResult(0), layout_after_matmul, output_layout));
+
+  op->getOpResult(0).replaceAllUsesWith(final_output);
+  op->erase();
+
+  return final_output.getDefiningOp();
+}
+
+StatusOr<Layout> MatMulSPMDExpander::OutputLayoutAndReducedDims(
+    bool allow_unknown_layouts, mlir::Operation* op,
+    absl::flat_hash_set<std::string>* reduced_dims, std::optional<Layout>* left,
+    std::optional<Layout>* right) {
+  // These layouts are 2d layouts for the non-batch dimensions.
+  Layout left_layout;
+  Layout right_layout;
+  bool left_transposed;
+  bool right_transposed;
+
+  // This will hold the batch layout for the output.
+  Layout batch_layout;
+
+  if (!*left || !*right) {
+    if (allow_unknown_layouts) return absl::OkStatus();
+    return absl::UnimplementedError(
+        absl::StrCat("failed to do SPMD expansion for ", OpName(op),
+                     " operand layouts "
+                     "unknown"));
+  }
+
+  if (mlir::isa<mlir::TF::BatchMatMulV2Op>(op)) {
+    mlir::TF::BatchMatMulV2Op mm = mlir::cast<mlir::TF::BatchMatMulV2Op>(op);
+
+    // Note that it doesn't matter if we pass the global or local shape to
+    // GetBroadcastLayoutForElementWise, it will return the same result.
+    TF_ASSIGN_OR_RETURN(const auto left_shape, GetShapeOfValue(mm.getX()));
+    TF_ASSIGN_OR_RETURN(const auto right_shape, GetShapeOfValue(mm.getY()));
+    std::vector<std::string> left_splits;
+    std::vector<std::string> right_splits;
+    TF_ASSIGN_OR_RETURN(
+        batch_layout,
+        GetBroadcastLayoutForElementWise(
+            left->value(), right->value(), left_shape, right_shape,
+            /*dims_to_ignore=*/2, left_splits, right_splits));
+
+    left_layout = (*left)->Truncate(left_shape.size() - 2, /*end=*/true);
+    right_layout = (*right)->Truncate(right_shape.size() - 2, /*end=*/true);
+  } else if (mlir::isa<mlir::TF::MatMulOp>(op)) {
+    // There are no batch dims for MatMul op, so get an 'empty' layout that
+    // we can concat later.
+    batch_layout = (*left)->Truncate(/*split_point=*/0, /*end=*/false);
+    left_layout = left->value();
+    right_layout = right->value();
+  } else {
+    return absl::InternalError(absl::StrCat("Unknown op ", OpName(op)));
+  }
+  GetTransposeSettings(op, &left_transposed, &right_transposed);
+
+  if (left_transposed) {
+    TF_ASSIGN_OR_RETURN(left_layout, Layout::Transposed2D(left_layout));
+  }
+  if (right_transposed) {
+    TF_ASSIGN_OR_RETURN(right_layout, Layout::Transposed2D(right_layout));
+  }
+
+  // Input layouts are [batch...],a,b;[batch...],b,c
+  // Output layout is [batch...],a,c
+  const auto& batch_sharding_specs = batch_layout.sharding_spec_strs();
+  std::vector<std::string> output_dims(batch_sharding_specs.begin(),
+                                       batch_sharding_specs.end());
+  if (Layout::IsShardedDimension(left_layout.sharding_spec(0)) &&
+      left_layout.sharding_spec(0) == right_layout.sharding_spec(1)) {
+    // If a and c above are the same and sharded, we should output a replicated
+    // layout during propagation. This is so we don't create an illegal layout.
+    output_dims.resize(output_dims.size() + 2);
+    output_dims[output_dims.size() - 2] = Layout::kUnshardedDim;
+    output_dims[output_dims.size() - 1] = Layout::kUnshardedDim;
+  } else {
+    output_dims.emplace_back(left_layout.sharding_spec(0));
+    output_dims.emplace_back(right_layout.sharding_spec(1));
+  }
+
+  return Layout::GetLayout(output_dims, left_layout.mesh());
+}
+
+// This function will take the left and right input, possibly slice or add
+// AllConcat along various mesh dimensions before the MatMul operation takes
+// place. This also returns:
+// * The mesh dimension, if any, that the output of the matmul should be
+//   summed along.
+// * The resulting layout of the matmul tensor, so we can insert an AllConcat/
+//   split to make the output have the desired layout.
+// * The left and right value for use as input to the matmul.
+absl::Status MatMulSPMDExpander::MaybeRelayoutInputs(
+    mlir::Operation* op, const Layout& left_layout, bool left_transposed,
+    const Layout& right_layout, bool right_transposed,
+    const Layout& output_layout, std::string& reduced_dim,
+    Layout& matmul_layout, mlir::Value& left, mlir::Value& right) {
+  // These two lists will contain the mesh dimensions desired for the left
+  // and right inputs before the matmul. Since a Layout is generally immutable,
+  // we use these vectors to store sharding for the layout and produce the
+  // final layout at the end via Layout::GetLayout.
+  std::vector<std::string> left_specs = left_layout.sharding_spec_strs();
+  std::vector<std::string> right_specs = right_layout.sharding_spec_strs();
+  // Specs for the layout of the matmul.
+  std::vector<std::string> matmul_specs(output_layout.rank());
+
+  TF_ASSIGN_OR_RETURN(const std::vector<int64_t> left_shape,
+                      GetShapeOfValue(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const std::vector<int64_t> right_shape,
+                      GetShapeOfValue(op->getOperand(1)));
+
+  const std::vector<int64_t> left_global_shape =
+      left_layout.GlobalShapeFromLocalShape(left_shape);
+  const std::vector<int64_t> right_global_shape =
+      right_layout.GlobalShapeFromLocalShape(right_shape);
+
+  // From this point on, we will use a short hand in the comments for the
+  // last two dimensions of the inputs and output:
+  // left: a,b right: c,d output: e,f after left and right are appropriately
+  // transposed.
+
+  std::string& a = left_specs[left_specs.size() - 2];
+  std::string& b = left_specs[left_specs.size() - 1];
+  std::string& c = right_specs[right_specs.size() - 2];
+  std::string& d = right_specs[right_specs.size() - 1];
+  const std::string& e = output_layout.sharding_spec(output_layout.rank() - 2);
+  const std::string& f = output_layout.sharding_spec(output_layout.rank() - 1);
+
+  if (left_transposed) std::swap(a, b);
+  if (right_transposed) std::swap(c, d);
+
+  // Set the mesh dimensions along the batch axis for the left and right side
+  // from the output. This is relatively simple choice and there are cases that
+  // we could improve:
+  // - Left and right input are both sharded on a dimension and the output
+  //   is not. With the current algorithm we will unshard the inputs. But it
+  //   would be more efficient to leave the inputs sharded and unshard the
+  //   output.
+  llvm::SmallSet<std::string, 4> used_mesh_dimensions;
+  used_mesh_dimensions.insert(a);
+  used_mesh_dimensions.insert(b);
+  used_mesh_dimensions.insert(c);
+  used_mesh_dimensions.insert(d);
+  for (int i = 0; i < matmul_specs.size() - 2; ++i) {
+    matmul_specs[i] = output_layout.sharding_spec(i);
+    if (used_mesh_dimensions.contains(matmul_specs[i]))
+      matmul_specs[i] = Layout::kUnshardedDim;
+    if (i >= matmul_specs.size() - left_specs.size()) {
+      const int64_t left_pos = left_specs.size() - matmul_specs.size() + i;
+      left_specs[left_pos] = matmul_specs[i];
+      // If the left global shape is 1, its broadcasted so just set the
+      // dimension to unsharded.
+      if (left_global_shape[left_pos] == 1)
+        left_specs[left_pos] = Layout::kUnshardedDim;
+    }
+    if (i >= matmul_specs.size() - right_specs.size()) {
+      const int64_t right_pos = right_specs.size() - matmul_specs.size() + i;
+      right_specs[right_pos] = matmul_specs[i];
+      // If the right global shape is 1, its broadcasted so just set the
+      // dimension to unsharded.
+      if (right_global_shape[right_pos] == 1)
+        right_specs[right_pos] = Layout::kUnshardedDim;
+    }
+  }
+
+  // Reject the cases that we don't yet support, namely the contracting
+  // dimensions are sharded not equal, or the input and output non-contracting
+  // dimensions are equal and are sharded. These would require more extensive
+  // relayout to solve.
+  if (b != c && Layout::IsShardedDimension(b) && Layout::IsShardedDimension(c))
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Contracting dimension for matmul has sharding dimension ", b,
+        " for the left input and ", c,
+        " for the right input which are not equal. This case is currently not "
+        "supported."));
+
+  if (a != e && Layout::IsShardedDimension(a) && Layout::IsShardedDimension(e))
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Non-contracting dimension for left argument of matmul has sharding "
+        "dimension ",
+        a,
+        " and the second to last dimension of the output has sharding "
+        "dimension ",
+        e, ", which are not equal. This case is currently not supported."));
+
+  if (d != f && Layout::IsShardedDimension(d) && Layout::IsShardedDimension(f))
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Non-contracting dimension for right argument of matmul has sharding "
+        "dimension ",
+        d, " and the last dimension of the output has sharding dimension ", f,
+        ", which are not equal. This case is currently not supported."));
+
+  // If the output is sharded and the corresponding non-contracting input is not
+  // sharded, then shard the input on that dim, to reduce the amount of work
+  // done. Note that this sharding spec can't be used anywhere in the batch
+  // dimensions due the agreement between the sharding specs of the batch
+  // dimensions for the input and output, so this is safe.
+  // This handles the *,x . x,* -> *,y case.
+  if (Layout::IsUnshardedDimension(a) && Layout::IsShardedDimension(e) &&
+      e != b && e != c && e != d)
+    a = e;
+  if (Layout::IsUnshardedDimension(d) && Layout::IsShardedDimension(f) &&
+      f != a && f != b && f != c)
+    d = f;
+
+  // Handle the case when the non-contracting dimensions have the same
+  // sharding spec. This can't happen if either of the previous two cases are
+  // true as it would imply that e and f have the same sharding spec. So, a and
+  // d are sharded in the input and we need to AllConcat one of them.
+  // This handles the y,x . x,y -> *,y case.
+  if (Layout::IsShardedDimension(a) && a == d) {
+    if (a == e)
+      d = Layout::kUnshardedDim;
+    else if (d == f)
+      a = Layout::kUnshardedDim;
+    else
+      // TODO(bfontain): Update this to pick a or d to AllConcat based on shape.
+      a = Layout::kUnshardedDim;
+  }
+
+  // Handle the case where a non-contracting and contracting dim have the same
+  // sharding spec. For now we always unshard the contracting axis. Note that
+  // this is safe since, e.g. if a = c and both are sharded, then b must be
+  // unsharded.
+  // This handles the case x,y . *,y -> x,y
+  // Beware:
+  // Consider *,y . *,y -> *,*, there are two choices, unshard b or unshard d.
+  // If we unshard d, then we will need to shard c and all reduce. This maybe
+  // a good idea for performance, but the current EmitAllGather cannot handle
+  // the transformation of *,y to y,*. Thus it is safer to always unshard b in
+  // this case.
+  if (Layout::IsShardedDimension(a) && a == c) c = Layout::kUnshardedDim;
+  if (Layout::IsShardedDimension(b) && b == d) b = Layout::kUnshardedDim;
+
+  // Finally, we make both contracting axes agree on a sharding.
+  // If b and c are sharded, we checked about that their sharding is equal.
+  // If c is sharded then a != c (if a == c, then the above case would set it
+  // to unsharded). And since c was never part of the batch dimensions (we
+  // specifically excluded it earlier), the dimension c is not used in the left
+  // input so it is always safe to set b = c.
+  if (b != c) {
+    if (Layout::IsShardedDimension(b))
+      c = b;
+    else
+      b = c;
+  }
+  reduced_dim = b;
+
+  // Generate the layout that will be the output of the matmul.
+  // This may be different from the final output layout in the last two
+  // dimensions.
+  matmul_specs[output_layout.rank() - 2] = a;
+  matmul_specs[output_layout.rank() - 1] = d;
+
+  TF_ASSIGN_OR_RETURN(matmul_layout,
+                      Layout::GetLayout(matmul_specs, output_layout.mesh()));
+  if (left_transposed) std::swap(a, b);
+  if (right_transposed) std::swap(c, d);
+
+  TF_ASSIGN_OR_RETURN(auto new_left_layout,
+                      Layout::GetLayout(left_specs, left_layout.mesh()));
+  TF_ASSIGN_OR_RETURN(auto new_right_layout,
+                      Layout::GetLayout(right_specs, right_layout.mesh()));
+
+  TF_ASSIGN_OR_RETURN(
+      left, EmitRelayout(op->getOperand(0), left_layout, new_left_layout));
+  TF_ASSIGN_OR_RETURN(
+      right, EmitRelayout(op->getOperand(1), right_layout, new_right_layout));
+
+  return absl::OkStatus();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> MatMulSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  if (input_layouts.empty()) return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  absl::flat_hash_set<std::string> reduced_dims;
+  TF_ASSIGN_OR_RETURN(const auto left_shape,
+                      GetShapeOfValue(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const auto right_shape,
+                      GetShapeOfValue(op->getOperand(1)));
+
+  // At least one input is set, calculate an output layout.
+  std::optional<Layout> left, right;
+  if (input_layouts.find(0) != input_layouts.end())
+    left.emplace(input_layouts.lookup(0));
+  else
+    left.emplace(Layout::ReplicatedOnMesh(mesh, left_shape.size()));
+  if (input_layouts.find(1) != input_layouts.end())
+    right.emplace(input_layouts.lookup(1));
+  else
+    right.emplace(Layout::ReplicatedOnMesh(mesh, right_shape.size()));
+
+  TF_ASSIGN_OR_RETURN(
+      const Layout output_layout,
+      OutputLayoutAndReducedDims(
+          /*allow_unknown_layouts=*/true, op, &reduced_dims, &left, &right));
+
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> MatMulSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout output_layout = output_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(const auto left_shape,
+                      GetShapeOfValue(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const auto right_shape,
+                      GetShapeOfValue(op->getOperand(1)));
+
+  // We take output layout and 'copy' it to the two input layouts, but with
+  // the contracting dimension set to replicated.
+  // Note that some complication are introduced by the possibility of
+  // broadcasting in the BatchMatMulV2 case.
+
+  // Truncate layout in case of broadcast. Note that since
+  // output->rank() == std::max(left_shape.size(), right_shape.size()) due to
+  // broadcasting one of these truncations is just a copy of output and the
+  // other may be shorter.
+  Layout left = output_layout.Truncate(output_layout.rank() - left_shape.size(),
+                                       /*end=*/true);
+  Layout right =
+      output_layout.Truncate(output_layout.rank() - right_shape.size(),
+                             /*end=*/true);
+
+  // Make sure necessary dimensions are replicated.
+  //
+  // Due to broadcasting, each of the batch dimensions (i.e. from dimension 0
+  // to dim - 2), one of the two inputs may have dimension 1 while the
+  // other has dimension > 1 and equal to the dim of the output. Since a
+  // tensor with dimension 1 cannot be sharded, we set this to unsharded.
+  auto specs_matmul_operands = [](const llvm::ArrayRef<int64_t>& tensor_shape,
+                                  const Layout& layout,
+                                  bool is_left_operand) -> StatusOr<Layout> {
+    int contracting_dim =
+        is_left_operand ? layout.rank() - 1 : layout.rank() - 2;
+    // Assign "any" to the contracting dim and "unsharded" to any tensor dim
+    // of length = 1.
+    std::vector<std::string> sharding_specs = layout.sharding_spec_strs();
+    for (size_t i = 0; i < layout.rank(); ++i)
+      if (i == contracting_dim) {
+        sharding_specs[i] = Layout::kAny;
+      } else if (tensor_shape[i] == 1) {
+        sharding_specs[i] = Layout::kUnshardedDim;
+      }
+    return Layout::GetLayout(sharding_specs, layout.mesh());
+  };
+
+  TF_ASSIGN_OR_RETURN(left, specs_matmul_operands(left_shape, left,
+                                                  /*is_left_operand=*/true));
+  TF_ASSIGN_OR_RETURN(right, specs_matmul_operands(right_shape, right,
+                                                   /*is_left_operand=*/false));
+
+  // Transpose the layouts if needed, as we just generated the non-transposed
+  // layouts.
+  bool left_transposed;
+  bool right_transposed;
+  GetTransposeSettings(op, &left_transposed, &right_transposed);
+  if (left_transposed) {
+    TF_ASSIGN_OR_RETURN(left, Layout::Transposed2D(left));
+  }
+  if (right_transposed) {
+    TF_ASSIGN_OR_RETURN(right, Layout::Transposed2D(right));
+  }
+
+  return llvm::DenseMap<int, Layout>({{0, left}, {1, right}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.h

@@ -0,0 +1,76 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_MATMUL_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_MATMUL_SPMD_EXPANDER_H_
+
+#include <optional>
+#include <string>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class MatMulSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ private:
+  StatusOr<Layout> OutputLayoutAndReducedDims(
+      bool allow_unknown_layouts, mlir::Operation* op,
+      absl::flat_hash_set<std::string>* reduced_dims,
+      std::optional<Layout>* left, std::optional<Layout>* right);
+
+  // This function prepares the inputs (x, y or a, b) to (Batch)MatMul by
+  // possibly computing a new layout for each input that allows us to simply
+  // emit a local (Batch)MatMul op. Once the layouts are computed, the function
+  // calls EmitRelayout to transform from left_layout, right_layout to the
+  // newly computed layouts.
+  // The left and right arguments are set to the mlir::Values representing the
+  // tensors with the possibly new layout.
+  // reduced_dim will contain the dim that must be reduced on after the local
+  // MatMul. It may be set to Layout::kUnsharded if no reduction is needed.
+  // matmul_layout will be set to the layout of the output of the local matmul
+  // (after the above reduction). This may be different from the desired output
+  // layout.
+  absl::Status MaybeRelayoutInputs(
+      mlir::Operation* op, const Layout& left_layout, bool left_transposed,
+      const Layout& right_layout, bool right_transposed,
+      const Layout& output_layout, std::string& reduced_dim,
+      Layout& matmul_layout, mlir::Value& left, mlir::Value& right);
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_MATMUL_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/meta_spmd_expander.h

@@ -0,0 +1,150 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_META_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_META_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Pack/Unpack (aka tf.stack/unstack)
+// For Pack, we verify input tensors have the same layout, and produce a new
+// tensor of rank N + 1 with an unsharded first dimension.
+class PackSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class UnpackSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class PadSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class TileSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// SPMD expansion for reshape.
+//
+// If an explicit layout is provided, reshape will adjust the output to
+// conform to the new layout. N.B. not all possible input/output shapes+layouts
+// are implemented.
+//
+// A fully general reshape involves arbitrary send/recv or collective
+// permutations, and may be inefficient.
+//
+// We provide special cases for a number of common cases.
+class ReshapeSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class TransposeSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class OneHotSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// SPMD expansion for shape/rank metadata operations.
+class ShapeSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_META_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.cc

@@ -0,0 +1,149 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.h"
+
+#include <cassert>
+#include <cstdint>
+#include <vector>
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> NullarySPMDExpander::ExpandOp(mlir::Operation* op) {
+  if (op->getNumResults() == 0) return op;
+
+  bool all_operands_fully_replicated = true;
+  TF_ASSIGN_OR_RETURN(auto op_layouts, ExtractLayoutFromOp(op));
+  for (const auto& op_layout : op_layouts) {
+    if (!op_layout)
+      return absl::InvalidArgumentError(
+          "Nullary op layouts must be known before SPMD expansion.");
+    all_operands_fully_replicated =
+        all_operands_fully_replicated && op_layout->IsFullyReplicated();
+  }
+
+  if (all_operands_fully_replicated) return op;
+
+  if (auto const_op = mlir::dyn_cast<mlir::TF::ConstOp>(op)) {
+    if (auto dense =
+            mlir::dyn_cast<mlir::DenseElementsAttr>(const_op.getValue())) {
+      if (dense.isSplat()) {
+        // A 'splat' value for a DenseElementsAttr, has a single value for
+        // all its elements. For these inputs, we don't need to slice. We just
+        // need to update the shape of the attribute given the requested
+        // sharding.
+        assert(dense.getType().getRank() == op_layouts[0]->rank());
+        auto shape = dense.getType().getShape();
+        std::vector<int64_t> new_shape(dense.getType().getRank());
+        for (int i = 0; i < op_layouts[0]->rank(); ++i) {
+          const int num_shards = op_layouts[0]->num_shards_for_dim(i);
+          if (shape[i] % num_shards != 0)
+            return absl::InvalidArgumentError(absl::StrCat(
+                "has output dimension size ", shape[i],
+                " which is not evenly divisible by the number of shards ",
+                num_shards, " in the layout for that dimension."));
+          new_shape[i] = shape[i] / num_shards;
+        }
+        const_op.setValueAttr(mlir::DenseElementsAttr::get(
+            mlir::RankedTensorType::get(new_shape,
+                                        dense.getType().getElementType()),
+            dense.getSplatValue<mlir::Attribute>()));
+        return InferSPMDExpandedLocalShape(op);
+      }
+    }
+  }
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  llvm::SmallVector<mlir::Value, 4> generated_outputs;
+  llvm::SmallVector<mlir::Type, 4> generated_types;
+
+  mlir::OpBuilder builder(op);
+  builder.setInsertionPointAfter(op);
+  for (int i = 0; i < op_layouts.size(); ++i) {
+    // Split each output to the correct layout by assuming the input is
+    // replicated.
+    TF_ASSIGN_OR_RETURN(
+        const mlir::Value output,
+        EmitAllScatter(builder, op->getOpResult(i),
+                       Layout::ReplicatedOnMesh(op_layouts[i]->mesh(),
+                                                op_layouts[i]->rank()),
+                       *op_layouts[i], &newly_created_ops));
+    generated_outputs.emplace_back(output);
+    generated_types.emplace_back(output.getType());
+  }
+
+  auto identity_op = mlir::TF::IdentityNOp::create(
+      builder, op->getLoc(), generated_types, generated_outputs);
+
+  newly_created_ops.insert(identity_op);
+  for (int i = 0; i < op_layouts.size(); ++i)
+    op->getOpResult(i).replaceAllUsesExcept(identity_op.getResult(i),
+                                            newly_created_ops);
+
+  return identity_op.getOperation();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> NullarySPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  auto enclosing_mesh = op->getParentOfType<mlir::tf_device::ClusterOp>();
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshFromOp(enclosing_mesh));
+  if (!mesh.has_value())
+    return absl::InternalError("Failure in extracting mesh from Nullary Op.");
+  llvm::DenseMap<int, Layout> output_layouts;
+  // Nullary ops always output replicated layout for output values.
+  for (auto i = 0; i < op->getNumResults(); ++i) {
+    auto output_ranked_type =
+        mlir::dyn_cast<mlir::RankedTensorType>(op->getResult(i).getType());
+    if (!output_ranked_type) {
+      return absl::InvalidArgumentError(
+          llvm::formatv("requires output type to have statically known rank, "
+                        "but got : {0}",
+                        output_ranked_type)
+              .str());
+    }
+    output_layouts[i] =
+        Layout::ReplicatedOnMesh(*mesh, output_ranked_type.getRank());
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+NullarySPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // No operand inputs.
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_NULLARY_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_NULLARY_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// SPMD Expander for handling ops without operands.
+class NullarySPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_NULLARY_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/optional_spmd_expander.cc

@@ -0,0 +1,112 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/optional_spmd_expander.h"
+
+#include <vector>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/dtensor_location.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> OptionalGetValueSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  auto original_op = mlir::cast<mlir::TF::OptionalGetValueOp>(op);
+  mlir::OpBuilder builder(op);
+
+  TF_ASSIGN_OR_RETURN(std::vector<Layout> output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+
+  llvm::SmallVector<mlir::Type, 4> local_types(original_op->getNumResults());
+
+  for (int i = 0; i < original_op->getNumResults(); ++i) {
+    mlir::TensorType global_output_type =
+        mlir::cast<mlir::TensorType>(original_op.getResult(i).getType());
+    TF_ASSIGN_OR_RETURN(
+        mlir::TensorType local_type,
+        LocalTypeFromGlobalType(output_layouts[i], global_output_type));
+    local_types[i] = local_type;
+  }
+
+  auto new_op = mlir::TF::OptionalGetValueOp::create(
+      builder, DT_LOC(op->getLoc()), local_types, original_op->getOperand(0));
+
+  for (int i = 0; i < original_op->getNumResults(); ++i) {
+    original_op.getResult(i).replaceAllUsesWith(new_op.getResult(i));
+  }
+  original_op.erase();
+  return InferSPMDExpandedLocalShape(new_op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+OptionalGetValueSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // Extract the output element layouts from some op in the input chain that has
+  // the `tf._element_layouts` attribute set.
+  TF_ASSIGN_OR_RETURN(const auto layouts,
+                      ExtractElementLayoutsFromOperand(op->getOpOperand(0)));
+
+  llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    output_layouts[i] = layouts[i];
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+OptionalGetValueSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)))}});
+}
+
+StatusOr<mlir::Operation*> OptionalHasValueSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+OptionalHasValueSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, /*rank=*/0)}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+OptionalHasValueSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(0)))}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/optional_spmd_expander.h

@@ -0,0 +1,57 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_OPTIONAL_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_OPTIONAL_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class OptionalGetValueSPMDExpander final : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class OptionalHasValueSPMDExpander final : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_OPTIONAL_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/qr_spmd_expander.cc

@@ -0,0 +1,118 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/qr_spmd_expander.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> QRSPMDExpander::ExpandOp(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const auto output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(const auto operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+
+  // Relayout all layouts to the first output layout with the last two
+  // dimensions replicated. We can do more optimization but this is fine
+  TF_ASSIGN_OR_RETURN(
+      Layout new_layout,
+      output_layouts[0].GetLayoutWithReducedDims({-1, -2}, /*keep_dims=*/true));
+
+  TF_ASSIGN_OR_RETURN(
+      const auto new_operand,
+      EmitRelayout(op->getOperand(0), operand_layouts[0], new_layout));
+  op->setOperand(0, new_operand);
+
+  mlir::OpBuilder builder(op);
+  op = InferSPMDExpandedLocalShape(op);
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  llvm::SmallVector<mlir::Value, 4> generated_outputs;
+  llvm::SmallVector<mlir::Type, 4> generated_types;
+
+  // Relayout outputs
+  for (auto i = 0; i < output_layouts.size(); i++) {
+    TF_ASSIGN_OR_RETURN(auto new_output,
+                        EmitRelayout(op->getOpResult(i), new_layout,
+                                     output_layouts[i], &newly_created_ops));
+    generated_outputs.emplace_back(new_output);
+    generated_types.emplace_back(new_output.getType());
+  }
+  if (generated_outputs[0].getDefiningOp()->isBeforeInBlock(
+          generated_outputs[1].getDefiningOp()))
+    builder.setInsertionPointAfterValue(generated_outputs[1]);
+  else
+    builder.setInsertionPointAfterValue(generated_outputs[0]);
+
+  // Tie the two outputs together with an identity op
+  auto identity_op = mlir::TF::IdentityNOp::create(
+      builder, op->getLoc(), generated_types, generated_outputs);
+  newly_created_ops.insert(identity_op);
+  for (int i = 0; i < output_layouts.size(); i++) {
+    op->getOpResult(i).replaceAllUsesExcept(identity_op.getResult(i),
+                                            newly_created_ops);
+  }
+  return identity_op.getOperation();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> QRSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If no input layout at index 0 is present then do not infer any output
+  // layout.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  // Set the output layouts as the copy of the input layouts with the last 2
+  // dimensions replicated.
+  TF_ASSIGN_OR_RETURN(Layout output_layout,
+                      input_layouts.lookup(0).GetLayoutWithReducedDims(
+                          {-1, -2}, /*keep_dims=*/true));
+  return llvm::DenseMap<int, Layout>({{0, output_layout}, {1, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> QRSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If no output layout is present then do not infer any operand layouts.
+  if (output_layouts.empty()) return llvm::DenseMap<int, Layout>();
+
+  // Set the operand layout as the copy of the output layouts with the last 2
+  // dimensions replicated.
+  Layout layout = output_layouts.find(0) != output_layouts.end()
+                      ? output_layouts.lookup(0)
+                      : output_layouts.lookup(1);
+  TF_ASSIGN_OR_RETURN(
+      Layout layout_reduced_dims,
+      layout.GetLayoutWithReducedDims({-1, -2}, /*keep_dims=*/true));
+  return llvm::DenseMap<int, Layout>({{0, layout_reduced_dims}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/qr_spmd_expander.h

@@ -0,0 +1,47 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_QR_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_QR_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// SPMD Expander for QR op. The First Rank-2 dimensions are the batch dimensions
+// and the last two dimensions must be replicated
+class QRSPMDExpander : public SPMDExpanderBase {
+ protected:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ private:
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_QR_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/random_op_spmd_expander.cc

@@ -0,0 +1,401 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/random_op_spmd_expander.h"
+
+#include <algorithm>
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/constants.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/dtensor_location.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+absl::Status CheckLayoutIsSupported(const Layout& layout) {
+  // Currently we support small mesh rank for arbitrary layout.
+  if (layout.mesh().rank() > 3)
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Large mesh rank size is not supported", layout.ToString()));
+
+  return absl::OkStatus();
+}
+
+absl::Status ValidateShapeAndGetNewShape(
+    const llvm::SmallVector<int64_t, 4>& op_shape, const Layout& layout,
+    llvm::SmallVectorImpl<int64_t>& new_random_shape) {
+  TF_RETURN_IF_ERROR(CheckLayoutIsSupported(layout));
+
+  // Validate that sharding of random op is compatible with it's user defined
+  // shape and calculate new shape of local random op.
+  const auto op_sharding = layout.num_shards();
+  new_random_shape.reserve(op_shape.size());
+
+  if (op_sharding.size() != op_shape.size())
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Sharding dimension of random op does not match rank of the "
+        "random op. Received sharding: ",
+        layout.ToString()));
+
+  for (int i = 0; i < op_sharding.size(); ++i) {
+    const auto dimension_sharding = op_sharding[i];
+    const auto op_dimension_size = op_shape[i];
+    if (op_dimension_size % dimension_sharding != 0) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Sharding of random op incompatible with shape. Received "
+          "sharding: ",
+          layout.ToString()));
+    }
+    new_random_shape.emplace_back(op_dimension_size / dimension_sharding);
+  }
+  return absl::OkStatus();
+}
+
+// Get a device seed for this layout and device_id.
+//
+// The computation will be inserted directly after the mesh coordinate
+// computation in the current cluster. First we search for a Squeeze with the
+// attribute kDeviceSeedForMeshDims = layout.mesh_dims
+// If it exists, we return that, otherwise we insert the ops to compute a device
+// seed.
+StatusOr<mlir::Value> GetDeviceSeed(const Layout& layout, mlir::Operation* op) {
+  // We need both a set, to check for membership and a vector that we sort
+  // to use as the attribute attached to the squeeze op.
+  llvm::SmallVector<int32_t, 4> layout_dims;
+  llvm::SmallSet<int32_t, 4> layout_dims_set;
+  for (const auto& spec : layout.sharding_spec_strs()) {
+    if (Layout::IsUnshardedDimension(spec)) continue;
+    layout_dims.emplace_back(layout.mesh().GetMeshDimIndexWithName(spec));
+    layout_dims_set.insert(layout_dims.back());
+  }
+  llvm::sort(layout_dims);
+
+  mlir::tf_device::ClusterOp cluster =
+      op->getParentOfType<mlir::tf_device::ClusterOp>();
+  if (!cluster)
+    return absl::InvalidArgumentError(
+        "random op not in ClusterOp when it should be");
+
+  for (mlir::TF::SqueezeOp squeeze : cluster.getOps<mlir::TF::SqueezeOp>())
+    if (squeeze->hasAttrOfType<mlir::DenseIntElementsAttr>(
+            kDeviceSeedForMeshDims) &&
+        std::equal(layout_dims.begin(), layout_dims.end(),
+                   squeeze
+                       ->getAttrOfType<mlir::DenseIntElementsAttr>(
+                           kDeviceSeedForMeshDims)
+                       .getValues<uint32_t>()
+                       .begin()))
+      return squeeze.getOutput();
+
+  TF_ASSIGN_OR_RETURN(mlir::Value mesh_coordinates,
+                      GetMeshCoordinatesFromCluster(cluster));
+
+  mlir::OpBuilder builder(cluster.getContext());
+  builder.setInsertionPointAfterValue(mesh_coordinates);
+
+  // mesh_coordinates is a [1, mesh.rank()] shaped tensor containing the current
+  // mesh coordinates of the device.
+  // If there are 4 mesh dimensions [w, x, y, z] and only [w, x, z] are used in
+  // this layout then one way of getting the device id would be
+  // w_coord + x_coord*size_w + z_coord*size_x*size_w
+  // Note that only the dims in layout_dims count.
+  llvm::SmallVector<uint32_t, 4> multipliers(layout.mesh().rank(), 0);
+
+  // By starting with 65536, we effective perform a left shift of the id by
+  // 16 bits.
+  int32_t running_product = 65536;
+  for (int i = 0; i < layout.mesh().rank(); ++i) {
+    if (layout_dims_set.contains(i)) {
+      multipliers[i] = running_product;
+      running_product = running_product * layout.mesh().dim_sizes()[i];
+    }
+  }
+
+  mlir::RankedTensorType const_type =
+      mlir::RankedTensorType::get({static_cast<int64_t>(multipliers.size()), 1},
+                                  builder.getIntegerType(32));
+  mlir::Attribute const_attr =
+      mlir::DenseIntElementsAttr::get(const_type, multipliers);
+  mlir::Value multiplier =
+      mlir::TF::ConstOp::create(builder, cluster.getLoc(), const_attr)
+          .getOutput();
+
+  const mlir::RankedTensorType one_by_one =
+      mlir::RankedTensorType::get({1, 1}, builder.getIntegerType(32));
+
+  mlir::Value seed = mlir::TF::MatMulOp::create(
+      builder, cluster.getLoc(), one_by_one, mesh_coordinates, multiplier);
+
+  // Largest prime in 16 bits.
+  mlir::Value prime = CreateIntScalarConst(
+      /*value=*/65521, builder, cluster.getLoc(), /*use_int64=*/false);
+
+  mlir::Value seed_plus_prime =
+      mlir::TF::AddV2Op::create(builder, cluster.getLoc(), one_by_one, seed,
+                                prime)
+          .getZ();
+
+  mlir::TF::SqueezeOp squeeze = mlir::TF::SqueezeOp::create(
+      builder, cluster.getLoc(),
+      mlir::RankedTensorType::get({}, builder.getIntegerType(32)),
+      seed_plus_prime, builder.getI64ArrayAttr({0, 1}));
+
+  squeeze->setAttr(kDeviceSeedForMeshDims,
+                   builder.getI32TensorAttr(layout_dims));
+
+  return squeeze.getOutput();
+}
+
+// Compute the new local shape for SPMD expansion and ensure it is valid.
+template <typename RandomOp>
+StatusOr<llvm::SmallVector<int64_t, 4>> GetNewLocalShape(mlir::Operation* op,
+                                                         const Layout& layout) {
+  auto random_op = llvm::cast<RandomOp>(op);
+  llvm::SmallVector<int64_t, 4> op_shape;
+  TF_RETURN_IF_ERROR(
+      ExtractConstVectorFromValue(random_op.getShape(), &op_shape));
+
+  // Validate that sharding of random op is compatible with it's user defined
+  // shape and calculate new shape of local random op.
+  llvm::SmallVector<int64_t, 4> new_random_shape;
+  TF_RETURN_IF_ERROR(
+      ValidateShapeAndGetNewShape(op_shape, layout, new_random_shape));
+  return new_random_shape;
+}
+
+// Calculate the new local seed
+template <typename RandomOp>
+StatusOr<mlir::Value> ComputeNewSeed(mlir::OpBuilder& builder,
+                                     mlir::Operation* op, const Layout& layout,
+                                     mlir::Location& location,
+                                     mlir::Value op_seed) {
+  TF_ASSIGN_OR_RETURN(auto device_id_seed, GetDeviceSeed(layout, op));
+  mlir::Type seed_type =
+      mlir::cast<mlir::TensorType>(op_seed.getType()).getElementType();
+
+  device_id_seed = mlir::TF::CastOp::create(
+      builder, location, mlir::RankedTensorType::get({}, seed_type),
+      device_id_seed);
+
+  mlir::Value seed_xor = mlir::TF::BitwiseXorOp::create(
+      builder, location, op_seed, device_id_seed);
+  return seed_xor;
+}
+
+template <typename RandomOp>
+StatusOr<mlir::Operation*> CreatedShardedLocalRandomOpV1(const Layout& layout,
+                                                         mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto new_random_shape,
+                      GetNewLocalShape<RandomOp>(op, layout));
+
+  // Create new seed using already existing seed and a device id.
+  mlir::OpBuilder builder(op);
+  auto location = DT_LOC(op);
+
+  auto random_op = llvm::cast<RandomOp>(op);
+  // Create device_id_seed for local RNG.
+  TF_ASSIGN_OR_RETURN(auto seed_xor,
+                      ComputeNewSeed<RandomOp>(builder, op, layout, location,
+                                               random_op.getSeed()));
+
+  // Create a new random op with new `local` shape and newly generated seed.
+  // StatelessRandom op is used to make random op SPMD expansion
+  // deterministic.
+  mlir::Type new_random_type = mlir::RankedTensorType::get(
+      new_random_shape, mlir::cast<mlir::TensorType>(op->getResult(0).getType())
+                            .getElementType());
+
+  auto new_shape_value = Int64Const(builder, location, new_random_shape);
+  // TODO(zhonglinhan) : check different input for StatelessRandomUniformInt
+  auto local_random = RandomOp::create(builder, location, new_random_type,
+                                       new_shape_value, seed_xor);
+  op->getResult(0).replaceAllUsesWith(local_random.getOutput());
+  op->erase();
+  return local_random.getOperation();
+}
+
+template <typename RandomOp>
+StatusOr<mlir::Operation*> CreatedShardedLocalRandomOpV2(const Layout& layout,
+                                                         mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto new_random_shape,
+                      GetNewLocalShape<RandomOp>(op, layout));
+
+  // Create new seed using already existing seed and a device id.
+  mlir::OpBuilder builder(op);
+  auto location = DT_LOC(op);
+
+  auto random_op = llvm::cast<RandomOp>(op);
+  // Create device_id_seed for local RNG.
+  TF_ASSIGN_OR_RETURN(auto seed_xor,
+                      ComputeNewSeed<RandomOp>(builder, op, layout, location,
+                                               random_op.getKey()));
+
+  // Create a new random op with new `local` shape and newly generated seed.
+  // StatelessRandom op is used to make random op SPMD expansion
+  // deterministic.
+  mlir::Type new_random_type = mlir::RankedTensorType::get(
+      new_random_shape, mlir::cast<mlir::TensorType>(op->getResult(0).getType())
+                            .getElementType());
+
+  auto new_shape_value = Int64Const(builder, location, new_random_shape);
+
+  auto local_random =
+      RandomOp::create(builder, location, new_random_type, new_shape_value,
+                       seed_xor, random_op.getCounter(), random_op.getAlg());
+  op->getResult(0).replaceAllUsesWith(local_random.getOutput());
+  op->erase();
+  return local_random.getOperation();
+}
+
+template <typename RandomOp>
+StatusOr<mlir::Operation*> CreatedShardedLocalRandomOpV2Range(
+    const Layout& layout, mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto new_random_shape,
+                      GetNewLocalShape<RandomOp>(op, layout));
+
+  // Create new seed using already existing seed and a device id.
+  mlir::OpBuilder builder(op);
+  auto location = DT_LOC(op);
+
+  auto random_op = llvm::cast<RandomOp>(op);
+  // Create device_id_seed for local RNG.
+  TF_ASSIGN_OR_RETURN(auto seed_xor,
+                      ComputeNewSeed<RandomOp>(builder, op, layout, location,
+                                               random_op.getKey()));
+
+  // Create a new random op with new `local` shape and newly generated seed.
+  // StatelessRandom op is used to make random op SPMD expansion
+  // deterministic.
+  mlir::Type new_random_type = mlir::RankedTensorType::get(
+      new_random_shape, mlir::cast<mlir::TensorType>(op->getResult(0).getType())
+                            .getElementType());
+
+  auto new_shape_value = Int64Const(builder, location, new_random_shape);
+
+  auto local_random =
+      RandomOp::create(builder, location, new_random_type, new_shape_value,
+                       seed_xor, random_op.getCounter(), random_op.getAlg(),
+                       random_op.getMinval(), random_op.getMaxval());
+  op->getResult(0).replaceAllUsesWith(local_random.getOutput());
+  op->erase();
+  return local_random.getOperation();
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> RandomOpSPMDExpander::ExpandOp(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto layout, ExtractSingleLayoutFromOp(op));
+
+  if (!layout)
+    return absl::InvalidArgumentError(
+        "layout of Random op must be known before SPMD expansion.");
+
+  // For fully replicated random ops, all devices have the same random
+  // value. As so, SPMD expansion is a no-op.
+  if (layout->IsFullyReplicated()) return op;
+  if (llvm::isa<mlir::TF::StatelessRandomUniformOp>(op)) {
+    return CreatedShardedLocalRandomOpV1<mlir::TF::StatelessRandomUniformOp>(
+        *layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessRandomUniformFullIntOp>(op)) {
+    return CreatedShardedLocalRandomOpV1<
+        mlir::TF::StatelessRandomUniformFullIntOp>(*layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessRandomNormalOp>(op)) {
+    return CreatedShardedLocalRandomOpV1<mlir::TF::StatelessRandomNormalOp>(
+        *layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessTruncatedNormalOp>(op)) {
+    return CreatedShardedLocalRandomOpV1<mlir::TF::StatelessTruncatedNormalOp>(
+        *layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessRandomUniformFullIntV2Op>(op)) {
+    return CreatedShardedLocalRandomOpV2<
+        mlir::TF::StatelessRandomUniformFullIntV2Op>(*layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessRandomNormalV2Op>(op)) {
+    return CreatedShardedLocalRandomOpV2<mlir::TF::StatelessRandomNormalV2Op>(
+        *layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessTruncatedNormalV2Op>(op)) {
+    return CreatedShardedLocalRandomOpV2<
+        mlir::TF::StatelessTruncatedNormalV2Op>(*layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessRandomUniformV2Op>(op)) {
+    return CreatedShardedLocalRandomOpV2<mlir::TF::StatelessRandomUniformV2Op>(
+        *layout, op);
+  }
+  if (llvm::isa<mlir::TF::StatelessRandomUniformIntV2Op>(op)) {
+    return CreatedShardedLocalRandomOpV2Range<
+        mlir::TF::StatelessRandomUniformIntV2Op>(*layout, op);
+  }
+  return absl::UnimplementedError(absl::StrCat(
+      "SPMD expansion for op : ", OpName(op), " is not implemented"));
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+RandomOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  llvm::DenseMap<int, Layout> output_layouts;
+
+  // For random op, input is always replicated and we always respect layouts
+  // from consumers.
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    output_layouts[i] =
+        Layout::ReplicatedOnMesh(mesh, /*rank=*/ValueRank(op->getResult(i)));
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+RandomOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  llvm::DenseMap<int, Layout> input_layouts;
+
+  // For random op, default the input layout as replicated layout.
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    input_layouts[i] =
+        Layout::ReplicatedOnMesh(mesh, /*rank=*/ValueRank(op->getOperand(i)));
+  }
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/random_op_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RANDOM_OP_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RANDOM_OP_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class RandomOpSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RANDOM_OP_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/range_spmd_expander.cc

@@ -0,0 +1,63 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/range_spmd_expander.h"
+
+#include "absl/status/status.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> RangeSPMDExpander::ExpandOp(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto layout, ExtractSingleLayoutFromOp(op));
+
+  if (!layout)
+    return absl::InvalidArgumentError(
+        "layout of RangeOp must be known before SPMD expansion.");
+
+  if (!layout->IsFullyReplicated())
+    return absl::InternalError("Shared RangeOp is not supported yet.");
+
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> RangeSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  // Always return a Replicated layout. This will always respect the consumer
+  // requested layouts.
+  return llvm::DenseMap<int, Layout>({{0, Layout::ReplicatedOnMesh(mesh, 1)}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> RangeSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  // Always assign a replicated layout to the operands.
+  llvm::DenseMap<int, Layout> input_layouts;
+  for (int i = 0; i < op->getNumOperands(); ++i)
+    input_layouts[i] =
+        Layout::ReplicatedOnMesh(mesh, /*rank=*/ValueRank(op->getOperand(i)));
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/range_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RANGE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RANGE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// SPMD Expander for RangeOp.
+class RangeSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RANGE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc

@@ -0,0 +1,327 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+absl::string_view StringRefToView(llvm::StringRef ref) {
+  return absl::string_view(ref.data(), ref.size());
+}
+
+absl::string_view DefiningOpName(mlir::Value operand) {
+  return StringRefToView(operand.getDefiningOp()->getName().getStringRef());
+}
+
+absl::Status AssertReplicated(mlir::Value operand) {
+  TF_ASSIGN_OR_RETURN(auto layout, ExtractLayoutFromOperand(operand));
+  if (!layout) return absl::OkStatus();
+
+  if (!layout->IsFullyReplicated()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Expected layout for ", DefiningOpName(operand),
+        " to be fully replicated, but found ", layout->ToString()));
+  }
+  return absl::OkStatus();
+}
+
+absl::flat_hash_set<std::string> ReducedMeshDimensions(
+    const dtensor::Layout& input, const dtensor::Layout& output) {
+  absl::flat_hash_set<std::string> mesh_dims;
+  for (const auto& dim : input.sharding_spec_strs()) {
+    mesh_dims.insert(dim);
+  }
+  for (const auto& dim : output.sharding_spec_strs()) {
+    mesh_dims.erase(dim);
+  }
+  return mesh_dims;
+}
+
+template <typename OpType>
+absl::Status ExtractDims(mlir::Operation* op,
+                         llvm::SmallVector<int64_t, 4>* reduced_dims,
+                         bool* keep_dims, bool* matched) {
+  if (!llvm::isa<OpType>(op)) return absl::OkStatus();
+  auto reduce_op = llvm::cast<OpType>(op);
+  *keep_dims = reduce_op.getKeepDims();
+  TF_RETURN_IF_ERROR(ExtractConstVectorFromValue(
+      reduce_op.getReductionIndices(), reduced_dims));
+  TF_RETURN_IF_ERROR(AssertReplicated(reduce_op.getReductionIndices()));
+  *matched = true;
+
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status ExtractDims<mlir::TF::L2LossOp>(
+    mlir::Operation* op, llvm::SmallVector<int64_t, 4>* reduced_dims,
+    bool* keep_dims, bool* matched) {
+  if (!llvm::isa<mlir::TF::L2LossOp>(op)) return absl::OkStatus();
+  auto loss_op = llvm::cast<mlir::TF::L2LossOp>(op);
+  *reduced_dims = llvm::SmallVector<int64_t, 4>{};
+  reduced_dims->resize(ValueRank(loss_op->getOperand(0)));
+  for (int i = 0; i < reduced_dims->size(); ++i) {
+    (*reduced_dims)[i] = i;
+  }
+  *keep_dims = false;
+  *matched = true;
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status ExtractDims<mlir::TF::BiasAddGradOp>(
+    mlir::Operation* op, llvm::SmallVector<int64_t, 4>* reduced_dims,
+    bool* keep_dims, bool* matched) {
+  if (!llvm::isa<mlir::TF::BiasAddGradOp>(op)) return absl::OkStatus();
+  auto bias_add_grad_op = llvm::cast<mlir::TF::BiasAddGradOp>(op);
+  auto data_format = bias_add_grad_op.getDataFormat();
+  // rank is at least 2 (required by BiasAddGrad).
+  int rank = ValueRank(bias_add_grad_op->getOperand(0));
+  if (data_format == "NHWC") {
+    for (int dim = 0; dim < rank - 1; ++dim) {
+      reduced_dims->push_back(dim);
+    }
+  } else if (data_format == "NCHW") {
+    for (int dim = 0; dim < rank; ++dim) {
+      if (dim == 1) continue;
+      reduced_dims->push_back(dim);
+    }
+  } else {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Unsupported data_format for BiasAddGrad: ",
+                     StringRefToView(data_format)));
+  }
+  *keep_dims = false;
+  *matched = true;
+  return absl::OkStatus();
+}
+
+template <>
+absl::Status ExtractDims<mlir::TF::EncodePngOp>(
+    mlir::Operation* op, llvm::SmallVector<int64_t, 4>* reduced_dims,
+    bool* keep_dims, bool* matched) {
+  if (!llvm::isa<mlir::TF::EncodePngOp>(op)) return absl::OkStatus();
+  *reduced_dims = {-3, -2, -1};
+  *keep_dims = false;
+  *matched = true;
+  return absl::OkStatus();
+}
+
+absl::Status ExtractReductionParameters(
+    mlir::Operation* op, absl::flat_hash_set<int>& reduced_dims_set,
+    bool& keep_dims) {
+  llvm::SmallVector<int64_t, 4> reduced_dims;
+  bool matched = false;
+  TF_RETURN_IF_ERROR(ExtractDims<mlir::TF::EncodePngOp>(op, &reduced_dims,
+                                                        &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::SumOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::AllOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::AnyOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::MaxOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::MinOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::MeanOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::ProdOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(
+      ExtractDims<mlir::TF::L2LossOp>(op, &reduced_dims, &keep_dims, &matched));
+  TF_RETURN_IF_ERROR(ExtractDims<mlir::TF::BiasAddGradOp>(
+      op, &reduced_dims, &keep_dims, &matched));
+
+  if (!matched)
+    return absl::UnimplementedError(
+        absl::StrCat("Op type: ", OpName(op), " not yet implemented."));
+
+  reduced_dims_set.insert(reduced_dims.begin(), reduced_dims.end());
+  return absl::OkStatus();
+}
+
+StatusOr<Layout> ComputeResultLayout(mlir::Operation* op,
+                                     const Layout& input_layout) {
+  absl::flat_hash_set<int> reduced_dims_set;
+  bool keep_dims;
+  TF_RETURN_IF_ERROR(
+      ExtractReductionParameters(op, reduced_dims_set, keep_dims));
+
+  return input_layout.GetLayoutWithReducedDims(reduced_dims_set, keep_dims);
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> ReduceSPMDExpander::ExpandOp(mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto requested_output_layout,
+                      ExtractSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(auto input_layout,
+                      ExtractLayoutFromOperand(op->getOperand(0)));
+
+  if (!input_layout || !requested_output_layout)
+    return absl::InvalidArgumentError("is missing input or output layouts.");
+
+  // Generate an error message for TPU int64.
+  if (input_layout->mesh().is_tpu_mesh()) {
+    if (auto tensor_type =
+            mlir::dyn_cast<mlir::TensorType>(op->getOperand(0).getType())) {
+      if (tensor_type.getElementType().isInteger(64)) {
+        return absl::InvalidArgumentError(
+            "ReduceOp on TPU does not support int64 as dtype.");
+      }
+    }
+  }
+
+  mlir::OpBuilder builder(op->getBlock(), ++mlir::Block::iterator(op));
+
+  TF_ASSIGN_OR_RETURN(auto output_layout,
+                      ComputeResultLayout(op, input_layout.value()));
+
+  absl::flat_hash_set<std::string> reduced_dims =
+      ReducedMeshDimensions(*input_layout, output_layout);
+  InferSPMDExpandedLocalShape(op);
+
+  mlir::Operation* reduce_op;
+  if (mlir::isa<mlir::TF::SumOp, mlir::TF::L2LossOp, mlir::TF::BiasAddGradOp,
+                mlir::TF::EncodePngOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpAdd));
+  } else if (mlir::isa<mlir::TF::AllOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpAll));
+  } else if (mlir::isa<mlir::TF::AnyOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpAny));
+  } else if (mlir::isa<mlir::TF::MaxOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpMax));
+  } else if (mlir::isa<mlir::TF::MinOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpMin));
+  } else if (mlir::isa<mlir::TF::ProdOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpMul));
+  } else if (mlir::isa<mlir::TF::MeanOp>(op)) {
+    TF_ASSIGN_OR_RETURN(
+        reduce_op,
+        EmitAllReduce(builder, output_layout, reduced_dims, op, kReduceOpMean));
+  } else {
+    return DT_CTX(absl::UnimplementedError(absl::StrCat(
+        "Failed to create AllReduce op during SPMD expansion. Op type: ",
+        OpName(op), " not yet implemented in DTensor SPMD pass.")));
+  }
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  TF_ASSIGN_OR_RETURN(
+      auto final_output,
+      EmitAllScatter(builder, reduce_op->getOpResult(0), output_layout,
+                     requested_output_layout.value(), &newly_created_ops));
+  reduce_op->getOpResult(0).replaceAllUsesExcept(final_output,
+                                                 newly_created_ops);
+  return final_output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> ReduceSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // Do not infer any output layout if no input layout exists.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(auto result_layout,
+                      ComputeResultLayout(op, input_layouts.lookup(0)));
+  return llvm::DenseMap<int, Layout>({{0, result_layout}});
+}
+
+// For Reduction op, we do not propagate consumer preferred layouts to operand
+// layouts as reduction operation explicitly converts tensor dimension of
+// reduced dimensions to replicated.
+StatusOr<llvm::DenseMap<int, Layout>> ReduceSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Do not infer any operand layouts if no output layouts exists.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  Layout output_layout = output_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(auto input_shape, GetShapeOfValue(op->getOperand(0)));
+
+  std::vector<std::string> inferred_operand_layout_str;
+
+  absl::flat_hash_set<int> reduced_dims_set;
+  bool keep_dims;
+  TF_RETURN_IF_ERROR(
+      ExtractReductionParameters(op, reduced_dims_set, keep_dims));
+
+  // For each dimension, if dimension is not reduced dimension, then propagate
+  // the sharding of output value to operand. Else, set input dimension as
+  // replicated.
+  int output_dim = 0;
+  for (int i = 0; i < input_shape.size(); ++i) {
+    // reduced_dims may contain negative values.
+    if (reduced_dims_set.contains(i) ||
+        reduced_dims_set.contains(i - input_shape.size())) {
+      inferred_operand_layout_str.push_back(Layout::kAny);
+      if (keep_dims) output_dim += 1;
+    } else {
+      inferred_operand_layout_str.push_back(
+          output_layout.sharding_spec(output_dim));
+      output_dim += 1;
+    }
+  }
+  TF_ASSIGN_OR_RETURN(
+      auto inferred_operand_layout,
+      Layout::GetLayout(inferred_operand_layout_str, output_layout.mesh()));
+  return llvm::DenseMap<int, Layout>({{0, inferred_operand_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.h

@@ -0,0 +1,47 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_REDUCE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_REDUCE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Emit SPMD expansion for reduction operations
+// (max, min, prod, sum, std, mean, variance)
+class ReduceSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_REDUCE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.cc

@@ -0,0 +1,144 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.h"
+
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Relayout all operands and outputs to replicated layout.
+StatusOr<mlir::Operation*>
+ReplicatedOpSPMDExpander::ReplicatedRelayoutOperandsAndOutputs(
+    mlir::Operation* op, const std::vector<Layout>& operand_layouts,
+    const std::vector<Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Relayout operands
+  for (auto i = 0; i < operand_layouts.size(); ++i) {
+    Layout new_layout =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(i)));
+    TF_ASSIGN_OR_RETURN(
+        const auto new_operand,
+        EmitRelayout(op->getOperand(i), operand_layouts[i], new_layout));
+    op->setOperand(i, new_operand);
+  }
+  // Expand to local shape
+  op = InferSPMDExpandedLocalShape(op);
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  llvm::SmallVector<mlir::Value, 4> generated_outputs;
+  llvm::SmallVector<mlir::Type, 4> generated_types;
+
+  // Track the op that comes last after splitting.
+  mlir::Operation* last_op_after_splitting = op;
+
+  // Relayout outputs
+  for (auto i = 0; i < output_layouts.size(); ++i) {
+    Layout new_layout =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getResult(i)));
+    TF_ASSIGN_OR_RETURN(auto new_output,
+                        EmitRelayout(op->getOpResult(i), new_layout,
+                                     output_layouts[i], &newly_created_ops));
+    generated_outputs.emplace_back(new_output);
+    generated_types.emplace_back(new_output.getType());
+    if (last_op_after_splitting->isBeforeInBlock(new_output.getDefiningOp())) {
+      last_op_after_splitting = new_output.getDefiningOp();
+    }
+  }
+  mlir::OpBuilder builder(op);
+  builder.setInsertionPointAfter(last_op_after_splitting);
+
+  // Tie all outputs together with identity_n
+  auto identity_op = mlir::TF::IdentityNOp::create(
+      builder, op->getLoc(), generated_types, generated_outputs);
+  newly_created_ops.insert(identity_op);
+  for (int i = 0; i < output_layouts.size(); ++i) {
+    op->getOpResult(i).replaceAllUsesExcept(identity_op.getResult(i),
+                                            newly_created_ops);
+  }
+
+  return identity_op.getOperation();
+}
+
+StatusOr<mlir::Operation*> ReplicatedOpSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const auto output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(const auto operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+  if (relayout_when_sharded_)
+    return ReplicatedRelayoutOperandsAndOutputs(op, operand_layouts,
+                                                output_layouts);
+  if (!AllReplicated(output_layouts) || !AllReplicated(operand_layouts)) {
+    return absl::InvalidArgumentError(
+        llvm::formatv("Expecting {0} to have input and output layouts to be "
+                      "fully replicated but was not. ",
+                      OpName(op))
+            .str());
+  }
+  return op;
+}
+
+// Always return a set of replicated layouts.
+StatusOr<llvm::DenseMap<int, Layout>>
+ReplicatedOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    output_layouts[i] =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getResult(i)));
+  }
+  return output_layouts;
+}
+
+// Always return a set of replicated layouts.
+StatusOr<llvm::DenseMap<int, Layout>>
+ReplicatedOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands());
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    input_layouts[i] =
+        Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(i)));
+  }
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.h

@@ -0,0 +1,70 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_REPLICATED_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_REPLICATED_SPMD_EXPANDER_H_
+
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// General SPMD Expander that enforces input(s) and output(s) are replicated.
+class ReplicatedOpSPMDExpander : public SPMDExpanderBase {
+ protected:
+  // If true, then during ExpandOp, relayouts all operands and outputs
+  // to be have replicated layout. If false, then will emit
+  // an error if not all operand and output layouts are replicated after layout
+  // propagation.
+  //
+  // This is needed because some ops like RngReadAndSkip need to enforce input
+  // and output are replicated, while some ops don't need to enforce it, so
+  // we can just relayout to replicated on those during ExpandOp.
+  bool relayout_when_sharded_;
+
+  // Expand the op to the local op after checking all layouts are replicated.
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  // Compute the layouts as always replicated.
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  // Compute the layouts as always replicated.
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ public:
+  explicit ReplicatedOpSPMDExpander(bool relayout_when_sharded = false) {
+    relayout_when_sharded_ = relayout_when_sharded;
+  }
+
+ private:
+  // Relayouts all operands and outputs to a replicated layout.
+  StatusOr<mlir::Operation*> ReplicatedRelayoutOperandsAndOutputs(
+      mlir::Operation* op, const std::vector<Layout>& operand_layouts,
+      const std::vector<Layout>& output_layouts);
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_REPLICATED_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/resource_spmd_expander.cc

@@ -0,0 +1,354 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/resource_spmd_expander.h"
+
+#include <algorithm>
+#include <cassert>
+#include <iterator>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/constants.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+template <typename AttrType>
+std::vector<AttrType> CreateOrGetMutableAttributeList(
+    mlir::tf_device::ClusterOp op, std::string attr_name) {
+  auto array_attribute = op->getAttrOfType<mlir::ArrayAttr>(attr_name);
+
+  std::vector<AttrType> output;
+  if (array_attribute) auto attr_list = array_attribute.getValue().vec();
+  return output;
+}
+
+StatusOr<mlir::Operation*> ExpandVarHandleOp(mlir::Operation* op) {
+  mlir::OpBuilder builder(op);
+  builder.setInsertionPointAfter(op);
+
+  // This is the layout of the value held by the resource.
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> resource_layout,
+                      ExtractSingleLayoutFromOp(op));
+
+  if (!resource_layout) {
+    // If resource does not have a layout, perform local SPMD expansion.
+    return InferSPMDExpandedLocalShape(op);
+  }
+
+  // If resource has a layout, create VarHandleOps with local shape.
+  // auto var_op = llvm::cast<mlir::TF::VarHandleOp>(op);
+  auto op_result = op->getOpResult(0);
+  TF_RETURN_IF_ERROR(InferSPMDExpandedLocalShapeForResourceOutput(
+      &op_result, resource_layout.value(), builder.getContext()));
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<mlir::Operation*> ExpandSummaryWriterOp(mlir::Operation* op) {
+  // This is the layout of the value held by the resource.
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> resource_layout,
+                      ExtractSingleLayoutFromOp(op));
+
+  if (!resource_layout) {
+    // If resource does not have a layout, perform local SPMD expansion.
+    return InferSPMDExpandedLocalShape(op);
+  } else if (!resource_layout->IsFullyReplicated()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("SummaryWriter op should have fully replicated layout. ",
+                     "Got: ", resource_layout->ToString()));
+  }
+
+  // For SummaryWriter op, we will expand as fully replicated, but only the
+  // replica 0 will be used to write the summary. This is implemented by all the
+  // other summary op expanders.
+  return InferSPMDExpandedLocalShape(op);
+}
+
+absl::Status ValidateAndAssignResourceInputLayout(
+    mlir::tf_device::ClusterOp op, const std::string& layout_string,
+    const int resource_arg_index, mlir::OpBuilder* builder) {
+  const auto add_layout_as_attributes =
+      [&](std::vector<mlir::StringRef> new_resource_layouts,
+          std::vector<int> new_resource_indices, int resource_arg_index,
+          std::string layout) {
+        new_resource_layouts.emplace_back(layout);
+        new_resource_indices.emplace_back(resource_arg_index);
+        op->setAttr(kNewResourceArgLayouts,
+                    builder->getStrArrayAttr(
+                        llvm::ArrayRef<mlir::StringRef>(new_resource_layouts)));
+        op->setAttr(kNewResourceLayoutIndices,
+                    builder->getI32VectorAttr(new_resource_indices));
+      };
+
+  auto resource_input_layouts_attrs =
+      CreateOrGetMutableAttributeList<mlir::StringAttr>(op,
+                                                        kNewResourceArgLayouts);
+  auto resource_input_indices_attrs =
+      CreateOrGetMutableAttributeList<mlir::IntegerAttr>(
+          op, kNewResourceLayoutIndices);
+  std::vector<llvm::StringRef> mutable_input_layouts;
+  std::vector<int> mutable_input_indices;
+  for (auto layout_index_pair :
+       llvm::zip(resource_input_indices_attrs, resource_input_layouts_attrs)) {
+    mutable_input_indices.emplace_back(std::get<0>(layout_index_pair).getInt());
+    mutable_input_layouts.emplace_back(
+        std::get<1>(layout_index_pair).getValue());
+  }
+
+  if (!mutable_input_indices.empty()) {
+    assert(mutable_input_indices.size() == mutable_input_layouts.size());
+
+    auto it = std::find(mutable_input_indices.begin(),
+                        mutable_input_indices.end(), resource_arg_index);
+
+    if (it != mutable_input_indices.end()) {
+      // Input layout for given resource was already inferred from previous
+      // SPMD expansions. Check that layouts of resource are consistent.
+      auto previous_layout = mutable_input_layouts[std::distance(
+          mutable_input_indices.begin(), it)];
+
+      // TODO(hongjunchoi): Implement relayout logic for resource ops.
+      if (layout_string != previous_layout.str())
+        return absl::InvalidArgumentError(
+            "Trying to assign a variable to a resource with a different "
+            "layout.");
+    } else {
+      add_layout_as_attributes(mutable_input_layouts, mutable_input_indices,
+                               resource_arg_index, layout_string);
+    }
+  } else {
+    add_layout_as_attributes(mutable_input_layouts, mutable_input_indices,
+                             resource_arg_index, layout_string);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> ResourceSPMDExpander::ExpandOp(mlir::Operation* op) {
+  // These ops need no special handling.
+  if (llvm::isa<mlir::TF::DestroyResourceOp, mlir::TF::VarIsInitializedOp>(op))
+    return InferSPMDExpandedLocalShape(op);
+
+  if (llvm::isa<mlir::TF::VarHandleOp>(op)) {
+    return ExpandVarHandleOp(op);
+  }
+  if (llvm::isa<mlir::TF::SummaryWriterOp>(op)) {
+    return ExpandSummaryWriterOp(op);
+  }
+
+  mlir::OpBuilder builder(op);
+
+  // Output of read variable may need to be sliced, so it needs to be treated
+  // specially.
+  if (llvm::isa<mlir::TF::ReadVariableOp>(op)) {
+    builder.setInsertionPointAfter(op);
+    TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
+    TF_ASSIGN_OR_RETURN(auto input_layout,
+                        ExtractLayoutFromOperand(op->getOperand(0)));
+    if (!output_layout)
+      TF_RETURN_WITH_CONTEXT(absl::InternalError("output layout is missing"));
+    if (!input_layout)
+      TF_RETURN_WITH_CONTEXT(absl::InternalError("input layout is missing"));
+
+    InferSPMDExpandedLocalShape(op);
+    llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+    TF_ASSIGN_OR_RETURN(
+        auto final_output,
+        EmitAllScatter(builder, op->getOpResult(0), input_layout.value(),
+                       output_layout.value(), &newly_created_ops));
+    op->getOpResult(0).replaceAllUsesExcept(final_output, newly_created_ops);
+    return final_output.getDefiningOp();
+  }
+
+  if (!llvm::isa<mlir::TF::AssignVariableOp, mlir::TF::AssignAddVariableOp,
+                 mlir::TF::AssignSubVariableOp>(op))
+    TF_RETURN_WITH_CONTEXT(absl::InternalError("unsupported resource op"));
+
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> output_layout,
+                      ExtractSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> resource_layout,
+                      ExtractLayoutFromOperand(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> value_layout,
+                      ExtractLayoutFromOperand(op->getOperand(1)));
+
+  // For assignment operations, the layout for the resource (first operand),
+  // when not present, is, inferred from the layout of the input value (second
+  // operand). We attach the inferred layout to the resource.
+  // Note that in the case that input_resource_value.getDefiningOp() exists, it
+  // is a DTensorLayout and this means that the corresponding block argument
+  // already has a layout set.
+  // If the resource is specified in the graph as an op (e.g. VarHandleOp), we
+  // attach the layout directly. Otherwise, the resource is an argument to the
+  // SPMD function, and we attach the layout to the appropriate argument.
+  auto input_resource_value = op->getOpOperand(0).get();
+  if (input_resource_value.getDefiningOp()) {
+    if (!resource_layout)
+      TF_RETURN_WITH_CONTEXT(absl::InternalError("missing layout on resource"));
+    if (!value_layout)
+      TF_RETURN_WITH_CONTEXT(absl::InternalError("missing layout on value"));
+    if (resource_layout != value_layout) {
+      TF_ASSIGN_OR_RETURN(auto new_value,
+                          EmitRelayout(op->getOperand(1), value_layout.value(),
+                                       resource_layout.value()));
+      op->setOperand(1, new_value);
+    }
+  } else {
+    if ((!resource_layout || resource_layout->IsEmpty()) && !value_layout)
+      TF_RETURN_WITH_CONTEXT(absl::InternalError(
+          "at least one of resource or value layout must be set"));
+    // This error should not happen: if resource_layout is set, then we expect
+    // a DTensorLayout op between the resource tensor and this op, so we should
+    // actaully be in the if case rather than the else case.
+    if (resource_layout && !resource_layout->IsEmpty() && value_layout &&
+        resource_layout != value_layout)
+      TF_RETURN_WITH_CONTEXT(absl::InternalError(
+          "if both resource and value layout are set they must be equal"));
+
+    auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(input_resource_value);
+    auto enclosing_device_cluster =
+        op->getParentOfType<mlir::tf_device::ClusterOp>();
+
+    if (!enclosing_device_cluster)
+      TF_RETURN_WITH_CONTEXT(
+          absl::InvalidArgumentError("op must be enclosed by a cluster"));
+
+    auto block_arg_index = block_arg.getArgNumber();
+
+    // If layout of resource already exists, then check that layouts are
+    // consistent. Otherwise, add newly inferred layout of resource argument
+    // as attributes to the enclosing cluster op to be propagated to custom
+    // device.
+    std::string layout_string;
+    if (resource_layout && !resource_layout->IsEmpty())
+      layout_string = resource_layout->ToString();
+    else
+      layout_string = value_layout->ToString();
+    TF_RETURN_IF_ERROR(ValidateAndAssignResourceInputLayout(
+        enclosing_device_cluster, layout_string, block_arg_index, &builder));
+  }
+
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+ResourceSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // VarHandle and VarIsInitialized have 0 rank outputs.
+  if (llvm::isa<mlir::TF::VarHandleOp, mlir::TF::VarIsInitializedOp>(op)) {
+    return llvm::DenseMap<int, Layout>({{0, Layout::Empty()}});
+  }
+
+  // SummaryWriterOp has a rank 0 output, but it need to be fully replicated.
+  if (llvm::isa<mlir::TF::SummaryWriterOp>(op)) {
+    TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+    return llvm::DenseMap<int, Layout>(
+        {{0, Layout::ReplicatedOnMesh(mesh, /*rank=*/0)}});
+  }
+
+  // Handling of resource destruction is no-op.
+  if (llvm::isa<mlir::TF::DestroyResourceOp>(op))
+    return llvm::DenseMap<int, Layout>();
+
+  // Read variable ops have one input so infer the output layout if input
+  // layout exists.
+  if (llvm::isa<mlir::TF::ReadVariableOp>(op)) {
+    return input_layouts;
+  }
+
+  // These ops do not have outputs, so do not infer any layout.
+  if (llvm::isa<mlir::TF::AssignVariableOp, mlir::TF::AssignAddVariableOp,
+                mlir::TF::AssignSubVariableOp>(op)) {
+    return llvm::DenseMap<int, Layout>();
+  }
+  // Return an error if not any of the ops above.
+  return absl::InvalidArgumentError(
+      llvm::formatv(
+          "Found unexpected resource op {0} during layout propagation.",
+          OpName(op))
+          .str());
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+ResourceSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts,
+    const llvm::DenseMap<int, Layout>& output_layouts) {
+  // For Assign* ops, propagate the resource tensor layout to the tensor if
+  // resource tensor layout exists.
+  if (llvm::isa<mlir::TF::AssignVariableOp, mlir::TF::AssignAddVariableOp,
+                mlir::TF::AssignSubVariableOp>(op)) {
+    if (input_layouts.find(0) != input_layouts.end()) {
+      auto resource_layout = input_layouts.lookup(0);
+      if (resource_layout.IsEmpty()) {
+        auto mesh = GetMeshOnParentCluster(op);
+        if (mesh.ok()) {
+          auto layout = Layout::ReplicatedOnMesh(
+              mesh.value(), ValueRank(op->getOpOperand(1).get()));
+
+          // Resource has an empty layout, propagate back replicated layout
+          // for the resource.
+          return llvm::DenseMap<int, Layout>({{0, layout}, {1, layout}});
+        }
+      }
+      return llvm::DenseMap<int, Layout>({{1, input_layouts.lookup(0)}});
+    }
+    return llvm::DenseMap<int, Layout>();
+  }
+  // Handling of these ops are no-ops.
+  if (llvm::isa<mlir::TF::DestroyResourceOp, mlir::TF::VarHandleOp,
+                mlir::TF::VarIsInitializedOp, mlir::TF::ReadVariableOp,
+                mlir::TF::SummaryWriterOp>(op)) {
+    return llvm::DenseMap<int, Layout>();
+  }
+
+  // Return an error if not any of the ops above.
+  return absl::InvalidArgumentError(
+      llvm::formatv(
+          "Found unexpected resource op {0} during layout propagation.",
+          OpName(op))
+          .str());
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/resource_spmd_expander.h

@@ -0,0 +1,46 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RESOURCE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RESOURCE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class ResourceSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  // Override the backward with both input and output layouts as arguments
+  // because we need to compute some operand layouts from existing operand
+  // layouts.
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_RESOURCE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/save_restore_spmd_expander.cc

@@ -0,0 +1,930 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/save_restore_spmd_expander.h"
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/save_restore_util.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/device_utils.h"
+#include "tensorflow/dtensor/mlir/dtensor_send_recv.h"
+#include "tensorflow/dtensor/mlir/ir/tf_dtensor.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+// Given a string tensor `prefix` of shape [k], produces a new string tensor
+// of shape [k*n] where n = number of devices in `mesh` by appending
+// device_id from [0, n) to `prefix`.
+//
+// For example:
+//   before:
+//     prefix = tf.Constant(["alice", "bob"])
+//     mesh.num_devices() = 2
+//   after =
+//     result = tf.Constant(["alice_device_0", "bob_device_0", "alice_device_1",
+//     "bob_device_1"])
+//
+// This is needed for DTensorCheckpointV2 tf.MergeV2Checkpoint SPMD expansion
+// to generate all candidate checkpoint prefix string that we generated
+// during tf.SaveV2 SPMD Expansion.
+StatusOr<mlir::Value> GetAllCandidateCheckpointPrefixes(
+    mlir::OpBuilder& builder, mlir::Value prefix, mlir::Value zero_const,
+    const Mesh& mesh) {
+  if (mesh.num_devices() == 0) return prefix;
+
+  mlir::Value new_prefix =
+      mlir::TF::AddOp::create(
+          builder, prefix.getLoc(),
+          mlir::dyn_cast<mlir::RankedTensorType>(prefix.getType()), prefix,
+          StringConst(builder, prefix.getLoc(),
+                      llvm::SmallVector<llvm::StringRef>(
+                          {DeviceSuffix(0, mesh.num_devices())})))
+          .getZ();
+
+  for (int64_t device_id = 1; device_id < mesh.num_devices(); ++device_id) {
+    mlir::Value prefix_plus_dtensor_suffix =
+        mlir::TF::AddOp::create(
+            builder, prefix.getLoc(),
+            mlir::dyn_cast<mlir::RankedTensorType>(prefix.getType()), prefix,
+            StringConst(builder, prefix.getLoc(),
+                        llvm::SmallVector<llvm::StringRef>(
+                            {DeviceSuffix(device_id, mesh.num_devices())})))
+            .getZ();
+
+    new_prefix =
+        mlir::TF::ConcatOp::create(builder, prefix.getLoc(),
+                                   /*output=*/prefix.getType(),
+                                   /*concat_dim=*/
+                                   zero_const,
+                                   llvm::SmallVector<mlir::Value, 4>{
+                                       new_prefix, prefix_plus_dtensor_suffix})
+            .getResult();
+  }
+  return new_prefix;
+}
+
+// Maps a device_id to a 0 based switch-case branch index.
+//
+// For Save/Restore ops, constructing a switch-case on all global devices is not
+// going to scale to larger slices as the function grows with the number of
+// devices. Instead, we only need to look at devices that are local to the
+// current host and generate SPMD for those. This allows the SPMD become
+// O(variables) since the local devices are constant for all device types.
+//
+// The challenge is that the switch-case op branch index is 0 based, meaning
+// that we can not use the device_id the same way in the global devices switch.
+// To deal with that, we will use this function to map the local_device_id on
+// the hosts into a 0 base, by constructing a 1D tensor with all local device
+// ids and using the index of the tensor as the branch index.
+//
+// A concrete example would be:
+//
+// local_device_ids = [1, 2, 4, 5, 6] -- We shouldn't assume continuity in
+// device_ids.
+//
+// switching device_id = [4]
+//
+// branch_index = idx_of(local_device_ids) = 2
+//
+// The tf op equivalent would be:
+// tf.reshape(tf.where(tf.equal(local_device_ids, device_id)), ())
+mlir::Value DeviceIdToLocalBranchIndex(
+    const mlir::Location& location,
+    const llvm::ArrayRef<int64_t>& local_device_ids, mlir::Value device_id,
+    mlir::OpBuilder& builder) {
+  mlir::Value local_device_id_tensors =
+      IntConst(builder, location,
+               llvm::SmallVector<int32_t>(local_device_ids.begin(),
+                                          local_device_ids.end()));
+  mlir::Value condition = mlir::TF::EqualOp::create(
+      builder, location, local_device_id_tensors, device_id,
+      /*incompatible_shape_error=*/builder.getBoolAttr(true));
+  auto where_op = mlir::TF::WhereOp::create(
+      builder, location,
+      mlir::RankedTensorType::get({1, 1}, builder.getI64Type()), condition);
+  // cast to int32 as where_op returns a int64 array.
+  auto cast_op = mlir::TF::CastOp::create(
+      builder, location,
+      mlir::RankedTensorType::get({1, 1}, builder.getI32Type()),
+      where_op.getResult());
+
+  // Reshape the output to i32 Scalar.
+  auto size_type = mlir::RankedTensorType::get({}, builder.getI32Type());
+  mlir::Value scalar_shape = mlir::TF::collection_ops_util::GetR1Const(
+      size_type.getShape(), builder, location);
+  auto branch_index_scalar = mlir::TF::ReshapeOp::create(
+      builder, location, mlir::ArrayRef<mlir::Type>{size_type},
+      mlir::ArrayRef<mlir::Value>{cast_op.getResult(), scalar_shape},
+      mlir::ArrayRef<mlir::NamedAttribute>{});
+
+  return branch_index_scalar.getResult();
+}
+
+// Builds a switch case function that only conditionally runs save with its
+// slice_specs on sharded tensors.
+//
+// Note that this would generate multiple prefixes for saving rather than the
+// single one passed in from the original op.
+// DTensor uses DTensorShardedPrefix to query the generated ones and use those
+// in MergeV2.
+StatusOr<mlir::TF::CaseOp> ConditionalSave(
+    mlir::TF::SaveV2Op original_save, const Mesh& mesh,
+    const absl::flat_hash_map<
+        int64_t, absl::flat_hash_map<int64_t, std::vector<std::string>>>&
+        saving_specs) {
+  mlir::ModuleOp module = original_save->getParentOfType<mlir::ModuleOp>();
+  if (!module)
+    return absl::InternalError(
+        "SaveV2 op isn't enclosed inside a mlir::ModuleOp");
+
+  mlir::SymbolTable symbol_table(module);
+
+  mlir::OpBuilder builder(original_save);
+  const auto& location = original_save.getLoc();
+
+  llvm::SmallVector<mlir::func::FuncOp, 8> branch_funs;
+
+  // Best effort extraction on shape_and_slices and verify they are empty. If
+  // the extraction failed to just ignore those values and work as if those are
+  // empty.
+  llvm::SmallVector<std::string, 4> original_shape_and_slices;
+  const absl::Status extraction_status = ExtractConstStringVectorFromValue(
+      original_save.getShapeAndSlices(), original_shape_and_slices);
+  if (extraction_status.ok()) {
+    for (const std::string& shape_and_slice : original_shape_and_slices) {
+      if (!shape_and_slice.empty())
+        return absl::InvalidArgumentError(
+            absl::StrCat("DTensor SaveV2 requires shape_and_slices() field to "
+                         "be empty for tensors, but get : ",
+                         shape_and_slice));
+    }
+  } else {
+    VLOG(2) << "Failed to extract and verify shape_and_slices() from "
+               "original SaveV2 op. SaveV2 SPMD would proceed as if "
+               "shape_and_slices are empty for all the tensors.";
+  }
+
+  // Branch functions have shared function type, where input is simply all the
+  // inputs from origial saveV2 and no outputs.
+  auto func_type = mlir::FunctionType::get(builder.getContext(),
+                                           original_save.getOperandTypes(),
+                                           /*results=*/{});
+  // Only generates save functions for devices that is local to the client.
+  // This would mean that we will run different functions on different client,
+  // but it would be fine as we're running on CPU for this.
+  for (int device_id : mesh.local_device_ids()) {
+    //  If saving_spec doesn't contain the device_id, then that device_id is a
+    //  no-op on the save.
+    const auto& it = saving_specs.find(device_id);
+    if (it == saving_specs.end()) {
+      // Builds place holder for the no_op function, which takes the exact same
+      // args as the original save op and returns nothing.
+      mlir::func::FuncOp no_op = mlir::func::FuncOp::create(
+          location,
+          llvm::formatv("{0}_no_op_on_device_{1}_{2}", OpName(original_save),
+                        device_id, OpHash(original_save))
+              .str(),
+          func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+      // Set function visibility to private to indicate that it is only used in
+      // this module.
+      no_op.setVisibility(mlir::SymbolTable::Visibility::Private);
+      symbol_table.insert(no_op);
+
+      mlir::Block* fn_block = no_op.addEntryBlock();
+      mlir::OpBuilder fn_builder = mlir::OpBuilder::atBlockBegin(fn_block);
+      mlir::TF::NoOp::create(fn_builder, location);
+      mlir::func::ReturnOp::create(fn_builder, location);
+
+      branch_funs.push_back(no_op);
+    } else {
+      const absl::flat_hash_map<int64_t, std::vector<std::string>>&
+          per_device_specs = it->second;
+
+      // Build the new SaveV2 that contains proper SliceSpec on this device.
+      // tensor_names and slice_spec would be concatted into a 1d string tensor.
+      mlir::func::FuncOp new_save = mlir::func::FuncOp::create(
+          location,
+          llvm::formatv("{0}_save_op_on_device_{1}_{2}", OpName(original_save),
+                        device_id, OpHash(original_save))
+              .str(),
+          func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+      // Set function visibility to private to indicate that it is only used in
+      // this module.
+      new_save.setVisibility(mlir::SymbolTable::Visibility::Private);
+      symbol_table.insert(new_save);
+
+      mlir::Block* fn_block = new_save.addEntryBlock();
+      mlir::OpBuilder fn_builder = mlir::OpBuilder::atBlockBegin(fn_block);
+
+      mlir::Value prefix = new_save.getArgument(0);
+      mlir::Value tensor_names = new_save.getArgument(1);
+      // It is currently unsupported if user passes in shape_and_slices.
+      // TODO(hthu): Implement this.
+      // mlir::Value shape_and_slices = new_save.getArgument(2);
+
+      // First run a split op on the tensor_names so that we can use the proper
+      // splitted output(one of the tensor_name) to reconstruct tensor_names
+      // field in the new SaveV2 op.
+      TF_ASSIGN_OR_RETURN(llvm::ArrayRef<int64_t> tensor_names_shape,
+                          GetGlobalShapeOfValueFromDTensorLayout(
+                              original_save.getTensorNames()));
+      if (tensor_names_shape.size() != 1)
+        return absl::InternalError(
+            llvm::formatv("SaveV2 op got `tensor_names` with rank {0}) but "
+                          "expects rank to be 1.",
+                          tensor_names_shape.size())
+                .str());
+      mlir::TF::SplitOp name_splits;
+      TF_RETURN_IF_ERROR(CreateSplitOp(/*num_split=*/tensor_names_shape[0],
+                                       /*split_dimension=*/0, location,
+                                       /*src_input=*/tensor_names, &fn_builder,
+                                       &name_splits));
+
+      // Builds the per device saving spec, that takes care of tensor_name
+      // uniqueness requirement. Each save op should use new_tensor_indices and
+      // new_specs to map the corresponding saving tensor and its slice spec.
+      SaveOpSpecs specs = BuildPerDeviceSave(
+          fn_builder, per_device_specs, device_id, prefix, mesh.num_devices());
+      const std::vector<std::vector<int>>& new_tensor_indices =
+          specs.tensor_indices;
+      const std::vector<std::vector<std::string>>& new_specs =
+          specs.shape_and_slice_spec;
+
+      // Prepare corresponding SaveOp arguments.
+      for (int save_op_index = 0; save_op_index < new_tensor_indices.size();
+           ++save_op_index) {
+        llvm::SmallVector<mlir::Value, 4> new_tensor_names;
+        llvm::SmallVector<std::string, 4> new_shape_and_slices;
+        llvm::SmallVector<mlir::Value, 4> new_tensors;
+
+        // Per_device_specs records the index of the tensor_names from the
+        // original save, and all slice_specs needed to save that tensor.
+        // The corresponding saving tensor can be found in the original save op
+        // by adding 3 to the index (as 0, 1, 2) are fixed inputs for prefix,
+        // tensor_names and shapes_and_slices.
+        for (int i = 0; i < new_tensor_indices[save_op_index].size(); ++i) {
+          int tensor_name_index = new_tensor_indices[save_op_index][i];
+          int tensor_index = 3 + tensor_name_index;
+          new_tensor_names.push_back(name_splits.getResult(tensor_name_index));
+          new_shape_and_slices.push_back(new_specs[save_op_index][i]);
+          new_tensors.push_back(new_save.getArgument(tensor_index));
+        }
+        // Build the new SaveV2 op.
+        mlir::Value tensor_names = new_tensor_names[0];
+        if (new_tensor_names.size() > 1) {
+          // For tensor_names that has more than 1 entry, we concat the list of
+          // names into a 1d vector.
+          tensor_names =
+              mlir::TF::ConcatOp::create(
+                  fn_builder, location,
+                  /*output=*/original_save.getTensorNames().getType(),
+                  /*concat_dim=*/
+                  IntConst(fn_builder, location, /*values=*/{0}),
+                  new_tensor_names)
+                  .getResult();
+        }
+
+        mlir::TF::SaveV2Op::create(
+            fn_builder, location, specs.new_prefixes[save_op_index],
+            /*tensor_names=*/tensor_names,
+            /*shape_and_slices=*/
+            StringConst(
+                fn_builder, location,
+                llvm::SmallVector<llvm::StringRef>(new_shape_and_slices.begin(),
+                                                   new_shape_and_slices.end())),
+            new_tensors);
+      }
+      branch_funs.push_back(new_save);
+      mlir::func::ReturnOp::create(fn_builder, location);
+    }
+  }
+
+  llvm::SmallVector<mlir::Attribute, 4> symbols;
+  for (auto& func : branch_funs)
+    symbols.push_back(mlir::SymbolRefAttr::get(func));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value device_id, DeviceId(original_save));
+  llvm::SmallVector<int64_t> local_device_ids(mesh.local_device_ids().begin(),
+                                              mesh.local_device_ids().end());
+  mlir::Value branch_index = DeviceIdToLocalBranchIndex(
+      location, local_device_ids, device_id, builder);
+
+  auto case_op =
+      mlir::TF::CaseOp::create(builder, location,
+                               // SaveV2 doesn't return a value.
+                               /*output=*/llvm::ArrayRef<mlir::Type>{},
+                               /*branch_index=*/branch_index,
+                               /*input=*/original_save.getOperands(),
+                               /*branches=*/builder.getArrayAttr(symbols),
+                               /*is_stateless=*/builder.getBoolAttr(false));
+
+  return case_op;
+}
+
+StatusOr<mlir::Operation*> ExpandSaveV2Op(mlir::Operation* op) {
+  if (!llvm::isa<mlir::TF::SaveV2Op>(op)) {
+    return absl::InvalidArgumentError(
+        llvm::formatv("Expecting SaveV2Op but got {0}", OpName(op)).str());
+  }
+
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto save_v2 = mlir::cast<mlir::TF::SaveV2Op>(op);
+
+  mlir::OpBuilder builder(save_v2);
+  std::vector<SavingTensorMetadata> metadata;
+  for (const auto& it : llvm::enumerate(save_v2.getTensors())) {
+    mlir::Value tensor = it.value();
+    // We use index to select the tensor names and shape_and_slices from the
+    // inputs. This is generic regardless whether the inputs are constants or
+    // just arguments.
+    int index = it.index();
+    TF_ASSIGN_OR_RETURN(std::optional<Layout> layout,
+                        ExtractLayoutFromOperand(tensor));
+    if (!layout)
+      return absl::InvalidArgumentError(
+          "layout is required when saving a DTensor but find no layout "
+          "attached");
+
+    TF_ASSIGN_OR_RETURN(llvm::ArrayRef<int64_t> tensor_shape,
+                        GetGlobalShapeOfValueFromDTensorLayout(it.value()));
+
+    metadata.push_back(SavingTensorMetadata(
+        index, std::vector<int64_t>(tensor_shape.begin(), tensor_shape.end()),
+        *layout));
+  }
+  TF_ASSIGN_OR_RETURN(auto saving_specs, BuildSavingSpec(metadata));
+
+  // Now we have a complete map on device_id and its saving tensors and specs.
+  // Build a switch case conditioned on device_id and do saves properly.
+  TF_ASSIGN_OR_RETURN(mlir::TF::CaseOp case_op,
+                      ConditionalSave(save_v2, mesh, saving_specs));
+
+  save_v2->replaceAllUsesWith(case_op);
+  save_v2->erase();
+
+  return case_op.getOperation();
+}
+
+// SPMD Expander for MergeV2.
+//
+// The op is expected to have one and only one of the prefix input, which is
+// used as a key to query all the saved shard prefixed generated in SaveV2 op
+// SPMD.
+//
+// The expanded MergeV2 contains all the shard_prefix generated, and only runs
+// on Device 0.
+StatusOr<mlir::Operation*> ExpandMergeV2Op(mlir::Operation* op) {
+  mlir::TF::MergeV2CheckpointsOp merge_v2 =
+      mlir::dyn_cast<mlir::TF::MergeV2CheckpointsOp>(op);
+  if (!merge_v2) {
+    return absl::InvalidArgumentError(
+        llvm::formatv("Expecting MergeV2CheckpointsOp but got {0}", OpName(op))
+            .str());
+  }
+
+  // Build an if op that only runs MergeV2 on device 0. Note that if condition
+  // is tested false when device_id == 0, so that the `then` branch will be
+  // no_op while the else branch will be the real MergeV2 op that is on device
+  // 0.
+  auto module = merge_v2->getParentOfType<mlir::ModuleOp>();
+  mlir::SymbolTable symbol_table(module);
+  auto location = merge_v2.getLoc();
+  mlir::OpBuilder builder(merge_v2);
+
+  TF_ASSIGN_OR_RETURN(auto device_id, DeviceId(merge_v2));
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Emit a barrier before every merge to ensure that every device's save
+  // finishes running before the merge is called and use this as the
+  // suffix for the new prefix string tensors so that it is not
+  // removed from dead node analysis.
+  //
+  // Note that this barrier is implemented with a DTensorAllReduce which
+  // gets lowered to a CollectiveReduce op, which is a SideEffect op. Since
+  // both Merge and Save ops are maximally conservative SideEffect ops,
+  // the execution order is guaranteed to be Save -> Barrier -> Merge, which
+  // is the desired execution order to guarantee correctness of IO.
+  TF_ASSIGN_OR_RETURN(
+      mlir::Operation * rank_1_zero,
+      EmitBarrierWithConstValue(builder, location, mesh, /*value=*/0));
+
+  // Add a zero value to the end of the function types, used for
+  // adding a barrier before every merge op.
+  llvm::SmallVector<mlir::Type> input_types(merge_v2.getOperandTypes());
+  input_types.push_back(rank_1_zero->getResult(0).getType());
+
+  auto func_type = mlir::FunctionType::get(builder.getContext(), input_types,
+                                           llvm::ArrayRef<mlir::Type>{});
+  // Build then_func that is the branch of device_id != 0, which only contains a
+  // single NoOp.
+  mlir::func::FuncOp then_func = mlir::func::FuncOp::create(
+      location,
+      llvm::formatv("{0}_then_func_{1}", OpName(merge_v2), OpHash(merge_v2))
+          .str(),
+      func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+  // Set function visibility to private to indicate that it is only used in
+  // this module.
+  then_func.setVisibility(mlir::SymbolTable::Visibility::Private);
+  mlir::Block* then_fn_block = then_func.addEntryBlock();
+  mlir::OpBuilder then_fn_builder =
+      mlir::OpBuilder::atBlockBegin(then_fn_block);
+  mlir::TF::NoOp::create(then_fn_builder, location);
+  mlir::func::ReturnOp::create(then_fn_builder, location);
+
+  // Build else_func that is the branch of device_id == 0.
+  // The else func is just the original MergeV2 itself.
+  mlir::func::FuncOp else_func = mlir::func::FuncOp::create(
+      location,
+      llvm::formatv("{0}_else_func_{1}", OpName(merge_v2), OpHash(merge_v2))
+          .str(),
+      func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+  // Set function visibility to private to indicate that it is only used in
+  // this module.
+  else_func.setVisibility(mlir::SymbolTable::Visibility::Private);
+
+  mlir::Block* else_fn_block = else_func.addEntryBlock();
+  mlir::OpBuilder else_fn_builder =
+      mlir::OpBuilder::atBlockBegin(else_fn_block);
+  mlir::Value checkpoint_prefixes = else_fn_block->getArgument(0);
+
+  mlir::Value zero_value =
+      else_fn_block->getArgument(else_fn_block->getNumArguments() - 1);
+  TF_ASSIGN_OR_RETURN(
+      checkpoint_prefixes,
+      GetAllCandidateCheckpointPrefixes(else_fn_builder, checkpoint_prefixes,
+                                        zero_value, mesh));
+
+  mlir::Value destination_prefixes = else_fn_block->getArgument(1);
+
+  mlir::TF::MergeV2CheckpointsOp::create(
+      else_fn_builder, location, checkpoint_prefixes, destination_prefixes,
+      /*delete_old_dirs=*/
+      else_fn_builder.getBoolAttr(merge_v2.getDeleteOldDirs()),
+      /*allow_missing_files=*/else_fn_builder.getBoolAttr(true));
+
+  mlir::func::ReturnOp::create(else_fn_builder, location);
+
+  symbol_table.insert(then_func);
+  symbol_table.insert(else_func);
+
+  auto new_operands = llvm::to_vector<4>(merge_v2.getOperands());
+  new_operands.push_back(rank_1_zero->getResult(0));
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value zero_scalar,
+      CreateZeroScalarConst(
+          builder, location,
+          mlir::cast<mlir::TensorType>(device_id.getType()).getElementType()));
+
+  mlir::TF::NotEqualOp not_equal = mlir::TF::NotEqualOp::create(
+      builder, location, device_id, zero_scalar,
+      /*incompatible_shape_error=*/builder.getBoolAttr(false));
+
+  auto if_op = mlir::TF::IfOp::create(
+      builder, location, then_func.getFunctionType().getResults(),
+      /*cond=*/not_equal.getResult(),
+      /*input=*/new_operands,
+      /*then_branch=*/then_func.getSymName(),
+      /*else_branch=*/else_func.getSymName(), /*is_stateless=*/false);
+
+  merge_v2->replaceAllUsesWith(if_op);
+  merge_v2.erase();
+  return if_op.getOperation();
+}
+
+// SPMD Expander for RestoreV2 op.
+//
+// Both tf.RestoreV2 and DTensorRestoreV2 op will be expanded the same way.
+// That is, they will be updated to only restore the slice for the
+// given device_id. For replicated tensors, that would be the full tensor slice.
+// For sharded tensors, we compute its slice using device coordinates and tensor
+// layout.
+//
+// `global_shapes` refers to the global shapes of the outputs of the op.
+// `layouts` refers to the output layouts of the op.
+StatusOr<mlir::Operation*> ExpandRestoreV2OpHelper(
+    mlir::Operation* op, std::vector<std::vector<int64_t>> global_shapes,
+    std::vector<Layout> layouts, std::vector<mlir::Type> output_types,
+    mlir::MutableOperandRange shapes_and_slices_mutable) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  // Prepare for building CaseOp.
+  mlir::ModuleOp module = op->template getParentOfType<mlir::ModuleOp>();
+  if (!module)
+    return absl::InternalError(
+        "DTensorRestoreV2 op isn't enclosed inside a mlir::ModuleOp");
+
+  mlir::SymbolTable symbol_table(module);
+
+  mlir::OpBuilder builder(op);
+  const auto& location = op->getLoc();
+
+  // Tracks case branch functions for each local_device_id.
+  llvm::SmallVector<mlir::func::FuncOp> branch_funcs(
+      mesh.local_device_ids().size());
+  // Stores restore ops for each device_id in a function, that is suitable for
+  // feeding into a CaseOp.
+  //
+  // Branch functions have shared function type as original restore_v2.
+  const auto func_type =
+      mlir::FunctionType::get(builder.getContext(), op->getOperandTypes(),
+                              mlir::TypeRange(output_types));
+
+  for (int local_device_idx = 0;
+       local_device_idx < mesh.local_device_ids().size(); ++local_device_idx) {
+    int device_id = mesh.local_device_ids()[local_device_idx];
+    TF_ASSIGN_OR_RETURN(const DeviceLocation& coords,
+                        mesh.device_location(device_id));
+
+    llvm::SmallVector<std::string> new_shapes_and_slices(op->getNumResults());
+
+    // For each tensor, build its restore shape_and_slice.
+    for (const auto& it : llvm::enumerate(llvm::zip(global_shapes, layouts))) {
+      std::vector<int64_t> global_shape = std::get<0>(it.value());
+      Layout layout = std::get<1>(it.value());
+      // Fully replicated tensor does not need a slice and spec field and we
+      // simply leave it as empty string. Note that Non-DTensor restore will
+      // use replicated layout from SaveSpec.
+      if (layout.IsFullyReplicated()) {
+        new_shapes_and_slices[it.index()] = "";
+        continue;
+      }
+
+      TF_ASSIGN_OR_RETURN(
+          std::vector<std::string> slice_specs,
+          SliceSpecOnDevice(layout, mesh, coords, global_shape));
+
+      // Concat shape and slice specs
+      new_shapes_and_slices[it.index()] =
+          llvm::formatv("{0} {1}", absl::StrJoin(global_shape, " "),
+                        absl::StrJoin(slice_specs, ":"))
+              .str();
+    }
+
+    // Builds the restore op on device_id.
+    mlir::OpBuilder builder(op);
+    shapes_and_slices_mutable.assign(StringConst(
+        builder, op->getLoc(),
+        llvm::SmallVector<llvm::StringRef>(new_shapes_and_slices.begin(),
+                                           new_shapes_and_slices.end())));
+    mlir::func::FuncOp device_restore_fn = mlir::func::FuncOp::create(
+        location,
+        llvm::formatv("{0}_on_device_{1}_{2}", OpName(op), device_id,
+                      OpHash(op))
+            .str(),
+        func_type, llvm::ArrayRef<mlir::NamedAttribute>{});
+    // Set function visibility to private to indicate that it is only used in
+    // this module.
+    device_restore_fn.setVisibility(mlir::SymbolTable::Visibility::Private);
+    symbol_table.insert(device_restore_fn);
+
+    mlir::Block* fn_block = device_restore_fn.addEntryBlock();
+    mlir::OpBuilder fn_builder = mlir::OpBuilder::atBlockBegin(fn_block);
+    mlir::Value prefix = device_restore_fn.getArgument(0);
+    mlir::Value tensor_names = device_restore_fn.getArgument(1);
+    // Constructs shapes and slices ourselves while reusing all other
+    // arguments.
+    auto new_restore_v2 = mlir::TF::RestoreV2Op::create(
+        fn_builder, location, mlir::TypeRange(output_types), prefix,
+        tensor_names,
+        StringConst(
+            fn_builder, location,
+            llvm::SmallVector<llvm::StringRef>(new_shapes_and_slices.begin(),
+                                               new_shapes_and_slices.end())));
+    mlir::func::ReturnOp::create(fn_builder, location,
+                                 new_restore_v2.getResults());
+
+    branch_funcs[local_device_idx] = device_restore_fn;
+  }
+
+  // Builds the final case op.
+  llvm::SmallVector<mlir::Attribute, 4> symbols;
+  for (auto& func : branch_funcs)
+    symbols.push_back(mlir::SymbolRefAttr::get(func));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value device_id, DeviceId(op));
+  llvm::SmallVector<int64_t> local_device_ids(mesh.local_device_ids().begin(),
+                                              mesh.local_device_ids().end());
+  mlir::Value branch_index = DeviceIdToLocalBranchIndex(
+      location, local_device_ids, device_id, builder);
+
+  auto case_op =
+      mlir::TF::CaseOp::create(builder, location,
+                               /*output=*/mlir::TypeRange(output_types),
+                               /*branch_index=*/branch_index,
+                               /*input=*/op->getOperands(),
+                               /*branches=*/builder.getArrayAttr(symbols),
+                               /*is_stateless=*/builder.getBoolAttr(false));
+
+  op->replaceAllUsesWith(case_op);
+  op->erase();
+
+  return case_op.getOperation();
+}
+
+// DTensorRestoreV2 op has layouts and shapes as the attribute of the op
+// itself. We extract those attributes and call the helper expander.
+StatusOr<mlir::Operation*> ExpandDTensorRestoreV2Op(mlir::Operation* op) {
+  mlir::TF::DTensorRestoreV2Op restore_v2 =
+      mlir::dyn_cast<mlir::TF::DTensorRestoreV2Op>(op);
+  if (!restore_v2) {
+    return absl::InvalidArgumentError(
+        llvm::formatv("Expecting DTensorRestoreV2Op but got {0}", OpName(op))
+            .str());
+  }
+
+  mlir::ArrayAttr input_shapes_attr =
+      restore_v2->getAttrOfType<mlir::ArrayAttr>("input_shapes");
+  if (!input_shapes_attr) {
+    return absl::InvalidArgumentError(
+        "DTensorRestoreV2Op requires input_shapes attributes.");
+  }
+
+  std::vector<std::vector<int64_t>> input_shapes;
+  input_shapes.reserve(input_shapes_attr.size());
+  for (const auto& shape : input_shapes_attr) {
+    mlir::TF::ShapeAttr shape_attr = mlir::cast<mlir::TF::ShapeAttr>(shape);
+    if (!shape_attr.hasStaticShape()) {
+      return absl::InvalidArgumentError(
+          llvm::formatv("DTensorRestoreV2Op requires statically known input "
+                        "shape, but got non-static shape: {0}.",
+                        shape_attr)
+              .str());
+    }
+    input_shapes.push_back(std::vector<int64_t>(shape_attr.getShape().begin(),
+                                                shape_attr.getShape().end()));
+  }
+
+  mlir::ArrayAttr input_layouts_attr = restore_v2.getInputLayouts();
+  if (!input_layouts_attr) {
+    return absl::InvalidArgumentError(
+        "DTensorRestoreV2Op requires input_layouts attributes.");
+  }
+  std::vector<Layout> input_layouts;
+  input_layouts.reserve(input_layouts_attr.size());
+  for (const auto& layout : input_layouts_attr.getValue().vec()) {
+    input_layouts.push_back(
+        Layout::FromString(
+            mlir::cast<mlir::StringAttr>(layout).getValue().str())
+            .value());
+  }
+
+  return ExpandRestoreV2OpHelper(
+      op, input_shapes, input_layouts,
+      std::vector<mlir::Type>(op->getResultTypes().begin(),
+                              op->getResultTypes().end()),
+      restore_v2.getShapeAndSlicesMutable());
+}
+
+// Extract the layout and shapes the normal way. By this time, we should
+// have all necessary DTensorLayout op as the outputs of each op
+// and the correct Type shapes and dtypes as the outputs of the tf.RestoreV2
+// op.
+//
+// Call the helper expander function with those shapes and layouts.
+StatusOr<mlir::Operation*> ExpandRestoreV2Op(mlir::Operation* op) {
+  // Fetch the shape of each output.
+  std::vector<std::vector<int64_t>> global_shapes;
+  global_shapes.reserve(op->getNumResults());
+
+  // This is subtle. For tf.train.Checkpoint.save_counter scalar variable,
+  // this variable may not yet be created by the time we call
+  // Checkpoint.restore.
+  //
+  // In this case, the tf.RestoreV2 is called eagerly, and thus there is no
+  // tf.AssignVariable op. This means that we cannot infer the shapes and layout
+  // from previous pass CreateDTensorInferShapesForRestoreV2Op.
+  //
+  // But for save_counter, we know this is always replicated, and we can just
+  // return the op itself. For now, we will do this hacky way, but eventually
+  // we need to generalize restoring variables that are not yet created.
+  //
+  // TODO(b/235373719) Generalize support for checkpoint restoration for
+  // variables that are not yet created.
+  if (op->getNumResults() == 1 && !GetShapeOfValue(op->getResult(0)).ok()) {
+    return op;
+  }
+
+  for (auto result : op->getResults()) {
+    global_shapes.push_back(GetShapeOfValue(result).value());
+  }
+
+  // Fetch the layout of each output.
+  TF_ASSIGN_OR_RETURN(std::vector<Layout> layouts,
+                      ExtractRequiredLayoutFromOp(op));
+
+  // Calculate the new local type range needed for the new RestoreV2Op we will
+  // emit.
+  std::vector<mlir::Type> new_types;
+  new_types.reserve(op->getNumResults());
+
+  for (const auto& it :
+       llvm::zip(op->getResultTypes(), global_shapes, layouts)) {
+    mlir::Type type = std::get<0>(it);
+    std::vector<int64_t>& shape = std::get<1>(it);
+    Layout& layout = std::get<2>(it);
+    new_types.push_back(mlir::RankedTensorType::get(
+        layout.LocalShapeFromGlobalShape(shape),
+        mlir::dyn_cast<mlir::RankedTensorType>(type).getElementType()));
+  }
+
+  return ExpandRestoreV2OpHelper(
+      op, global_shapes, layouts, new_types,
+      mlir::dyn_cast<mlir::TF::RestoreV2Op>(op).getShapeAndSlicesMutable());
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> SaveRestoreSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  if (llvm::isa<mlir::TF::SaveV2Op>(op)) {
+    return ExpandSaveV2Op(op);
+  }
+  if (llvm::isa<mlir::TF::MergeV2CheckpointsOp>(op)) {
+    return ExpandMergeV2Op(op);
+  }
+  if (llvm::isa<mlir::TF::DTensorRestoreV2Op>(op)) {
+    return ExpandDTensorRestoreV2Op(op);
+  }
+  if (llvm::isa<mlir::TF::RestoreV2Op>(op)) {
+    return ExpandRestoreV2Op(op);
+  }
+
+  return absl::UnimplementedError(
+      llvm::formatv("SPMD for op : {0} is not implemented ", OpName(op)).str());
+}
+
+// Find all the resource tensor layouts attached to the AssignVariableOp
+// that `restore_op` is restoring to.
+StatusOr<llvm::SmallVector<Layout>> GetLayoutsFromAssignVariableOps(
+    mlir::ModuleOp module, mlir::TF::RestoreV2Op* restore_op) {
+  llvm::SmallVector<Layout> layouts(restore_op->getNumResults());
+
+  for (auto result : restore_op->getResults()) {
+    // Find the AssignVariableOp connected to this output. There should only
+    // be at most one IdentityOp and one DTensorSend between this result
+    // and the AssignVariableOp.
+    for (auto consuming_op : result.getUsers()) {
+      // To get to the AssignVariableOp that consumes `result`, we expect
+      // an IdentityOp, CastOp, or a DTensorSend op on the path. So, skip past
+      // these ops first.
+      while (llvm::isa<mlir::TF::CastOp, mlir::TF::IdentityOp,
+                       mlir::TF::RelayoutOp, mlir::TF::DTensorLayout,
+                       mlir::TF::DTensorSend>(consuming_op)) {
+        if (auto send_op =
+                mlir::dyn_cast_or_null<mlir::TF::DTensorSend>(consuming_op)) {
+          TF_ASSIGN_OR_RETURN(
+              consuming_op, GetCorrespondingDTensorSendRecvOp(module, send_op));
+        }
+        auto next_op = consuming_op->getResult(0).getUsers();
+        if (next_op.empty()) {
+          return absl::InternalError(
+              "Expected a result of an identity op to be consumed by another "
+              "op, but was empty during RestoreV2 Expansion.");
+        }
+        consuming_op = *next_op.begin();
+      }
+      // We skipped past ops like Identity and Send's. There might be an
+      // AssignVariableOp now.
+      if (auto assign_op = llvm::dyn_cast_or_null<mlir::TF::AssignVariableOp>(
+              consuming_op)) {
+        TF_ASSIGN_OR_RETURN(auto layout, ExtractRequiredLayoutFromOperand(
+                                             assign_op.getResource()));
+        layouts[result.getResultNumber()] = layout;
+        break;
+      }
+    }
+  }
+  return layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SaveRestoreSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts,
+    const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Save op doesn't have return values.
+  if (llvm::isa<mlir::TF::SaveV2Op, mlir::TF::MergeV2CheckpointsOp>(op)) {
+    return llvm::DenseMap<int, Layout>();
+  }
+  if (llvm::isa<mlir::TF::RestoreV2Op>(op)) {
+    // If there are already output layouts specified, this means that
+    // we are in the Late Variable Creation restoration. For this path,
+    // the output layout is already specified, through the default layout
+    // scope. So just return that layout.
+    if (!output_layouts.empty()) return output_layouts;
+
+    mlir::ModuleOp module_op = op->getParentOfType<mlir::ModuleOp>();
+    mlir::TF::RestoreV2Op restore_v2 = mlir::cast<mlir::TF::RestoreV2Op>(op);
+    TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+    if (!mesh.is_cpu_mesh()) {
+      return absl::InvalidArgumentError(
+          llvm::formatv(
+              "RestoreV2Op must run on a CPU mesh, but was running on: {0}",
+              mesh.ToString())
+              .str());
+    }
+    // Extract the layout of each resource tensor from the AssignVariableOp
+    // consuming each result. This layout sharding will be used as the
+    // output layout for each result tensor.
+    TF_ASSIGN_OR_RETURN(
+        auto layouts, GetLayoutsFromAssignVariableOps(module_op, &restore_v2));
+    if (layouts.size() != restore_v2.getNumResults()) {
+      return absl::InternalError(
+          llvm::formatv("Failed to get {0} output layouts "
+                        "for RestoreV2Op. Got {1} layouts.",
+                        restore_v2.getNumResults(), layouts.size())
+              .str());
+    }
+    llvm::DenseMap<int, Layout> output_layouts(restore_v2.getNumResults());
+
+    // Change the mesh of each layout to `mesh` since RestoreOp always runs on
+    // the CPU.
+    for (int i = 0; i < layouts.size(); ++i) {
+      TF_ASSIGN_OR_RETURN(
+          Layout host_mesh_layout,
+          Layout::GetLayout(Layout::LayoutType::kStatic,
+                            layouts[i].sharding_spec_strs(), mesh));
+      output_layouts[i] = host_mesh_layout;
+    }
+    return output_layouts;
+  }
+  if (llvm::isa<mlir::TF::DTensorRestoreV2Op>(op)) {
+    mlir::TF::DTensorRestoreV2Op restore_v2 =
+        mlir::cast<mlir::TF::DTensorRestoreV2Op>(op);
+    llvm::DenseMap<int, Layout> output_layouts(restore_v2.getNumResults());
+    // Output layout is simply the layout from the arguments.
+    for (const auto& it : llvm::enumerate(restore_v2.getInputLayouts())) {
+      TF_ASSIGN_OR_RETURN(
+          Layout layout,
+          Layout::FromString(
+              mlir::cast<mlir::StringAttr>(it.value()).getValue().str()));
+      output_layouts[it.index()] = layout;
+    }
+    return output_layouts;
+  }
+  return absl::UnimplementedError(
+      llvm::formatv("Layout propagation for op : {0} is not implemented",
+                    OpName(op))
+          .str());
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SaveRestoreSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/save_restore_spmd_expander.h

@@ -0,0 +1,68 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SAVE_RESTORE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SAVE_RESTORE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class SaveRestoreSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// SPMD expansion for DTensorShardPrefix op.
+// The expansion merely sets the output layout as a rank 1 replicated tensor.
+class DTensorShardPrefixSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override {
+    return op;
+  };
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override {
+    TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+    return llvm::DenseMap<int, Layout>(
+        {{0, Layout::ReplicatedOnMesh(mesh, /*rank=*/1)}});
+  }
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override {
+    return llvm::DenseMap<int, Layout>();
+  }
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SAVE_RESTORE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.cc

@@ -0,0 +1,457 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.h"
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/optional.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/collection_ops_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+StatusOr<Layout> GetOutputLayout(const absl::optional<Layout>& tensor_layout,
+                                 int tensor_rank,
+                                 const absl::optional<Layout>& updates_layout,
+                                 int updates_rank, const Mesh& mesh) {
+  // The first tensor_rank - update_rank dimensions of the output should be set
+  // to replicated. The remainder are set from tensor_layout and updates_layout
+  // with tensor_layout taking priority, as it is generally larger than updates
+  // (as unsharding updates is faster).
+  std::vector<std::string> output_specs(tensor_rank);
+
+  // The number of dimensions at the start of the tensor input that are used
+  // for the index, also the size of the second dimension of the indices tensor.
+  const int index_dimensions = tensor_rank - (updates_rank - 1);
+
+  for (int i = 0; i < tensor_rank; ++i) output_specs[i] = Layout::kUnshardedDim;
+
+  absl::flat_hash_set<std::string> used_mesh_dims;
+
+  if (tensor_layout) {
+    for (int i = index_dimensions; i < tensor_rank; ++i) {
+      output_specs[i] = tensor_layout->sharding_spec(i);
+      if (Layout::IsShardedDimension(output_specs[i]))
+        used_mesh_dims.emplace(output_specs[i]);
+    }
+  }
+
+  if (updates_layout) {
+    for (int i = index_dimensions; i < tensor_rank; ++i) {
+      const auto& update_spec =
+          updates_layout->sharding_spec(i - index_dimensions + 1);
+
+      if (Layout::IsUnshardedDimension(output_specs[i]) &&
+          Layout::IsShardedDimension(update_spec) &&
+          !used_mesh_dims.contains(update_spec))
+        output_specs[i] = update_spec;
+    }
+  }
+
+  return Layout::GetLayout(output_specs, mesh);
+}
+
+template <typename OpType>
+StatusOr<mlir::Operation*> TensorScatterOpExpand(mlir::Operation* op) {
+  auto scatter_op = llvm::cast<OpType>(op);
+  TF_ASSIGN_OR_RETURN(auto tensor_layout,
+                      ExtractLayoutFromOperand(scatter_op.getTensor()));
+  TF_ASSIGN_OR_RETURN(auto indices_layout,
+                      ExtractLayoutFromOperand(scatter_op.getIndices()));
+  TF_ASSIGN_OR_RETURN(auto updates_layout,
+                      ExtractLayoutFromOperand(scatter_op.getUpdates()));
+  TF_ASSIGN_OR_RETURN(auto output_layout,
+                      ExtractSingleLayoutFromOp(scatter_op));
+
+  const int tensor_rank = ValueRank(scatter_op.getTensor());
+  const int updates_rank = ValueRank(scatter_op.getUpdates());
+
+  if (tensor_rank == -1 || updates_rank == -1)
+    return absl::InvalidArgumentError("all inputs must have valid rank.");
+
+  // Get the global shape of all inputs as we need them for the Relayout
+  // operations.
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> tensor_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(scatter_op.getTensor()));
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> indices_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(scatter_op.getIndices()));
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> updates_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(scatter_op.getUpdates()));
+
+  // Start by relaying out the inputs. Indices should replicated.
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value new_indices,
+      EmitRelayout(scatter_op.getIndices(), *indices_layout,
+                   Layout::ReplicatedOnMesh(indices_layout->mesh(),
+                                            indices_shape.size())));
+
+  // Create intermediate layouts for tensors and updates. Since the layout of
+  // tensor and the output of the local tensor-scatter are the same we can reuse
+  // GetOutputLayout.
+  // If the true output layout is even more sharded, we could forward those
+  // shardings here for even better performance.
+  TF_ASSIGN_OR_RETURN(
+      Layout pre_output_layout,
+      GetOutputLayout(tensor_layout, tensor_rank, updates_layout, updates_rank,
+                      tensor_layout->mesh()));
+
+  std::vector<std::string> updates_specs(updates_rank);
+  updates_specs[0] = Layout::kUnshardedDim;
+
+  const int index_dimensions = tensor_rank - (updates_rank - 1);
+
+  for (int i = 0; i < updates_rank - 1; ++i)
+    updates_specs[i + 1] =
+        pre_output_layout.sharding_spec(index_dimensions + i);
+
+  TF_ASSIGN_OR_RETURN(Layout new_updates_layout,
+                      Layout::GetLayout(updates_specs, updates_layout->mesh()));
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value new_tensor,
+      EmitRelayout(scatter_op.getTensor(), *tensor_layout, pre_output_layout));
+  TF_ASSIGN_OR_RETURN(mlir::Value new_updates,
+                      EmitRelayout(scatter_op.getUpdates(), *updates_layout,
+                                   new_updates_layout));
+
+  mlir::OpBuilder builder(op);
+  OpType new_scatter =
+      OpType::create(builder, op->getLoc(), new_tensor.getType(), new_tensor,
+                     new_indices, new_updates);
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value new_output,
+      EmitRelayout(new_scatter.getOutput(), pre_output_layout, *output_layout));
+
+  op->getResult(0).replaceAllUsesWith(new_output);
+  op->erase();
+
+  return new_output.getDefiningOp();
+}
+
+template <typename OpType>
+StatusOr<llvm::DenseMap<int, Layout>> TensorScatterOpComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto scatter_op = llvm::cast<OpType>(op);
+
+  const int tensor_rank = ValueRank(scatter_op.getTensor());
+  const int updates_rank = ValueRank(scatter_op.getUpdates());
+  if (tensor_rank == -1 || updates_rank == -1)
+    return absl::InvalidArgumentError("all inputs must have valid rank.");
+
+  std::optional<Layout> tensor_layout;
+  if (input_layouts.find(0) != input_layouts.end())
+    tensor_layout.emplace(input_layouts.lookup(0));
+  std::optional<Layout> updates_layout;
+  if (input_layouts.find(2) != input_layouts.end())
+    updates_layout.emplace(input_layouts.lookup(2));
+
+  if (tensor_layout || updates_layout) {
+    TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                        GetOutputLayout(tensor_layout, tensor_rank,
+                                        updates_layout, updates_rank, mesh));
+    return llvm::DenseMap<int, Layout>({{0, output_layout}});
+  }
+
+  return llvm::DenseMap<int, Layout>();
+}
+
+template <typename OpType>
+StatusOr<llvm::DenseMap<int, Layout>> TensorScatterOpComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto scatter_op = llvm::cast<OpType>(op);
+
+  const int tensor_rank = ValueRank(scatter_op.getTensor());
+  const int indices_rank = ValueRank(scatter_op.getIndices());
+  const int updates_rank = ValueRank(scatter_op.getUpdates());
+  if (tensor_rank == -1 || indices_rank == -1 || updates_rank == -1)
+    return absl::InvalidArgumentError("all inputs must have valid rank.");
+
+  // The number of dimensions at the start of the tensor input that are used
+  // for the index, also the size of the second dimension of the indices tensor.
+  const int index_dimensions = tensor_rank - (updates_rank - 1);
+
+  llvm::DenseMap<int, Layout> input_layouts(scatter_op.getNumOperands());
+
+  // Always set indices layout to replicated.
+  const Layout indices_layout = Layout::ReplicatedOnMesh(mesh, indices_rank);
+  input_layouts[1] = indices_layout;
+
+  if (output_layouts.find(0) != output_layouts.end()) {
+    const Layout output_layout = output_layouts.lookup(0);
+
+    std::vector<std::string> tensor_sharding_specs(tensor_rank);
+    std::vector<std::string> updates_sharding_specs(updates_rank);
+
+    for (int i = 0; i < index_dimensions; ++i)
+      tensor_sharding_specs[i] = Layout::kUnshardedDim;
+    updates_sharding_specs[0] = Layout::kUnshardedDim;
+
+    for (int i = index_dimensions; i < tensor_rank; ++i) {
+      tensor_sharding_specs[i] = output_layout.sharding_spec(i);
+      updates_sharding_specs[i - index_dimensions + 1] =
+          output_layout.sharding_spec(i);
+    }
+
+    TF_ASSIGN_OR_RETURN(const Layout tensor_layout,
+                        Layout::GetLayout(tensor_sharding_specs, mesh));
+    TF_ASSIGN_OR_RETURN(const Layout updates_layout,
+                        Layout::GetLayout(updates_sharding_specs, mesh));
+    input_layouts[0] = tensor_layout;
+    input_layouts[2] = updates_layout;
+  }
+
+  return input_layouts;
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> TensorScatterOpSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  if (llvm::isa<mlir::TF::TensorScatterUpdateOp>(op)) {
+    return TensorScatterOpExpand<mlir::TF::TensorScatterUpdateOp>(op);
+  }
+  if (llvm::isa<mlir::TF::TensorScatterAddOp>(op)) {
+    return TensorScatterOpExpand<mlir::TF::TensorScatterAddOp>(op);
+  }
+  return absl::UnimplementedError(absl::StrCat(
+      "SPMD expansion for op : ", OpName(op), " is not implemented"));
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorScatterOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  if (llvm::isa<mlir::TF::TensorScatterUpdateOp>(op)) {
+    return TensorScatterOpComputeLayoutForward<mlir::TF::TensorScatterUpdateOp>(
+        op, input_layouts);
+  }
+  if (llvm::isa<mlir::TF::TensorScatterAddOp>(op)) {
+    return TensorScatterOpComputeLayoutForward<mlir::TF::TensorScatterAddOp>(
+        op, input_layouts);
+  }
+  return absl::UnimplementedError(absl::StrCat(
+      "Layout propagation for op : ", OpName(op), " is not implemented"));
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorScatterOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (llvm::isa<mlir::TF::TensorScatterUpdateOp>(op)) {
+    return TensorScatterOpComputeLayoutBackward<
+        mlir::TF::TensorScatterUpdateOp>(op, output_layouts);
+  }
+  if (llvm::isa<mlir::TF::TensorScatterAddOp>(op)) {
+    return TensorScatterOpComputeLayoutBackward<mlir::TF::TensorScatterAddOp>(
+        op, output_layouts);
+  }
+  return absl::UnimplementedError(absl::StrCat(
+      "Layout propagation for op : ", OpName(op), " is not implemented"));
+}
+
+StatusOr<mlir::Operation*> ScatterNdOpSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const std::vector<Layout>& operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+  TF_ASSIGN_OR_RETURN(const Layout& output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+
+  const Layout& indices_layout = operand_layouts[0];
+  const Layout& updates_layout = operand_layouts[1];
+
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto scatter_op = llvm::cast<mlir::TF::ScatterNdOp>(op);
+
+  const int output_rank = ValueRank(scatter_op.getResult());
+  const int updates_rank = ValueRank(scatter_op.getUpdates());
+  const int indices_rank = ValueRank(scatter_op.getIndices());
+
+  if (output_rank == -1 || updates_rank == -1) {
+    return absl::InvalidArgumentError(
+        "Dynamic shaped inputs are not supported. Please file a feature "
+        "request to TF DTensor: component id: 8333864");
+  }
+
+  llvm::SmallVector<int64_t, 4> global_shape;
+  if (!ExtractConstVectorFromValue(scatter_op.getShape(), &global_shape).ok()) {
+    return absl::InvalidArgumentError(
+        "Failed in extracting constant vector from shape tensor. Please file "
+        "a bug to TF DTensor: component id: 833864");
+  }
+
+  // Only do computation after replicating indices tensor.
+  // The expansion will work as the following:
+  // Let N be the rank of the output tensor.
+  // Let K be the rank of each update tensor.
+  // Then the size of each tensor of indices must be size N-K+1.
+  // We will enforce the indices tensor to be replicated, which means
+  // that also the first N-K+1 dimensions of the output must be replicated as
+  // well.
+  // We will shard the updates tensor however much we can, which also means
+  // we will shard the last K dimension tensors of the output tensor to be
+  // as sharded as the updates tensor.
+  //
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value new_indices,
+      EmitRelayout(scatter_op.getIndices(), indices_layout,
+                   Layout::ReplicatedOnMesh(mesh, indices_rank)));
+
+  // Create intermediate layouts for tensors and updates. Since the layout of
+  // tensor and the output of the local tensor-scatter are the same we can reuse
+  // GetOutputLayout. This intermediate layout will be the layout that
+  // both the output and the updates tensor will agree upon. The updates
+  // intermediate layout will also be computed from the last K dimension of
+  // this.
+  TF_ASSIGN_OR_RETURN(Layout output_intermediate_layout,
+                      GetOutputLayout(output_layout, output_rank,
+                                      updates_layout, updates_rank, mesh));
+
+  std::vector<std::string> updates_specs(updates_rank);
+  if (updates_rank == 0) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Expected updates_rank to be greater than zero, but got: ",
+                     updates_rank));
+  }
+  updates_specs[0] = Layout::kUnshardedDim;
+
+  for (int i = 1; i < updates_rank; ++i) {
+    updates_specs[updates_rank - i] =
+        output_intermediate_layout.sharding_spec(output_rank - i);
+  }
+
+  TF_ASSIGN_OR_RETURN(Layout new_updates_layout,
+                      Layout::GetLayout(updates_specs, mesh));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_updates,
+                      EmitRelayout(scatter_op.getUpdates(), updates_layout,
+                                   new_updates_layout));
+
+  const std::vector<int64_t>& local_shape =
+      output_layout.LocalShapeFromGlobalShape(global_shape);
+
+  mlir::OpBuilder builder(op);
+  mlir::Operation* new_scatter = mlir::TF::ScatterNdOp::create(
+      builder, op->getLoc(), op->getResult(0).getType(), new_indices,
+      new_updates,
+      /*shape=*/
+      ::mlir::TF::collection_ops_util::GetR1Const(local_shape, builder,
+                                                  op->getLoc()));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_output,
+                      EmitRelayout(new_scatter->getResult(0),
+                                   output_intermediate_layout, output_layout));
+
+  op->getResult(0).replaceAllUsesWith(new_output);
+  op->erase();
+
+  return InferSPMDExpandedLocalShape(new_output.getDefiningOp());
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+ScatterNdOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto scatter_op = llvm::cast<mlir::TF::ScatterNdOp>(op);
+
+  const int output_rank = ValueRank(scatter_op.getResult());
+  const int updates_rank = ValueRank(scatter_op.getUpdates());
+  if (output_rank == -1 || updates_rank == -1)
+    return absl::InvalidArgumentError(
+        "Dynamic shaped inputs are not supported. Please file a feature "
+        "request to TF DTensor: component id: 8333864");
+
+  std::optional<Layout> updates_layout;
+  auto iter = input_layouts.find(1);
+  if (iter == input_layouts.end()) {
+    return llvm::DenseMap<int, Layout>();
+  }
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      GetOutputLayout(std::nullopt, output_rank,
+                                      iter->getSecond(), updates_rank, mesh));
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+ScatterNdOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto scatter_op = llvm::cast<mlir::TF::ScatterNdOp>(op);
+
+  const int output_rank = ValueRank(scatter_op.getResult());
+  const int indices_rank = ValueRank(scatter_op.getIndices());
+  const int updates_rank = ValueRank(scatter_op.getUpdates());
+
+  if (output_rank == -1 || indices_rank == -1 || updates_rank == -1)
+    return absl::InvalidArgumentError(
+        "Dynamic shaped inputs are not supported. Please file a feature "
+        "request to TF DTensor: component id: 8333864");
+
+  llvm::DenseMap<int, Layout> input_layouts(scatter_op.getNumOperands());
+
+  // Always set `indices` tensor and 'shape' tensor to replicated.
+  input_layouts[0] = Layout::ReplicatedOnMesh(mesh, /*rank=*/indices_rank);
+  input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+
+  auto iter = output_layouts.find(0);
+  if (iter == output_layouts.end()) {
+    return input_layouts;
+  }
+  // Compute the updates layout.
+  const Layout& output_layout = iter->getSecond();
+  std::vector<std::string> updates_sharding_specs(updates_rank);
+
+  // Replicate the first dimension. This is the number of update tensors.
+  // Set the rest of the dimensions equal to the output's corresponding
+  // sharding.
+  updates_sharding_specs[0] = Layout::kUnshardedDim;
+  for (int i = 1; i < updates_rank; ++i) {
+    updates_sharding_specs[i] = output_layout.sharding_spec(output_rank - i);
+  }
+
+  TF_ASSIGN_OR_RETURN(const Layout updates_layout,
+                      Layout::GetLayout(updates_sharding_specs, mesh));
+  input_layouts[1] = updates_layout;
+
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.h

@@ -0,0 +1,56 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SCATTER_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SCATTER_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class TensorScatterOpSPMDExpander : public SPMDExpanderBase {
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class ScatterNdOpSPMDExpander : public SPMDExpanderBase {
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SCATTER_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/segmentation_spmd_expander.cc

@@ -0,0 +1,153 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/segmentation_spmd_expander.h"
+
+#include <string>
+
+#include "absl/container/flat_hash_set.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<llvm::DenseMap<int, Layout>>
+UnsortedSegmentSumSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto unsorted_segmented_sum = llvm::cast<mlir::TF::UnsortedSegmentSumOp>(op);
+  const int output_rank = ValueRank(unsorted_segmented_sum.getOutput());
+  if (input_layouts.find(0) != input_layouts.end()) {
+    // If the data layout exists, we can use it to forward propagate a layout
+    // to the output.
+    const int segment_ids_rank =
+        ValueRank(unsorted_segmented_sum.getSegmentIds());
+
+    Layout input_layout_truncated =
+        input_layouts.lookup(0).Truncate(segment_ids_rank, /*end=*/true);
+    return llvm::DenseMap<int, Layout>(
+        {{0, input_layout_truncated.LeftPad(output_rank)}});
+  }
+
+  // When we don't have a data layout we can only output a replicated layout.
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, output_rank)}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+UnsortedSegmentSumSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+  auto unsorted_segmented_sum = llvm::cast<mlir::TF::UnsortedSegmentSumOp>(op);
+
+  Layout segment_ids_layout =
+      Layout::ReplicatedOnMesh(mesh, ValueRank(op->getOperand(1)));
+  Layout num_segments_layout = Layout::ReplicatedOnMesh(mesh, /*rank=*/0);
+  if (!output_layouts.empty()) {
+    // If we have an output layout, we can send it backwards to the last few
+    // dimension
+    const int data_rank = ValueRank(unsorted_segmented_sum.getData());
+    Layout output_layout_truncated =
+        output_layouts.lookup(0).Truncate(1, /*end=*/true);
+    return llvm::DenseMap<int, Layout>(
+        {{0, output_layout_truncated.LeftPad(data_rank)},
+         {1, segment_ids_layout},
+         {2, num_segments_layout}});
+  }
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(
+               mesh, /*rank=*/ValueRank(unsorted_segmented_sum.getData()))},
+       {1, segment_ids_layout},
+       {2, num_segments_layout}});
+}
+
+StatusOr<mlir::Operation*> UnsortedSegmentSumSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  // The algorithm is simple
+  //
+  // 1. Relayout segment_ids to match the layout of data[:rank(segment_ids)].
+  //    An improved version of this would merge the layouts of segment_ids and
+  //    data into a common layout (e.g. data with layout [x,*,*,*] and
+  //    segment_ids with layout [*,y] would result in [x,y,*,*] and [x,y].
+  // 2. Emit a local UnsortedSegmentSum.
+  // 3. Emit an AllReduce on the output.
+  // 4. Emit a Relayout to the output layout.
+
+  auto sum_op = mlir::cast<mlir::TF::UnsortedSegmentSumOp>(op);
+  mlir::Value data = sum_op.getData();
+  mlir::Value segment_ids = sum_op.getSegmentIds();
+
+  TF_ASSIGN_OR_RETURN(Layout data_layout,
+                      ExtractRequiredLayoutFromOperand(data));
+  TF_ASSIGN_OR_RETURN(Layout segment_ids_layout,
+                      ExtractRequiredLayoutFromOperand(segment_ids));
+
+  const int data_rank = data_layout.rank();
+  const int segment_ids_rank = segment_ids_layout.rank();
+
+  Layout new_segment_ids_layout = data_layout.Truncate(segment_ids_rank);
+
+  absl::flat_hash_set<std::string> reduce_dimensions;
+  for (int i = 0; i < segment_ids_rank; i++)
+    if (new_segment_ids_layout.sharding_spec(i) != Layout::kUnshardedDim)
+      reduce_dimensions.insert(new_segment_ids_layout.sharding_spec(i));
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value new_segment_ids,
+      EmitRelayout(segment_ids, segment_ids_layout, new_segment_ids_layout));
+
+  mlir::OpBuilder builder(op);
+  mlir::Operation* new_sum_op = mlir::TF::UnsortedSegmentSumOp::create(
+      builder, op->getLoc(), sum_op.getOutput().getType(), data,
+      new_segment_ids, sum_op.getNumSegments());
+
+  InferSPMDExpandedLocalShape(new_sum_op);
+
+  Layout data_layout_truncated =
+      data_layout.Truncate(segment_ids_rank, /*end=*/true);
+  Layout result_output_layout = data_layout_truncated.LeftPad(
+      data_rank - segment_ids_rank + 1);  // This is output rank.
+
+  TF_ASSIGN_OR_RETURN(
+      new_sum_op, EmitAllReduce(builder, result_output_layout,
+                                reduce_dimensions, new_sum_op, kReduceOpAdd));
+
+  // Transform the result to the expected output_layout, if necessary.
+  TF_ASSIGN_OR_RETURN(Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(mlir::Value final_output,
+                      EmitRelayout(new_sum_op->getResult(0),
+                                   result_output_layout, output_layout));
+  op->getResult(0).replaceAllUsesWith(final_output);
+  op->erase();
+
+  return final_output.getDefiningOp();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/segmentation_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SEGMENTATION_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SEGMENTATION_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// TODO(b/171079751): Implement SPMD logic for non-trivial layouts.
+class UnsortedSegmentSumSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SEGMENTATION_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/slice_spmd_expander.cc

@@ -0,0 +1,245 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/slice_spmd_expander.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+absl::Status GetSliceOpArguments(mlir::TF::SliceOp slice_op,
+                                 llvm::SmallVector<int64_t, 4>& begins,
+                                 bool& dynamic_begins,
+                                 llvm::SmallVector<int64_t, 4>& sizes) {
+  absl::Status begins_result =
+      ExtractConstVectorFromValue(slice_op.getBegin(), &begins);
+  dynamic_begins = !begins_result.ok();
+
+  TF_RETURN_WITH_CONTEXT(
+      ExtractConstVectorFromValue(slice_op.getSize(), &sizes),
+      "expected constant argument for SliceOp::size()");
+
+  return absl::OkStatus();
+}
+
+StatusOr<Layout> VerifySliceLayout(
+    mlir::Operation* slice_op, mlir::Value value, const Layout& layout,
+    llvm::ArrayRef<int64_t>* global_shape = nullptr) {
+  if (layout.IsFullyReplicated()) return layout;
+
+  TF_ASSIGN_OR_RETURN(llvm::ArrayRef<int64_t> shape,
+                      GetShapeOfValue(value, /*fail_on_dynamic=*/true));
+  const int64_t rank = shape.size();
+  if (global_shape != nullptr) {
+    // In ExpandOp, tensor shape is local shape. So, call site needs to provide
+    // global shape expliclity.
+    shape = *global_shape;
+  }
+
+  llvm::SmallVector<int64_t, 4> begins, sizes;
+  bool dynamic_begins = false;
+  begins.reserve(rank);
+  sizes.reserve(rank);
+
+  TF_RETURN_IF_ERROR(GetSliceOpArguments(
+      llvm::cast<mlir::TF::SliceOp>(slice_op), begins, dynamic_begins, sizes))
+
+  auto num_shards = layout.num_shards();
+
+  std::vector<std::string> sharding_specs;
+
+  for (int64_t i = 0; i < rank; ++i) {
+    const bool begins_starts_at_zero =
+        (sizes[i] == shape[i]) || (!dynamic_begins && begins[i] == 0);
+    const bool ends_at_full_size =
+        (sizes[i] == shape[i]) || (!dynamic_begins && sizes[i] == -1);
+
+    if (begins_starts_at_zero && ends_at_full_size) {
+      // We support slicing with dynamic begins when the sharded dimensions are
+      // getting a full slice. Since we don't know the begins in this case, we
+      // need to rely in the sizes being static and equal to the global shape.
+      // In particular sizes[i] == shape[i] implies begins[i] == 0.
+      // A full slice over the any dimension can be performed locally.
+      sharding_specs.push_back(layout.sharding_spec(i));
+    } else {
+      // Slicing on sharded dim is not trivial. Propose an unsharded dim for
+      // that.
+      sharding_specs.push_back(Layout::kUnshardedDim);
+    }
+  }
+  return Layout::GetLayout(sharding_specs, layout.mesh());
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> SliceSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto slice_op = mlir::cast<mlir::TF::SliceOp>(op);
+  TF_ASSIGN_OR_RETURN(auto input_layout,
+                      ExtractLayoutFromOperand(slice_op.getInput()));
+  TF_ASSIGN_OR_RETURN(auto output_layout, ExtractSingleLayoutFromOp(op));
+
+  if (!output_layout || !input_layout)
+    return absl::UnimplementedError(
+        "layout of Slice op must be known before SPMD expansion.");
+
+  // The dyn_cast will never be nullptr as it is checked in
+  // GetLayoutFromOperands.
+  auto input_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(slice_op.getInput().getType());
+  if (!input_type)
+    return absl::InvalidArgumentError(
+        "rank of input tensor must be statically known for slice op.");
+
+  TF_ASSIGN_OR_RETURN(auto global_shape,
+                      ExtractGlobalInputShape(op->getOpOperand(0)));
+  const int64_t input_rank = input_type.getRank();
+
+  llvm::SmallVector<int64_t, 4> begins, sizes;
+  bool dynamic_begins = false;
+  begins.reserve(input_rank);
+  sizes.reserve(input_rank);
+
+  TF_RETURN_IF_ERROR(
+      GetSliceOpArguments(slice_op, begins, dynamic_begins, sizes));
+
+  TF_ASSIGN_OR_RETURN(auto proposed_layout,
+                      VerifySliceLayout(slice_op, slice_op.getInput(),
+                                        *input_layout, &global_shape));
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+  TF_ASSIGN_OR_RETURN(auto relayout_input,
+                      EmitRelayout(op->getOperand(0), *input_layout,
+                                   proposed_layout, &newly_created_ops));
+  {
+    // Adjusts the sizes when it is full slicing on sharded dimension.
+    // Note that proposed layout is unsharded in the cases that:
+    // 1) We can't determine the begins and sizes != global shape
+    // 2) begins != 0
+    // 3) sizes != global shape or -1
+    const std::vector<int> num_shards = proposed_layout.num_shards();
+    for (int64_t i = 0; i < input_rank; ++i) {
+      if (num_shards[i] == 1) continue;
+
+      if (sizes[i] == -1 && !dynamic_begins && begins[i] == 0) continue;
+
+      if (sizes[i] == global_shape[i]) {
+        // Set the correct output size. If the input dynamic and this is -1,
+        // then shape inference can't tell the output shape.
+        sizes[i] = global_shape[i] / num_shards[i];
+        continue;
+      }
+
+      return absl::InvalidArgumentError(
+          "Non-full-slicing on the sharded dimension is not allowed. "
+          "internal bug.");
+    }
+  }
+
+  mlir::OpBuilder builder(op);
+  mlir::Value new_size;
+  auto loc = op->getLoc();
+  // Both begin and size need to be the same type, so we must match the new
+  // size input with the type of begin.
+  if (!mlir::isa<mlir::ShapedType>(slice_op.getBegin().getType()))
+    return absl::InternalError("type of begin is not a ShapedType");
+  mlir::ShapedType type =
+      mlir::cast<mlir::ShapedType>(slice_op.getBegin().getType());
+  if (type.getElementType().isInteger(32))
+    new_size =
+        IntConst(builder, loc,
+                 llvm::SmallVector<int32_t, 4>(sizes.begin(), sizes.end()));
+  else
+    new_size = Int64Const(builder, loc, sizes);
+
+  auto new_op =
+      mlir::TF::SliceOp::create(builder, loc, slice_op.getOutput().getType(),
+                                relayout_input, slice_op.getBegin(), new_size)
+          .getOperation();
+  new_op = InferSPMDExpandedLocalShape(new_op);
+
+  TF_ASSIGN_OR_RETURN(auto relayout_output,
+                      EmitRelayout(new_op->getResult(0), proposed_layout,
+                                   *output_layout, &newly_created_ops));
+
+  op->getOpResult(0).replaceAllUsesExcept(relayout_output, newly_created_ops);
+  op->erase();
+  return relayout_output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> SliceSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If the input layout is missing, don't return an output layout.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  auto slice_op = mlir::cast<mlir::TF::SliceOp>(op);
+
+  const Layout& input_layout = input_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(
+      auto proposed_layout,
+      VerifySliceLayout(slice_op, slice_op.getInput(), input_layout));
+  return llvm::DenseMap<int, Layout>({{0, proposed_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> SliceSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  auto slice_op = mlir::cast<mlir::TF::SliceOp>(op);
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(slice_op.getNumOperands());
+  // Set replicated layout for begin and size operands.
+  input_layouts[1] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+
+  // input
+  if (output_layouts.find(0) != output_layouts.end()) {
+    const Layout& output_layout = output_layouts.lookup(0);
+    TF_ASSIGN_OR_RETURN(
+        auto proposed_layout,
+        VerifySliceLayout(slice_op, slice_op.getOutput(), output_layout));
+    input_layouts[0] = proposed_layout;
+  }
+
+  return input_layouts;
+}
+
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/slice_spmd_expander.h

@@ -0,0 +1,45 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SLICE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SLICE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class SliceSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SLICE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/softmax_spmd_expander.cc

@@ -0,0 +1,764 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/softmax_spmd_expander.h"
+
+#include <cassert>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/op_utils.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander_common.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+// Computes a local reduce followed by an EmitAllReduce. This performs a global
+// reduction, output will have global shape 1 on the reduced axes if keep dims
+// is true otherwise the axes will be removed.
+// Assumes builder's insertion point is after input.
+StatusOr<mlir::Value> ComputeGlobalReduce(
+    mlir::OpBuilder& builder, const mlir::Value& input,
+    const Layout& input_layout, const absl::flat_hash_set<int>& reduced_dims,
+    absl::string_view reduce_op, bool keep_dims) {
+  TF_ASSIGN_OR_RETURN(
+      const Layout reduction_layout,
+      input_layout.GetLayoutWithReducedDims(reduced_dims,
+                                            /*keep_dims=*/true));
+  std::vector<int32_t> reduce_dim_array(reduced_dims.begin(),
+                                        reduced_dims.end());
+  const mlir::Value reduction_indices =
+      IntConst(builder, input.getLoc(), reduce_dim_array);
+  mlir::Operation* local_reduce;
+
+  // First compute a local reduce
+  if (reduce_op == kReduceOpAdd) {
+    local_reduce = mlir::TF::SumOp::create(
+        builder, input.getLoc(), input, reduction_indices,
+        /*keep_dims=*/builder.getBoolAttr(true));
+  } else if (reduce_op == kReduceOpMax) {
+    local_reduce = mlir::TF::MaxOp::create(
+        builder, input.getLoc(), input, reduction_indices,
+        /*keep_dims=*/builder.getBoolAttr(true));
+  } else {
+    return absl::UnimplementedError(
+        absl::StrCat("reduction ", reduce_op, " not implemented"));
+  }
+
+  // Then an all reduce.
+  absl::flat_hash_set<std::string> reduced_sharding_specs;
+  for (const int dim : reduced_dims)
+    if (Layout::IsShardedDimension(input_layout.sharding_spec(dim)))
+      reduced_sharding_specs.emplace(input_layout.sharding_spec(dim));
+  TF_ASSIGN_OR_RETURN(
+      mlir::Operation * global_reduce,
+      EmitAllReduce(builder, reduction_layout, reduced_sharding_specs,
+                    local_reduce, reduce_op));
+
+  if (!keep_dims) {
+    mlir::RankedTensorType output_type = mlir::dyn_cast<mlir::RankedTensorType>(
+        global_reduce->getResult(0).getType());
+    if (!output_type)
+      return absl::InternalError(
+          "output of EmitAllReduce is not a RankedTensorType");
+    std::vector<int64_t> new_shape;
+    for (int i = 0; i < output_type.getRank(); ++i)
+      if (!reduced_dims.contains(i))
+        new_shape.emplace_back(output_type.getDimSize(i));
+    mlir::RankedTensorType new_type =
+        mlir::RankedTensorType::get(new_shape, output_type.getElementType());
+    // Upcast the dimensions to int64_t as SqueezeOp requires this for its
+    // dimension attribute type. Everything else is OK with int32_t dimensions.
+    std::vector<int64_t> reduce_dim_array_64(reduced_dims.begin(),
+                                             reduced_dims.end());
+    global_reduce = mlir::TF::SqueezeOp::create(
+        builder, input.getLoc(), new_type, global_reduce->getResult(0),
+        builder.getI64ArrayAttr(reduce_dim_array_64));
+  }
+  return global_reduce->getResult(0);
+}
+
+// Takes a sharded logits and compute both the shifted exponentiation of the
+// logits and its sum. Assumes that builder's insertion point is after logits.
+absl::Status ComputeExpAndSum(mlir::OpBuilder& builder,
+                              const mlir::Value& logits,
+                              const Layout& logits_layout,
+                              mlir::Value& shifted_logits,
+                              mlir::Value& exp_of_shifted_logits,
+                              mlir::Value& sum_of_exp) {
+  auto loc = logits.getLoc();
+
+  if (logits_layout.rank() == 0)
+    return absl::UnimplementedError(
+        "softmax not supported for rank 0 tensors.");
+
+  const int64_t class_dimension = logits_layout.rank() - 1;
+
+  // Softmax is exp(input)/sum(exp(input)) and LogSoftmax is
+  // logits - log(sum(exp(input)) where the sum takes place on the
+  // last axis.
+  // For numerical stability, we shift the logits by the max (along
+  // the last axis) before doing the above calculation.
+
+  // Construct the max.
+  TF_ASSIGN_OR_RETURN(
+      const mlir::Value max_logits,
+      ComputeGlobalReduce(builder, logits, logits_layout, {class_dimension},
+                          kReduceOpMax, /*keep_dims=*/true));
+
+  // Subtract max from local copy of logits.
+  shifted_logits =
+      mlir::TF::SubOp::create(builder, loc, logits, max_logits).getResult();
+  exp_of_shifted_logits =
+      mlir::TF::ExpOp::create(builder, loc, shifted_logits).getResult();
+
+  // Sum the exponential.
+  TF_ASSIGN_OR_RETURN(
+      sum_of_exp,
+      ComputeGlobalReduce(builder, exp_of_shifted_logits, logits_layout,
+                          {class_dimension}, kReduceOpAdd,
+                          /*keep_dims=*/true));
+  return absl::OkStatus();
+}
+
+// Computes softmax from its components. Assumes that builder's insertion point
+// is after sum_of_exp and exp_of_shifted_logits.
+mlir::Value ComputeSoftmax(mlir::OpBuilder& builder,
+                           const mlir::Value& exp_of_shifted_logits,
+                           const mlir::Value& sum_of_exp) {
+  // For Softmax, we compute exp(shifted_logits)/sum(exp(shifted_logits))
+  auto softmax =
+      mlir::TF::DivOp::create(builder, exp_of_shifted_logits.getLoc(),
+                              exp_of_shifted_logits, sum_of_exp);
+  return softmax.getResult();
+}
+
+// Computes softmax from its components. Assumes that builder's insertion point
+// is after shifted_logits and sum_of_exp.
+mlir::Value ComputeLogSoftmax(mlir::OpBuilder& builder,
+                              const mlir::Value& shifted_logits,
+                              const mlir::Value& sum_of_exp) {
+  // For LogSoftmax, we compute shifted_logs - log(sum(exp(shifted_logits)))
+  auto log_of_sum =
+      mlir::TF::LogOp::create(builder, shifted_logits.getLoc(), sum_of_exp);
+  auto log_softmax = mlir::TF::SubOp::create(
+      builder, shifted_logits.getLoc(), shifted_logits, log_of_sum.getResult());
+  return log_softmax.getResult();
+}
+
+// Computes the softmax of the input along the last axis, assuming that the
+// input is sharded along that axis.
+StatusOr<mlir::Value> ComputeShardedSoftmax(mlir::OpBuilder& builder,
+                                            const mlir::Value& logits,
+                                            const Layout& logits_layout,
+                                            bool log_softmax) {
+  mlir::Value shifted_logits;
+  mlir::Value exp_of_shifted_logits;
+  mlir::Value sum_of_exp;
+  TF_RETURN_IF_ERROR(ComputeExpAndSum(builder, logits, logits_layout,
+                                      shifted_logits, exp_of_shifted_logits,
+                                      sum_of_exp));
+
+  if (log_softmax) {
+    return ComputeLogSoftmax(builder, shifted_logits, sum_of_exp);
+  } else {
+    return ComputeSoftmax(builder, exp_of_shifted_logits, sum_of_exp);
+  }
+}
+
+// Creates a layout from specs which is
+// 1) Left truncated to match the size of global_shape.
+// 2) Has unsharded dimensions where ever global_shape is 1.
+StatusOr<Layout> GetBroadcastedLayout(llvm::ArrayRef<int64_t> global_shape,
+                                      const std::vector<std::string>& specs,
+                                      const Mesh& mesh) {
+  std::vector<std::string> new_specs(global_shape.size());
+  for (int i = 0; i < global_shape.size(); ++i) {
+    if (global_shape[i] == 1)
+      new_specs[i] = Layout::kUnshardedDim;
+    else
+      new_specs[i] = specs[i + specs.size() - global_shape.size()];
+  }
+
+  return Layout::GetLayout(new_specs, mesh);
+}
+
+// Gets a scalar floating point constant with the same element type as the input
+// value. Assumes builder's insertion point is after input.
+StatusOr<mlir::Value> GetFPConstOfType(mlir::OpBuilder& builder,
+                                       const mlir::Value& input, float value) {
+  if (mlir::TensorType type =
+          mlir::dyn_cast<mlir::TensorType>(input.getType())) {
+    return mlir::TF::ConstOp::create(
+               builder, input.getLoc(),
+               mlir::DenseFPElementsAttr::get<float>(
+                   mlir::RankedTensorType::get({}, type.getElementType()),
+                   {value}))
+        .getOutput();
+  } else {
+    return absl::UnimplementedError(
+        "non tensor type for labels is not supported");
+  }
+}
+
+// Takes input, which has layout agreeing with the truncation of desired_layout
+// and runs OneHot on it to make it 2 dimensions.
+// Assumes builder's insertion point is after input and desired_layout is rank
+// 2.
+//
+// OneHot's element type matches that of features and the number of class is
+// derived from features last dimension and the number of shards in the last
+// dimension of desired layout.
+//
+// TODO(bfontain): Extract and share with OneHotSPMDExpander
+StatusOr<mlir::Value> ComputeOneHot(mlir::OpBuilder& builder,
+                                    const mlir::Value& input,
+                                    const mlir::Value& features,
+                                    const Layout& desired_layout) {
+  // Get the number of classes for this onehot. The number of classes is the
+  // global size of the last dimension of features.
+  mlir::RankedTensorType features_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(features.getType());
+  if (!features_type)
+    return absl::InvalidArgumentError(
+        "feature input shape must be statically known");
+  if (features_type.getRank() == 0)
+    return absl::InvalidArgumentError(
+        "expected feature input to have at least rank 1, but found rank 0");
+
+  const int64_t local_classes = features_type.getShape().back();
+  const int64_t classes = local_classes * desired_layout.num_shards_for_dim(
+                                              desired_layout.rank() - 1);
+
+  int64_t num_shards = desired_layout.num_shards_for_dim(1);
+  if (classes % num_shards)
+    return absl::InvalidArgumentError(
+        absl::StrCat("unable to shard onehot with size ", classes,
+                     " over dimension with ", num_shards, " shards"));
+  const mlir::Location& loc = input.getLoc();
+
+  mlir::Value depth = CreateIntScalarConst(classes / num_shards, builder, loc,
+                                           /*use_int64=*/false);
+
+  // TODO(bfontain): Extract this block (upto and including the SqueezeOp) to
+  // a common function.
+  mlir::tf_device::ClusterOp cluster =
+      depth.getDefiningOp()->getParentOfType<mlir::tf_device::ClusterOp>();
+
+  // `mesh_coordinates` is tensor of size [1, mesh_size] where each
+  // element in the tensor refers to shard id for the specified mesh
+  // dimension.
+  TF_ASSIGN_OR_RETURN(mlir::Value mesh_coordinates,
+                      GetMeshCoordinatesFromCluster(cluster));
+
+  const int mesh_dim_index = desired_layout.mesh().GetMeshDimIndexWithName(
+      desired_layout.sharding_spec(/*idx=*/1));
+
+  // Slice out the [1,1] for mesh_dim_index.
+  mlir::Value shard_id =
+      mlir::TF::SliceOp::create(
+          builder, loc,
+          mlir::RankedTensorType::get({1, 1}, builder.getI32Type()),
+          mesh_coordinates,
+          IntConst(builder, input.getLoc(), {0, mesh_dim_index}),
+          IntConst(builder, input.getLoc(), {1, 1}))
+          .getOutput();
+
+  shard_id =
+      mlir::TF::SqueezeOp::create(
+          builder, loc, mlir::RankedTensorType::get({}, builder.getI32Type()),
+          shard_id, builder.getI64ArrayAttr({0, 1}))
+          .getOutput();
+
+  // `new_indices` = `input` - `shard_id` * (classes/num_shards)
+  mlir::Value id_offset =
+      mlir::TF::MulOp::create(builder, loc, shard_id, depth).getZ();
+
+  // Note that the type of id_offset (int32) may not match the type of input.
+  // So we insert a cast in this case.
+  mlir::TensorType input_type =
+      mlir::dyn_cast<mlir::TensorType>(input.getType());
+  if (!input_type)
+    return absl::InvalidArgumentError("input is not a TensorType");
+  if (!input_type.getElementType().isInteger(32))
+    id_offset = mlir::TF::CastOp::create(builder, loc,
+                                         mlir::RankedTensorType::get(
+                                             {}, input_type.getElementType()),
+                                         id_offset)
+                    .getY();
+
+  mlir::Value indices =
+      mlir::TF::SubOp::create(builder, loc, input, id_offset).getZ();
+
+  TF_ASSIGN_OR_RETURN(mlir::Value on_value,
+                      GetFPConstOfType(builder, features, 1.0));
+  TF_ASSIGN_OR_RETURN(mlir::Value off_value,
+                      GetFPConstOfType(builder, features, 0.0));
+
+  return mlir::TF::OneHotOp::create(builder, input.getLoc(), indices, depth,
+                                    on_value, off_value,
+                                    builder.getI64IntegerAttr(1))
+      .getOutput();
+}
+
+}  // namespace
+
+// Expander for Softmax and LogSoftmax ops.
+StatusOr<mlir::Operation*> SoftmaxOpSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(auto logits_layout,
+                      ExtractLayoutFromOperand(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+
+  if (!logits_layout) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Failed during SPMD expansion of ", OpName(op),
+                     ". Layout of logits input must be known."));
+  }
+
+  // (Log)Softmax's logits are a rank >= 1 tensor. We reduce over the last
+  // dimension. If this is replicated, we don't need any cross-replica
+  // operations and can just emit the op as is.
+  mlir::Value softmax_result;
+  if (logits_layout->IsLastDimReplicated()) {
+    InferSPMDExpandedLocalShape(op);
+    softmax_result = op->getOpResult(0);
+  } else {
+    mlir::OpBuilder builder(op);
+    builder.setInsertionPointAfter(op);
+
+    TF_ASSIGN_OR_RETURN(
+        softmax_result,
+        ComputeShardedSoftmax(builder, op->getOperand(0), *logits_layout,
+                              mlir::isa<mlir::TF::LogSoftmaxOp>(op)));
+    op->getOpResult(0).replaceAllUsesWith(softmax_result);
+    op->erase();
+  }
+
+  // Add a final Relayout in case the output layout is not the same as the
+  // layout of input logits.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  auto final_result = EmitRelayout(softmax_result, *logits_layout,
+                                   output_layout, &newly_created_ops);
+  if (!final_result.ok()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Failed to relayout the output of softmax: ",
+                     final_result.status().message()));
+  }
+  softmax_result.replaceAllUsesExcept(*final_result, newly_created_ops);
+  return final_result->getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SoftmaxOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // We want to use the same layout for the output.
+  return input_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SoftmaxOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // We want to use the same layout for the input.
+  return output_layouts;
+}
+
+// Takes the input and output layouts and
+// 1) Selects a batch and class sharding from the layouts
+// 2) Applies relayout to the input
+// 3) Sets the new features and loss layout. Takes into account broadcasting.
+// 4) Returns the full layout for backprop/loss.
+StatusOr<Layout> SoftmaxLossOpSPMDExpander::MaybeRelayoutInputs(
+    mlir::Operation* op, bool is_sparse, const Layout& features_layout,
+    const Layout& labels_layout, const Layout& loss_layout,
+    const Layout& backprop_layout, Layout& new_features_layout,
+    Layout& new_labels_layout) {
+  // This layout represents the 'internal layout' that the softmax will be
+  // operating on. Inputs will be relayout'ed to this layout and outputs will be
+  // relayout'ed from this layout to their desired layout.
+  std::vector<std::string> internal_layout(2);
+  internal_layout[0] = Layout::kUnshardedDim;
+  internal_layout[1] = Layout::kUnshardedDim;
+
+  // Choose an internal layout, ideally this layout would be chosen so that
+  // the relayout costs for the inputs (from features_layout/labels_layout to
+  // internal_layout) and the outputs (from internal_layout to
+  // loss_layout/backprop_layout) are minimized, but we will do something more
+  // naive for now.
+
+  // Pick a batch sharding, first from features, then labels, loss and backprop.
+  // Due to possible broadcasting on features and labels, they will only
+  // have a batch dim if they are rank 2.
+  if (features_layout.rank() == 2)
+    internal_layout[0] = features_layout.sharding_spec(0);
+  if (((labels_layout.rank() == 2) ||
+       (is_sparse && labels_layout.rank() == 1)) &&
+      Layout::IsUnshardedDimension(internal_layout[0]))
+    internal_layout[0] = labels_layout.sharding_spec(0);
+  if (Layout::IsUnshardedDimension(internal_layout[0]))
+    internal_layout[0] = loss_layout.sharding_spec(0);
+  if (Layout::IsUnshardedDimension(internal_layout[0]))
+    internal_layout[0] = backprop_layout.sharding_spec(0);
+
+  // Pick a class sharding, first from features, then labels and backprop.
+  // The class dim for features and labels is always the last dim if it exists.
+  // Note that loss and backprop have fixed ranks 1 and 2 respectively where as
+  // ranks of features and labels may involved broadcasting.
+  if (features_layout.rank() > 0 &&
+      (internal_layout[0] !=
+       features_layout.sharding_spec(features_layout.rank() - 1)))
+    internal_layout[1] =
+        features_layout.sharding_spec(features_layout.rank() - 1);
+  if (!is_sparse && labels_layout.rank() > 0 &&
+      Layout::IsUnshardedDimension(internal_layout[1]) &&
+      (internal_layout[0] !=
+       labels_layout.sharding_spec(labels_layout.rank() - 1)))
+    internal_layout[1] = labels_layout.sharding_spec(labels_layout.rank() - 1);
+  if (Layout::IsUnshardedDimension(internal_layout[1]) &&
+      (internal_layout[0] != backprop_layout.sharding_spec(1)))
+    internal_layout[1] = backprop_layout.sharding_spec(1);
+
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> features_global_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(op->getOperand(0)));
+
+  // At this point we need to compute the new layout of features and labels.
+  // Broadcasting makes this more complicated: First we truncate the correct
+  // rank and then set any dimensions where the global shape is size 1 to
+  // unsharded.
+  TF_ASSIGN_OR_RETURN(
+      new_features_layout,
+      GetBroadcastedLayout(features_global_shape, internal_layout,
+                           features_layout.mesh()));
+
+  TF_ASSIGN_OR_RETURN(
+      const mlir::Value new_features,
+      EmitRelayout(op->getOperand(0), features_layout, new_features_layout));
+
+  op->setOperand(0, new_features);
+
+  TF_ASSIGN_OR_RETURN(
+      llvm::ArrayRef<int64_t> labels_global_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(op->getOperand(1)));
+
+  if (is_sparse) {
+    // If we are sparse, then the only possible dimension is the batch_dim.
+    std::vector<std::string> sparse_specs = {internal_layout[0]};
+    TF_ASSIGN_OR_RETURN(new_labels_layout,
+                        GetBroadcastedLayout(labels_global_shape, sparse_specs,
+                                             labels_layout.mesh()));
+  } else {
+    TF_ASSIGN_OR_RETURN(
+        new_labels_layout,
+        GetBroadcastedLayout(labels_global_shape, internal_layout,
+                             labels_layout.mesh()));
+  }
+
+  TF_ASSIGN_OR_RETURN(
+      const mlir::Value new_labels,
+      EmitRelayout(op->getOperand(1), labels_layout, new_labels_layout));
+
+  op->setOperand(1, new_labels);
+
+  return Layout::GetLayout(internal_layout, features_layout.mesh());
+}
+
+// Takes the given loss, backprop values and relayouts them out to the required
+// layouts and pass them through an IdentityN op.
+// This assumes that the input have local shape in their type.
+StatusOr<mlir::Operation*> SoftmaxLossOpSPMDExpander::MaybeRelayoutOutputs(
+    mlir::Operation* op, const mlir::Value& loss, const mlir::Value& backprop,
+    const Layout& output_layout, const Layout& loss_layout,
+    const Layout& backprop_layout) {
+  const Layout current_loss_layout = output_layout.Truncate(/*split_point=*/1);
+  const Layout& current_backprop_layout = output_layout;
+
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+  TF_ASSIGN_OR_RETURN(
+      const mlir::Value new_loss,
+      EmitRelayout(loss, current_loss_layout, loss_layout, &newly_created_ops));
+
+  TF_ASSIGN_OR_RETURN(const mlir::Value new_backprop,
+                      EmitRelayout(backprop, current_backprop_layout,
+                                   backprop_layout, &newly_created_ops));
+
+  mlir::OpBuilder builder(loss.getContext());
+
+  if (new_loss.getDefiningOp()->isBeforeInBlock(new_backprop.getDefiningOp()))
+    builder.setInsertionPointAfterValue(new_backprop);
+  else
+    builder.setInsertionPointAfterValue(new_loss);
+
+  llvm::SmallVector<mlir::Type, 4> types = {new_loss.getType(),
+                                            new_backprop.getType()};
+  llvm::SmallVector<mlir::Value, 4> values = {new_loss, new_backprop};
+
+  mlir::TF::IdentityNOp identity_op =
+      mlir::TF::IdentityNOp::create(builder, loss.getLoc(), types, values);
+
+  newly_created_ops.insert(identity_op);
+
+  op->getResult(0).replaceAllUsesExcept(identity_op.getResult(0),
+                                        newly_created_ops);
+  op->getResult(1).replaceAllUsesExcept(identity_op.getResult(1),
+                                        newly_created_ops);
+
+  // If the op we are expanding isn't being used any more, erase it from the
+  // program.
+  if (op->getResult(0).use_empty() && op->getResult(1).use_empty()) op->erase();
+
+  return identity_op.getOperation();
+}
+
+StatusOr<mlir::Operation*> SoftmaxLossOpSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  if (!mlir::isa<mlir::TF::SoftmaxCrossEntropyWithLogitsOp>(op) &&
+      !mlir::isa<mlir::TF::SparseSoftmaxCrossEntropyWithLogitsOp>(op))
+    return absl::InvalidArgumentError(
+        "unsupported op for in SoftmaxLossOpSPMDExpander");
+
+  TF_ASSIGN_OR_RETURN(const Layout& features_layout,
+                      ExtractRequiredLayoutFromOperand(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const Layout& labels_layout,
+                      ExtractRequiredLayoutFromOperand(op->getOperand(1)));
+  TF_ASSIGN_OR_RETURN(const std::vector<Layout>& output_layouts,
+                      ExtractRequiredLayoutFromOp(op));
+
+  const bool is_sparse =
+      mlir::isa<mlir::TF::SparseSoftmaxCrossEntropyWithLogitsOp>(op);
+
+  Layout new_features_layout;
+  Layout new_labels_layout;
+
+  TF_ASSIGN_OR_RETURN(
+      const Layout internal_layout,
+      MaybeRelayoutInputs(op, is_sparse, features_layout, labels_layout,
+                          output_layouts[0], output_layouts[1],
+                          new_features_layout, new_labels_layout));
+
+  assert(internal_layout.rank() == 2);
+
+  // If the class dim is unshared, we can emit a local op.
+  if (Layout::IsUnshardedDimension(internal_layout.sharding_spec(1))) {
+    op = InferSPMDExpandedLocalShape(op);
+    return MaybeRelayoutOutputs(op, op->getResult(0), op->getResult(1),
+                                internal_layout, output_layouts[0],
+                                output_layouts[1]);
+  }
+
+  mlir::OpBuilder builder(op);
+  builder.setInsertionPointAfter(op);
+
+  mlir::Value features = op->getOperand(0);
+  mlir::Value labels = op->getOperand(1);
+  if (is_sparse) {
+    // SparseSoftmaxCrossEntropyWithLogits(features, labels) can be rewritten
+    // as SoftmaxCrossEntropyWithLogits(features, OneHot(labels)).
+    // Note that this is what is done in the XLA kernel for this op.
+    TF_ASSIGN_OR_RETURN(
+        labels, ComputeOneHot(builder, labels, features, internal_layout));
+  }
+
+  if (features_layout.rank() == 0)
+    return absl::UnimplementedError(
+        "scalar values features is not currently supported");
+
+  // SoftmaxCrossEntropyWithLogitsOp is the same as:
+  // loss = -tf.reduce_sum(labels*tf.LogSoftmax(features), class_dim)
+  // backprop = tf.Softmax(features) - labels
+
+  mlir::Value shifted_logits;
+  mlir::Value exp_of_shifted_logits;
+  mlir::Value sum_of_exp;
+
+  // Note that its possible that features is shape [x, 1] and is broadcasted
+  // to match labels. In this case we are doing a bunch of extra work, since
+  // softmax is 1 and log_softmax is 0.
+  TF_RETURN_IF_ERROR(ComputeExpAndSum(builder, features, new_features_layout,
+                                      shifted_logits, exp_of_shifted_logits,
+                                      sum_of_exp));
+
+  const mlir::Value log_softmax =
+      ComputeLogSoftmax(builder, shifted_logits, sum_of_exp);
+  const mlir::Value softmax =
+      ComputeSoftmax(builder, exp_of_shifted_logits, sum_of_exp);
+
+  // Mimic the XLA, which uses where/select to ensure that sub is zero when
+  // labels are zero.
+  TF_ASSIGN_OR_RETURN(const mlir::Value features_zero,
+                      GetFPConstOfType(builder, features, 0.0));
+  TF_ASSIGN_OR_RETURN(const mlir::Value labels_zero,
+                      GetFPConstOfType(builder, labels, 0.0));
+
+  const mlir::Value is_labels_zero =
+      mlir::TF::EqualOp::create(builder, op->getLoc(), labels, labels_zero,
+                                builder.getBoolAttr(true))
+          .getZ();
+  const mlir::Value safe_softmax =
+      mlir::TF::SelectV2Op::create(builder, op->getLoc(), is_labels_zero,
+                                   features_zero, log_softmax)
+          .getOutput();
+  const mlir::Value prod =
+      mlir::TF::MulOp::create(builder, op->getLoc(), labels, safe_softmax)
+          .getZ();
+
+  // Compute the reduce sum
+  TF_ASSIGN_OR_RETURN(
+      mlir::Value positive_loss,
+      ComputeGlobalReduce(builder, prod, internal_layout, /*reduced_dims=*/{1},
+                          kReduceOpAdd, /*keep_dims=*/false));
+
+  builder.setInsertionPointAfterValue(positive_loss);
+  mlir::Value loss =
+      mlir::TF::NegOp::create(builder, op->getLoc(), positive_loss).getY();
+
+  mlir::Value backprop =
+      mlir::TF::SubOp::create(builder, op->getLoc(), softmax, labels);
+
+  return MaybeRelayoutOutputs(op, loss, backprop, internal_layout,
+                              output_layouts[0], output_layouts[1]);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SoftmaxLossOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  const bool is_sparse =
+      mlir::isa<mlir::TF::SparseSoftmaxCrossEntropyWithLogitsOp>(op);
+
+  // loss is sum(-labels * logsoftmax(features)), so the layout is batch
+  // sharded if labels and features are batch sharded on the same mesh dim or
+  // if one is replicated.
+  // backprop is softmax(features) - labels
+
+  std::optional<Layout> features_layout;
+  if (input_layouts.find(0) != input_layouts.end())
+    features_layout.emplace(input_layouts.lookup(0));
+  std::optional<Layout> labels_layout;
+  if (input_layouts.find(1) != input_layouts.end())
+    labels_layout.emplace(input_layouts.lookup(1));
+
+  // We need to compute shardings for two dimensions: batch and class.
+  std::vector<std::string> layout_specs(2);
+  layout_specs[0] = Layout::kUnshardedDim;
+  layout_specs[1] = Layout::kUnshardedDim;
+
+  // First pick the batch dimension, set it to the batch dimension of features
+  // if it exists otherwise to the batch dimesion of labels.
+  if (features_layout && (features_layout->rank() == 2))
+    layout_specs[0] = features_layout->sharding_spec(0);
+  if (labels_layout &&
+      (labels_layout->rank() == 2 ||
+       (is_sparse && labels_layout->rank() == 1)) &&
+      Layout::IsUnshardedDimension(layout_specs[0]))
+    layout_specs[0] = labels_layout->sharding_spec(0);
+
+  // The class sharding_spec for features and labels is always the last
+  // sharding_spec if it exists.
+  if (features_layout && (features_layout->rank() > 0) &&
+      (layout_specs[0] !=
+       features_layout->sharding_spec(features_layout->rank() - 1)))
+    layout_specs[1] =
+        features_layout->sharding_spec(features_layout->rank() - 1);
+  if (!is_sparse && labels_layout && (labels_layout->rank() > 0) &&
+      Layout::IsUnshardedDimension(layout_specs[1]) &&
+      (layout_specs[0] !=
+       labels_layout->sharding_spec(labels_layout->rank() - 1)))
+    layout_specs[1] = labels_layout->sharding_spec(labels_layout->rank() - 1);
+
+  TF_ASSIGN_OR_RETURN(const Layout backprop_layout,
+                      Layout::GetLayout(layout_specs, mesh));
+  const Layout loss_layout = backprop_layout.Truncate(/*split_point=*/1);
+
+  return llvm::DenseMap<int, Layout>({{0, loss_layout}, {1, backprop_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SoftmaxLossOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  const bool is_sparse =
+      mlir::isa<mlir::TF::SparseSoftmaxCrossEntropyWithLogitsOp>(op);
+
+  std::optional<Layout> loss_layout;
+  if (output_layouts.find(0) != output_layouts.end())
+    loss_layout.emplace(output_layouts.lookup(0));
+  std::optional<Layout> backprop_layout;
+  if (output_layouts.find(1) != output_layouts.end())
+    backprop_layout.emplace(output_layouts.lookup(1));
+
+  // We need to compute two possible shardings:
+  // One for the batch dimension and one for the class dimension.
+  std::vector<std::string> layout_specs(2);
+  layout_specs[0] = Layout::kUnshardedDim;
+  layout_specs[1] = Layout::kUnshardedDim;
+
+  // Respect the loss layout if it is set, otherwise use the backprop
+  // layout for the batch_dim.
+  if (loss_layout) layout_specs[0] = loss_layout->sharding_spec(0);
+  if (backprop_layout && Layout::IsUnshardedDimension(layout_specs[0]))
+    layout_specs[0] = backprop_layout->sharding_spec(0);
+
+  // Only backprop has class dim so use that if it is available.
+  if (backprop_layout && backprop_layout->sharding_spec(1) != layout_specs[0])
+    layout_specs[1] = backprop_layout->sharding_spec(1);
+
+  TF_ASSIGN_OR_RETURN(const auto features_shape,
+                      GetShapeOfValue(op->getOperand(0)));
+  TF_ASSIGN_OR_RETURN(const auto labels_shape,
+                      GetShapeOfValue(op->getOperand(1)));
+  TF_ASSIGN_OR_RETURN(const Layout features_layout,
+                      GetBroadcastedLayout(features_shape, layout_specs, mesh));
+  if (is_sparse) {
+    // Drop the class sharding as the labels don't have class dimension in
+    // the sparse version.
+    layout_specs.resize(1);
+  }
+  TF_ASSIGN_OR_RETURN(const Layout labels_layout,
+                      GetBroadcastedLayout(labels_shape, layout_specs, mesh));
+
+  return llvm::DenseMap<int, Layout>(
+      {{0, features_layout}, {1, labels_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/softmax_spmd_expander.h

@@ -0,0 +1,80 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SOFTMAX_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SOFTMAX_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Expander for Softmax and LogSoftmax ops.
+class SoftmaxOpSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// Expander for SoftmaxCrossEntropyWithLogits ops.
+class SoftmaxLossOpSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ private:
+  // Computes the relayouts of the inputs of the softmax loss op. Returns the
+  // internal layout of the softmax loss in new_features_layout and
+  // new_labels_layout.
+  StatusOr<Layout> MaybeRelayoutInputs(mlir::Operation* op, bool is_sparse,
+                                       const Layout& features_layout,
+                                       const Layout& labels_layout,
+                                       const Layout& loss_layout,
+                                       const Layout& backprop_layout,
+                                       Layout& new_features_layout,
+                                       Layout& new_labels_layout);
+
+  // Computes relayouts of the outputs, returning an IdentityN op that ties
+  // together the loss and backprop returns.
+  StatusOr<mlir::Operation*> MaybeRelayoutOutputs(
+      mlir::Operation* op, const mlir::Value& loss, const mlir::Value& backprop,
+      const Layout& output_layout, const Layout& loss_layout,
+      const Layout& backprop_layout);
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SOFTMAX_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.cc

@@ -0,0 +1,72 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.h"
+
+#include <optional>
+
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> SparseToDenseSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  // Set the op's shape as the local shape of the input tensors from the
+  // layouts.
+  TF_ASSIGN_OR_RETURN(std::optional<Layout> computed_layout,
+                      ExtractSingleLayoutFromOp(op));
+  auto local_shape = computed_layout->LocalShapeFromGlobalShape(
+      ExtractGlobalOutputShape(op->getResult(0)).value());
+  auto op_result = op->getResult(0);
+
+  const auto element_type =
+      mlir::cast<mlir::TensorType>(op_result.getType()).getElementType();
+  op_result.setType(mlir::RankedTensorType::get(local_shape, element_type));
+  // No-op
+  return op;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SparseToDenseSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SparseToDenseSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If there is no output layout present then do not infer any operand layouts.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  Layout output_layout = output_layouts.lookup(0);
+  if (output_layout.mesh().is_tpu_mesh()) {
+    return absl::InvalidArgumentError(
+        "Layout for SparseToDenseOp must not be on TPU Mesh.");
+  }
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.h

@@ -0,0 +1,46 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SPARSE_TO_DENSE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SPARSE_TO_DENSE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Implements SPMD Expansion for TF::SparseToDenseOp See base class
+// SPMDExpanderBase for function documentation.
+class SparseToDenseSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SPARSE_TO_DENSE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/split_spmd_expander.cc

@@ -0,0 +1,194 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/split_spmd_expander.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+// Merges the output layouts to a common layout for Split/SplitV.
+//
+// Split has multiple outputs, so we first calculate a common output layout that
+// we can use for passing backwards and then merge the remaining layouts.
+StatusOr<Layout> MergeLayoutsForSplitOutput(
+    int64_t split_dim, const llvm::DenseMap<int, Layout>& layouts) {
+  assert(!layouts.empty());
+  const Layout& first_layout = layouts.begin()->getSecond();
+  std::vector<std::string> sharding_specs = first_layout.sharding_spec_strs();
+
+  // Merge remaining layouts. If there is a conflicting sharding, then set the
+  // dim to replicated.
+  for (auto it = layouts.begin(); it != layouts.end(); ++it) {
+    const Layout& output_layout = it->getSecond();
+    for (int dim = 0; dim < output_layout.rank(); ++dim) {
+      if (Layout::IsShardedDimension(output_layout.sharding_spec(dim)) &&
+          Layout::IsShardedDimension(sharding_specs[dim]) &&
+          output_layout.sharding_spec(dim) != sharding_specs[dim]) {
+        sharding_specs[dim] = Layout::kUnshardedDim;
+      }
+    }
+  }
+  // Force the split_dim to be unsharded.
+  sharding_specs[split_dim] = Layout::kUnshardedDim;
+  return Layout::GetLayout(sharding_specs, first_layout.mesh());
+}
+
+// Retrieves the value of the split_dim operand adjusted based on the input
+// rank. The split_dim operand's value can be [-rank(input), rank(input)), which
+// is adjusted to a positive value.
+StatusOr<int64_t> GetAdjustedSplitDim(mlir::Value split_dim_value,
+                                      mlir::Value input_value) {
+  TF_ASSIGN_OR_RETURN(int64_t split_dim,
+                      ExtractConstIntFromValue(split_dim_value));
+  if (split_dim < 0) {
+    int rank = ValueRank(input_value);
+    if (rank == -1) {
+      return absl::InvalidArgumentError("Input operand has rank -1.");
+    }
+    split_dim += rank;
+  }
+  return split_dim;
+}
+
+}  // namespace
+
+StatusOr<mlir::Operation*> SplitSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto split_op = mlir::cast<mlir::TF::SplitOp>(op);
+  TF_ASSIGN_OR_RETURN(const Layout input_layout,
+                      ExtractRequiredLayoutFromOperand(split_op.getValue()));
+  TF_ASSIGN_OR_RETURN(
+      const int64_t split_dim,
+      GetAdjustedSplitDim(split_op.getSplitDim(), split_op.getValue()));
+
+  if (Layout::IsShardedDimension(input_layout.sharding_spec(split_dim))) {
+    return absl::InvalidArgumentError(
+        "Spliting over sharded dimension is not supported.");
+  }
+
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> SplitSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  auto split_op = mlir::cast<mlir::TF::SplitOp>(op);
+  llvm::DenseMap<int, Layout> output_layouts(split_op.getNumResults());
+  if (input_layouts.find(1) != input_layouts.end()) {
+    const Layout& suggested_layout = input_layouts.lookup(1);
+    for (int i = 0; i < split_op.getNumResults(); ++i) {
+      output_layouts[i] = suggested_layout;
+    }
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> SplitSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  auto split_op = mlir::cast<mlir::TF::SplitOp>(op);
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(split_op.getNumOperands());
+  // axis
+  input_layouts[0] = Layout::ReplicatedOnMesh(mesh, /*rank=*/0);
+
+  if (!output_layouts.empty()) {
+    // Split has multiple outputs, first calculate a common output layout that
+    // we can use for passing backwards.
+    TF_ASSIGN_OR_RETURN(
+        const int64_t split_dim,
+        GetAdjustedSplitDim(split_op.getSplitDim(), split_op.getValue()));
+    TF_ASSIGN_OR_RETURN(const Layout common_output_layout,
+                        MergeLayoutsForSplitOutput(split_dim, output_layouts));
+    // value
+    input_layouts[1] = common_output_layout;
+  }
+
+  return input_layouts;
+}
+
+StatusOr<mlir::Operation*> SplitVSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto split_v_op = mlir::cast<mlir::TF::SplitVOp>(op);
+  TF_ASSIGN_OR_RETURN(const Layout input_layout,
+                      ExtractRequiredLayoutFromOperand(split_v_op.getValue()));
+  TF_ASSIGN_OR_RETURN(
+      const int64_t split_dim,
+      GetAdjustedSplitDim(split_v_op.getSplitDim(), split_v_op.getValue()));
+
+  if (Layout::IsShardedDimension(input_layout.sharding_spec(split_dim))) {
+    return absl::InvalidArgumentError(
+        "Spliting over sharded dimension is not supported.");
+  }
+
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> SplitVSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  auto split_v_op = mlir::cast<mlir::TF::SplitVOp>(op);
+  llvm::DenseMap<int, Layout> output_layouts(split_v_op.getNumResults());
+  if (input_layouts.find(0) != input_layouts.end()) {
+    const Layout& suggested_layout = input_layouts.lookup(0);
+    for (int i = 0; i < split_v_op.getNumResults(); ++i) {
+      output_layouts[i] = suggested_layout;
+    }
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> SplitVSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  auto split_v_op = mlir::cast<mlir::TF::SplitVOp>(op);
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(split_v_op.getNumOperands());
+  // size_splits
+  input_layouts[1] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  // axis
+  input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/0);
+
+  if (!output_layouts.empty()) {
+    // Split has multiple outputs, first calculate a common output layout that
+    // we can use for passing backwards.
+    TF_ASSIGN_OR_RETURN(
+        const int64_t split_dim,
+        GetAdjustedSplitDim(split_v_op.getSplitDim(), split_v_op.getValue()));
+    TF_ASSIGN_OR_RETURN(const Layout common_output_layout,
+                        MergeLayoutsForSplitOutput(split_dim, output_layouts));
+    // value
+    input_layouts[0] = common_output_layout;
+  }
+
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/split_spmd_expander.h

@@ -0,0 +1,61 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SPLIT_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SPLIT_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Implements SPMD Expansion for the TF::Split op. See base class
+// SPMDExpanderBase for function documentation.
+class SplitSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+// Implements SPMD Expansion for the TF::SplitV op. See base class
+// SPMDExpanderBase for function documentation.
+class SplitVSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SPLIT_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.cc

@@ -0,0 +1,157 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+namespace {
+
+std::set<int64_t> GetSqueezeDims(mlir::Operation* op, int64_t rank) {
+  auto array_attribute = op->getAttrOfType<mlir::ArrayAttr>("squeeze_dims");
+  std::set<int64_t> squeeze_dims;
+  if (array_attribute) {
+    auto attr_list = array_attribute.getValue().vec();
+    for (const auto& attr : attr_list) {
+      int64_t dim =
+          mlir::cast<mlir::IntegerAttr>(attr).getValue().getSExtValue();
+      // Offset the negative indices to positive range.
+      squeeze_dims.insert((dim + rank) % rank);
+    }
+  }
+  return squeeze_dims;
+}
+
+}  // namespace
+
+StatusOr<llvm::DenseMap<int, Layout>> SqueezeSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If there is no tensor layout then do not infer any output layouts.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout& input_layout = input_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(auto shape, ExtractGlobalInputShape(op->getOpOperand(0)));
+  std::set<int64_t> squeeze_dims = GetSqueezeDims(op, /*rank=*/shape.size());
+
+  std::vector<std::string> layout_specs;
+  layout_specs.reserve(input_layout.rank());
+  for (int64_t i = 0; i < input_layout.rank(); ++i) {
+    if (squeeze_dims.empty()) {
+      if (shape[i] > 1) {
+        layout_specs.push_back(input_layout.sharding_spec(i));
+      }
+    } else {
+      if (squeeze_dims.find(i) == squeeze_dims.end()) {
+        layout_specs.push_back(input_layout.sharding_spec(i));
+      }
+    }
+  }
+
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      Layout::GetLayout(layout_specs, input_layout.mesh()));
+  return llvm::DenseMap<int, Layout>({{0, output_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+SqueezeSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If there is no output layout present then do not infer any operand layouts.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  const Layout& output_layout = output_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(auto shape, ExtractGlobalInputShape(op->getOpOperand(0)));
+  std::set<int64_t> squeeze_dims = GetSqueezeDims(op, /*rank=*/shape.size());
+
+  std::vector<std::string> layout_specs;
+  layout_specs.reserve(output_layout.rank());
+  size_t j = 0;
+  for (size_t i = 0; i < shape.size(); ++i) {
+    if (squeeze_dims.empty()) {
+      if (shape[i] > 1) {
+        layout_specs.push_back(output_layout.sharding_spec(j++));
+      } else {
+        layout_specs.push_back(Layout::kUnshardedDim);
+      }
+    } else {
+      if (squeeze_dims.find(i) == squeeze_dims.end()) {
+        layout_specs.push_back(output_layout.sharding_spec(j++));
+      } else {
+        layout_specs.push_back(Layout::kUnshardedDim);
+      }
+    }
+  }
+
+  TF_ASSIGN_OR_RETURN(const Layout input_layout,
+                      Layout::GetLayout(layout_specs, output_layout.mesh()));
+  return llvm::DenseMap<int, Layout>({{0, input_layout}});
+}
+
+StatusOr<mlir::Operation*> SqueezeSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto squeeze_op = mlir::cast<mlir::TF::SqueezeOp>(op);
+  TF_ASSIGN_OR_RETURN(auto layout, ExtractSingleLayoutFromOp(op));
+
+  if (!layout) {
+    return absl::InvalidArgumentError(
+        "layout of SqueezeOp must be known before SPMD expansion.");
+  }
+
+  TF_ASSIGN_OR_RETURN(auto input_shape,
+                      ExtractGlobalInputShape(op->getOpOperand(0)));
+  std::set<int64_t> squeeze_dims =
+      GetSqueezeDims(op, /*rank=*/input_shape.size());
+
+  if (squeeze_dims.empty()) {
+    // If the squeeze dim is empty, make sure the squeeze only happens on the
+    // dims that is not sharded and has global_shape is 1. Otherwise if the
+    // local shape happens to have size 1 on the dim, it got squeezed
+    // unexpected.
+    for (int i = 0; i < input_shape.size(); ++i) {
+      // Global shape 1 implies the dim cannot be sharded -- worst case it can
+      // be sharded over a mesh with dim size 1, and we would squeeze it as is.
+      if (input_shape[i] == 1) {
+        squeeze_dims.insert(i);
+      }
+    }
+    mlir::OpBuilder builder(squeeze_op);
+    squeeze_op->setAttr("squeeze_dims",
+                        builder.getI64ArrayAttr(llvm::SmallVector<int64_t>(
+                            squeeze_dims.begin(), squeeze_dims.end())));
+  }
+
+  return InferSPMDExpandedLocalShape(squeeze_op);
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SQUEEZE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SQUEEZE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class SqueezeSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_SQUEEZE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/strided_slice_spmd_expander.cc

@@ -0,0 +1,502 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/strided_slice_spmd_expander.h"
+
+#include <cassert>
+#include <cstdint>
+#include <vector>
+
+#include "absl/numeric/bits.h"
+#include "absl/status/status.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/slice_util.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+namespace {
+
+// Tokenizes the arguments of a StridedSlice Op to a vector of Tokens.
+// Most arguments are converted directly. If begin, end, or strides are dynamic
+// shaped then the converted Tokens will have dynamic_mask set to true.
+// Fails if the rank is dynamic.
+template <typename T>
+StatusOr<std::vector<slice_util::Token>> TokenizeOp(T strided_slice) {
+  std::vector<slice_util::Token> tokens;
+
+  llvm::SmallVector<int64_t, 4> spec_begin;
+  llvm::SmallVector<int64_t, 4> spec_end;
+  llvm::SmallVector<int64_t, 4> spec_strides;
+
+  TF_ASSIGN_OR_RETURN(llvm::ArrayRef<int64_t> strides_shape,
+                      GetShapeOfValue(strided_slice.getStrides(),
+                                      /*fail_on_dynamic=*/true));
+  if (strides_shape.size() != 1)
+    return absl::InvalidArgumentError(
+        "strides input to strided operation is not rank 1");
+
+  int64_t spec_rank = strides_shape[0];
+  bool dynamic = false;
+  tokens.reserve(spec_rank);
+
+  if (!ExtractConstVectorFromValue(strided_slice.getStrides(), &spec_strides)
+           .ok()) {
+    spec_strides.resize(spec_rank, 0);
+    dynamic = true;
+  }
+  if (ExtractConstVectorFromValue(strided_slice.getBegin(), &spec_begin).ok()) {
+    if (spec_begin.size() != spec_rank)
+      return absl::InvalidArgumentError(
+          "rank of begin input to strided operation does not equal rank of "
+          "strides input");
+  } else {
+    spec_begin.resize(spec_rank, 0);
+    dynamic = true;
+  }
+  if (ExtractConstVectorFromValue(strided_slice.getEnd(), &spec_end).ok()) {
+    if (spec_end.size() != spec_rank)
+      return absl::InvalidArgumentError(
+          "rank of end input to strided operation does not equal rank of "
+          "strides input");
+  } else {
+    spec_end.resize(spec_rank, 0);
+    dynamic = true;
+  }
+  const uint64_t new_axis_mask = strided_slice.getNewAxisMask();
+  const uint64_t shrink_axis_mask = strided_slice.getShrinkAxisMask();
+  const uint64_t spec_begin_mask = strided_slice.getBeginMask();
+  const uint64_t spec_end_mask = strided_slice.getEndMask();
+  uint64_t ellipsis_mask = strided_slice.getEllipsisMask();
+
+  if (absl::popcount(ellipsis_mask) > 1)
+    return absl::InvalidArgumentError(
+        "strided slice only supports at most one ellipsis");
+
+  for (int64_t token_index = 0; token_index < spec_rank; ++token_index) {
+    uint64_t bit = 1 << token_index;
+    slice_util::Token::TokenType token_type = slice_util::Token::REGULAR;
+    if (bit & new_axis_mask) {
+      token_type = slice_util::Token::NEW_AXIS;
+    } else if (bit & shrink_axis_mask) {
+      token_type = slice_util::Token::SHRINK_AXIS;
+    } else if (bit & ellipsis_mask) {
+      token_type = slice_util::Token::ELLIPSIS;
+    }
+    tokens.emplace_back(token_type,
+                        /*begin=*/spec_begin[token_index],
+                        /*end=*/spec_end[token_index],
+                        /*stride=*/spec_strides[token_index],
+                        /*dynamic_mask=*/dynamic,
+                        /*begin_mask=*/bit & spec_begin_mask,
+                        /*end_mask=*/bit & spec_end_mask);
+  }
+
+  return tokens;
+}
+
+// Updates an Op's inputs and attributes using the Token vector.
+// NOTE(feyu): This function only updates the end argument because currently
+// this is the only meaningful change when a global Token vector is converted
+// to the local Token vector.
+template <typename T>
+absl::Status UpdateOpFromTokens(T strided_slice,
+                                const std::vector<slice_util::Token>& tokens) {
+  mlir::OpBuilder builder(strided_slice);
+  llvm::SmallVector<int64_t, 4> end;
+  end.reserve(tokens.size());
+  for (const auto& token : tokens) {
+    end.push_back(token.end);
+  }
+  assert(end.size() == tokens.size());
+  mlir::Value new_end = IntConstWithMatchingType(
+      builder, strided_slice.getLoc(), end, strided_slice.getBegin().getType());
+  strided_slice.getEndMutable().assign(new_end);
+  return absl::OkStatus();
+}
+}  // namespace
+
+//
+// StridedSlice
+//
+StatusOr<mlir::Operation*> StridedSliceSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  mlir::OpBuilder builder(op);
+
+  auto strided_slice_op = mlir::cast<mlir::TF::StridedSliceOp>(op);
+  TF_ASSIGN_OR_RETURN(Layout input_layout, ExtractRequiredLayoutFromOperand(
+                                               strided_slice_op.getInput()));
+  TF_ASSIGN_OR_RETURN(Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  TF_ASSIGN_OR_RETURN(
+      const llvm::ArrayRef<int64_t> global_input_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(strided_slice_op.getInput()));
+
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_op));
+
+  TF_ASSIGN_OR_RETURN(
+      auto forward,
+      slice_util::CreateAndRun<slice_util::ForwardLayoutInference>(
+          tokens, input_layout, global_input_shape));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_input,
+                      EmitRelayout(strided_slice_op.getInput(), input_layout,
+                                   forward.expander_input_layout()));
+
+  TF_RETURN_IF_ERROR(
+      UpdateOpFromTokens(strided_slice_op, forward.local_tokens()));
+
+  strided_slice_op.getInputMutable().assign(new_input);
+
+  op = InferSPMDExpandedLocalShape(op);
+
+  // Do a final relayout to the correct output layout in case there are any
+  // differences between intermediate_output_layout and output_layout.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+  TF_ASSIGN_OR_RETURN(mlir::Value output,
+                      EmitRelayout(strided_slice_op.getOutput(),
+                                   forward.expander_value_layout(),
+                                   output_layout, &newly_created_ops));
+
+  strided_slice_op.getOutput().replaceAllUsesExcept(output, newly_created_ops);
+
+  return output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+StridedSliceSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  mlir::TF::StridedSliceOp strided_slice_op =
+      mlir::cast<mlir::TF::StridedSliceOp>(op);
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_input_shape,
+                      GetShapeOfValue(strided_slice_op.getInput(),
+                                      /*fail_on_dynamic=*/true));
+
+  if (input_layouts.find(0) == input_layouts.end()) {
+    return llvm::DenseMap<int, Layout>();
+  }
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_op));
+  TF_ASSIGN_OR_RETURN(
+      auto forward,
+      slice_util::CreateAndRun<slice_util::ForwardLayoutInference>(
+          tokens, input_layouts.lookup(0), global_input_shape));
+
+  return llvm::DenseMap<int, Layout>({{0, forward.expander_value_layout()}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+StridedSliceSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  mlir::TF::StridedSliceOp strided_slice_op =
+      mlir::cast<mlir::TF::StridedSliceOp>(op);
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_input_shape,
+                      GetShapeOfValue(strided_slice_op.getInput(),
+                                      /*fail_on_dynamic=*/true));
+
+  llvm::DenseMap<int, Layout> input_layouts(strided_slice_op.getNumOperands());
+  // Set replicated layout for begin, end, and strides operands.
+  input_layouts[1] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[3] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+
+  // input
+  if (output_layouts.find(0) == output_layouts.end()) {
+    return input_layouts;
+  }
+
+  const Layout& output_layout = output_layouts.lookup(0);
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_op));
+  TF_ASSIGN_OR_RETURN(
+      auto backward,
+      slice_util::CreateAndRun<slice_util::BackwardLayoutInference>(
+          tokens, output_layout, global_input_shape));
+
+  input_layouts[0] = backward.expander_input_layout();
+  return input_layouts;
+}
+
+//
+//  TensorStridedSliceUpdate
+//
+StatusOr<mlir::Operation*> TensorStridedSliceUpdateSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  mlir::TF::TensorStridedSliceUpdateOp strided_slice_op =
+      llvm::cast<mlir::TF::TensorStridedSliceUpdateOp>(op);
+  TF_ASSIGN_OR_RETURN(
+      const Layout input_layout,
+      ExtractRequiredLayoutFromOperand(strided_slice_op.getInput()));
+  TF_ASSIGN_OR_RETURN(
+      const Layout value_layout,
+      ExtractRequiredLayoutFromOperand(strided_slice_op.getValue()));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+
+  TF_ASSIGN_OR_RETURN(
+      const llvm::ArrayRef<int64_t> global_input_shape,
+      GetGlobalShapeOfValueFromDTensorLayout(strided_slice_op.getInput()));
+
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_op));
+
+  TF_ASSIGN_OR_RETURN(
+      auto forward,
+      slice_util::CreateAndRun<slice_util::ForwardLayoutInference>(
+          tokens, input_layout, global_input_shape));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_input,
+                      EmitRelayout(strided_slice_op.getInput(), input_layout,
+                                   forward.expander_input_layout()));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_value,
+                      EmitRelayout(strided_slice_op.getValue(), value_layout,
+                                   forward.expander_value_layout()));
+
+  strided_slice_op.getInputMutable().assign(new_input);
+  strided_slice_op.getValueMutable().assign(new_value);
+
+  TF_RETURN_IF_ERROR(
+      UpdateOpFromTokens(strided_slice_op, forward.local_tokens()));
+
+  op = InferSPMDExpandedLocalShape(op);
+
+  mlir::OpBuilder builder(op);
+
+  // Do a final relayout to the correct output layout in case there are any
+  // differences between intermediate_output_layout and output_layout.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+  TF_ASSIGN_OR_RETURN(mlir::Value output,
+                      EmitRelayout(strided_slice_op.getOutput(),
+                                   forward.expander_input_layout(),
+                                   output_layout, &newly_created_ops));
+
+  strided_slice_op.getOutput().replaceAllUsesExcept(output, newly_created_ops);
+
+  return output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorStridedSliceUpdateSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  mlir::TF::TensorStridedSliceUpdateOp strided_slice_op =
+      mlir::cast<mlir::TF::TensorStridedSliceUpdateOp>(op);
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_input_shape,
+                      GetShapeOfValue(strided_slice_op.getInput(),
+                                      /*fail_on_dynamic=*/true));
+
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_op));
+  // We have a choice to determine the output layout, we will default to use
+  // input_layout if available, otherwise we will expand value_layout and use
+  // that.
+  std::vector<Layout> candidates;
+  if (input_layouts.find(0) != input_layouts.end()) {
+    // If we have an input_layout, prefer to keep it.
+    candidates.push_back(input_layouts.lookup(0));
+  }
+
+  if (input_layouts.find(4) != input_layouts.end()) {
+    // When we don't have the input layout, use value layout to 'create' the
+    // input layout. We do this by applying the backward inference.
+    // This is because in the case of a normal strided slice the layout of
+    // value would be output layout.
+    const Layout& value_layout = input_layouts.lookup(4);
+    TF_ASSIGN_OR_RETURN(
+        auto backward,
+        slice_util::CreateAndRun<slice_util::BackwardLayoutInference>(
+            tokens, value_layout, global_input_shape));
+    candidates.push_back(backward.expander_input_layout());
+  }
+  if (candidates.empty()) {
+    return llvm::DenseMap<int, Layout>();
+  }
+
+  TF_ASSIGN_OR_RETURN(auto input_layout, GetLeastShardedLayout(candidates));
+
+  return llvm::DenseMap<int, Layout>({{0, input_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorStridedSliceUpdateSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  mlir::TF::TensorStridedSliceUpdateOp strided_slice_op =
+      mlir::cast<mlir::TF::TensorStridedSliceUpdateOp>(op);
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_input_shape,
+                      GetShapeOfValue(strided_slice_op.getInput(),
+                                      /*fail_on_dynamic=*/true));
+
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_op));
+
+  llvm::DenseMap<int, Layout> input_layouts(strided_slice_op.getNumOperands());
+  // Set replicated layout for begin, end, and strides operands.
+  input_layouts[1] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[3] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+
+  // input and value layouts
+  if (output_layouts.find(0) != output_layouts.end()) {
+    const Layout& output_layout = output_layouts.lookup(0);
+    input_layouts[0] = output_layout;
+
+    // We also need a layout for value as well, and for that we just take the
+    // forward inference of the input layout.
+    TF_ASSIGN_OR_RETURN(
+        auto forward,
+        slice_util::CreateAndRun<slice_util::ForwardLayoutInference>(
+            tokens, output_layout, global_input_shape));
+    input_layouts[4] = forward.expander_value_layout();
+  }
+
+  return input_layouts;
+}
+
+//
+// StridedSliceGrad
+//
+StatusOr<mlir::Operation*> StridedSliceGradSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  auto strided_slice_grad_op = llvm::cast<mlir::TF::StridedSliceGradOp>(op);
+  TF_ASSIGN_OR_RETURN(
+      const Layout input_layout,
+      ExtractRequiredLayoutFromOperand(strided_slice_grad_op.getDy()));
+  TF_ASSIGN_OR_RETURN(const Layout output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_output_shape,
+                      GetGlobalShapeOfValueFromDTensorLayout(
+                          strided_slice_grad_op.getOutput()));
+
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_grad_op));
+
+  TF_ASSIGN_OR_RETURN(
+      auto backward,
+      slice_util::CreateAndRun<slice_util::BackwardLayoutInference>(
+          tokens, input_layout, global_output_shape));
+
+  TF_ASSIGN_OR_RETURN(mlir::Value new_dy,
+                      EmitRelayout(strided_slice_grad_op.getDy(), input_layout,
+                                   backward.expander_value_layout()));
+
+  TF_RETURN_IF_ERROR(
+      UpdateOpFromTokens(strided_slice_grad_op, backward.local_tokens()));
+  strided_slice_grad_op.getDyMutable().assign(new_dy);
+
+  mlir::OpBuilder builder(op);
+
+  // The shape input to StridedSliceGrad will still be global, so we need to
+  // compute the local shape update it.
+  std::vector<int64_t> computed_output_shape =
+      backward.expander_input_layout().LocalShapeFromGlobalShape(
+          global_output_shape);
+  mlir::Value new_shape = IntConstWithMatchingType(
+      builder, strided_slice_grad_op.getLoc(), computed_output_shape,
+      strided_slice_grad_op.getBegin().getType());
+  strided_slice_grad_op.getShapeMutable().assign(new_shape);
+
+  op = InferSPMDExpandedLocalShape(op);
+
+  // Do a final relayout to the correct output layout in case there are any
+  // differences between intermediate_output_layout and output_layout.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+  TF_ASSIGN_OR_RETURN(mlir::Value output,
+                      EmitRelayout(strided_slice_grad_op.getOutput(),
+                                   backward.expander_input_layout(),
+                                   output_layout, &newly_created_ops));
+
+  strided_slice_grad_op.getOutput().replaceAllUsesExcept(output,
+                                                         newly_created_ops);
+
+  return output.getDefiningOp();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+StridedSliceGradSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If the input layout is missing, don't return an output layout.
+  if (input_layouts.find(4) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  mlir::TF::StridedSliceGradOp strided_slice_grad_op =
+      mlir::cast<mlir::TF::StridedSliceGradOp>(op);
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_output_shape,
+                      GetShapeOfValue(strided_slice_grad_op.getOutput(),
+                                      /*fail_on_dynamic=*/true));
+  const Layout& input_layout = input_layouts.lookup(4);
+
+  TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_grad_op));
+  TF_ASSIGN_OR_RETURN(
+      auto backward,
+      slice_util::CreateAndRun<slice_util::BackwardLayoutInference>(
+          tokens, input_layout, global_output_shape));
+
+  return llvm::DenseMap<int, Layout>({{0, backward.expander_input_layout()}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+StridedSliceGradSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  mlir::TF::StridedSliceGradOp strided_slice_grad_op =
+      mlir::cast<mlir::TF::StridedSliceGradOp>(op);
+  TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  TF_ASSIGN_OR_RETURN(const llvm::ArrayRef<int64_t> global_output_shape,
+                      GetShapeOfValue(strided_slice_grad_op.getOutput(),
+                                      /*fail_on_dynamic=*/true));
+  llvm::DenseMap<int, Layout> input_layouts(
+      strided_slice_grad_op.getNumOperands());
+  // Set replicated layout for shape, begin, end, stride operands.
+  input_layouts[0] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[1] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[2] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+  input_layouts[3] = Layout::ReplicatedOnMesh(mesh, /*rank=*/1);
+
+  // dy
+  if (output_layouts.find(0) != output_layouts.end()) {
+    const Layout& output_layout = output_layouts.lookup(0);
+    TF_ASSIGN_OR_RETURN(auto tokens, TokenizeOp(strided_slice_grad_op));
+
+    TF_ASSIGN_OR_RETURN(
+        auto forward,
+        slice_util::CreateAndRun<slice_util::ForwardLayoutInference>(
+            tokens, output_layout, global_output_shape));
+    input_layouts[4] = forward.expander_value_layout();
+  }
+
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/strided_slice_spmd_expander.h

@@ -0,0 +1,70 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_STRIDED_SLICE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_STRIDED_SLICE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class StridedSliceSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class TensorStridedSliceUpdateSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+class StridedSliceGradSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_STRIDED_SLICE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/tensorlist_getitem_spmd_expander.cc

@@ -0,0 +1,68 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/tensorlist_getitem_spmd_expander.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> TensorListGetItemSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const auto operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+  TF_ASSIGN_OR_RETURN(const auto output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+
+  // Do a final relayout to the correct output layout in case there are any
+  // differences between layout of `handle` and output_layout.
+  llvm::SmallPtrSet<mlir::Operation*, 4> newly_created_ops;
+
+  TF_ASSIGN_OR_RETURN(mlir::Value output,
+                      EmitRelayout(op->getOpResult(0), operand_layouts[0],
+                                   output_layout, &newly_created_ops));
+
+  op->getOpResult(0).replaceAllUsesExcept(output, newly_created_ops);
+  return InferSPMDExpandedLocalShape(output.getDefiningOp());
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorListGetItemSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // Prefer the output to be the layout of the list handle.
+  if (input_layouts.find(0) != input_layouts.end()) {
+    return llvm::DenseMap<int, Layout>({{0, input_layouts.lookup(0)}});
+  }
+  return llvm::DenseMap<int, Layout>();
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorListGetItemSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Do not infer any layout to the operands.
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/tensorlist_getitem_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_GETITEM_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_GETITEM_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class TensorListGetItemSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_GETITEM_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/tensorlist_reserve_spmd_expander.cc

@@ -0,0 +1,88 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/tensorlist_reserve_spmd_expander.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/dtensor_location.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> TensorListReserveSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  // Recompute the new local shape from the Layout and the global shape.
+  TF_ASSIGN_OR_RETURN(auto global_shape, GetShapeOfValue(op->getOpResult(0)));
+  TF_ASSIGN_OR_RETURN(Layout layout, ExtractRequiredSingleLayoutFromOp(op));
+  std::vector<int64_t> local_shape =
+      layout.LocalShapeFromGlobalShape(global_shape);
+
+  mlir::TF::TensorListReserveOp tensorlist_op =
+      llvm::dyn_cast<mlir::TF::TensorListReserveOp>(op);
+  mlir::OpBuilder builder(op);
+
+  mlir::Type element_type =
+      mlir::cast<mlir::TensorType>(GetSubtypeOrSelf(op->getOpResult(0)))
+          .getElementType();
+
+  mlir::RankedTensorType new_output_type = mlir::RankedTensorType::get(
+      {}, mlir::TF::VariantType::get(
+              mlir::RankedTensorType::get(local_shape, element_type),
+              builder.getContext()));
+  mlir::Value new_shape_value = Int64Const(builder, DT_LOC(op), local_shape);
+  mlir::TF::TensorListReserveOp new_op = mlir::TF::TensorListReserveOp::create(
+      builder, DT_LOC(op), new_output_type, new_shape_value,
+      tensorlist_op.getNumElements());
+
+  op->getResult(0).replaceAllUsesWith(new_op.getResult());
+  op->erase();
+  return new_op.getOperation();
+}
+
+// Always prefer the output to be replicated.
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorListReserveSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  TF_ASSIGN_OR_RETURN(auto global_shape, GetShapeOfValue(op->getOpResult(0)));
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, global_shape.size())}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorListReserveSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Do not infer any operand layouts.
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/tensorlist_reserve_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_RESERVE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_RESERVE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class TensorListReserveSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_RESERVE_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.cc

@@ -0,0 +1,83 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> TensorListSetItemSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  TF_ASSIGN_OR_RETURN(const auto operand_layouts,
+                      ExtractRequiredLayoutFromOperands(op));
+  TF_ASSIGN_OR_RETURN(const auto output_layout,
+                      ExtractRequiredSingleLayoutFromOp(op));
+  // If the layout of handle and layout of output handle is not equal,
+  // then return an error. This is probably not possible to have different
+  // handles have differing layouts since this means we are changing the shapes
+  // and have to relayout all the items that this points to.
+  if (operand_layouts[0] != output_layout) {
+    return absl::InternalError(
+        "Differing layouts for variant tensor input and variant tensor output "
+        "is not yet allowed.");
+  }
+
+  // Relayout 'item' tensor to the layout of the 'handle',
+  TF_ASSIGN_OR_RETURN(
+      const auto new_item,
+      EmitRelayout(op->getOperand(2), operand_layouts[2], operand_layouts[0]));
+
+  op->setOperand(2, new_item);
+  // We do not need to relayout the output since above we checked that
+  // the input list handle and output list handle is the same.
+  return op;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorListSetItemSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op));
+  TF_ASSIGN_OR_RETURN(auto global_shape, GetShapeOfValue(op->getOpResult(0)));
+  // Prefer the output list handle layout to be the same layout
+  // as the input list handle layout.
+  if (input_layouts.find(0) != input_layouts.end())
+    return llvm::DenseMap<int, Layout>({{0, input_layouts.lookup(0)}});
+  // Assume replicated layout if the list handle layout does not exist.
+  return llvm::DenseMap<int, Layout>(
+      {{0, Layout::ReplicatedOnMesh(mesh, global_shape.size())}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+TensorListSetItemSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // Prefer the layout of `item` tensor to be the layout of the output list
+  // handle.
+  if (output_layouts.find(0) != output_layouts.end()) {
+    return llvm::DenseMap<int, Layout>({{2, output_layouts.lookup(0)}});
+  }
+  return llvm::DenseMap<int, Layout>();
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_SETITEM_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_SETITEM_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class TensorListSetItemSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TENSORLIST_SETITEM_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.cc

@@ -0,0 +1,117 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.h"
+
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/collectives.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/proto/layout.pb.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// layout -> layout[:-1] + unsharded
+StatusOr<Layout> GetSuggestedLayout(const Layout& input_layout) {
+  std::vector<std::string> layout_specs(input_layout.rank());
+
+  for (int i = 0; i < input_layout.rank() - 1; ++i) {
+    layout_specs[i] = input_layout.sharding_spec(i);
+  }
+  layout_specs[input_layout.rank() - 1] = Layout::kUnshardedDim;
+  return Layout::GetLayout(layout_specs, input_layout.mesh());
+}
+
+StatusOr<mlir::Operation*> TopKSPMDExpander::ExpandOp(mlir::Operation* op) {
+  auto top_k_op = mlir::cast<mlir::TF::TopKV2Op>(op);
+  mlir::Value input = top_k_op.getInput();
+  TF_ASSIGN_OR_RETURN(auto input_layout, ExtractLayoutFromOperand(input));
+
+  if (!input_layout)
+    return absl::InvalidArgumentError(
+        "layout of TopKV2Op must be known before SPMD expansion.");
+
+  TF_ASSIGN_OR_RETURN(auto layouts, ExtractLayoutFromOp(op));
+  for (const auto& layout : layouts) {
+    if (layout.has_value() && !layout->IsLastDimReplicated()) {
+      return absl::InvalidArgumentError(
+          "The last dimensions of TopKV2Op outputs should be UNSHARDED.");
+    }
+  }
+  mlir::OpBuilder builder(op);
+  if (!input_layout->IsLastDimReplicated()) {
+    TF_ASSIGN_OR_RETURN(Layout new_layout, GetSuggestedLayout(*input_layout));
+    TF_ASSIGN_OR_RETURN(
+        input, EmitAllGather(builder, input, *input_layout, new_layout));
+    mlir::IRMapping mapping;
+    mapping.map(op->getOperand(0), input);
+    mlir::Operation* new_op = builder.clone(*op, mapping);
+    new_op = InferSPMDExpandedLocalShape(new_op);
+
+    op->getResult(0).replaceAllUsesWith(new_op->getResult(0));
+    op->getResult(1).replaceAllUsesWith(new_op->getResult(1));
+    op->erase();
+    return new_op;
+  }
+
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> TopKSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  // If the input layout is missing, don't return an output layout.
+  if (input_layouts.find(0) == input_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(Layout output_layout,
+                      GetSuggestedLayout(input_layouts.lookup(0)));
+
+  return llvm::DenseMap<int, Layout>({
+      {0, output_layout},  // values
+      {1, output_layout},  // indices
+  });
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> TopKSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  // If the output values layout is missing, don't return an input layout.
+  if (output_layouts.find(0) == output_layouts.end())
+    return llvm::DenseMap<int, Layout>();
+
+  TF_ASSIGN_OR_RETURN(Layout input_layout,
+                      GetSuggestedLayout(output_layouts.lookup(0)));
+  const Mesh& mesh = input_layout.mesh();
+
+  return llvm::DenseMap<int, Layout>({
+      {0, input_layout},                                // input
+      {1, Layout::ReplicatedOnMesh(mesh, /*rank=*/0)},  // k
+  });
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.h

@@ -0,0 +1,44 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TOP_K_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TOP_K_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class TopKSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TOP_K_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.cc

@@ -0,0 +1,113 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.h"
+
+#include <cassert>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/layout_parsing.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/value_utils.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> TerminatorSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  auto terminator_op = llvm::cast<mlir::tf_device::ReturnOp>(op);
+  auto parent_op = op->getParentOp();
+  auto output_types = llvm::to_vector<8>(terminator_op.getOperandTypes());
+  assert(output_types.size() == parent_op->getNumResults());
+
+  for (const auto& output_type_and_index : llvm::enumerate(output_types)) {
+    const int index = output_type_and_index.index();
+    const auto& type = output_type_and_index.value();
+    parent_op->getResult(index).setType(type);
+  }
+  return op;
+}
+
+StatusOr<mlir::Operation*> MetadataSPMDExpander::ExpandOp(mlir::Operation* op) {
+  for (auto operand : op->getOperands()) {
+    TF_ASSIGN_OR_RETURN(auto input_layout, ExtractLayoutFromOperand(operand));
+    if (!input_layout.has_value())
+      return absl::InternalError(
+          "All input layouts to Metadata op must be specified at SPMD "
+          "expansion.");
+
+    if (!input_layout->IsFullyReplicated())
+      return absl::InvalidArgumentError(
+          "Metadata ops like tf.BroadcastGradientArgs op must have replicated "
+          "input layouts.");
+  }
+
+  TF_ASSIGN_OR_RETURN(auto result_layouts, ExtractLayoutFromOp(op));
+  for (const auto& layout : result_layouts) {
+    if (!layout.has_value())
+      return absl::InternalError(
+          "All op result layouts of Metadata op must be specified for SPMD "
+          "expansion.");
+
+    if (!layout->IsFullyReplicated()) {
+      return absl::InvalidArgumentError(
+          "Metadata ops like tf.BroadcastGradientArgs op must have replicated "
+          "output layouts.");
+    }
+  }
+  return InferSPMDExpandedLocalShape(op);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+MetadataSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> output_layouts(op->getNumResults());
+  for (const auto& result_and_index : llvm::enumerate(op->getOpResults())) {
+    const int index = result_and_index.index();
+    auto result = result_and_index.value();
+    output_layouts.insert(
+        {index, Layout::ReplicatedOnMesh(mesh, ValueRank(result))});
+  }
+  return output_layouts;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+MetadataSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op));
+
+  llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands());
+  for (const auto& operand_and_index : llvm::enumerate(op->getOperands())) {
+    const int index = operand_and_index.index();
+    auto operand = operand_and_index.value();
+    input_layouts.insert(
+        {index, Layout::ReplicatedOnMesh(mesh, ValueRank(operand))});
+  }
+  return input_layouts;
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.h

@@ -0,0 +1,92 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TRIVIAL_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TRIVIAL_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/shape_utils.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class NoOpSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override {
+    return InferSPMDExpandedLocalShape(op);
+  }
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override {
+    return llvm::DenseMap<int, Layout>();
+  }
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override {
+    return llvm::DenseMap<int, Layout>();
+  }
+};
+
+class TerminatorSPMDExpander : public SPMDExpanderBase {
+ private:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override {
+    return llvm::DenseMap<int, Layout>();
+  }
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override {
+    return llvm::DenseMap<int, Layout>();
+  }
+};
+
+// Expansion for metadata operations (like BroadcastGradientArgs) which always
+// take replicated inputs and emit a fully replicated output.
+class MetadataSPMDExpander : public SPMDExpanderBase {
+ public:
+  // BroadcastGradientArgs accepts 2 shape tensors and returns 2 values, which
+  // indicate reduction dimensions that should be used a part of gradient
+  // computation. These dimensions should be computed against the global shape,
+  // as there are various specializations that occur when a dimension is of
+  // size 1.
+  //
+  // Since shapes are passed in as value (as opposed to being pulled directly
+  // from the input shape), the operation will be performed on the global shape
+  // by default.
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_TRIVIAL_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.cc

@@ -0,0 +1,51 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.h"
+
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+UnsupportedOpSPMDExpander::UnsupportedOpSPMDExpander(
+    const absl::string_view error_message) {
+  error_message_ = error_message;
+}
+
+StatusOr<mlir::Operation*> UnsupportedOpSPMDExpander::ExpandOp(
+    mlir::Operation* op) {
+  return absl::UnimplementedError(error_message_);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+UnsupportedOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  return absl::UnimplementedError(error_message_);
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+UnsupportedOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  return absl::UnimplementedError(error_message_);
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.h

@@ -0,0 +1,50 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_UNSUPPORTED_OP_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_UNSUPPORTED_OP_SPMD_EXPANDER_H_
+
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+class UnsupportedOpSPMDExpander : public SPMDExpanderBase {
+ public:
+  explicit UnsupportedOpSPMDExpander(absl::string_view error_message);
+
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+
+ private:
+  absl::string_view error_message_;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_UNSUPPORTED_OP_SPMD_EXPANDER_H_

tensorflow/dtensor/mlir/expansions/where_spmd_expander.cc

@@ -0,0 +1,82 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/dtensor/mlir/expansions/where_spmd_expander.h"
+
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+StatusOr<mlir::Operation*> WhereOpSPMDExpander::ExpandOp(mlir::Operation* op) {
+  // Where op does not do anything during SPMD expansion.
+  // Because the output layout is parted layout and it follows local semantic.
+  return op;
+}
+
+StatusOr<llvm::DenseMap<int, Layout>> WhereOpSPMDExpander::ComputeLayoutForward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) {
+  if (input_layouts.find(0) == input_layouts.end()) {
+    return llvm::DenseMap<int, Layout>();
+  }
+  // Currently Where op only supports 1D input.
+  Layout layout = input_layouts.lookup(0);
+  if (layout.rank() != 1) {
+    return llvm::DenseMap<int, Layout>();
+  }
+
+  // Append an unsharded sharding spec for the index dimension generated by the
+  // Where op.
+  std::vector<std::string> layout_specs;
+  layout_specs.push_back(layout.sharding_spec(0));
+  layout_specs.push_back(Layout::kUnshardedDim);
+  // The output of Where Op contains dynamic shape and has parted layout.
+  // This is the source of the parted layout and it is propagated to descendent
+  // ops.
+  TF_ASSIGN_OR_RETURN(Layout new_layout,
+                      Layout::GetLayout(Layout::LayoutType::kParted,
+                                        layout_specs, layout.mesh()));
+  return llvm::DenseMap<int, Layout>({{0, new_layout}});
+}
+
+StatusOr<llvm::DenseMap<int, Layout>>
+WhereOpSPMDExpander::ComputeLayoutBackward(
+    mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) {
+  if (output_layouts.find(0) == output_layouts.end()) {
+    return llvm::DenseMap<int, Layout>();
+  }
+
+  // Remove the unsharded sharding spec generated by the Where op.
+  Layout layout = output_layouts.lookup(0);
+  std::vector<std::string> layout_specs;
+  layout_specs.reserve(layout.rank() - 1);
+  for (int i = 0; i < layout.rank() - 1; i++) {
+    layout_specs.push_back(layout.sharding_spec(i));
+  }
+  TF_ASSIGN_OR_RETURN(Layout new_layout,
+                      Layout::GetLayout(layout_specs, layout.mesh()));
+  return llvm::DenseMap<int, Layout>({{0, new_layout}});
+}
+
+}  // namespace dtensor
+}  // namespace tensorflow

tensorflow/dtensor/mlir/expansions/where_spmd_expander.h

@@ -0,0 +1,46 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_WHERE_SPMD_EXPANDER_H_
+#define TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_WHERE_SPMD_EXPANDER_H_
+
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/dtensor/cc/dstatus.h"
+#include "tensorflow/dtensor/cc/tensor_layout.h"
+#include "tensorflow/dtensor/mlir/spmd_expander.h"
+
+namespace tensorflow {
+namespace dtensor {
+
+// Expander class for Where op.
+// Where op returns the indices of the non-zero elements in the input tensor.
+class WhereOpSPMDExpander : public SPMDExpanderBase {
+ public:
+  StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutForward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& input_layouts) override;
+
+  StatusOr<llvm::DenseMap<int, Layout>> ComputeLayoutBackward(
+      mlir::Operation* op,
+      const llvm::DenseMap<int, Layout>& output_layouts) override;
+};
+
+}  // namespace dtensor
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_DTENSOR_MLIR_EXPANSIONS_WHERE_SPMD_EXPANDER_H_