ops->read_only_memory_region_ops = static_cast<TF_ReadOnlyMemoryRegionOps*>(
      plugin_memory_allocate(TF_READ_ONLY_MEMORY_REGION_OPS_SIZE));
  ops->read_only_memory_region_ops->cleanup =
      tf_read_only_memory_region::Cleanup;
  ops->read_only_memory_region_ops->data = tf_read_only_memory_region::Data;
  ops->read_only_memory_region_ops->length = tf_read_only_memory_region::Length;

  }
  ~MatMulGradientFunction() override {
    for (auto input : forward_inputs_) {
      if (input) {
        input->Unref();
      }
    }
  }

 private:
  // TODO(b/174778737): Only hold needed inputs.
  vector<AbstractTensorHandle*> forward_inputs_;
  AttrBuilder forward_attrs_;
};

class NegGradientFunction : public GradientFunction {
 public:
  Status Compute(AbstractContext* ctx,

                      /*use_function=*/UseFunction());
  ASSERT_EQ(errors::OK, s.code()) << s.message();
}

TEST_P(UnifiedAPI, TestUnknownShapeTracing) {
  if (!UseFunction()) {
    GTEST_SKIP() << "Tracing only test.";
  }
  if (UseMlir()) {
    // TODO(b/173074167): Remove this.
    GTEST_SKIP() << "MlirTensor::Shape is not implemented yet.";
  }
  AbstractContextPtr ctx(BuildFunction("test_fn"));

    core::RefCountPtr<AbstractFunction> scoped_func;
    // Returning null tensors from a tf.function is not supported, so we keep
    // track of indices in the model's outputs are nullptr in this set.
    // The FunctionDef only outputs the non-null tensors. We later pad the
    // function op outputs to have nullptrs at the `null_indices`.
    absl::flat_hash_set<int> null_indices;
    {
      AbstractContextPtr func_ctx(BuildFunction(fn_name));

  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,

   *             |   |  |  |   |
   *             |   v  |  v   |
   *             |      |      |
   *             +------+------+
   *                    |
   *                    v
   *
   *  Only the second output from split is used as function output
   */
  // Define
  TF_Operation* feed = Placeholder(func_graph_, s_);
  TF_Operation* split = Split3(feed, func_graph_, s_);

  TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
      ctx, inputs[0], inputs[1], &loss, &backprop,
      "SparseSoftmaxCrossEntropyWithLogits"));
  // `gradient_checker` only works with model that returns only 1 tensor.
  // Although, `ops::SparseSoftmaxCrossEntropyWithLogits` returns 2 tensors, the
  // second tensor isn't needed for computing gradient so we could safely drop
  // it.
  outputs[0] = loss;

  std::string translated_name = TranslateName(fname);
  ops_->new_read_only_memory_region_from_file(
      filesystem_.get(), translated_name.c_str(), region.get(),
      plugin_status.get());

  if (TF_GetCode(plugin_status.get()) == TF_OK)
    *result = std::make_unique<ModularReadOnlyMemoryRegion>(
        std::move(region), read_only_memory_region_ops_.get());

  return StatusFromTF_Status(plugin_status.get());
}

}

// This test verifies the following:
// 1. Start the GRPC server for worker 1 using single host server def A with
//    only worker 1.
// 2. Create a context B using A.
// 3. Create the variable in B.
// 4. Create another single host server def C with only worker 0.
// 5. Start the GRPC server for worker 0 using C.
// 6. Create a context D using the full cluster server def E.
// 7. Read the variable in D.