// Transpose conv supports hybrid computation with quantized weights.
def FoldQuantWeightsIntoTposeConv : Pat<
  (TFL_TransposeConvOp
    $output_shape,
    (TFL_DequantizeOp $quant_weights),
    $quant_input,
    $bias, $padding, $stride_h, $stride_w, $faf),
  (TFL_TransposeConvOp $output_shape, $quant_weights,
    $quant_input, $bias, $padding, $stride_h, $stride_w, $faf), 

  PartialTensorShape output_shape =
      ConvertTypeToTensorShape(mlir::UnrankedTensorType::get(b.getF32Type()));
  EXPECT_TRUE(output_shape.IsIdenticalTo(PartialTensorShape()));
}

TEST(ConvertTypeToTensorTypeTest, NonFullyDefinedRankedTensorType) {
  mlir::MLIRContext context;
  RegisterDialects(context);
  mlir::Builder b(&context);

  PartialTensorShape output_shape = ConvertTypeToTensorShape(

// cases, the additional attribute should be dropped.

// CHECK: name: "StatefulIf"
// CHECK-NOT: name:
// CHECK: op: "If"
// CHECK-NOT: is_stateless
// CHECK:  attr {
// CHECK:    key: "output_shapes"
// CHECK:    value {
// CHECK:      list {
// CHECK:        shape {
// CHECK:          dim {
// CHECK:            size: 4
// CHECK:          }
// CHECK:        }
// CHECK:      }
// CHECK:    }

  } : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>

  func.return %0 : tensor<!tf_type.variant>
// CHECK: "tf.MapDataset"(
// CHECK-SAME: <{f = @{{.*}}, metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], preserve_cardinality = true, use_inter_op_parallelism = true}> {Targuments = []}
}

// CHECK-LABEL: TfTakeWhileDataset

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op
     (TF_MapDatasetOp $input_dataset, $other_arguments, $f, $output_types,
        $output_shapes, $use_inter_op_parallelism, $preserve_cardinality,
        $force_synchronous, $map_dataset_metadata),
     $batch_size, $drop_remainder, $parallel_copy, $batch_output_types,

    return failure();

  // Build the lstm op.
  SmallVector<int64_t, 3> output_shape;
  if (time_majored) {
    output_shape = {time, batch, n_output};
  } else {
    output_shape = {batch, time, n_output};
  }
  auto result_type = mlir::RankedTensorType::get(
      output_shape,
      mlir::cast<RankedTensorType>(final_inputs.getType()).getElementType());

    SmallVector<int64_t, 4> output_shape;
    for (int i = 0; i < num_dimensions; ++i) {
      perm.push_back(perm_tensor.getValues<IntegerAttr>()[i].getInt());
      output_shape.push_back(input_shape[perm[i]]);

      // Check that the derived output shape matches the static shape.
      assert(!output_type.hasStaticShape() ||
             output_type.getShape()[i] == output_shape[i]);
    }

  attr {
    key: "branches"
    value {
      list {
        func {
          name: "foo"
        }
        func {
          name: "bar"
        }
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
      }
    }
  }
  experimental_debug_info {
  }
}
library {
  function {
    signature {
      name: "bar"
      output_arg {
        name: "bar"

inline DenseElementsAttr GetShape(Value output_val, bool truncate = false) {
  auto output_shape = output_val.getType().dyn_cast<ShapedType>().getShape();

  SmallVector<int32_t> shape;
  shape.reserve(output_shape.size());

  bool needs_truncation = true;
  for (size_t dim_idx = 0; dim_idx < output_shape.size(); ++dim_idx) {
    int64_t dim = output_shape[dim_idx];
    if (truncate && needs_truncation && dim == 1) {
      continue;

  // expected-remark@below {{Result #0, ID 0 : 0, 1}}
  %iter_handle1 = "tf.IteratorV2"() {container = "c", shared_name = "v0", output_shapes = [#tf_type.shape<>], output_types = [!tf_res]} : () -> !tf_res
  // expected-remark@below {{Result #0, ID 1 : 0, 1}}
  %iter_handle2 = "tf.IteratorV2"() {container = "c", shared_name = "v1", output_shapes = [#tf_type.shape<>], output_types = [!tf_res]} : () -> !tf_res
  func.return
}

// -----