Implement np.mgrid, np.ogrid, & np.r_, np.c_

[Conchylicultor]

Numpy has a very convenient np.mgrid which allow to easily create coordinate grids: https://numpy.org/doc/stable/reference/generated/numpy.mgrid.html.

I was expecting jax to have the same, but it seems to be missing:

jnp.mgrid  # raise AttributeError

---

Tracking implementation of these:

• jnp.mgrid Implement jnp.mgrid #6248
• jnp.ogrid Implement jnp.ogrid #6342
• jnp.r_ Implement jnp.r_ and jnp.c_ #6593
• jnp.c_ Implement jnp.r_ and jnp.c_ #6593
• jnp.ix_
• jnp.s_

[jakevdp]

Thanks! Agreed it would be useful to have a jax-native version of this. ogrid and r_ are similarly missing.

[VikasHanumegowda]

I'd love to solve this. Anyone with a good direction suggestion for me to point to, to solve this issue? @jakevdp @Conchylicultor

[Conchylicultor]

To add to the list, np.c_ is missing too: https://numpy.org/doc/stable/reference/generated/numpy.c_.html#numpy.c_

[jakevdp]

The way to implement these would be to create a private class, and override the __getitem__ method to mimic what numpy does. Something like this:

import jax.numpy as jnp

class _Mgrid:
  def __getitem__(self, key):
    if isinstance(key, slice):
      return jnp.arange(key.start or 0, key.stop, key.step or 1)
    else:
      raise NotImplementedError()

mgrid = _Mgrid()
print(mgrid[:5])
# [0 1 2 3 4]

The challenge will be to make certain it handles all the input cases that np.mgrid does.

[Conchylicultor]

You could also look at the numpy implementation to see how they are implementing it

[jakevdp]

You could also look at the numpy implementation to see how they are implementing it

But note that if you copy from numpy's implementation rather than writing your own version, the code should go in https://github.com/google/jax/tree/master/jax/_src/third_party for licensing reasons.

[hawkinsp]

And in this case, where there's probably no need to look at NumPy's implementation to implement the API faithfully and well, it's much better to implement from scratch so we don't have to split up the code based on license. In general:

• please try not to use code from NumPy
• if you do use code from NumPy, it must go in third_party/, and ultimately may need to be rewritten anyway.

[minoring]

Is there an any update for this issue? If not, I would like to progress it :)

[VikasHanumegowda]

I didn't find time for this @minoring
Please go ahead.

[minoring]

Thanks for your response! @VikasHanumegowda

[SondosAhmed]

Hi
I'd like to help with this issue if possible
any update?

[jakevdp]

mgrid has been implemented in #6248, but I don't know of any other work on this

[minoring]

Hi!
I was working on implementing jnp.ogrid.
Maybe you can implement jnp.r_ so that we do not conflict to each other :)

[SondosAhmed]

@jakevdp @minoring Thanks for informing me
I'm gonna try implementing jnp.r_

[arv-77]

Can I take up jnp.s_ if no one else is working on this.

[jakevdp]

I don't know of anyone working on jnp.s_ – feel free to dive in!

[lgeiger]

Looks like this issue can be closed since all function have been implemented 🚀
np.s_ was the last remaining function which was added in #7114.

[jakevdp]

Thanks!