Calculating element-to-element differences and sums in Nx

Many years ago, I used a system called Yorick for doing array calculations. One feature that was really useful was a series of functions for taking differences and sums within a vector. Ref: Yorick User Manual

The simplest operation to explain is “difference”; when applied to a dimension, it produces a result that’s one element shorter and has successive differences.

[0, 1, 4, 9, 16, 25] (length 6)

has differences

[1, 3, 5, 7, 9] (length 5)

This can be really useful for calculus-y things like “given a list of x values and a list of corresponding y values, estimate the slope as a function of x”.

The almost-inverse of “difference” is “cumulative sum”; when applied to a dimension, it produces a result that’s one element longer and has a running total:

[1, 3, 5, 7, 9] (length 5)

has cumulative sums

[0, 1, 4, 9, 16, 25] (length 6)

It’s an “almost” inverse because the first element of the cumulative sum is always zero. If this reminds you of “plus C” in integrals, you’re not wrong :stuck_out_tongue:

My root question: what’s the best way to write these kinds of calculations in Nx?

It’s certainly possible to produce tensors (via Nx.make_diagonal etc) that can be multiplied by the input vector to produce these results, but it seems inefficient given the simple nature of the calculations.

1 Like

There is Nx.diff and Nx.cumulative_sum for what you want

As an aside, I suggest you take a look into some of the source code for Scholar, as it uses quite a few different Nx functions for calculus stuff