Calling breakpoint_if_nonfinite on TPU

[lucaslingle]

The function breakpoint_if_nonfinite, defined in the Jax page on the debug module, and shown below, does not work on TPU.

import jax
import jax.numpy as jnp


def breakpoint_if_nonfinite(x):
    is_finite = jnp.isfinite(x).all()

    def true_fn(x):
        pass

    def false_fn(x):
        jax.debug.breakpoint()

    jax.lax.cond(is_finite, true_fn, false_fn, x)

The error is

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Host callback not supported for runtime type: stream_executor

However, I am able to call jax.debug.breakpoint() directly within my code on TPU. In light of this, I was wondering what the error above means, and how I can perform conditional breakpointing on TPU.

Thanks!

[mattjj]

Thanks for bringing this up! I think this may be a Colab issue; Colab is running on an older TPU runtime, which is missing several features, though I don't think this particular one was on our radar.

Can you confirm that you're using Colab? If so, we should just add a mention of this issue to the docs. When Colab eventually gets the new TPU runtime ("Cloud TPU VM" I think it's called), this will be fixed.

cc @sharadmv @skye

[lucaslingle]

Hi, thanks for your reply.

I am using a Cloud TPU VM (v3-8), with the latest versions of jax, jaxlib, and the nightly libtpu. For some reason, it still didn't work for me.

[sharadmv]

Hm I suspect this is an instance of #12260 (lack of TFRT support on cloud tpu) though it seems to be only in control flow.

I'll investigate this more this week.

[lucaslingle]

Any updates on this? The ability to perform conditional breakpointing on TPUs would be a great feature!