Support expand_dims

[asi1024]

🚀 Feature

Support torch.expand_dims.
(cc: @mruberry, @rgommers, @emcastillo and @kmaehashi)

Motivation and Pitch

To enhance the compatibility with NumPy’s interface (c.f. #38349).

Currently PyTorch supports the same functionality with a different function name: torch.unsqueeze, but I would like to support an alias to improve compatibility with NumPy and Python array API standard (#54581).

References

• NumPy numpy.expand_dims
• Python Data API expand_dims(x, axis, /)

cc @mruberry @rgommers @heitorschueroff @pmeier

[rgommers]

Thanks @asi1024. +1 for adding expand_dims as an alias to unsqueeze. It's a nicer name too imho, since unsqueeze isn't actually the inverse of squeeze.

The one question I have is whether it should be a Tensor method too.

[mruberry]

Thanks @asi1024. +1 for adding expand_dims as an alias to unsqueeze. It's a nicer name too imho, since unsqueeze isn't actually the inverse of squeeze.

Adding an alias is straightforward, too. See the note in native_functions.yaml. If unsqueeze doesn't already have an OpInfo I'd first submit a PR adding an OpInfo for it, since OpInfos are how we test aliasing now.

The one question I have is whether it should be a Tensor method too.

Most functions in the torch namespace have both a function and method variant. unsqueeze has a method variant, too, so for consistency I suggest we make this both a function and a method.

[asi1024]

@mruberry

If unsqueeze doesn't already have an OpInfo I'd first submit a PR adding an OpInfo for it, since OpInfos are how we test aliasing now.

unsqueeze doesn't have an OnInfo currently. I can submit a PR for it if I just have to add ShapeFuncInfo to op_db!

[mruberry]

unsqueeze doesn't have an OnInfo currently. I can submit a PR for it if I just have to add ShapeFuncInfo to op_db!

You can just add it as a regular OpInfo ;)