DimensionVector out_dimensions_;
};

// Calculates the flattened shapes for dynamic shaped operands in
// mhlo.dot_general:
//   1. flattened_out_dim = UnsortedSegmentProdOp(operand_shape, out_axes)
//   2. flattened_contracting_dim = UnsortedSegmentProdOp(operand_shape,
//   contracting_axes)
//   3. batch_dimensions = Gather(operand_shape, batch_axes)
//   4. flattened_shape = Concat(batch_dimensions, flattened_out_dim,

  // Build UnsortedSegmentProdOp
  Type segProdresultType =
      RankedTensorType::get(num_reshape_segids, rewriter->getIntegerType(32));
  auto segids_tensor = createI32ConstantOp(reshape_segids, loc, rewriter);
  auto num_reshape_segids_tensor =
      createI32ConstOp(num_reshape_segids, loc, rewriter);
  auto segprod = rewriter->create<TF::UnsortedSegmentProdOp>(

    TypeID::get<TF::TopKV2Op>(),
    TypeID::get<TF::_UnaryOpsCompositionOp>(),
    TypeID::get<TF::UnsortedSegmentMaxOp>(),
    TypeID::get<TF::UnsortedSegmentMinOp>(),
    TypeID::get<TF::UnsortedSegmentProdOp>(),
    TypeID::get<TF::UnsortedSegmentSumOp>(),
    TypeID::get<TF::XdivyOp>(),
    TypeID::get<TF::XlaSendTPUEmbeddingGradientsOp>(),
    TypeID::get<TF::XlaAllReduceOp>(),
    TypeID::get<TF::XlaGatherOp>(),

  DimensionVector out_dimensions_;
};

// Calculates the flattened shapes for dynamic shaped operands in
// mhlo.dot_general:
//   1. flattened_out_dim = UnsortedSegmentProdOp(operand_shape, out_axes)
//   2. flattened_contracting_dim = UnsortedSegmentProdOp(operand_shape,
//   contracting_axes)
//   3. batch_dimensions = Gather(operand_shape, batch_axes)
//   4. flattened_shape = Concat(batch_dimensions, flattened_out_dim,

LogicalResult UnsortedSegmentMaxOp::verify() {
  return VerifyUnsortedSegmentReduction(*this);
}
LogicalResult UnsortedSegmentMinOp::verify() {
  return VerifyUnsortedSegmentReduction(*this);
}
LogicalResult UnsortedSegmentProdOp::verify() {
  return VerifyUnsortedSegmentReduction(*this);
}
LogicalResult UnsortedSegmentSumOp::verify() {
  return VerifyUnsortedSegmentReduction(*this);
}

                                     rewriter);
  }
};

class ConvertUnsortedSegmentProdOp
    : public GenericConvertUnsortedSegmentReductionOp<
          ConvertUnsortedSegmentProdOp, TF::UnsortedSegmentProdOp, MulOp> {
 public:
  using GenericConvertUnsortedSegmentReductionOp::
      GenericConvertUnsortedSegmentReductionOp;

  static Value GetInitialValue(Type reduce_element_type, Location loc,