ENH: add kron and expand_dims

[lucascolley]

x-ref data-apis/array-api#837

[lucascolley]

@mdhaber I've taken your implementations from scipy/scipy#20077 (comment), made some modifications for edge-cases, and added type annotations as well as xp-compatible tests and docs based on NumPy's. Any feedback appreciated if you'd like to take a look!

[lucascolley]

@izaid

I've used kron. It's valuable. I'd have to double-check if it can be implemented just with array operations (probably can), and if it can I think it should be.

Indeed it can!

[mdhaber]

I was just zombie-translating from NumPy to array API with that, so I'm afraid I don't have any insight here.

[izaid]

Is this expand_dims for a tuple of axes?!?!?!? You have just made my day.

I still think this should be in array API properly, but happy for it to live here until / if that happens.

[lucascolley]

Is this expand_dims for a tuple of axes?!?!?!? You have just made my day.

yes 😁

It basically just sorts the tuple (my fix) and applies xp.expand_dims for each item (Matt's idea).

[izaid]

It basically just sorts the tuple (my fix) and applies xp.expand_dims for each item (Matt's idea).

Fantastic!!! I specifically asked for this change to the array API (going from supporting a single axis to a tuple of axes), but it hasn't happened yet. See: data-apis/array-api#760

I use such a thing literally all the time.

Edit: Also, some of the existing libraries (NumPy and CuPy, but not PyTorch) do support expand_dims with a tuple of axes already. Maybe you want to dispatch directly to them in those cases?

[lucascolley]

aha, reading @asmeurer's comments in that issue I see the concern for ambiguity now. This implementation chooses to always add the same number of dimensions as the length of the tuple, and resolves negatives by taking the positions modulo the number of dimensions in the output. I should check that this matches what NumPy does, and add a note to the docs.

Edit: Also, some of the existing libraries (NumPy and CuPy, but not PyTorch) do support expand_dims with a tuple of axes already. Maybe you want to dispatch directly to them in those cases?

My intention with this package is to target just the standard - it isn't concerned with which real-world implementations exist (other than array-api-strict for testing).

It's quite simple to write small wrappers which delegate to known array libraries around array-agnostic functions like this, but I would rather keep that separate for now. Perhaps a separate package could do that and this will just be a backend for that, or perhaps that will take over this package eventually. I would rather that consumer libraries deal with which array libraries they would like to delegate to in their own code for now.

[izaid]

aha, reading @asmeurer's comments in that issue I see the concern for ambiguity now. This implementation chooses to always add the same number of dimensions as the length of the tuple, and resolves negatives by taking the positions modulo the number of dimensions in the output. I should check that this matches what NumPy does, and add a note to the docs.

Edit: Also, some of the existing libraries (NumPy and CuPy, but not PyTorch) do support expand_dims with a tuple of axes already. Maybe you want to dispatch directly to them in those cases?

My intention with this package is to target just the standard - it isn't concerned with which real-world implementations exist (other than array-api-strict for testing).

It's quite simple to write small wrappers which delegate to known array libraries around array-agnostic functions like this, but I would rather keep that separate for now. Perhaps a separate package could do that and this will just be a backend for that, or perhaps that will take over this package eventually. I would rather that consumer libraries deal with which array libraries they would like to delegate to in their own code for now.

Yeah, I get the idea for keeping it simple! Makes sense, just wanted to point it out if you weren't aware. I'm just happy to not have to maintain an expand_dims anymore.

[lucascolley]

This PR is ready from my side. I'll leave it open for a while for any comments.

I've documented the resolution to the ambiguity in data-apis/array-api#760, which is to follow NumPy in disallowing a tuple containing both positive and negative indices referring to the same position in the expanded shape.

[lucascolley]

kron

Implementing a new function that is not in the standard.

• Signatures match between NumPy, PyTorch and JAX (PyTorch has an out= kwarg too).
• No particular design choices here, just following the existing implementations.
• No limitations to address as far as I am aware.

---

expand_dims

Adding the ability for axis to be a tuple of ints on top of the standard.

• NumPy and JAX support providing axis as a tuple. PyTorch (which calls it unsqueeze) does not.
• Design choices:
  • @asmeurer and @Micky774 discussed some ambiguities relating to simultaneously providing negative and positive indices which (could) refer to the same position in the expanded shape in RFC: add support for a tuple of axes in expand_dims array-api#760
    • To resolve the first point from RFC: add support for a tuple of axes in expand_dims array-api#760 (comment), I independently came up with the same argument mentioned there that len(axes) dims should be added
    • To resolve the second point there, I follow NumPy and JAX in taking the axes modulo the number of dimensions in the output, and raising an error if duplicates arise from that (or otherwise exist)
• One may hope that the apparent limitation described in RFC: add support for a tuple of axes in expand_dims array-api#760 (comment) could be solved, but as @asmeurer showed with the included proof there, that isn't possible in some instances.

---

@kgryte good to proceed?

[lucascolley]

Gentle bump here! If I don't hear back, I'll probably merge this soon and circle back before a new release is cut.