Jax tracers and adding scalars to arrays

[alexander-held]

Version of Awkward Array

ce63bf2 (main from 2 weeks ago)

Description and code to reproduce

The following currently does not work. I am trying to differentiate through an operation that adds a scalar to an array:

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 correct_jets(jets, alpha):
    jets = ak.Array(jets)
    new_pt = jets["pt"] + 25*alpha
    jets["pt"] = new_pt
    return jets

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(correct_jets, argnums=1)(jets, 0.1)

Result:

TypeError: Jax tracers cannot be used with `ak.from_arraylib`

This error occurred while calling

    numpy.add.__call__(
        <Array [[17.921875, 15.734375], ..., [...]] type='140 * var * float32'>
        JVPTracer-instance
    )

Am I setting this up wrong, or is this a bug?

[alexander-held]

Another thing I've been wondering about in this context: the trace points to an API ak.from_arraylib that I cannot find in the documentation. It lives in the internal ak._layout as ak._layout.from_arraylib, which explains why it is not documented. Is there a way to make this show up with its correct full path? I had a look at trying to understand the way this gets imported and am guessing that the ak.from_jax (which comes in via ak.operations) makes it have this name via

awkward/src/awkward/operations/ak_from_jax.py

Line 5 in 8ba3e30

from awkward._layout import from_arraylib, wrap_layout

but then it is unclear to me why I can't call ak.from_arraylib myself too.

[jpivarski]

Another autodiff issue to self-assign, @Saransh-cpp. Thanks!

[Saransh-cpp]

Linking #2556 (comment) because the error is originating from the exact same place.