std::swap(rhs_shape[rhs_dims - 1], rhs_shape[rhs_dims - 2]);
  }

  const int64_t rows = lhs_shape[lhs_dims - 2];
  const int64_t cols = rhs_shape[rhs_dims - 1];

  if (lhs_shape[lhs_dims - 1] != rhs_shape[rhs_dims - 2]) {
    // Input dimensions must be compatible for multiplication.
    return failure();
  }

  const auto matmul_type = RankedTensorType::get({rows, cols}, element_type);

  SmallVector<int64_t> new_shape;
  for (auto idx : llvm::reverse(llvm::seq<int32_t>(0, rhs_shape.getRank()))) {
    const int64_t rhs_dim = rhs_shape.getDimSize(idx);
    if (num_elements % rhs_dim != 0) {
      return {};
    }
    new_shape.push_back(rhs_dim);
    num_elements = num_elements / rhs_dim;
    if (num_elements == 1) break;
  }
  absl::c_reverse(new_shape);

  auto reshape_op = builder.create<TF::ReshapeOp>(