mlir::StringAttr::get(op->getContext(), hlosharding->ToString()));
}

mlir::LogicalResult ExtractInputsForLogicalDevices(
    const int num_cores_per_replica,
    mlir::tf_device::ClusterFuncOp cluster_func, mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list) {
  // Initialize the input list for each logical devices.
  input_list->reserve(num_cores_per_replica);

      builder->getStringAttr(kFuncAttr),
      mlir::SymbolRefAttr::get(builder->getContext(), outlined_func.getName()));

  builder->setInsertionPoint(cluster_op);
  auto cluster_func_op = builder->create<tf_device::ClusterFuncOp>(
      cluster_op.getLoc(), outlined_func.getFunctionType().getResults(),
      live_ins.getArrayRef(), cluster_op->getAttrs());
  cluster_op.replaceAllUsesWith(cluster_func_op);
  cluster_op.erase();
}

  // Returns whether it's legal to inline a call to a function.
  bool isLegalToInline(Operation* call, Operation* callable,
                       bool wouldBeCloned) const final {
    if (isa<ClusterFuncOp>(call)) return false;
    return true;
  }

  // Returns if its legal to inline 'src' region into the 'dest' region
  // attached to a TF Device operation.

    func.return %arg0: tensor<3x4xf32>
  }
}

// -----

// Test that unsupported input sharding type of TPUPartitionedInputV2Op inputs of
// ClusterFuncOp result in error.