Unimplemented NumPy core functions

[alexbw]

Remaining functions to be implemented:

• [x] np.argpartition (tricky because XLA does not provide any partial sort functionality)
• [x] 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:

• [x] np.linalg: https://github.com/google/jax/issues/1999
• [x] np.fft: https://github.com/google/jax/issues/1877
• [x] searching: https://github.com/google/jax/issues/2080
• [x] sorting: https://github.com/google/jax/issues/2079
• [x] sets: https://github.com/google/jax/issues/2078
• [x] NaN functions: https://github.com/google/jax/issues/2077
• [x] Index tricks: https://github.com/google/jax/issues/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 https://github.com/google/jax/pull/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.

[navneet-nmk]

@mattjj I am looking to implement np.unique. Is this something which is a priority at the moment?

[mattjj]

Sounds like a great idea! The core team isn't working on any numpy operations at the moment (we're mostly focused on improving performance in the XLA runtime, parallel computation, and Cloud TPU support) so that's a great place to make contributions.

np.unique is tricky though because AIUI the shape of the output depends on the values of the input. Since XLA can only express computations with static shapes (meaning independent of the input values), we'd have to re-compile the XLA computation for every new value of the input, and moreover it might be tricky to express in terms of XLA HLO (i.e. in terms of lax library calls).

@hawkinsp any thoughts on implementing np.unique?

[navneet-nmk]

@mattjj since this would be my first hand at this, I should probably look at numpy functions that are less tricky compared to unique?

Any suggestions would be really helpful.

[navneet-nmk]

@mattjj Apologies for the constant posts. I have decided to implement np.diff which seems fairly easier. I was just wondering whether there is a guide to developing the testcases in the lax_numpy_test.py file. Thanks!

[murphyk]

+1 to adding np.cov

[murphyk]

+1 to np.percentile

[mattjj]

@murphyk np.cov coming in #983

[hawkinsp]

@murphyk quantile and percentile are now present, at least for interpolation='linear' and floating-point types.

[adarob]

fyi, np.rank is deprecated in favor of np.ndim

[jessebett]

Anyone interested in np.convolve I've commented on possibly using lax.conv. It would require reshaping and flipping the inputs. I've described it in #1561.

[michiboo]

@jessebett I just created a PR for it #1831

[fedden]

I'd also like to get involved in this process of completing the remaining numpy ops. As a first pass, I am looking at implementing something really simple, to dip my toes in the water so I can be more familiar with the codebase before tackling something more serious!

I will look at np.alen unless there are any issues with that :)

[fedden]

Created a PR for alen and isscalar in #1924

[shoyer]

I took another pass at deleting irrelevant functions:

• anything currently deprecated by numpy
• anything that relies upon details of NumPy's internal array storage (strides, etc)
• anything related to np.matrix.
• anything that modifies arguments in place

I think most of the rest should be relevant, though some will be tricky (e.g., with output dependent shapes).

[StephenHogg]

Hi - I'm interested in getting nanmean done. Am I right in suggesting that this could be done by somehow piggybacking off nansum? If so, are there any tips that may be worth heeding in this regard?

[shoyer]

NumPy should be a pretty good model for most NaN functions, all of its NaN functions are written in pure Python.

[shoyer]

To make things a little easier to keep track of, I split some functions off into a handful of sub-issues:

• np.linalg: #1999
• np.fft: #1877
• searching: #2080
• sorting: #2079
• sets: #2078
• NaN functions: #2077

[oliverastrand]

Hi, should np.asanyarray be removed since it simply becomes np.asarray when there is no subclassing?

[jameskirkpatrick]

np.cast is not implemented in Jax. this is a dict of lambdas that allows casting to certain types.

[hawkinsp]

@jameskirkpatrick I agree that np.cast exists in my copy of NumPy, but it is not a documented API as far as I can tell. Are we sure it's an intentional, non-deprecated NumPy API? @shoyer any thoughts?

[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