(LiftAsTFXlaCallModule<"composite_gather_fn">
    (ArgumentList $operand, $start_indices),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"dimension_numbers"> $dimension_numbers),
      (NamedAttr<"slice_sizes"> $slice_sizes),
      (NamedAttr<"indices_are_sorted"> (DefaultOrNullAttr $indices_are_sorted)))),

    %2 = "tf.XlaGather"(%1, %arg0, %cst) {dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true} : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xf32>
    return %2 : tensor<1x300x10xf32>
  }

// }
func.func @xla_gather(%arg0: tensor<?x2xf32>, %arg1: tensor<1xi32>, %arg2: tensor<2xi32>) -> tensor<*xf32> {
  %0 = "tf.XlaGather"(%arg0, %arg1, %arg2) {device = "", dimension_numbers = "\0A\01\00\12\01\01\1A\01\01", indices_are_sorted = true} : (tensor<?x2xf32>, tensor<1xi32>, tensor<2xi32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

// CHECK: func @xla_gather

    %5 = stablehlo.reshape %4 : (tensor<100xi64>) -> tensor<100x1xi64>
    %6 = "stablehlo.scatter"(%arg0, %5, %arg3) ({
    ^bb0(%arg5: tensor<f32>, %arg6: tensor<f32>):
      stablehlo.return %arg6 : tensor<f32>

// Convert XlaGather op without batch to Slice op with above two functions.
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),

                                                         op->unique_indices)));
    attributes.emplace_back(builder.getNamedAttr(
        "indices_are_sorted",
        mlir::vhlo::BooleanV1Attr::get(builder.getContext(),
                                       op->indices_are_sorted)));
    attributes.emplace_back(builder.getNamedAttr(
        "index_vector_dim", BuildVhloIntV1Attr(op->index_vector_dim, builder)));
    return;
  }

    start_index_map = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>,
    index_vector_dim = #vhlo.integer_v1<2 : i64>,
    slice_sizes = #vhlo.tensor_v1<dense<[1, 2, 2]> : tensor<3xi64>>,
    indices_are_sorted = #vhlo.bool_v1<false>
  }> : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
  return %result : tensor<2x3x2x2xi32>
}

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

            offset_dims = [2, 3],
            collapsed_slice_dims = [0],
            start_index_map = [1, 0],
            index_vector_dim = 2>,
          slice_sizes = array<i64: 1, 2, 2>,
          indices_are_sorted = false
        } : (tensor<3x4x2xf32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32>
        func.return %2 : tensor<2x3x2x2xf32>
      }
    }
  )mlir";

  %0 = "mhlo.gather"(%arg0, %arg1) {
    dimension_numbers = #mhlo.gather<
      collapsed_slice_dims = [0],
      index_vector_dim = 2,
			start_index_map = [0],
    >,
    indices_are_sorted = false,
    slice_sizes = dense<1> : tensor<1xi64>
  } : (tensor<147456xf16>, tensor<192x256x1xi32>) -> tensor<192x256xf16>
  func.return %0 : tensor<192x256xf16>
}