Add .at[].subtract

[carlosgmartin]

Feature request: Add a subtract method to jax.numpy.ndarray.at.

[superbobry]

This sounds reasonable, even though you can do .add(-x) for signed types.

@jakevdp wdyt?

[jakevdp]

I think this is reasonable. We could do it the easy way, with something like this:

def subtract(self, x):
  return self.add(-x)

but as @superbobry points out, this would fail for unsigned types.

The full solution is a bit more work: it would involve defining a lax.scatter_sub primitive with appropriate autodiff rules, and dispatching to that. The implementation would look very similar to scatter_add:

jax/jax/_src/lax/slicing.py

Line 382 in 3c25da2

def scatter_add(

The only really involved pieces there would be the scatter_sub_jvp and scatter_sub_transpose rules, but those could probably be implemented in terms of the existing scatter_add_jvp and scatter_add_transpose.