limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_
#define TENSORFLOW_C_EAGER_CUSTOM_DEVICE_TESTUTIL_H_

// A simple logging device to test custom device registration.
#include <memory>

#include "tensorflow/c/c_api.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"

-   type: input id: tfversion attributes: label: Tensorflow Version description:
    placeholder: ex,. tf 2.8 validations: required: true

-   type: dropdown id: Code attributes: label: Custom Code description:
    options: - "Yes" - "No" validations: required: true

-   type: input id: OS attributes: label: OS Platform and Distribution
    description: placeholder: e.g., Linux Ubuntu 16.04 validations: required:

<li><a href="http://sponge2/$KOKORO_BUILD_ID">Sponge2</a></li>
<li><a href="http://sponge/target:$KOKORO_JOB_NAME">Sponge - recent jobs</a></li>
<li><a href="http://fusion2/ci/kokoro/prod:$(echo "$KOKORO_JOB_NAME" | sed 's!/!%2F!g')">Test Fusion - recent jobs</a></li>
<li><a href="http://cs/f:devtools/kokoro/config/prod/$KOKORO_JOB_NAME">Codesearch - job definition</a></li>

  const char* first_device = "/job:worker/replica:0/task:1/device:CPU:0";
  const char* second_device = "/job:worker/replica:0/task:2/device:CPU:0";
  const char* device_name = "/job:localhost/replica:0/task:0/device:CUSTOM:0";
  std::array<const char*, 2> underlying_devices{first_device, second_device};
  RegisterParallelDevice(context.get(), device_name, underlying_devices,
                         status.get());

# setup.python.sh: Install a specific Python version and packages for it.
# Usage: setup.python.sh <pyversion> <requirements.txt>

# Sets up custom apt sources for our TF images.

# Prevent apt install tzinfo from asking our location (assumes UTC)
export DEBIAN_FRONTEND=noninteractive

# Set up shared custom sources
apt-get update
apt-get install -y gnupg ca-certificates

# Deadsnakes: https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa

  virtual ImmediateExecutionContext* GetContext() const = 0;

  // Following two methods are used to support custom device.
  // Return true if the inputs contain custom device tensor handle. It means
  // that the argument need to be handled by a custom device.
  virtual bool HasCustomDeviceInput() const = 0;

  virtual const tensorflow::OpDef* OpDef() const = 0;

- Provide links to your updated versions of the above two colab notebooks.
- Provide links to your TensorFlow model and (optionally) TensorFlow Lite Model.
```

#### Option B: Paste your code here or provide a link to a custom end-to-end colab

```
(You can paste links or attach files by dragging & dropping them below)
- Include code to invoke the TFLite Converter Python API and the errors.

    CHECK_EQ(grad_inputs.size(), 1);
    grad_inputs[0] = grad_outputs[0];
    if (grad_inputs[0]) {
      grad_inputs[0]->Ref();
    }
    return absl::OkStatus();
  }
};

// Computes:
//
// @tf.custom_gradient
// def f(input):
//   def grad(grads):
//     return grads[0]
//   return tf.exp(input), grad
// outputs = [f(inputs[0])]
Status ExpWithPassThroughGrad(AbstractContext* ctx,

        ":not_differentiable",
        "//tensorflow/c/eager:abstract_context",
        "//tensorflow/c/eager:gradients_internal",
    ],
)

tf_cuda_cc_test(
    name = "custom_gradient_test",
    size = "small",
    srcs = [
        "custom_gradient_test.cc",
    ],
    args = ["--heap_check="],  # TODO(b/174752220): Remove
    tags = tf_cuda_tests_tags(),
    deps = [
        "//tensorflow/c:tf_status_helper",

computation, shared state, and the operations that mutate that state. It maps the nodes of a dataflow graph across many machines in a cluster, and within a machine across multiple computational devices, including multicore CPUs, general purpose GPUs, and custom-designed ASICs known as Tensor Processing Units (TPUs). This architecture gives flexibility to the application developer, whereas in previous “parameter server” designs the management of shared state is built into the system, TensorFlow enables developers...