Python tf.types.experimental.GenericFunction.experimental_get_compiler_ir用法及代码示例

参数

• *args 用于编译的参数；与用于调用函数的参数相同。需要即刻的张量。
• **kwargs 用于编译的关键字参数。

返回

• 可使用以下 kwargs 调用的函数：
  • stage应该序列化编译器 IR。允许的值为：
    • hlo:HLO 从 TF 转换后的输出 (https://www.tensorflow.org/xla/operation_semantics)。
    • hlo_serialized :Like stage= hlo ，但输出是一个序列化的 HLO 模块原型(一个字节对象)。
    • optimized_hlo：编译器优化后的HLO。
    • optimized_hlo_serialized :Like stage= optimized_hlo ，但输出是一个序列化的 HLO 模块原型(一个字节对象)。
    • optimized_hlo_dot ：优化的 HLO 以适合 Graphviz 的 DOT 格式。
  • device_name 可以是 None，在这种情况下，首选设备用于编译，也可以是设备名称。它可以是完整的设备名称，也可以是部分名称，例如 /device:CPU:0 。

  例如，对于

  @tf.function(jit_compile=True)
  def f(x):
    return x + 1
  
  f.experimental_get_compiler_ir(tf.random.normal([10, 10])(stage='hlo')

  输出是：

  HloModule a_inference_f_13__.9
  
  ENTRY %a_inference_f_13__.9 (arg0.1:f32[10,10]) -> f32[10,10] {
    %arg0.1 = f32[10,10]{1,0} parameter(0), parameter_replication={false}
    %reshape.2 = f32[10,10]{1,0} reshape(f32[10,10]{1,0} %arg0.1)
    %constant.3 = f32[] constant(1)
    %broadcast.4 = f32[10,10]{1,0} broadcast(f32[] %constant.3)
    %add.5 = f32[10,10]{1,0} add(f32[10,10]{1,0} %reshape.2,
                                 f32[10,10]{1,0} %broadcast.4)
    %reshape.6 = f32[10,10]{1,0} reshape(f32[10,10]{1,0} %add.5)
    %tuple.7 = (f32[10,10]{1,0}) tuple(f32[10,10]{1,0} %reshape.6)
    ROOT %get-tuple-element.8 = f32[10,10]{1,0}
      get-tuple-element((f32[10,10]{1,0}) %tuple.7), index=0
  }

抛出

• ValueError 如果选择了无效的 stage 或应用于未编译的函数(未设置 jit_compile=True)。
• TypeError 在图形模式下使用输入调用时。