assertEquals(3, mmHeap.size());
    assertTrue("Heap is not intact after remove()", mmHeap.isIntact());
    assertTrue("Heap contains two 1's", mmHeap.contains(1));
    assertTrue("Heap contains two 1's", mmHeap.remove(1));
    assertTrue("Heap contains 1", mmHeap.contains(1));
    assertTrue("Heap contains 1", mmHeap.remove(1));
    assertFalse("Heap does not contain 1", mmHeap.contains(1));

    assertEquals(3, mmHeap.size());
    assertTrue("Heap is not intact after remove()", mmHeap.isIntact());
    assertTrue("Heap contains two 1's", mmHeap.contains(1));
    assertTrue("Heap contains two 1's", mmHeap.remove(1));
    assertTrue("Heap contains 1", mmHeap.contains(1));
    assertTrue("Heap contains 1", mmHeap.remove(1));
    assertFalse("Heap does not contain 1", mmHeap.contains(1));

      Heap heap;
      if (crossOver == index) {
        heap = this;
      } else {
        index = crossOver;
        heap = otherHeap;
      }
      heap.bubbleUpAlternatingLevels(index, x);
    }

    /**
     * Bubbles a value from {@code index} up the levels of this heap, and returns the index the
     * element ended up at.
     */
    @CanIgnoreReturnValue

    ],
)

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",
        "//tensorflow/c/eager:abstract_context",
        "//tensorflow/c/eager:c_api",

*   Fixes a heap OOB read in `tf.ragged.cross`
    ([CVE-2021-41212](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-41212))
*   Fixes a heap OOB in shape inference for `QuantizeV2`
    ([CVE-2021-41211](https://cve.mitre.org/cgi-bin/cvename.cgi?name=CVE-2021-41211))
*   Fixes a heap OOB read in all `tf.raw_ops.QuantizeAndDequantizeV*` ops

    ],
)

tf_cc_test(
    name = "parallel_device_remote_test",
    srcs = ["parallel_device_remote_test.cc"],
    # TODO(b/136478427): Enable global heap checking when servers shut down
    # cleanly.
    args = ["--heap_check="],
    deps = [
        ":parallel_device",
        ":parallel_device_testlib",
        "//tensorflow/c:c_api",
        "//tensorflow/c:c_api_experimental",

        final JvmMemoryObj jvmMemoryObj = new JvmMemoryObj();
        jvmObj.memory = jvmMemoryObj;
        final JvmMemoryHeapObj jvmMemoryHeapObj = new JvmMemoryHeapObj();
        jvmMemoryObj.heap = jvmMemoryHeapObj;
        jvmMemoryHeapObj.used = mem.getHeapUsed().getBytes();
        jvmMemoryHeapObj.committed = mem.getHeapCommitted().getBytes();
        jvmMemoryHeapObj.max = mem.getHeapMax().getBytes();

public final class GcFinalization {
  private GcFinalization() {}

  /**
   * 10 seconds ought to be long enough for any object to be GC'ed and finalized. Unless we have a
   * gigantic heap, in which case we scale by heap size.
   */
  private static long timeoutSeconds() {
    // This class can make no hard guarantees.  The methods in this class are inherently flaky, but

        buf.append("\"heap\":{");
        append(buf, "used", () -> mem.getHeapUsed().getBytes()).append(',');
        append(buf, "committed", () -> mem.getHeapCommitted().getBytes()).append(',');
        append(buf, "max", () -> mem.getHeapMax().getBytes()).append(',');
        append(buf, "percent", () -> mem.getHeapUsedPercent());
        buf.append("},");
        buf.append("\"non_heap\":{");

public final class GcFinalization {
  private GcFinalization() {}

  /**
   * 10 seconds ought to be long enough for any object to be GC'ed and finalized. Unless we have a
   * gigantic heap, in which case we scale by heap size.
   */
  private static long timeoutSeconds() {
    // This class can make no hard guarantees.  The methods in this class are inherently flaky, but