ENH: create_diagonal: support broadcasting

[NeilGirdhar]

Fixes #136

[NeilGirdhar]

I'm not entirely sure what I should do about the numpy dependency. Do I re-implement ndindex somewhere?

Also, I'm getting a strange error when trying to run the tests about a bad device.

[lucascolley]

yes, we will want to remove the NumPy dependency. I don't know how easy it is to make ndindex agnostic.

Also, I'm getting a strange error when trying to run the tests about a bad device.

np.ndindex will be producing arrays with the default NumPy device, so you need to be careful to coerce to the correct device when that is used (things are probably going wrong in the at call). But this problem should be subsumed in removing the NumPy dependency.

[NeilGirdhar]

Ready for your review 😄

[lucascolley]

@mdhaber has ndindex come up in any of your work on vectorising things?

[NeilGirdhar]

What do you mean? I just need it for this create_diagonal.

[lucascolley]

I'm just wondering whether this is an instance of a problem that has appeared elsewhere

[NeilGirdhar]

Whenever I need ndindex, I use the one from Numpy. I've never been reluctant to import Numpy, so I would not use a version from xpx.

[lucascolley]

Sure, I was just asking Matt if anything like this had come up in the context of array-agnostic code yet for him

[NeilGirdhar]

Oh sorry! I didn't see the tag. I haven't been able to sleep tonight and I'm just tired.

[lucascolley]

no worries!

[mdhaber]

Yes, I used it in the decorator that adds batch support in linalg. For the decorator used in stats, I considered using ndindex before going with apply_along_axis, but if I were to rewrite it I'd just use ndindex. Both of these are just for NumPy right now, though.

I've thought about having a heuristic in marray that decides whether to loop over slices and compress (remove masked elements from) each slice before performing the operation. This can be faster than the usual thing (say, data[mask] = 0 in sum) when there are fewer slices and they are very big. But for that, I assumed I would reshape to have just one batch dimension at the front, loop over that single axis, and then reshape the result. But by "reshape" I really mean moveaxis + reshape, and ndindex would save the trouble of reshape.

[lucascolley]

Makes sense, thanks! In this case, does the approach of making ndindex agnostic seem good to you?

[NeilGirdhar]

@lucascolley Would you like me to do anything else with this PR?

[lucascolley]

I just need to find time to review this. @j-bowhay would you like to take a look?

[NeilGirdhar]

Sounds good. I added a lot more tests.

[NeilGirdhar]

(I renamed the test functions to make it a bit clearer why we have each test.)

[j-bowhay]

LGTM modulo the one query for @lucascolley

[NeilGirdhar]

@j-bowhay Thank you for much for you thorough and speedy review. We'll wait on Lucas's decision about the term to use.

[lucascolley]

looks like there is a merge conflict

[NeilGirdhar]

Should be fixed now

[lucascolley]

thanks for this @NeilGirdhar, LGTM! And thanks for the review @j-bowhay !

I will merge this after putting out a micro release.

[NeilGirdhar]

I will merge this after putting out a micro release.

Awesome thanks! Just curious, why not merge then release?

[lucascolley]

While it is very unlikely that anyone is actually relying on the behaviour before this PR, it's technically a change to an existing API, so I'm (conservatively) following https://jacobtomlinson.dev/effver/. There have been some new features and bug fixes which I would like to be available in a 0.6.x version.

Again, this doesn't really matter in this situation, and it may be that something else I've merged already is more likely to break things—it's more just trying to apply a principled way of doing things.

[lucascolley]

thanks again @NeilGirdhar !