def LiftGather : Pat<
  (StableHLO_GatherOp:$res
      $operand, $start_indices, $dimension_numbers, $slice_sizes, $indices_are_sorted),
  (LiftAsTFXlaCallModule<"composite_gather_fn">
    (ArgumentList $operand, $start_indices),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"dimension_numbers"> $dimension_numbers),
      (NamedAttr<"slice_sizes"> $slice_sizes),

      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
      (NamedAttr<"dimension_numbers"> $dimension_numbers),
      (NamedAttr<"feature_group_count"> $feature_group_count),
      (NamedAttr<"batch_group_count"> $batch_group_count),
      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),

  // Matches input dimensions corresponding to: [b, f, 0, 1].
  bool MatchInputDimensionNumbers(
      const ConvDimensionNumbersAttr dimension_numbers) const {
    return dimension_numbers.getInputBatchDimension() == 0 &&
           dimension_numbers.getInputFeatureDimension() == 1 &&
           dimension_numbers.getInputSpatialDimensions() ==
               ArrayRef<int64_t>{2, 3};
  }

    %1 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xbf16>
    %2 = "tf.XlaDotV2"(%arg0, %1) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xbf16>) -> tensor<1x2x2x1024xbf16>
    %3 = "tf.Identity"(%2) : (tensor<1x2x2x1024xbf16>) -> tensor<1x2x2x1024xbf16>

// CHECK: return %[[einsum]] : tensor<?x2x4x5xf32>

// -----

// dimension_numbers: {
//   offset_dims: 0
//   collapsed_slice_dims: 1
//   start_index_map: 1
// }
func.func @xla_gather(%arg0: tensor<?x2xf32>, %arg1: tensor<1xi32>, %arg2: tensor<2xi32>) -> tensor<*xf32> {

def ConvertXlaGatherOpWithoutBatch : Pat<
  (TF_XlaGatherOp:$gather $operand,
    $start_indices, $slice_sizes, $dimension_numbers, $indices_are_sorted),
  (CreateSliceAndReshapeOpFromXlaGatherOpWithoutBatch $operand,
    $start_indices, $slice_sizes, $gather, $dimension_numbers),

    OpBuilder& builder, const Location loc, Value operand, Value start_indices,
    Value slice_sizes, Value output, StringAttr dimension_numbers_str) {
  // Reads dimension numbers.
  xla::GatherDimensionNumbers dimension_numbers;
  dimension_numbers.ParseFromString(dimension_numbers_str.str());

  // Construct full start_indices with given start_indices and
  // start_index_map.
  const ArrayRef<int64_t> operand_shape =

// CHECK: %[[MATMUL_1:.*]] = "tf.XlaDotV2"(%arg3, %arg4) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> tensor<1x1024xbf16>
// CHECK: %[[IDENTITY_2:.*]] = "tf.Identity"(%arg4) {device = ""} : (tensor<1024x1024xbf16>) -> tensor<1024x1024xbf16>

    return stablehlo::ConvDimensionNumbersAttr::get(
        context, dimension_numbers.getInputBatchDimension(),
        dimension_numbers.getInputFeatureDimension(),
        dimension_numbers.getInputSpatialDimensions(), kernel_input_feature_dim,
        kernel_output_feature_dim, kernel_spatial_dims,
        dimension_numbers.getOutputBatchDimension(),
        dimension_numbers.getOutputFeatureDimension(),

                          $dimension_numbers, $indices_are_sorted),
          (MHLO_GatherOp $operand, $start_indices,
                        (ToGatherDimNumsAttr $dimension_numbers),
                        (CastElementsToI64Elements $slice_sizes),
                        $indices_are_sorted),
          [(HasValidGatherDims $dimension_numbers)]>;