return mapping;
}

// Adds a new replicated input to the replicate op.
tf_device::ReplicateOp AddInputsToReplicateOp(
    tf_device::ReplicateOp replicate,
    MutableArrayRef<TF::VarHandleOp> new_inputs,
    const llvm::SmallDenseMap<llvm::StringRef, llvm::SmallVector<StringRef, 4>>&
        devices) {
  int64_t num_replicas = replicate.getN();
  assert(new_inputs.size() == num_replicas);

// tf.StridedSlice op semantics.
static void CalculateSlicedShapeFromDenseIndices(
    MutableArrayRef<int64_t> input_shape, int32_t begin_mask, int32_t end_mask,
    int32_t shrink_axis_mask, MutableArrayRef<int64_t> begin,
    MutableArrayRef<int64_t> end, MutableArrayRef<int64_t> stride) {
  assert(input_shape.size() <= 32);  // Only 32-bit masks are supported.

                                tf_device::ReplicateOp replicate,
                                int32_t block_size) {
  // We need to know the devices to copy to.
  if (!replicate.getDevices()) return false;

  MutableArrayRef<OpOperand> inputs =
      replicate.GetOperandsForBlockArgument(block_arg);
  for (auto& input : inputs) {
    auto input_op = input.get().getDefiningOp();

    LLVM_DEBUG(llvm::dbgs() << GemmStyleOp::getOperationName()
                            << " op must have ranked tensor type.\n");
    return failure();
  }

  MutableArrayRef<BlockArgument> operands =
      entry_func_op.getBody().getArguments();
  // Function must have input, filter, and optionally bias.
  if (operands.size() != 2 && operands.size() != 3) {

                        const char* raw_input, const int element_byte_size,
                        const int64_t current_axis, char*& output_element_addr,
                        MutableArrayRef<uint64_t> current_input_index,
                        ShapedType input_shape_type) {
  const int64_t input_axis = perms[current_axis];
  const bool is_last_axis = current_axis == output_shape.size() - 1;