Ссылка: Jax: Визуализация вычислительного графа программы jax

Запустите следующий скрипт Python для создания вычислительного графа:

import jax
import jax.lax as lax
from jax.lib import xla_bridge
import jax.numpy as np
from jaxlib.xla_extension import HloPrintOptions

# jax.config.update('jax_platform_name', 'cpu')

backend = xla_bridge.get_backend()
print(backend.platform_version)

option = HloPrintOptions()
# option.print_metadata = False
# option.include_layout_in_shapes = False
# option.print_extra_attributes = False

@jax.xla_computation
def f(x, y):
    a = np.einsum('pqrs,tuqvr->pstuv', x, y)
    return lax.sin(a)

c = f(np.ones((3,4,5,6)), np.ones((7,8,4,9,5)))
module = backend.compile(c).hlo_modules()[0]
hlo_text = module.to_string(option)
print(hlo_text)

with open('/tmp/hlo.txt', 'w') as f:
    print(hlo_text, file=f)

Установите инструмент сборки Bazel:

sudo mkdir -p /usr/local/lib/bazel/bin
sudo wget -P /usr/local/lib/bazel/bin https://github.com/bazelbuild/bazel/releases/download/5.0.0/bazel-5.0.0-linux-x86_64
sudo chmod +x /usr/local/lib/bazel/bin/bazel-5.0.0-linux-x86_64

Скомпилируйте interactive_graphviz:

git clone --depth=1 https://github.com/tensorflow/tensorflow.git
cd tensorflow
bazel build tensorflow/compiler/xla/tools:interactive_graphviz

Компиляция занимает около 7 минут на виртуальной машине Cloud TPU (с 96 процессорами).

Выполнить interactive_graphviz:

bazel-bin/tensorflow/compiler/xla/tools/interactive_graphviz --hlo_text=/tmp/hlo.txt

Вход list computations:

command: list computations
Entry computation:
  xla_computation_f.13
Subcomputations:
  fused_computation.1.clone.clone
  fused_computation.4
  fused_computation.3

Введите вычисление входа в соответствии с выходом. В данном примере это xla_computation_f.13:

command: xla_computation_f.13
file:///tmp/interactive_graphviz.1645348689190091.html

Откройте file:///tmp/interactive_graphviz.1645348689190091.html для просмотра вычислительного графика:

Если вы хотите узнать больше о TPU и JAX, не забудьте ознакомиться с поддержанным сообществом подробным руководством по TPU: https://github.com/ayaka14732/tpu-starter.