tensorflow - 使用 TPUClusterResolver 时为 TPU 转储 HLO IR

Question

我正在通过 Google Colab 和 GCP 使用 TPU，并想转储 XLA IR。但是我在 github xla index中看到了 xla doc ，它只显示了后端是 CPU 或 GPU 的方式。

我曾尝试使用XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/content/iir/" TF_XLA_FLAGS=--tf_xla_cpu_global_jit运行以 CPU 为目标的程序并获取转储的 hlo 文件。我也尝试过capture_tpu_file并且只能在“op_profile”页面中为每个操作员获取 ir。那么当后端是 TPU 时，有没有办法为整个程序转储 XLA IR？

谢谢！

周杰伦

标签： tensorflowgoogle-colaboratorytpugoogle-cloud-tpu

Answer 1

不幸的是，目前没有办法在 Cloud TPU 上转储/访问 XLA IR，因为需要在 TPU 服务器上设置 XLA_FLAGS。