メインコンテンツまでスキップ

問題のステートメント

Issueのタイトル: ak.meanを通じた微分の実行

Awkward Arrayのバージョン​

main ブランチ

説明と再現コード​

これは、構成を少し簡略化した #2591 のフォローアップです。概念的には、平均を取ることで微分が可能になるはずです。現在、これは機能していません。

再現コード:

import awkward as ak
import jax
import uproot

ak.jax.register_and_check()

ttbar_file = "https://github.com/scikit-hep/scikit-hep-testdata/"\
    "raw/main/src/skhep_testdata/data/nanoAOD_2015_CMS_Open_Data_ttbar.root"

def mean_jet_pt(jets):
    return ak.mean(jets.pt)

with uproot.open(ttbar_file) as f:
    arr = f["Events"].arrays(["Jet_pt","Jet_eta", "Jet_phi", "Jet_mass"])
    evtfilter = ak.num(arr["Jet_pt"]) >= 2
    jets = ak.zip(dict(zip(["pt","eta", "phi", "mass"], ak.unzip(arr))), with_name="Momentum4D")[evtfilter]
    jets = ak.to_backend(jets, "jax")

jax.value_and_grad(mean_jet_pt, argnums=0)(jets)

結果:

RuntimeError: Cannot differentiate through count_zero

This error occurred while calling

    ak.mean(
        <Array [[...], [...], ..., [...], [...]] type='140 * var * float32'>
    )

平均を取るスタンドアロンのjaxバージョンは正常に機能します:

import jax.numpy as jnp

def mean(j):
    return jnp.mean(j)

data = jnp.array([1, 7, 3, 5],dtype=float)

jax.value_and_grad(mean, argnums=0)(data)

コード差分

src/awkward/operations/ak_covar.py​

diff --git a/src/awkward/operations/ak_covar.py b/src/awkward/operations/ak_covar.py
index a070ac68..f0decdeb 100644
--- a/src/awkward/operations/ak_covar.py
+++ b/src/awkward/operations/ak_covar.py
@@ -102,52 +102,52 @@ def _impl(x, y, weight, axis, keepdims, mask_identity, highlevel, behavior, attr
     y = ctx.wrap(y_layout)
     weight = ctx.wrap(weight_layout, allow_other=True)
 
     with np.errstate(invalid="ignore", divide="ignore"):
         xmean = ak.operations.ak_mean._impl(
             x, weight, axis, False, mask_identity,
             highlevel=True, behavior=None, attrs=None,
         )
         ymean = ak.operations.ak_mean._impl(
             y, weight, axis, False, mask_identity,
             highlevel=True, behavior=None, attrs=None,
         )
         if weight is None:
-            sumw = ak.operations.ak_count._impl(
-                x,
+            sumw = ak.operations.ak_sum._impl(
+                x * 0 + 1,
                 axis, keepdims, mask_identity,
                 highlevel=True, behavior=None, attrs=None,
             )
             sumwxy = ak.operations.ak_sum._impl(
                 (x - xmean) * (y - ymean),
                 axis, keepdims, mask_identity,
                 highlevel=True, behavior=None, attrs=None,
             )

コミットメッセージ

fix: ak.meanをJAXで微分可能にする

`ak.mean`をトレースすると、`jax.value_and_grad`は
`RuntimeError: Cannot differentiate through count_zero`をスローします。
根本的な原因は、重みの合計を計算するために内部で使用される`ak.count`に、
JAXの微分ルールがないことです。

`ak.count(x, ...)`を`ak.sum(x * 0 + 1, ...)`に置き換えます。
これにより、同じ数値結果が生成されますが、JAXで完全に
微分可能になります。同じ置換が`ak.var`、`ak.covar`、
`ak.moment`、および`ak.linear_fit`に適用されます。

Fixes #2595.

プルリクエスト​

概要​

ak.meanを呼び出す関数でjax.value_and_gradを使用すると、エラーメッセージがスローされる問題を解決します。(Fixes #2595、#2591のフォローアップ。)

問題​

JAXを使用してak.meanを微分しようとすると、次のエラーで失敗します。

RuntimeError: Cannot differentiate through count_zero

このエラーは、以下を呼び出したときにトリガーされます。

    ak.mean(
        <Array [[...], [...], ..., [...], [...]] type='140 * var * float32'>
    )

weightが指定されていない場合、ak.meanはak.countレデューサーを介して重みの合計を計算します。count（count_zero/count_nonzeroカーネルの上に実装されている）にはJAX微分ルールがないため、ak.mean、および同じパターンに従う他の統計操作を通る勾配は、エラーをトリガーします。

解決策​

この修正により、微分不可能なak.count(x, ...)の呼び出しが、数学的に等価なak.sum(x * 0 + 1, ...)に置き換えられます。

  • x * 0 + 1は、リスト構造と欠損値を保持しながら、xのすべての要素に1の重みをブロードキャストします。
  • ak.sumはJAXの下で微分可能です（x * 0による勾配への寄与は恒等的にゼロです）。
  • 数値結果は、以前のak.countベースの実装と同じです。

重みなしの重み合計にak.countを使用していたすべての統計操作に、同じ修正が適用されます。

  • src/awkward/operations/ak_mean.py
  • src/awkward/operations/ak_var.py（ak.varの上に構築されているak.stdも修正します）
  • src/awkward/operations/ak_covar.py（ak.corrも修正します）
  • src/awkward/operations/ak_moment.py
  • src/awkward/operations/ak_linear_fit.py

テスト​

  • 問題の再現パターンが機能するようになりました。
  arr = ak.Array([[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]], backend="jax")
  jax.value_and_grad(lambda x: ak.mean(x))(arr)
  # (Array(3.5, dtype=float32), <Array [[0.1667, 0.1667, 0.1667], ...]>)

正しい値と、要素あたり1/Nの予想される勾配を返します。

  • デフォルト（NumPy）バックエンドの既存のすべてのテストは、axis=None、axis=-1、および不規則（ragged）な入力を含め、ak.mean、ak.var、ak.std、ak.corr、ak.covar、ak.linear_fit、およびak.momentについて変更なく合格します。

ノート​

  • パブリックAPIは変更されていません。既存の呼び出しシグネチャは保持されます。
  • 新しい依存関係は導入されていません。
  • 置換x * 0 + 1は、実行時にはゼロコストの恒等式であり、測定可能なオーバーヘッドを追加しません。

テスト方法​

  1. GitHubのIssueにある再現コードを実行します。jax.value_and_grad(lambda x: ak.mean(x))(arr)は、RuntimeErrorをスローすることなく、値と勾配を返すはずです。
  2. 完全なテストスイート（pytest tests/）を実行して、ak.mean、ak.var、ak.std、ak.covar、ak.corr、ak.moment、およびak.linear_fitがNumPyバックエンドで同じように動作することを確認します。