auto output_types = fused_lstm_func_.getFunctionType().getResults();
  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),
            output_shape.size());
  for (int i = 0; i < output_shape.size(); i++) {
    EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getDimSize(i),
              output_shape[i]);
  }
}

namespace mlir {
namespace TF {

// Populate patterns to unroll tf.BatchMatMulV2 op into a sequence of TF ops.
// Since TFLite does not support BatchMatMul operation, it unrolls a BatchMatMul
// op into tf.Reshape, tf.Slice, tf.MatMul, tf.Pack, and tf.Reshape ops.
void PopulateUnrollTfBatchMatMul(MLIRContext* context,
                                 RewritePatternSet& patterns);

}  // namespace TF
}  // namespace mlir

  // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%[[ADD]], %[[CST_2]]) {device = ""} : (tensor<64x8x64xf32>, tensor<2xi32>) -> tensor<512x64xf32>
  // CHECK: return %[[RESHAPE]] : tensor<512x64xf32>
}

// CHECK-LABEL: dont_move_transpose_different_ranks

                                   ConversionPatternRewriter& rewriter) {
  if (index_vector_dim == indices_type.getRank()) {
    llvm::SmallVector<int64_t, 4> new_start_indices_shape(
        indices_type.getShape().begin(), indices_type.getShape().end());
    new_start_indices_shape.push_back(1);
    indices_type = RankedTensorType::get(new_start_indices_shape,
                                         indices_type.getElementType());

// CHECK:         }
func.func @tfl_wrapped_jax_random_normal(%arg0: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
  // This is a fake jax random normal body.
  %0 = stablehlo.constant dense<0.0> : tensor<12xf32>
  %1 = "stablehlo.reshape"(%0) : (tensor<12xf32>) -> tensor<3x4xf32>
  %2 = "stablehlo.tuple"(%1) : (tensor<3x4xf32>) -> tuple<tensor<3x4xf32>>
  func.return %2 : tuple<tensor<3x4xf32>>
}


// CHECK-LABEL:   func @tfl_wrapped_jax_random_uniform(

    auto rhs_type = mlir::dyn_cast<RankedTensorType>(input_rhs.getType());

    if (!lhs_type || !rhs_type) return failure();

    auto lhs_shape = lhs_type.getShape();
    auto rhs_shape = rhs_type.getShape();

    // Ensure that input ranks are at least 2.
    const int dims_a = lhs_shape.size();
    const int dims_b = rhs_shape.size();
    if (dims_a < 2 || dims_b < 2) {
      return failure();

  }
  ApiConverter::ToC(old_shape, &old_shape_c);
  executor->TpuTransferManager_GetInfeedLayoutFn(&old_shape_c, &new_shape_c);
  xla::Shape new_shape = ApiConverter::FromC(&new_shape_c);
  ApiConverter::Destroy(&old_shape_c);
  ApiConverter::Destroy(&new_shape_c);

  auto minor_to_major = new_shape.layout().minor_to_major();
  return std::vector<int64_t>(minor_to_major.begin(), minor_to_major.end());
}

    Value bcast_op_result = (*bcast_op)->getResult(0);
    auto bcast_op_result_type =
        mlir::cast<RankedTensorType>(bcast_op_result.getType());
    const ArrayRef<int64_t> bcast_shape = bcast_op_result_type.getShape();
    const TensorType new_bcast_op_result_type = bcast_op_result_type.cloneWith(
        bcast_shape, accumulation_quantized_element_type);
    bcast_op_result.setType(new_bcast_op_result_type);
  }

  Type dtype = getElementTypeOrSelf(type);
  if (RankedTensorType ty = type.dyn_cast<RankedTensorType>()) {
    llvm::SmallVector<int64_t, 4> shape = {dim};
    shape.append(ty.getShape().begin(), ty.getShape().end());
    return tensorflow::GetTypeFromTFTensorShape(shape, dtype);
  }
  return type;
}

Type GetTensorTypeForTensorList(Type element_type, TF::VariantType handle_dtype,

  %0 = tfr.call @Huu(%arg0)  : (tensor<?xf32>) -> !tfr.tensor<K>
  func.return %0 : !tfr.tensor<K>
}

// -----

// CHECK-LABEL: get_shape
func.func @get_shape(%arg0: !tfr.tensor) -> (!shape.shape, !shape.shape) {
  %0 = tfr.get_shape %arg0 -> !shape.shape
  %1 = "tfr.get_shape"(%arg0) : (!tfr.tensor) -> !shape.shape
  func.return %0, %1 : !shape.shape, !shape.shape
}

// -----

// CHECK-LABEL: get_real_shape