namespace {

#if !defined(IS_MOBILE_PLATFORM) && !defined(IS_SLIM_BUILD)

// Creates a placeholder representing an input to the cond or body graph.
// TODO(skyewm): remove these from final graph
bool CreateInput(const TF_Output& parent_input, TF_Graph* g, const char* name,
                 TF_Output* input, TF_Status* status) {

     *
     *    dA = U
     *    dbias = reduceSum(U, dims = channel_dim)
     */

    AbstractTensorHandle* upstream_grad = grad_outputs[0];
    DCHECK(upstream_grad);

    // Recover data format from forward pass for gradient.
    std::string data_format;
    TF_RETURN_IF_ERROR(forward_attrs_.Get("data_format", &data_format));

    // Grad for A
    grad_inputs[0] = upstream_grad;

      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
  auto collective_executor_handle = context->GetCollectiveExecutorHandle();
  tensorflow::Notification done;
  collective_executor_handle->get()->remote_access()->CheckPeerHealth(
      task, timeout_in_ms, [&done, status](const Status& s) {
        status->status = s;
        done.Notify();
      });
  done.WaitForNotification();
}

  // and not for correctness. The only downside of keeping this 1 seems to be
  // that the gradient accumulation is unbounded and we will never
  // aggressively aggregate accumulated gradients to recover memory.
  // Revisit and fix.
  return 1;
}

// Consumes references to the tensors in the gradient_tensors list and returns
// a tensor with the result.
AbstractTensorHandle* TapeVSpace::AggregateGradients(

        "Tensor expected to have scalar shape found rank: ", shape.dims());
  }
  return absl::OkStatus();
}

TEST_P(UnifiedAPI, TestTensorShapeScalar) {
  if (UseFunction() && UseMlir()) {
    // TODO(b/173074167): Remove this.
    GTEST_SKIP() << "MlirTensor::Shape is not implemented yet.";
  }
  AbstractContextPtr ctx;
  {
    AbstractContext* ctx_raw = nullptr;
    Status s =

  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  // Pack 3 variable handles into one TFE_TensorHandle.
  // When remote is false, function device is placed on task0. Handle types are
  // REMOTE, REMOTE, LOCAL on task0. When remote is true, function device is
  // placed on task1, Handle types are LOCAL, REMOTE, LOCAL on task1.
  int num_replicas = 3;
  std::vector<TFE_TensorHandle*> handles = {h0, h1, h2};

    for (int j = 0; j < underlying_devices_.size(); ++j) {
      components.push_back(std::move(per_device_output_tensors[j][i]));
    }
    if (expected_output_shapes[i].IsFullyDefined()) {
      per_device_outputs.push_back(ParallelTensor::FromTensorHandles(
          *this, std::move(components),
          absl::Span<const int64_t>(expected_output_shapes[i].dim_sizes()),
          status));
    } else {

  // Remove the device that has the host device name since host device is alreay
  // in an initialized context.
  for (auto d = devices.begin(); d != devices.end();) {
    if (absl::StrContains(d->get()->name(), "CPU:0")) {
      d = devices.erase(d);
    } else {
      ++d;
    }
  }

  status->status = tensorflow::unwrap(ctx)->AddDevices(std::move(devices));
}

  TF_Tensor* o = TF_TensorMaybeMove(t);
  ASSERT_TRUE(o == nullptr);  // It is unsafe to move memory TF might not own.
  TF_DeleteTensor(t);
  EXPECT_TRUE(deallocator_called);
}

TEST(CAPI, LibraryLoadFunctions) {
  // TODO(b/73318067): Fix linking for the GPU test generated by the
  // tf_cuda_cc_test() bazel rule and remove the next line.
  if (!GPUDeviceName().empty()) return;

  // Remote device on `worker1`.
  const char remote_device[] = "/job:localhost/replica:0/task:1/device:CPU:0";
  // `ctx_0`, `ctx_1`, `ctx_2` contains `remote_device`.
  {
    const std::vector<std::string>& device_names = ListDeviceNames(ctx_0);
    ASSERT_TRUE(std::find(device_names.begin(), device_names.end(),
                          remote_device) != device_names.end());
  }

  {