Faster `jnp.trapezoid` when `dx` is a scalar

[jurasic-pf]

TLDR

There's a faster and just as accurate way to compute jnp.trapezoid when dx is scalar.
Is there any reason against specializing in that case?

Motivation

For uniform x (when x=None and dx is passed), the cumulative trapezoid rule boils down to a sum + endpoint handling. This can be implemented more efficiently than the general trapezoid rule.
jnp.trapezoid is used in lots of scientific code, and this change translates to measureable improvements in overall execution time of simulations on both CPU and GPU (in one of our examples 1.42s -> 1.18s), in micro benchmarks of just jnp.trapezoid it can be up to 5x faster.

Current:

0.5 * (dx_array * (y_arr[..., 1:] + y_arr[..., :-1])).sum(-1)

Suggested:

if dx_array.ndim == 0:
    dx_array * (y_arr.sum(-1) - 0.5 * (y_arr[..., 0] + y_arr[..., -1]))

Please:

• Check for duplicate requests.
• Describe your goal, and if possible provide a code snippet with a motivating example.

[jurasic-pf]

Microbenchmark before:
jnp.trapezoid have the same performance when using x or dx. For comparison jnp.sum * dx is significantly faster.

Microbenchmark after:
jnp.trapezoid is faster when using dx and in line with the performance of jnp.sum * dx.

https://github.com/jurasic-pf/jax/tree/jurasic/faster-trapezoid Implements the change, if interested I'll open a PR.

[jakevdp]

Hi - thanks so much for raising this issue! I think computing via the more direct sum makes sense when dx is a scalar.

I wonder if it's worth also generalizing this in the N-dimensional case where dx is length 1 along the specified axis?

Would you be interested in submitting a pull request with this change?

[jurasic-pf]

Thanks! I'll give it a try.

[rosspeili]

hey @jurasic-pf how did this go?

[jurasic-pf]

Had a few busy days, so I only now addressed the review comments. I think it's ready?