value_str.copy(static_cast<char*>(data), value_str.length(), 0);
  result->data = data;
  result->length = value_str.length();
  result->data_deallocator = [](void* data, size_t length) { free(data); };
  return result;
}

TF_Buffer* TF_CoordinationServiceGetKeyValue(const char* key, int64_t key_size,
                                             TF_CoordinationServiceAgent* agent,

  // The resulting TensorHandle owns an opaque pointer to "device memory", which
  // for a ParallelDevice is really a ParallelTensor. When the TensorHandle is
  // deleted, it will call ParallelTensorDeallocator to free the struct.
  ParallelTensor* t_released = t.release();
  TFE_CustomDeviceTensorHandleMethods handle_methods;
  handle_methods.num_dims = &ParallelTensorNumDims;
  handle_methods.dim = &ParallelTensorDim;

      [gcs_file](const std::string& dir, GcsFileStat* stat, TF_Status* status) {
        auto children =
            GetChildrenBounded(gcs_file, dir, 1, true, true, status);
        if (TF_GetCode(status) != TF_OK) return;
        if (!children.empty()) {
          stat->base = {0, 0, true};
          return TF_SetStatus(status, TF_OK, "");
        } else {
          return TF_SetStatus(status, TF_INVALID_ARGUMENT, "Not a directory!");

  // Try to register as many filesystems as possible.
  // Free memory once we no longer need it
  Status status;
  for (int i = 0; i < info->num_schemes; i++) {
    status.Update(ValidateAndRegisterFilesystems(info, i));
    info->plugin_memory_free(info->ops[i].scheme);
    info->plugin_memory_free(info->ops[i].filesystem_ops);
    info->plugin_memory_free(info->ops[i].random_access_file_ops);

    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "op_to_reset should not be nullptr");
  }
}

void TFE_ContextEnableGraphCollection(TFE_Context* ctx) {
  tensorflow::unwrap(ctx)->SetShouldStoreGraphs(true);
}

void TFE_ContextDisableGraphCollection(TFE_Context* ctx) {
  tensorflow::unwrap(ctx)->SetShouldStoreGraphs(false);
}

uint64_t TFE_GetContextId(TFE_Context* ctx) {
  tensorflow::EagerContext* context =

  EXPECT_EQ(op_def.input_arg_size(), 0);
  EXPECT_EQ(op_def.output_arg_size(), 1);
  EXPECT_TRUE(op_def.is_stateful());

  TF_DeleteBuffer(buffer);
}

void AssertEqual(TF_Function* f1, TF_Function* f2) {
  string s1, s2;
  tensorflow::FunctionDef fdef1, fdef2;
  ASSERT_TRUE(GetFunctionDef(f1, &fdef1));
  ASSERT_TRUE(GetFunctionDef(f2, &fdef2));
  SerializeToStringDeterministic(fdef1, &s1);

                                /*transpose_a = */ true,
                                /*transpose_b = */ true, name_grad_A.c_str()));

      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, A.get(), &matmul_B_output,
                                /*transpose_a = */ true,
                                /*transpose_b = */ true, name_grad_B.c_str()));
    }

    // Gradient for A

    ASSERT_EQ(errors::OK, s.code()) << s.message();
    x.reset(x_raw);
  }

  PartialTensorShape shape;
  Status s = x->Shape(&shape);
  ASSERT_EQ(errors::OK, s.code()) << s.message();
  ASSERT_TRUE(shape.unknown_rank());
}

TEST_P(UnifiedAPI, TestPartialShapeTracing) {
  if (!UseFunction()) {
    GTEST_SKIP() << "Tracing only test.";
  }
  if (UseMlir()) {

                       /*use_function*/ ::testing::Values(true, false)));
#else
INSTANTIATE_TEST_SUITE_P(
    UnifiedCAPI, GradientCheckerTest,
    ::testing::Combine(::testing::Values("graphdef"),
                       /*tfrt*/ ::testing::Values(false),
                       /*use_function*/ ::testing::Values(true, false)));
#endif
}  // namespace
}  // namespace internal

      opts->session_options.options,
      static_cast<tensorflow::ContextDevicePlacementPolicy>(
          opts->device_placement_policy),
      opts->async, device_mgr.release(),
      /*device_mgr_owned*/ true, std::move(r),
      /*cluster_flr=*/nullptr,
      /*collective_executor_mgr=*/nullptr,
      /*run_eager_op_as_function=*/opts->run_eager_op_as_function,
      /*jit_compile_rewrite=*/opts->jit_compile_rewrite);