L-BFGS loss_gradient: is there real headroom in the Array API helpers on NumPy inputs? #34797

Your instinct is right, but the reason is more specific than "profiler overhead", and it changes what a useful benchmark would look like.

The short-circuit is real. get_namespace bails out before touching anything array-API related:

array_api_dispatch = get_config()["array_api_dispatch"]
if not array_api_dispatch:
    if xp is not None:
        return xp, False
    else:
        return np_compat, False

So with dispatch off there is no namespace inspection, no _remove_non_arrays, nothing. What is left is that first line, and that is where the time actually goes:

    # Return a copy of the threadlocal configuration so that users will
    # not be able to modify the configuration with the returned dict.
    return _get_threadlocal_config().copy()

Every call copies the whole config dict — nine keys today — reads one boolean out of it, and throws the copy away. So what you profiled is not Array API dispatch cost at all; it is a dict copy per helper call, wearing an Array API name in the trace.

That also explains the shape of what you saw. Counting the helpers reachable from one loss_gradient: get_namespace(X, y, sample_weight) at the top, get_namespace_and_device(X) inside weight_intercept_raw, get_namespace(coef) inside weight_intercept, plus one or two move_to calls depending on the branch. Call it four or five config copies per iteration, so a few hundred across a whole L-BFGS fit. Microseconds in total, and — the important part — completely independent of n_samples. On iris the BLAS work is tiny so that fixed cost is visible; on realistic X it is unchanged in absolute terms while everything around it grows, so it disappears. Both of your observations are correct simultaneously.

Which is why I would not write the asv benchmark you were considering. A LogisticRegression fit on realistic data will report "negligible" by construction, because you already know the numerator is fixed and the denominator scales. It confirms something you have proven analytically and tells nobody anything new. If you want a number, the honest quantity is the per-call cost of get_namespace itself, which is a microbenchmark, not an estimator fit.

And if it did turn out worth fixing, I would not do it the way you sketched. Special-casing move_to and get_namespace_and_device inside loss_gradient buys one call site and reintroduces exactly the local divergence you said you wanted to avoid. The copy lives in get_config, so reading the flag without copying — inside get_namespace, once — removes it for every caller in the library with no parallel NumPy branch anywhere. That is a much smaller diff than anything touching the dtype and Fortran-order handling in _linear_loss.py.

Whether maintainers want to bypass that deliberate defensive copy on a hot path is their call, and it is worth asking as a question about get_namespace rather than about LogisticRegression — the linear-model framing is what made this look like a solver problem when it is really a config-access one.

This sounds like a good candidate for an asv benchmark rather than an optimization based on the iris profile. Benchmark small, medium, and realistic dimensions, float32 and float64, dispatch enabled and disabled, and a few L-BFGS iterations. Measure full fit as well as isolated loss_gradient calls, because a helper can look large in inclusive time on tiny data while being negligible end to end. If the NumPy and dispatch-off path is consistently measurable, an early return that preserves existing matmul and dtype logic seems safer than duplicating the algorithm. Use benchmark evidence before changing this hot path.