// CHECK-LABEL: testDotToDotGeneralVectorVector
func.func @testDotToDotGeneralVectorVector(%arg0: tensor<3072xf32>, %arg1: tensor<3072xf32>) -> tensor<f32> {
  %0 = "mhlo.dot"(%arg0, %arg1) : (tensor<3072xf32>, tensor<3072xf32>) -> tensor<f32>
  func.return %0 : tensor<f32>

// CHECK:      %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME:   dot_dimension_numbers = #mhlo.dot<

  // CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<2.000000e+00>
  // CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
  // CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // CHECK: mhlo.add
  // CHECK: mhlo.return
  // CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
  // CHECK-NEXT: return %[[RESULT]]

  // CHECK-DAG: %[[EPS_BCAST:.+]] = mhlo.constant dense<1.001000e-05> : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
  // CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>

    if (failed(output_type)) {
      return failure();
    }

    auto input_quant = rewriter.create<mhlo::BitcastConvertOp>(
        op->getLoc(), *input_quant_type, input);
    auto result = rewriter.create<mhlo::UniformQuantizeOp>(
        op->getLoc(), *output_type, input_quant);
    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
        op,
        output_type->clone(

    lhs_flattend = rewriter.create<mhlo::ReshapeOp>(
        loc,
        RankedTensorType::get(lhs_flattened_shape, lhs_type.getElementType()),
        lhs_transposed.getResult());
  } else {
    auto lhs_flattend_shape_op = BuildDotOperandFlattenedShapeOp(
        lhs, lhs_dot_dimensions_info, builder, /*is_lhs=*/true);
    lhs_flattend = rewriter.create<mhlo::DynamicReshapeOp>(
        loc,

tensor<64x1001xf32> {mhlo.is_same_data_across_replicas = true, mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg4: tensor<1001xf32> {mhlo.is_same_data_across_replicas = true, mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg5: tensor<f32> {mhlo.is_same_data_across_replicas = true, mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg6: tensor<f32> {mhlo.is_same_data_across_replicas = true, mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg7: tensor<f32> {mhlo.is_same_data_across_replicas = true, mhlo.sharding = "\08\0...

#include "xla/client/sharding_builder.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/primitive_util.h"
#include "xla/side_effect_util.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {

using func::FuncOp;

namespace mhlo {

namespace {
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kFrontendAttributesAttr[] = "mhlo.frontend_attributes";

      rewriter.replaceOpWithNewOp<mhlo::TransposeOp>(
          scatter_op, scatter_op.getResult(0).getType(), tf_scatter_op,
          inverse_permutation);

      return success();
    }
  }
};

using ConvertScatterAddOp =
    ConvertScatterOp<mhlo::AddOp, TF::TensorScatterAddOp>;
using ConvertScatterMaxOp =
    ConvertScatterOp<mhlo::MaxOp, TF::TensorScatterMaxOp>;
using ConvertScatterMinOp =

#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

// Convert mhlo.dot to mhlo.dot_general.
LogicalResult ConvertDotToDotGeneral(mhlo::DotOp op,
                                     PatternRewriter &rewriter) {
  auto lhs_type = mlir::cast<ShapedType>(op.getLhs().getType());

  // StableHLO -> MHLO legalization for MHLO optimization.
  pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());
  // Rewrites legacy StableHLO ops.
  AddUnfuseMhloOpsPasses(pm);
  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
  // MHLO -> StableHLO legalization.
  pm.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());
}