Метод experimental_get_compiler_ir

Метод experimental_get_compiler_ir принадлежит классу Function и позволяет получить промежуточное представление (Intermediate Representation, IR) функции после ее компиляции. Метод применяется к объекту функции, полученному через get_concrete_function или декоратор tf.function. Первым параметром передаются позиционные аргументы, соответствующие сигнатуре функции. Вторым параметром можно указать тип представления через аргумент stage. Третьим параметром передается список device_names с именами устройств для компиляции.

Метод возвращает строку с IR-представлением функции. Это может быть граф TensorFlow, HLO-модуль XLA или другой формат в зависимости от указанной стадии компиляции. Метод помечен как экспериментальный и может измениться в будущих версиях TensorFlow.

Синтаксис

Function.experimental_get_compiler_ir(*args, **kwargs)

Параметры

Метод принимает следующие параметры:

*args - позиционные аргументы, соответствующие сигнатуре функции
stage - стадия компиляции: 'graph', 'hlo', 'hlo_serialized', 'optimized_hlo' и другие
device_names - список имен устройств для компиляции

Пример

Давайте создадим простую функцию и получим ее IR-представление на стадии графа:

import tensorflow as tf

@tf.function
def add_fn(x, y):
	return x + y

concrete = add_fn.get_concrete_function(
	tf.TensorSpec(shape=(2,), dtype=tf.float32),
	tf.TensorSpec(shape=(2,), dtype=tf.float32)
)

ir = concrete.experimental_get_compiler_ir(
	tf.constant([1.0, 2.0]),
	tf.constant([3.0, 4.0])
)(stage='graph')

print(type(ir))
print(ir)

Результат выполнения кода:

<class 'str'>
"node {
  name: 'add'
  op: 'AddV2'
  input: 'x'
  input: 'y'
  ...
}"

Пример

Давайте получим HLO-представление функции на стадии hlo:

import tensorflow as tf

@tf.function
def mul_fn(x, y):
	return x * y

concrete = mul_fn.get_concrete_function(
	tf.TensorSpec(shape=(2,), dtype=tf.float32),
	tf.TensorSpec(shape=(2,), dtype=tf.float32)
)

ir = concrete.experimental_get_compiler_ir(
	tf.constant([1.0, 2.0]),
	tf.constant([3.0, 4.0])
)(stage='hlo')

print(type(ir))
print(ir[:200])

Результат выполнения кода:

<class 'str'>
"HloModule mul_fn, entry_computation_layout={(f32[2]{0}, f32[2]{0})->f32[2]{0}}"

Пример

Давайте проверим, как метод работает с функцией, использующей матричное умножение:

import tensorflow as tf

@tf.function
def matmul_fn(x, y):
	return tf.matmul(x, y)

concrete = matmul_fn.get_concrete_function(
	tf.TensorSpec(shape=(2, 3), dtype=tf.float32),
	tf.TensorSpec(shape=(3, 2), dtype=tf.float32)
)

ir = concrete.experimental_get_compiler_ir(
	tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
	tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
)(stage='graph')

print(type(ir))
print('MatMul' in ir)

Результат выполнения кода:

<class 'str'>
True

Смотрите также

класс Function,
который представляет скомпилированную функцию TensorFlow
метод get_concrete_function,
который создает конкретную функцию с фиксированной сигнатурой
метод experimental_get_compiler_ir,
который возвращает IR-представление скомпилированной функции
класс Function,
который позволяет работать с графами и компиляцией TensorFlow