//   rhs: [batch, contract dims, other dims]
  auto dim_nums = dot.getDotDimensionNumbers();
  int batch_dims_count = dim_nums.getLhsBatchingDimensions().size();
  int contracting_dims_count = dim_nums.getLhsContractingDimensions().size();
  if (dim_nums.getLhsBatchingDimensions() !=
          ArrayRef<int64_t>(range(0, batch_dims_count)) ||
      dim_nums.getRhsBatchingDimensions() !=

  ImplicitLocOpBuilder builder(loc, rewriter);

  // Collects lhs and rhs dimensions information.
  DotDimensionsInfo lhs_dot_dimensions_info(
      lhs_type, dot_dimension_numbers.getLhsBatchingDimensions(),
      dot_dimension_numbers.getLhsContractingDimensions());
  DotDimensionsInfo rhs_dot_dimensions_info(
      rhs_type, dot_dimension_numbers.getRhsBatchingDimensions(),