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);

  // and concatenate scalars into the vector.
  if (is_all_tensors(0, axis) && is_all_scalars(1)) {
    std::array<int64_t, 1> rhs_dims{
        static_cast<int64_t>(op.getValues().size())};
    auto rhs_type =
        tensorflow::GetTypeFromTFTensorShape(rhs_dims, ranked.getElementType());
    return HoistParams{args(0),
                       args(1),
                       axis,
                       0,

  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>(

  llvm::SmallVector<int64_t, 4> dims;
  dims.reserve(lhs.getRank());
  for (auto dim : llvm::zip(lhs.getShape(), rhs.getShape())) {
    int64_t lhs_dim = std::get<0>(dim);
    if (lhs_dim == std::get<1>(dim)) {
      dims.push_back(lhs_dim);
    } else {
      dims.push_back(ShapedType::kDynamic);
    }
  }
  return tensorflow::GetTypeFromTFTensorShape(