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?