}
  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());
}

  ASSERT_EQ(TF_OK, TF_GetCode(status));
  TF_ShapeHandle* new_shape = TF_ShapeInferenceContextVectorFromSize(
      ctx, TF_ShapeInferenceContextRank(ctx, handle));
  TF_ShapeInferenceContextSetOutput(ctx, 0, new_shape, status);
  TF_DeleteShapeHandle(handle);
  TF_DeleteShapeHandle(new_shape);
}

TEST(OpsTest, TestShapeInference_VectorizeFunction) {

  auto input_type = mlir::cast<mlir::TensorType>(src_input.getType());

  if (input_type.hasRank()) {
    if (input_type.getShape()[split_dimension] == mlir::ShapedType::kDynamic) {
      output_type = input_type;
    } else {
      auto shape = llvm::to_vector<4>(input_type.getShape());
      if (shape[split_dimension] % num_split != 0) {
        return mlir::emitError(
            location,
            llvm::formatv(

  }
  return std::nullopt;
}

ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op) {
  auto composite_result_shape =
      mlir::cast<ShapedType>(old_op->getResults().front().getType()).getShape();
  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW
  output_shape[0] = composite_result_shape[0];
  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];

        type, builder.getFloatAttr(element_ty, unique_index));

  if (auto qtype = mlir::dyn_cast<QuantizedType>(element_ty)) {
    mlir::RankedTensorType new_type = tensorflow::GetTypeFromTFTensorShape(
        type.getShape(), qtype.getStorageType());
    return DenseElementsAttr::get(
        new_type, builder.getIntegerAttr(qtype.getStorageType(), unique_index));
  }
  llvm_unreachable("unhandled element type");
}

    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);
  }

                tensorflow::GetTypeFromTFTensorShape(
                    {splat_elements_attr.getType().getRank()},
                    op_builder->getI64Type()),
                splat_elements_attr.getType().getShape()));
    mlir::arith::ConstantOp fill_value =
        op_builder->create<mlir::arith::ConstantOp>(
            const_op->getLoc(),
            DenseElementsAttr::get(
                tensorflow::GetTypeFromTFTensorShape(

    /**
     * 
     * @return the domain this referral is for
     */
    String getDomain ();


    /**
     * @return the share this referral points to
     */
    String getShare ();


    /**
     * @return the number of characters from the unc path that were consumed by this referral
     */
    int getPathConsumed ();


    /**

      if (v_int > 7 || v_int < -8) {
        return failure();
      }
    }

    Builder builder(op.getContext());
    auto shaped_type =
        mlir::RankedTensorType::get(const_type.getShape(), builder.getI4Type());
    auto newAttr = DenseElementsAttr::getFromRawBuffer(
        shaped_type, mlir::cast<DenseElementsAttr>(op.getValue()).getRawData());
    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);

    return PartialTensorShape();
  }

  if (auto tensor_type = mlir::dyn_cast<mlir::RankedTensorType>(type)) {
    TensorShapeProto tensor_shape_proto;
    ConvertToTensorShapeProto(tensor_type.getShape(), &tensor_shape_proto);
    return PartialTensorShape(tensor_shape_proto);
  }

  // If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
  // Empty TensorShape indicates a scalar.