gtl::ArraySlice<AbstractTensorHandle*> output_gradients,
      absl::Span<AbstractTensorHandle*> result) const override;

  // Builds a tensor filled with ones with the same shape and dtype as `t`.
  Status BuildOnesLike(const TapeTensor& t,
                       AbstractTensorHandle** result) const override;

  // Looks up the ID of a Gradient.

                    /*input_tensors*/ {},
                    /*expected_shape*/ make_shape({3, kUnknownDim}));

  // Infer shape when some dimensions are unknown.
  CheckOutputShapes(
      matmul_op,
      /*input_shapes*/ {make_shape({kUnknownDim, 2}), make_shape({2, 4})},
      /*input_tensors*/ {},
      /*expected_shape*/ make_shape({kUnknownDim, 4}));

  // Infer shape when everything is unknown.

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all

    TF_RETURN_IF_ERROR(operation->SetAttrType("dtype", dtype));
    if (!shape.unknown_rank()) {
      TF_RETURN_IF_ERROR(operation->SetAttrShape(
          "shape", reinterpret_cast<int64_t*>(shape.dim_sizes().data()),
          shape.dims()));
    }
    int num_outputs = 1;
    std::vector<AbstractTensorHandle*> outputs(num_outputs);
    TF_RETURN_IF_ERROR(operation->Execute(

  // Create the variable handle.
  TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
  TFE_OpSetAttrString(op, "container", "localhost", 0);
  TFE_OpSetAttrString(op, "shared_name", "", 0);
  if (!device_name.empty()) {
    TFE_OpSetDevice(op, device_name.c_str(), status);
  }

}

TFE_MonitoringBuckets* TFE_MonitoringNewExponentialBuckets(double scale,
                                                           double growth_factor,
                                                           int bucket_count) {
  return new TFE_MonitoringBuckets([scale, growth_factor, bucket_count]() {
    return tensorflow::monitoring::Buckets::Exponential(scale, growth_factor,

  }

  Status Compute(AbstractContext* ctx,
                 absl::Span<AbstractTensorHandle* const> grad_outputs,
                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    // TODO(vnvo2409): Add shape broadcasting
    /* Given upstream grad U and a Div op: Z = X/Y, the gradients are:
     *
     *    dX = U / Y
     *    dY = -U*X / Y^2 = (X/Y) * -U / Y = -U*Z / Y
     *
     */

    ShapeHandle shape_handle = c.output(i);
    TF_ShapeAndType& shape = output_shapes_result->items[i];
    shape.num_dims = c.Rank(shape_handle);
    if (shape.num_dims == InferenceContext::kUnknownRank) {
      shape.dims = nullptr;
      continue;
    }
    shape.dims = new int64_t[shape.num_dims];
    for (size_t j = 0; j < shape.num_dims; ++j) {
      shape.dims[j] = c.Value(c.Dim(shape_handle, j));

  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  const std::vector<std::unique_ptr<ParallelTensor>>& handles = *outputs;
  const std::vector<int64_t>* shape;
  Status s = handles[0]->Shape(&shape);
  ASSERT_TRUE(s.ok());
  EXPECT_EQ(0, shape->size());
}

TEST(PARALLEL_DEVICE_LIB, TestCancelOnError) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);

TF_CAPI_EXPORT extern void TF_DeleteShapeAndTypeListArray(
    TF_ShapeAndTypeList** shape_list_array, int num_items);

// Infer shapes for the given `op`. The arguments mimic the arguments of the
// `shape_inference::InferenceContext` constructor. Note the following:
//   - The inputs of the `op` are not used for shape inference. So, it is
//     OK to not have the inputs properly set in `op`. See `input_tensors`