buf->Unref();
  return tensorflow::wrap(tensorflow::TensorHandle::CreateLocalHandle(
      std::move(t), device, device, context));
}

// This function will block till the operation that produces `h` has
// completed. This is only valid on local TFE_TensorHandles. Returns the size in
// bytes of the memory pointed to by the device pointer returned above.
size_t TFE_TensorHandleDeviceMemorySize(TFE_TensorHandle* h,

      immediate_execution_ctx_.get(), {X.get(), Y.get()}, UseFunction()));
}

TEST_P(CppGradients, TestBiasAddGrad) {
  if (UseFunction() && UseMlir()) {
    GTEST_SKIP() << "SetAttrString has not been implemented yet.\n";
  }

  // A
  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f};
  int64_t A_dims[] = {2, 2};
  AbstractTensorHandlePtr A;
  {
    AbstractTensorHandle* A_raw;

  EXPECT_EQ(attr_value.i(), 2);

  // Placeholder oper now has a consumer.
  ASSERT_EQ(1, TF_OperationOutputNumConsumers(TF_Output{feed, 0}));
  TF_Input feed_port;
  EXPECT_EQ(1, TF_OperationOutputConsumers(TF_Output{feed, 0}, &feed_port, 1));
  EXPECT_EQ(add, feed_port.oper);
  EXPECT_EQ(0, feed_port.index);

  // The scalar const oper also has a consumer.
  ASSERT_EQ(1, TF_OperationOutputNumConsumers(TF_Output{three, 0}));