// .device of shape is GPU since the op is executed on GPU
    device_name = TFE_TensorHandleDeviceName(retvals[0], status.get());
    ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
    ASSERT_TRUE(absl::StrContains(device_name, "GPU:0")) << device_name;

    // .backing_device of shape is CPU since the tensor is backed by CPU
    backing_device_name =

  def test_ptq_selective_denylist(self):
    """Tests that the op is not quantized when no quantization is enabled."""

    rng = np.random.default_rng(1230)
    random_tensor_gen_fn = lambda shape: rng.uniform(
        low=-1.0, high=1.0, size=shape
    ).astype(np.float32)

    class TwoMatmulModel(module.Module):
      """A model with two matmul ops."""

      @def_function.function

                                                           TF_DeleteStatus);

  TF_Tensor* tensor_shape = Int32Tensor({37, 1});
  TF_Operation* shape = Const(tensor_shape, func_graph.get(), s.get(), "shape");
  TF_Operation* random =
      RandomUniform(shape, TF_FLOAT, func_graph.get(), s.get());

  TF_Output outputs[] = {{random, 0}};
  *func = TF_GraphToFunction(func_graph.get(), name,

    // If weight_shape[2] != 1, it means weight shape was already restored.
    if (weight_shape[2] != 1) return failure();

    // Weight was reshaped into [H, W, 1, InxMul].
    // Since we know in_channels from input_shape, we can derive multiplier.
    int64_t in_channels = input_shape[3];
    // If in_channels is 1, there is no need to restore weight shape.
    if (in_channels == 1) return failure();

    builder.setInsertionPoint(op);
    if (!const_true)
      const_true = builder.create<TF::ConstOp>(
          op.getLoc(),
          DenseIntElementsAttr::get(
              RankedTensorType::get(/*shape=*/{}, builder.getI1Type()), true));

    op.getIsInitialized().replaceAllUsesWith(const_true);
    op.erase();
  }
}

// Performs store-load forwarding. This effectively removes

            output_zero_point_value);

    Value conv_output_value = op.getResult();
    auto output_uniform_quantized_tensor_type = RankedTensorType::getChecked(
        rewriter.getUnknownLoc(),
        /*shape=*/
        mlir::cast<TensorType>(conv_output_value.getType()).getShape(),
        output_uniform_quantized_type);

    SmallVector<Type> new_conv_output_types = {
        output_uniform_quantized_tensor_type};

// It returns three pieces so it doesn't need an allocation for string
// concatenation.
func funcNamePiecesForPrint(name string) (string, string, string) {
	// Replace the shape name in generic function with "...".
	i := bytealg.IndexByteString(name, '[')
	if i < 0 {
		return name, "", ""
	}
	j := len(name) - 1
	for name[j] != ']' {
		j--
	}
	if j <= i {

  func.return %1 : tensor<*xi32>
}

// -----

// Tests that resource ops with both load and store are hoisted
// but input to load and output from store have mixed defined/undefined shapes.

// CHECK-LABEL: func @same_resource_load_and_store_cast
func.func @same_resource_load_and_store_cast() -> tensor<1xi32> {

  // CHECK: %[[RES_HANDLE:[0-9]*]] = "tf.VarHandleOp"