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 在圖形模式下使用輸入調用時。