Unimplemented NumPy core functions

[alexbw]

Remaining functions to be implemented:

• np.argpartition (tricky because XLA does not provide any partial sort functionality)
• np.partition (tricky because XLA does not provide any partial sort functionality)

The list above was made by inspecting jnp._NOT_IMPLEMENTED and excluding deprecated functions (such as np.alen, np.ipmt, etc.), functions not relevant to JAX (such as np.setbufsize, np.ascontiguousarray, etc), and functions that modify buffers in-place (np.put, np.place, etc.):

Bugs for high-level categories:

• np.linalg: Implement all of numpy.linalg #1999
• np.fft: Implement all numpy.fft operations #1877
• searching: Implement NumPy search routines #2080
• sorting: Implement NumPy sorting routines #2079
• sets: Implement NumPy set routines #2078
• NaN functions: Implement NaN aggregation functions from NumPy #2077
• Index tricks: Implement np.mgrid, np.ogrid, & np.r_, np.c_ #5850

[mattjj]

The list needs a bit of refinement, since some things were already implemented and other things are out of scope, like to do with loading/saving/string manipulation. I deleted a few of the out-of-scope ones, and marked as done some things that were already included and some things that are coming in #96.

[rainwoodman]

Is there a guide on how to write customized primitive functions (or is that even possible?). I am looking for the equivalent to autograd's defvjp.

A quick glance into the code suggests everything is delegated to XLA -- if an operator is nontrivial then it has to be first implemented in XLA?

[mattjj]

Great question. It's possible, and JAX's internals are actually very similar to Autograd's in this respect, but the API is a bit different. We need to write up how to do this, and maybe add a convenience layer to the API.

There are a few different use cases that we'd want to cover. One is that you just want to define a custom VJP for a function that is otherwise implemented in terms of NumPy code, like in the Autograd tutorial section. But another use case is to define a custom primitive and VJP for some external routine, like a Cython or Fortran function that isn't implemented in terms of NumPy. If you have an intended use case, does it fit into one of those categories, or is it another?

Let's track this in #116. If you have a simplified example of what you want to do, post it there and we'll use it can help guide our convenience API and/or explanation.

[kovasb]

Looking at contributing a function or two from this list.

What is the approach to functions with nonstatically-sized return values? For example, setxor1d.

My understanding is XLA needs to compile for each specific shape. Is there a pattern for avoiding worse-possible-case behavior, or should functions like setxor1d be removed from this list?

[souravsingh]

@mattjj I am looking to take a stab at the issue. How do I start?

[hawkinsp]

@souravsingh Awesome! In general, lots of these are fairly easy to do. Pick one you like off the list, and take a look at previous PRs implementing numpy ops. For example, I just sent out #298 adding np.cumprod/cumsum support. Usually, it's just a question of implementing a numpy op in terms of the primitives in the lax.py library.

[bhushan-borole]

Hello @hawkinsp @jakevdp
I am willing to take up np.insert, np. min_scalar_type.
Is someone else working on it?

[jakevdp]

I don't know of anyone working on these functions, although I think min_scalar_type cannot be made compatible with JIT because of its semantics. Please let us know if you have questions!

[jakevdp]

I updated the list to only include remaining functions that would be good candidates for implementation in JAX.

[avani17101]

np.poly, np.polyfit and np.poly1d have transitioned as mentioned in https://numpy.org/doc/stable/reference/routines.polynomials.html.
Shall we transition it as well?

[jakevdp]

I'm not sure whether we'd want to try to replicate numpy's polynomial interface in JAX. The nice thing about np.poly* functions is that they basically just operate on arrays, which fits JAX's model more readily than does manipulation of custom polynomial objects.

[ntlm1686]

Hi @jakevdp @hawkinsp!
I am interested in implementing np.polydiv. Is someone else working on it?

[jakevdp]

There was a start on it in #7729, but that PR seems to have stalled. Feel free to work on it!

[Gairick52]

@alexbw Hello sir i want to contribute can you point me to some resources

[jakevdp]

Hi @Gairick52 - I removed the "good first issue" label because this issue is nearly complete, and the remaining TODOs are pretty difficult (they'd likely involve updates to XLA)

As far as contributing to JAX, the best approach I think would be to find something that overlaps with your own areas of expertise. What kinds of things do you use JAX for?

[alexbw]

https://youtu.be/xP8tFFJrtXU