output_uniform_quantized_tensor_type};
    auto new_conv_op_with_output_type =
        rewriter.create<stablehlo::ConvolutionOp>(
            op.getLoc(), new_conv_output_types, op.getOperands(),
            op->getAttrs());

    rewriter.replaceAllUsesWith(op.getResult(),
                                new_conv_op_with_output_type.getResult());

    auto new_output_dequant_op =

      // make it to the return array.
      auto segment_sizes = getOperation()->getAttrOfType<DenseI32ArrayAttr>(getOperandSegmentSizeAttr());
      return OperandRangeRange(getOperands(), segment_sizes).join();
    }
  }];
}

def TF_BatchMatMulOp : TF_Op<"BatchMatMul", [Pure, TF_SameOperandsAndResultElementTypeResolveRef]> {
  let summary = "Multiplies slices of two tensors in batches.";