No `/device: TPU:x` When Runing TensorBoard

[Danc2050]

Environment information (required)

diagnose_tensorboard.txt Chrome Version: 98.0.4758.102

Issue description

TensorBoard does not seem to show any TPU device, unlike it does with a GPU. Is this not supported? I ran this on my own code as well as a public notebook with JAX (just adding in the JAX code to trace). The issue is also present when I use TensorFlow.

• Screenshot of issue:

Minimal Code Example

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()

key = random.PRNGKey(0)
x = random.normal(key, (10,))
print(x)

size = 3000

jax.profiler.start_trace("/tmp/tensorboard")
x = random.normal(key, (size, size), dtype=jnp.float32)
%timeit jnp.dot(x, x.T).block_until_ready()  # runs on the GPU, but I also think would use TPU as it's a matrix multiply

jax.profiler.stop_trace()

!pip3 install tensorflow tbp-nightly

# # Load the TensorBoard notebook extension.
%load_ext tensorboard

# Launch TensorBoard and navigate to the Profile tab to view performance profile
%tensorboard --logdir=/tmp/tensorboard

[dmmolitor]

A couple of things to check:

• It looks like a number of threads were captured, but only a few are shown in this screenshot. If you scroll down in the Trace View, are all threads running on CPU and none on TPU?

• Are you certain that the code is indeed running on TPU? Could you add the logs?