// 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,
//   flattened_contracting_dim)

  // Construct full start_indices with given start_indices and
  // start_index_map.
  const ArrayRef<int64_t> operand_shape =
      mlir::cast<ShapedType>(operand.getType()).getShape();
  const int64_t operand_rank = operand_shape.size();

  // Fills zeros if start_index is not given in start_indices.
  Value empty_start_indices = builder.create<TF::FillOp>(

    // Shape after dilation.
    SmallVector<int64_t> dilated_shape(rank);
    ArrayRef<int64_t> operand_shape = operand_type.getShape();
    for (int i = 0; i < rank; ++i) {
      dilated_shape[i] =
          operand_shape[i] + interior_padding_i64[i] * (operand_shape[i] - 1);
    }
    TensorType output_type = op.getResult().getType().cast<TensorType>();

// 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,
//   flattened_contracting_dim)

    XlaSelectAndScatterOp op) {
  DCOMMENT_OP(op, "Inferring shape for XlaSelectAndScatterOp");

  auto operand_shape = mlir::cast<ShapedType>(op.getOperand().getType());
  auto source_shape = mlir::cast<ShapedType>(op.getSource().getType());
  DenseElementsAttr window_dimensions, window_strides, padding;
  if (operand_shape.hasRank() && source_shape.hasRank() &&
      matchPattern(op.getWindowDimensions(), m_Constant(&window_dimensions)) &&

    if (shaped_type.hasStaticShape()) {
      has_static_operand_shape = true;
      operand_shape = shaped_type.getShape();
    }
  }

  SmallVector<int64_t, 4> broadcastedShape;
  if (has_static_cond_shape && has_static_operand_shape &&
      !OpTrait::util::getBroadcastedShape(cond_shape, operand_shape,
                                          broadcastedShape)) {

    std::array<llvm::ArrayRef<int64_t>, 2> operand_shapes;
    operand_shapes[i] = broadcast_arg_type.getShape();
    operand_shapes[1 - i] = argument_type.getShape();

    // Check that the input of the broadcast and the other operand is broadcast
    // compatible.
    llvm::SmallVector<int64_t, 4> broadcasted_shape;
    if (!get_broadcasted_shape(operand_shapes[0], operand_shapes[1],
                               broadcasted_shape))

    TensorType bn_train_input_type_tensor =
        mlir::cast<TensorType>(bn_train_input.getType());

    if (op.getIsTraining()) {
      // Training case.
      auto operand_shape = bn_train_input_type_tensor.getShape();
      // The mean and variance are each 1 dimensional arrays the size of the
      // feature dimension, with the same element type as the operand (x).