Support numpy.mgrid

[99991]

>>> import autograd.numpy as np
>>> np.mgrid[:3, :4]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'autograd.numpy' has no attribute 'mgrid'

Support for numpy.mgrid would be nice.

[Karthik-Dulam]

@agriyakhetarpal

I'm interested in contributing if someone hasn't already started working on it.
However since I am not familiar with the project (I'm actively trying to understand how it works) I might take some time.
If that is okay, please point me towards resources or files in the project that would be helpful.

As I understand (correct me if I am wrong), most of the work should lie in figuring out whether mgrid should be a primitive and if it is write a vjp for it, otherwise copy the numpy implementation.

[fjosw]

Hey @99991 and @Karthik-Dulam, autograd actually already supports numpy.msort, but the method was removed with numpy 2, see the release notes:

np.msort has been removed. For a replacement, np.sort(a, axis=0) should be used instead.

The error you reported only occured because autograd was not compatible with numpy 2. In the meantime, we released autograd 1.7 which fully supports numpy 2 and the AttributeError should not occur anymore. Please reopen the issue in case you still experience problems with the new version.

[99991]

@fjosw I do not have permission to reopen this issue. Could you reopen?

This issue is about np.mgrid, not np.msort. The issue still occurs with the latest versions of NumPy and autograd on PyPI:

$ pip freeze | grep -P 'autograd|numpy'
autograd==1.7.0
numpy==2.1.1
$ python3
Python 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import autograd.numpy as np
>>> np.mgrid[:3, :4]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'autograd.numpy' has no attribute 'mgrid'

NumPy 2 still supports np.mgrid:

>>> import numpy
>>> numpy.mgrid[:3, :4]
array([[[0, 0, 0, 0],
        [1, 1, 1, 1],
        [2, 2, 2, 2]],

       [[0, 1, 2, 3],
        [0, 1, 2, 3],
        [0, 1, 2, 3]]])

[fjosw]

Oh, I'm sorry for the confusion. i reopened the issue.

[99991]

No worries. Thanks for reopening.

[fjosw]

One more question: Do you really need np.mgrid to be differentiable? At first glance I would not know how to define the derivative. Would it work for your use-case to create an array via

import numpy as onp
onp.mgrid[:3, :4]

and then use this array as input for function you would like to derive with autograd.numpy?

[99991]

Personally, I do not need the derivative. I just do not want to change my code when switching between numpy and autograd.numpy.

I do not know whether there are applications where a derivative with mgrid is possible or useful.

[fjosw]

Thanks for clarifying. I create a draft PR for adding a few index tricks to autograd.numpy (#639).