Nx バージョン 0.6.0 の破壊的変更

はじめに

Nx は Elixir で行列演算をするためのモジュールです

2023/8/15 にバージョン 0.6.0 がリリースされました

このバージョンでの破壊的変更について説明します

変更内容

公式の CHANGELOG から引用します

Deprecations

• [Nx.Serving] The post-processing function must now be a two-arity function that receives the {output, metadata} as a pair or the stream

Breaking changes

• [Nx.Serving] The nx.serving.postprocessing telemetry event no longer receives the serving output or serving metadata as event metadata

影響を受けるのは Nx.Serving.client_postprocessing の後処理関数の引数です

• v0.5.3: (Nx.Container.t(), metadata(), client_info() -> term())
• v0.6.0: ({Nx.Container.t(), metadata()}, client_info() -> term())

今まで３つの引数だったものが、2つの引数に変わります

Nx.Serving の後処理を使っていない場合は影響を受けません

必要な修正

具体的には以下のような修正を行います

serving =
  fn opts -> Nx.Defn.jit(&Nx.multiply(&1, 2), opts) end
  |> Nx.Serving.new(compiler: EXLA)
  |> Nx.Serving.client_preprocessing(fn input ->
    # 前処理
    {Nx.Batch.stack([input]), :client_info}
  end)
- |> Nx.Serving.client_postprocessing(fn output, _metadata, _client_info ->
+ |> Nx.Serving.client_postprocessing(fn {output, _metadata}, _client_info ->
    # 後処理
    Nx.squeeze(output, axes: [0])
  end)
  |> Nx.Serving.distributed_postprocessing(fn output ->
    # 分散している場合の後処理
    Nx.backend_transfer(output, Nx.BinaryBackend)
  end)

単純に　output, _metadata を {} で囲うだけです

まとめ

定期的に以下のリポジトリーを更新しているのですが、軽微で単純な変更だったので、修正も簡単でした