Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

NumPy should warn (eventually raise an error?) on comparisons between complex numbers #15981

Open
shoyer opened this issue Apr 15, 2020 · 14 comments

Comments

@shoyer
Copy link
Member

shoyer commented Apr 15, 2020

Python (version 3, at least) raises an error, but NumPy uses lexicographic order (see #11505 for discussion).

We should clean this up to match Python. This behavior is fine for sorting, but I don't think it makes sense for comparisons between numbers. I think most comparisons between complex numbers are probably latent bugs (as was the case where I discovered this).

My suggestion would be to start by issuing a DeprecationWarning, and then maybe make it a true error at some point in the distance future.

Reproducing code example:

>>> x = 1 + 1j
>>> y = 2 + 1j
>>> x < y
TypeError: '<' not supported between instances of 'complex' and 'complex'
>>> np.array(x) < np.array(y)
True

Numpy/Python version information:

1.18.2 3.6.9 (default, Nov 7 2019, 10:44:02)
[GCC 8.3.0]

@seberg
Copy link
Member

seberg commented Apr 15, 2020

Sounds like a reasonable thing to me to try, pending no bigger push-backs issues (which I do not suspect). A (maybe) similar sort being different from comparisons is that we sort NaN to the end of an array.

@shoyer
Copy link
Member Author

shoyer commented Apr 15, 2020

I think it would also be fair to require setting a special option in np.sort to handle complex values, but I can imagine that might get more push-back.

shoyer added a commit to shoyer/jax that referenced this issue Apr 15, 2020
The tests for CG were failing on TPUs:

- `test_cg_pytree` is fixed by requiring slightly less precision than the
  unit-test default.
- `test_cg_against_scipy` is fixed for complex values in two independent ways:
  1. We don't set both `tol=0` and `atol=0`, which made the termination
     behavior of CG (convergence or NaN) dependent on exactly how XLA handles
     arithmetic with denormals.
  2. We make use of *real valued* inner products inside `cg`, even for complex
     values. It turns that all these inner products are mathematically
     guaranteed to yield a real number anyways, so we can save some flops and
     avoid ill-defined comparisons of complex-values (see
     numpy/numpy#15981) by ignoring the complex part
     of the result from `jnp.vdot`. (Real numbers also happen to have the
     desired rounding behavior for denormals on TPUs, so this on its own would
     also fix these failures.)
@seberg
Copy link
Member

seberg commented Apr 15, 2020

True, although maybe a second step. Deprecating comparison (and min/max) this would also fix gh-15630. In many ways, I guess that is the same issue, since minimum and maximum should stay in sync with both comparisons and clip.

shoyer added a commit to google/jax that referenced this issue Apr 15, 2020
…2717)

* Fixup complex values and tol in tests for jax.scipy.linalg.sparse.cg

The tests for CG were failing on TPUs:

- `test_cg_pytree` is fixed by requiring slightly less precision than the
  unit-test default.
- `test_cg_against_scipy` is fixed for complex values in two independent ways:
  1. We don't set both `tol=0` and `atol=0`, which made the termination
     behavior of CG (convergence or NaN) dependent on exactly how XLA handles
     arithmetic with denormals.
  2. We make use of *real valued* inner products inside `cg`, even for complex
     values. It turns that all these inner products are mathematically
     guaranteed to yield a real number anyways, so we can save some flops and
     avoid ill-defined comparisons of complex-values (see
     numpy/numpy#15981) by ignoring the complex part
     of the result from `jnp.vdot`. (Real numbers also happen to have the
     desired rounding behavior for denormals on TPUs, so this on its own would
     also fix these failures.)

* comment fixup

* fix my comment
@charris
Copy link
Member

charris commented Apr 15, 2020

We do define "sorting order", which may not be what people expect for things like complex. Perhaps we should be clearer in the documentation. It is reasonable to raise errors for comparison operators and such, searchsorted may be a special case. I think the resolution should be driven by use cases, if any. It would also be good to know if changing things breaks anyone's code.

@eric-wieser
Copy link
Member

eric-wieser commented Apr 15, 2020

If we wanted to deprecate sorting complex numbers, the motivation would typically be explained as "the request is ambiguous, refuse the temptation to guess, make the user write what they want explicitly".

In python, we can do exactly that with:

sorted(list_of_complex, key=lambda c: (c.real, c.imag))

Do we have an equivalent spelling in numpy today? At the very least, we need to come up with such a spelling to include in the message of such a deprecation. Preferably, it would be concise.

@eric-wieser
Copy link
Member

eric-wieser commented Apr 15, 2020

Ah, I think that spelling is

np.take_along_axis(c_arr, np.lexsort((c_arr.real, c_arr.imag), axis), axis)

or in the simplest 1d case

c_arr[np.lexsort((c_arr.real, c_arr.imag))]

which isn't too bad.

@seberg
Copy link
Member

seberg commented Apr 15, 2020

Indeed, lexsort is the way to spell it out, and probably fine in most cases, since it is not very common hopefully.
IT could be tedious in code that is generic to float or complex but requires this sort order. But I wonder if there is an actual likely usecase. And the fact that I doubt it, means that we might actually get away with it.

In any case, I think that comparisons and min/max have a higher priority than the sort order change, although if we consider it, doing it at the same time may be best.

@shoyer
Copy link
Member Author

shoyer commented Apr 15, 2020

I could also imagine np.sort(c_arr, key=(c_arr.real, c_arr.imag)) as sugar for np.take_along_axis(c_arr, np.lexsort((c_arr.real, c_arr.imag), axis), axis)

@shoyer
Copy link
Member Author

shoyer commented Apr 15, 2020

Or maybe np.sortby(c_arr, key=(c_arr.real, c_arr.imag))

@eric-wieser
Copy link
Member

eric-wieser commented Apr 15, 2020

Which also would presumably bring about the existence of np.max_by, np.min_by, np.search_sorted_by, etc.

Edit: Let's move this tangent discussion to gh-16004

@Qiyu8
Copy link
Member

Qiyu8 commented Apr 16, 2020

Is np.isclose a special complex comparison case? as @mruberry mentioned, a well-defined and widely accepted notion of saying "a is close to b" even when both a and b are complex, That is what Python's cmath.isclose did.

@mruberry
Copy link

That's true, in PyTorch we do plan to support complex isclose, following NumPy's and cmath's approach. We also plan to error on complex functions that logically depend on comparison operations (clamp, min, max, sort...), and on complex inputs to torch.sign. We do have requests for a complex sign function (x/|x|) and sorting complex numbers by absolute value, however, so we'll likely create new functions (e.g. torch.sgn, torch.sortabs) or add kwargs (e.g. torch.sort(use_abs=True...)) to let our users specify the behavior they'd like.

@seberg
Copy link
Member

seberg commented Apr 17, 2020

We also have the very strange np.sort_complex function already ;). Whatever its oddness, it might actually help a bit with migrating the hopefully very very few existing users of complex sorting. Assuming they know their array is complex.

jacobjinkelly pushed a commit to jacobjinkelly/jax that referenced this issue Apr 21, 2020
…oogle#2717)

* Fixup complex values and tol in tests for jax.scipy.linalg.sparse.cg

The tests for CG were failing on TPUs:

- `test_cg_pytree` is fixed by requiring slightly less precision than the
  unit-test default.
- `test_cg_against_scipy` is fixed for complex values in two independent ways:
  1. We don't set both `tol=0` and `atol=0`, which made the termination
     behavior of CG (convergence or NaN) dependent on exactly how XLA handles
     arithmetic with denormals.
  2. We make use of *real valued* inner products inside `cg`, even for complex
     values. It turns that all these inner products are mathematically
     guaranteed to yield a real number anyways, so we can save some flops and
     avoid ill-defined comparisons of complex-values (see
     numpy/numpy#15981) by ignoring the complex part
     of the result from `jnp.vdot`. (Real numbers also happen to have the
     desired rounding behavior for denormals on TPUs, so this on its own would
     also fix these failures.)

* comment fixup

* fix my comment
gnecula added a commit to google/jax that referenced this issue May 11, 2020
* Implement jax.ops.index_mul. (#2696)

* Implement jax.ops.index_mul.

* Add index_mul to documentation.

* Fix RHS JVP rule for scatter_mul, fix test bug that meant it was not tested.

* Fix typo in docstring.

* Add missing functions to autodoc

* Update XLA. (#2703)

* Fix packbits/unpackbits tests (#2702)

* Make type of value_and_grad slightly more precise. (#2704)

* Fix minor typo in cell (#2692)

* Fix minor typo in cell

One of the arguments to `hvp` wasn't being used, which made the example slightly confusing.

* Fix both definitions of hvp in the autodiff cookbook.

Co-authored-by: Peter Hawkins <phawkins@google.com>

* Update semantics of to_dlpack. (#2707)

to_dlpack now takes ownership of the original buffer, leaving it in an invalid state.

* Add type annotations to optix. (#2687)

* Add type annotations to optix.

* Fix function signature for chain() and remove unused collections import.

* Include Sequence[OptState] as possible output of Init.

* Update np.linalg docs with missing funcitons (#2710)

* Update np.linalg docs with missing funcitons

* Update np.linalg docs with missing funcitons

* Implement numpy fmin() & fmax() (#2711)

* Implement numpy fmin() & fmax()

* Use Tuple[int, ...] rather than Sequence[int] in jnp.ndarray shape annotation.

* Fix some test failures. (#2713)

* Instantiate RNG in testcase instead of test harness in a few more places. (#2706)

* Plumb precision argument into convolution in signal.py (#2715)

* Modify syntax to `x.at[idx].set(y)` and similar.

* Add support for `mul`

* Remove unused textwrap

* Fixup complex values and tol in tests for jax.scipy.linalg.sparse.cg (#2717)

* Fixup complex values and tol in tests for jax.scipy.linalg.sparse.cg

The tests for CG were failing on TPUs:

- `test_cg_pytree` is fixed by requiring slightly less precision than the
  unit-test default.
- `test_cg_against_scipy` is fixed for complex values in two independent ways:
  1. We don't set both `tol=0` and `atol=0`, which made the termination
     behavior of CG (convergence or NaN) dependent on exactly how XLA handles
     arithmetic with denormals.
  2. We make use of *real valued* inner products inside `cg`, even for complex
     values. It turns that all these inner products are mathematically
     guaranteed to yield a real number anyways, so we can save some flops and
     avoid ill-defined comparisons of complex-values (see
     numpy/numpy#15981) by ignoring the complex part
     of the result from `jnp.vdot`. (Real numbers also happen to have the
     desired rounding behavior for denormals on TPUs, so this on its own would
     also fix these failures.)

* comment fixup

* fix my comment

* Explicitly build specific CUDA capabilities. (#2722)

We choose the same set as TensorFlow (minus 3.7, which TF is apparently considering dropping anyway).

This avoids a slow PTX -> SASS compilation on first time startup.

* Add a dynamic type check that the value returned by an XLA translation rule is an XlaOp. (#2723)

Helps give a more understandable error on erroneous translation rules.

* Allow ShardedDeviceArrays to represent arbitrary data shardings. (#2142)

This change introduces ShardingSpec, a struct describing how an array should be sharded. This is integrated into ShardedDeviceArray to allow more flexible sharding. It supports partitioning (both "pmap-style", where an entire axis is decomposed into separate shards and doesn't appear in the on-device shape at all, and "sharded_jit-style", where an axis is chunked into shards but remains in the on-device shape) and replication.

This removes the need for ChunkedDeviceArray, since a ShardedDeviceArray can now represent chunks.

Here are pmap_benchmark times showing that the overall effect of this change neutral to positive (integer indexing is much faster!).

**pmap_shard_args**
```
---------Benchmark summary for pmap_shard_args---------
  nargs    nshards       mean       %std    relative    mean/baseline
-------  ---------  ---------  ---------  ----------  ---------------
     10          8  0.041855    4.15223      1               1.01466
    100          8  0.129884    4.85321      3.1032          0.988543
    101          8  0.136347    6.20233      3.2576          0.967138
    500          8  0.533207    3.6815      12.7394          1.0294
   1000          8  1.10338     0.525193    26.362           0.960435
   5000          8  5.33911     0          127.562           0.963319
    100          2  0.0638619  10.7069       1.52579         1.0362
    100          4  0.0868253   6.76701      2.07443         0.967323
    100          8  0.128151    6.46004      3.06177         0.979742
    100        100  1.22631     1.94885     29.299           1.00371
    100        500  6.60746     0          157.865           0.956657
```
**pmap_shard_outputs**
```
  nouts    nshards        mean       %std    relative    mean/baseline
-------  ---------  ----------  ---------  ----------  ---------------
     10          8   0.0664526   9.49251      1               0.938466
    100          8   0.195711    2.19429      2.94512         1.04239
    500          8   0.82577     0.330864    12.4265          0.994669
   1000          8   1.68323     1.0516      25.3298          0.966915
   5000          8   8.89032     0          133.784           0.998038
    100          2   0.074806   10.1734       1.12571         0.980254
    100          4   0.121334    5.76774      1.82588         1.02033
    100          8   0.185253    5.45068      2.78775         1.01666
    100        100   2.37076     0           35.6759          1.08629
    100        500  17.0832      0          257.074           0.976879
```
**ShardedDeviceArray_indexing**
```
indices_fn                mean     %std    relative    mean/baseline
------------------  ----------  -------  ----------  ---------------
integer_indices      0.0603473  8.29159       1             0.359496
integer_2D_indices  18.0241     0           298.672         1.00583
```

This is how I ran the benchmark:
```
TARGET_TOTAL_SECS=2 CUDA_VISIBLE_DEVICES= XLA_FLAGS=--xla_force_host_platform_device_count=500 python3 benchmarks/pmap_benchmark.py --baseline_dir=<results as of a3cc9a7>
```

* Add type hint to fix pytype error. (#2727)

Without this, pytype (correctly) points out that AbstractValues do not have shape/type information.

* Implement numpy.linalg.multi_dot (#2726)

* Implement numpy.linalg.multi_dot

* Thread precision through multi_dot

* Update XLA. (#2733)

* Temporarily make ShardedDeviceArray.__init__ optionally accept old si… (#2730)

This allows us to incrementally update ShardedDeviceArray creators to the new constructor introduced in 07571ae.

* Fix some bugs in _reshape_sharded_device_array (#2732)

* Fix copy-paste error

It looks as though `_device_put_scalar` should be used here. If not, `_device_put_scalar` should be removed, as it is otherwise unused.

* Release jaxlib 0.1.44. (#2740)

* Add numpy.rint to lax numpy (#2724)

* Add numpy.rint to lax numpy

* Use round_to_nearest_even for numpy.rint

* Add rint to jax.numpy docs

* Fix np.rint float promotion

* Some cleanup and reformatting in `xla.py`.

- Make creation of a few dictionaries more readable.
- Use f-strings where possible.
- Remove unused imports and function parameters.
- Don't format string before passing to `log` function.

* Add a regression test that runs the same computation on all devices that are present. (#2741)

* fix scipy_signal_test convolve failures

* Update names and documentation.

* Fix typo

* Thread precision through np.convolve & np.correlate

* set precision=HIGHEST only for TPU test

* Add FIXMEs for AD type errors

* Add a simple form of partial evaluation for while_loop. (#2497)

The issue that I wanted to fix was that when running grad(while_loop),
the error was a cryptic assertion failure (that all primals are known
after linearization, in ad.py:linearize). I could not figure out
how to detect before that assertion that we are doing a reverse AD
for while_loop. So, I implemented a simple form of partial evaluation,
to allow the primals after linearization to be known, so that the
code proceeds and can then fail gracefully when trying to transpose the
while.

This is not a proper implementation of partial evaluation. The known
outputs are computed early, properly. But the unknown outputs
are computed by a *whole* computation of, including the known
parts.

Fixes issue: #2129

* Assert that reduction computations don't have constants. (#2754)

This case wouldn't work anyway, because there's no good way to pass constants to an XLA reducer.

* Raise an error if stop_gradient is called on non-arrays (#2750)

* Raise an error if stop_gradient is called on non-arrays

* Fix incorrect usage of stop_gradient in solve()

* fix *other* misuse of stop_gradient

* skip two unreliable tests

* Update automatic jaxlib install command to use nvidia-smi instead of nvcc. (#2758)

This is just to get the CUDA version number, and nvidia-smi is more
commonly available.

* Fixes in the FAQ for RST (#2761)

* add adamax

* add adamax test

* Added FAQ entry about relationship between VJP and JVP (#2762)

* Delete cotangent references on their last use (#2719)

* Delete cotangent references on their last use

Current implementation of transposition may add a factor of 2x to
peak memory usage in real cases and _potentially an unbounded factor_
in pathological programs. The reason why this happens is because the
cotangents computed by the `backward_pass` are never evicted from the
environment until the whole transposition is complete. Other systems
(e.g. PyTorch) generally make use of refcounting or liveness analysis
to remove unnecessary references as soon as they are known to no
longer be needed.

A simple example that showcases this issue is this:
```python
def f(x):
  for i in range(1000):
    x = x * 4
  return x

x = np.ones(4)
vjp(f, x)[1](x)
```

Adding `print(len(ct_env))` at the end of `backward_pass` reveals that
the dictionary actually holds a thousand `DeviceArray`s, while both the
forward and backward can be computed in constant memory. Of course this
is the pathological example I mentioned above, but one can easily see
that keeping the cotangents alive for the whole duration of differentiation
causes the memory cost to be approximately `fwd_coefs + all_fwd_intermediates`
instead of `fwd_memory + program_pathwidth` where:
* `fwd_coefs` is the amount of memory necessary to store all constant
  coefficients of the linearized function
* `all_fwd_intermediates` is the amount of memory necessary to
  store _all intermediates appearing in the forward program_.
* `program_pathwidth` is the maximum over amounts of memory necessary
  to store the live values over all transposed program locations

Note that usually we have that
`all_fwd_intermediates > fwd_coefs >> program_pathwidth`
(`>>` meaning that the RHS is usually significantly smaller).

* Import Set

* Use a list instead of a dict

* Type annotation

* Import List

* Fix confusing documentation typo. (#2773)

* Implement np.unique (#2760)

* Implement np.unique

This is an implementation of np.unique
It follows the original numpy implementation of sorting.
While unique it self is intrinsically hard to make compatible
with jit, a helper function has been added which is compatible.
This function could for example be used for jit-compatible
computation of number of unique elements.

* Add test for np.unique

This test tests all possible combinations of inputs for np.unique
with the standard generated array inputs

* Fix return type of the inverse

* Remove complex arrays from np.unique

Since xla can not do size comparisons between complex numbers, and
np.unique depends on np.sort they are removed as possible input.

* Add jit wrap to _unique1d_sorted_mask

* Add pmean to lax documentation (#2778)

* Fix distribution name in docstring (#2764)

* Skip jnp.unique test on GPU (#2780)

Broken due to use of unstable sort (#2779).

* also skip jax.numpy.unique test on tpu

* update version and changelog for pypi

* Bump jaxlib version to 0.1.45 and update WORKSPACE and CHANGELOG. (#2785)

* autodiff cookbook: assume continuous second derivatives

fixes #2772

* attempt to fix changelog formatting bugs

* try optimize=True with einsum

closes #2583

can revert if this ends up problematic for some reason!

* Bump jaxlib version in README to 0.1.45

* apply is_stable=True to sort translation rules (#2789)

fixes #2779

* factor out process_map / post_process_map (#2788)

* factor out process_map / post_process_map

Also fix a bug from reusing post_process_call for pmap. Fixes #2787

* consolidate call_bind / map_bind code

* changelog fixes

* adjust test tolerance for tpu

* bump min jaxlib version (thanks @hawkinsp)

* update travis to match min jaxlib version

* Added clearer error message for tracers in numpy.split (#2508)

* Added clearer error message for tracers in numpy.split

Now we print:

ConcretizationTypeError: Abstract tracer value where concrete value is expected (in
jax.numpy.split argument 1).
Use transformation parameters such as `static_argnums` for `jit` to avoid
tracing input values.
See `https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-where-concrete-value-is-expected-error`.
Encountered value: Traced<ShapedArray>

* Fixed tests, slight change to the error message

* Expanded the FAQ entry about abstract tracers for higher-order primitives

* Added clarification for tracers inside jit of grad

* Updated FAQ language in response to reviews

* attempt to fix failing travis (numerical issues)

* Switch jaxlib Python code to use the lower-level xla.ops API when building XLA ops. (#2798)

Change in preparation for deleting xla_client.ComputationBuilder.

* small typo fix (#2799)

* loosen scipy convolve test tolerance (GPU flaky)

* Add explicit derivative for jax.numpy.linalg.pinv. (#2794)

* Add explicit derivative for jax.numpy.linalg.pinv.

* Fix type confusion problems in the JVP rule for SVD that meant it produced 64-bit tangents for 32-bit primals.

* add tanh rule (#2653)

change expit taylor rule

add manual expit check, check stability of expit and tanh

* use _max instead of max, fix #2795 (#2803)

* use _max instead of max, fix #2795

* revert xla.py python scalar casting='safe' check

* skip pinv test on tpu because no svd

* in custom_jvp/vjp stop_gradient on nondiff_argnums (#2804)

fixes #2784

* disable mypy checks causing new errors

* stop_gradient_p -> ad_util.py, re-enable some mypy (#2806)

* rewrite axis_index implementation, use custom bind (#2807)

* rewrite axis_index implementation, use custom bind

fixes #2716

Co-authored-by: Trevor Cai <tycai@google.com>

* add test for #2716

Co-authored-by: Trevor Cai <tycai@google.com>

* Remove usage of xla_client.{Computation,ComputationBuilder}. (#2808)

* Remove usage of xla_client.{Computation,ComputationBuilder}.

ComputationBuilder is a fairly pointless wrapper class that mimics an outdated version of the the C++ XLA API. It dates back from when we used to have SWIG bindings and needed to write a non-trivial Python shim to keep the interface pleasant to use. Now that we have pybind11-based bindings that are reasonably ergonomic by themselves, we don't need the wrapper class. Instead, we can simply call the pybind11-wrapped C++ API directly, removing the impedance mismatch between the C++ and Python APIs and allowing us to delete the Python ComputationBuilder class.

Similarly we can delete xla_client.Computation for the same reasons; it doesn't do anything useful on top of the C++ API.

* Make ShardedDeviceArray._value threadsafe again. (#2810)

Fixes #2759

* Pin mypy version in .travis.yml. (#2811)

This is recommended in https://mypy.readthedocs.io/en/stable/existing_code.html#continuous-integration, to avoid unexpected upgrades introducing new type errors.

* use static_argnums in xla_computation (#2812)

* use static_argnums in xla_computation

fixes #1017

* add static_argnums to make_jaxpr

* fix type error: handle int case

* implement jet rules by lowering to other primitives (#2816)

merge jet_test

add jet rules

use lax.square

* Feature/permutation (#1568)

* added test for random.permutation

* added permutation that wraps shuffle with behaviour of np.random.permutation

* update docstring

* need to shuffle also the integer range input

* fixed test for permutation with integer

* tweak handling of random.permutation scalar case

* NotImplementedError for random.permutation on >1d

pending resolution to #2066

* address reviewer comments: improve tests

Co-authored-by: Matthew Johnson <mattjj@google.com>

* Fix time issues in odeint reverse mode (#2817)

* Fix time issues in odeint reverse mode

* Add regression test

* add ode test file (#2818)

* add ode test file

* control test tolerances based on precision

* add more ode tests (#2819)

* Remove platform canonicalization from xla_bridge.py (#2815)

* Fix lax.rng_uniform. (#2830)

* only maximally stage out for some call primitives (#2834)

fixes #2833

* handle mapped_invars correctly in more places (#2828)

fixes #2822

We didn't handle `pmap`'s `mapped_invars` correctly in all places in #1959. (I'm actually not sure if #1959 introduced the bug where things were working before, or just refactored it in terms of `mapped_invars`, though my guess is that because the information now contained in `mapped_invars` was implicitly contained in the pmapped jaxpr's `constvars` and `env_vars` that it was working correctly before #1959.) In particular, in #1959 we:
  1. assumed the `mapped_invars` parameter of xla_pmap_p was only populated after partial_eval and set to None otherwise (i.e. staging out for a jit or a control flow primitive),
  2. didn't update it correctly in JVPTrace.process_map (which adds new inputs corresponding to nonzero tangents, and hence `mapped_invars` must be grown),
  3. didn't update it correctly in JaxprTrace.process_map (which adds residual inputs to the staged-out version of the primitive),
  4. didn't forward it correctly in JaxprTrace.process_map anyway (we were setting it to all-true for the staged out eqn for all tracers regardless of what the original `mapped_invars` said),
  5. removed the leading axes of all pvs in JaxprTrace.process_map regardless of whether the corresponding entry of `mapped_invars` was True or False.

The reason we didn't notice 2 and 3 was that they only arise when doing control flow (e.g. scan or remat) of pmap involving closed-over tracers (apparently a rare case), since that's the case where we first form a jaxpr (populating `mapped_invars`) and then later have to apply transformations like AD and further partial eval (thus engaging JVPTrace.process_map and JaxprTrace.process_map with a populated `mapped_invars` parameter). It worked in other cases, e.g. when the pmap was not inside control flow or a remat, because in those cases we left `mapped_invars` set to None, indicating all-true of any length (so it didn't matter if we add inputs).

This commit fixes those issues by
  1. making `mapped_invars` non-optional,
  2. handling `mapped_invars` correctly in
    * JaxprTrace.process_map
    * JVPTrace.process_map
    * ad.map_transpose (since having symbolic-zero cotangents effectively prunes inputs, and having undefined-primal args also prunes inputs)
    * ad._eval_subjaxpr_primals (since having undefined-primal args prunes inputs)
  3. making the separate cases of calls and maps handled more explicitly by adding a new Primitive.map_primitive boolean attribute (analogous to Primitive.call_primitive), to be revised further in #2829.

This is begging for a more coherent cleanup. For example, we reuse the same Primitive class but tag it with `call_primitive` or `map_primitive` (only one of which can be True); we should instead just have a separate Primitive class for these cases and track the type tag with built-in
Python mechanisms. Moreover, when `call_primitive=True` or `map_primitive=True` implies things about what `params` must be present (`call_jaxpr` and `mapped_invars`). I plan to follow up with those cleanups in #2829, but I wanted to get something working first.

* Update deprecated API usages in lapack.pyx. (#2838)

* Replace uses of xla_client.Buffer.from_pyval() with backend.buffer_from_pyval(). (#2839)

Change in preparation for deleting xla_client.Buffer.

* Remove some tests for Jaxlib versions older than the minimum. (#2840)

* Enable some tests that now pass. (#2841)

* Custom derivative for np.linalg.det (#2809)

* Add vjp and jvp rules for jnp.linalg.det

* Add tests for new determinant gradients

* Replace index_update with concatenate in cofactor_solve

This avoids issues with index_update not having a transpose rule, removing one bug in the way of automatically converting the JVP into a VJP (still need to deal with the np.where).

* Changes to cofactor_solve so it can be transposed

This allows a single JVP rule to give both forward and backward derivatives

* Update det grad tests

All tests pass now - however second derivatives still do not work for nonsingular matrices.

* Add explanation to docstring for _cofactor_solve

* Fixed comment

* Fix typo in docstring for _cofactor_solve (#2844)

Found a small typo in the description of _cofactor_solve

* split testDetGradOfSingularMatrix into corank=1,2 (#2845)

* Add precision only arguments (#2850)

* Make precision argument keyword only in jax.numpy

* Fix private functions

* Simplify _odeint_rev (#2832)

* Update jaxpr.rst (#2859)

* Update jaxpr doc

* Make jaxpr.rst doctestable

* check step size is greater than zero (#2857)

loosen tols for grad test

set tol only for float64

* Add pmap_shard_device_array_benchmark. (#2864)

Also renames pmap_shard_args_benchmark to pmap_shard_sharded_device_array_benchmark.

* Fix chi-squared tests in random_test.py (#2847)

As far as I can tell, the previous implementation of the chi-squared test
for samples from discrete probability distributions was broken. It should have
been asserting that the p-value was greater 0.01, e.g., as illustrated here:
http://hamelg.blogspot.com/2015/11/python-for-data-analysis-part-25-chi.html

This hid a few other bugs, such a miscalculation of expected frequencies.

Fortunately, the existing random tests for Bernoulli and Categorical *mostly*
still pass, which the exception of multi-dimensional logits for Categorical.
Those tests are disabled by this PR.

* refactor ode tests, add scipy benchmark (#2824)

* refactor ode tests, add scipy benchmark

remove double import

rename to scipy merge vmap test properly

* clean up more global trace state after errors

Co-authored-by: Matthew Johnson <mattjj@google.com>

* Remove unused `ispure` method (#2781)

* Add population_count primitive to lax (#2753)

* add population_count primitive (needs new jaxlib)

fixes #2263

* Add popcount docs

* Add population_count to lax_reference

* Use int prng (since we're only testing uints)

Co-authored-by: Matthew Johnson <mattjj@google.com>

* Correct the order of .format arguments in vjp wrapper (#2866)

* Change isinstance test in xla_bridge.py to not explicitly name xla_client.Backend. (#2868)

Change in preparation for removing xla_client.Backend in favor of the underlying C++ classes.

* Add top_k jvp and batching rules

* Clarify that `grad` requires arguments to be differentiated to be of inexact type. (#2712)

* Fix definition of qr primitive to return only the upper triangular part of r. (#2870)

Issue #2863.

* Qr complex jvp fix (#2872)

* Fix qr jvp for complex input

* Fix qr jvp for complex64 inputs when jax_enable_x64=True

* Reenable complex jvp test for qr

* Updated README wrt. new features for Stax. (#2862)

* Updated README wrt. new features for Stax.

* fix sort_key_val return type annotation, docstring

* Document how jax.hessian and pytrees interact. (#2705)

* Document how jax.hessian and pytrees interact.

* add spacing to numpy.gradient (#2545)

* Implement nanargmin-max and add tests (#2398)

Co-authored-by: vlad <veryfakemail@ya.ru>

* Make dlpack code robust against upcoming XLA Python binding change. (#2876)

* Fix bug in ShardedDeviceArrayTest.testThreadsafeIndexing (#2875)

* Add nanargmin and nanargmax to documentation. (#2877)

* iterate on jax.hessian docs (#2873)

* iterate on jax.hessian docs

* tweaks

* add back note about block structure

* Add ReLU6, Hard sigmoid, swish (#2709)

* Fix slices in Gated Linear Unit activation (#2341)

* Check for unsupported dtypes and issue a helpful error. (#2885)

* Reset parameter replication default (#2880)

* Reset parameter replication default

* add tests

* Add relu6, hard_swish, and hard_sigmoid to docs. (#2886)

* Fix lax_reference implementation of round() to match lax. (#2894)

lax.round() is documented to round half away from zero, but np.round() rounds to nearest even.

* Make sure gather/scatter indices in lax gradient tests aren't out of bounds. (#2895)

Out-of-bounds gathers are clamped to be in bounds, but out-of-bounds scatters are dropped entirely. This can cause gradient tests to fail because the two operations aren't duals of one another, as the gradient rules expect.

* add jets for sines fns (#2892)

refactor

remove duplicate

* Fix jit with device placement (#2883)

In setups with multiple backends, a jit happens on the default
backend, unless we give a `backend` parameter. This is true
even if the inputs are committed to a device on the non-default
backend, or if we pass a `device` parameter to jit.

* Fix typo in tests; caught on GPU and TPU (#2902)

* err on empty operand in numpy argmin and argmax

fixes #2899

* Remove assert from ShardedDeviceArray staging. (#2908)

This would erroneously fail on Cloud TPU because the TPU client has its own buffer type.

* Update jax version to 0.1.65 (#2909)

* update changelog

* err on empty operand dimension in numpy argmin and argmax

see #2899

* revise xla.device_put device logic (#2907)

* revise xla.device_put device logic, fixes #2905

* remove test of behavior we don't want

Previously, we were testing that for a DeviceArray x, writing
jax.device_put(x) would evaluate to a DeviceArray *on the default
device*. Instead, we should be happy with just returning the same
DeviceArray without any movement.

* Add flag to enable checking, and turn on checking in tests. (#2900)

Fix an error in check_jaxpr.

* Fixed a few places where device sticky-ness was lost. Added FAQ  (#2882)

* Fixed a few places where device sitckyness was lost. Added FAQ for device
placement.

I have also added a new test (multi_device_test.test_computation_follows_data),
written more as part of the documentation. It is shorted than the
old test_computation_follows_data (which is still there, renamed
as test_computation_follows_data_old). I believe there is no
extra coverage in test_computation_follows_data_old w.r.t. all the
other tests we have.

* Fix mypy annotations and updates based on comments

* Undid some changes, will make another PR

* Relax some test tolerances. (#2917)

* Avoid tuple allreduce lowering of psum on TPUs (#2914)

Tuple-shaped allreduces aren't supported in an XLA:TPU optimization pass (see internal bug), but since our use of them on GPU is due to compiler nondeterminism that isn't present on TPU, it should be fine to avoid this bug by disabling tuple psum on TPU.

* Explicitly broadcast values in nn.one_hot and nn.initializers.orthogonal. (#2901)

At head the following fails:

```python
>>> import jax
>>> import jax.numpy as jnp
>>> jax.config.update('jax_numpy_rank_promotion', 'raise')
>>> jax.nn.one_hot(jnp.ones([8]), 512)
...
ValueError: Operands could not be broadcast together for equal on shapes (8, 1) (512,) and with the config option jax_numpy_rank_promotion='raise'. For more information, see https://jax.readthedocs.io/en/latest/rank_promotion_warning.html.
```

* Fix test flakiness in autodiff tests for min/max type functions (#2918)

* Fix test flakiness in autodiff tests for clamp, reduce, and reduce-window.

We change the tests to avoid computing numerical gradients in the neighborhood of nondifferentiable points where, for example, the maximum element in a reduce-max changes. The autodiff approximation is only valid within an epsilon ball around a point, and close to an inflection point the approximation may not be valid.

* Only test reduce-grad-mul for float types.

* Reapply #2017 (Allow shapecheck of PixelCNN++), fixing #2245 (#2800)

* Unrevert "Allow shapecheck of PixelCNN++ (#2017)"

This reverts commit ceab1e3.

* Fix out-of-bound slices (#2245)

* Minor

* Add type annotations

* Fix Poly.__rsub__

* any -> _any

* tweaks, mostly comments/whitespace

* separate polymorphic code path, patch _slice_sizes

* put back some logic for handling Poly sizes

* improve test_slice_indices

* Remove to_index, replace with canonicalize_shape

* Fix slicing with polymorphic start/stop

* Test negative step for polymorphic slicing

* Refactor polymorphic slicing

* Simplify diff

* Fix shapecheck(iota)

Co-authored-by: Matthew Johnson <mattjj@google.com>

* skip failing shapecheck tests

cc @juliuskunze

* Add support for in_axes=None (but not out_axes, or in_axes>0) to pmap (#2896)

* allow in_axes=None for pmap in api.py

* wire in_axes=None through parallel_callable

* add test

* fix error string

* fixes

* fixes

* add test for nested pmap with in_axes

* test pmap still defaults to (implicit) out_axes=0

* replace accidental use of jax.numpy.min w/ builtin

* revert previous change

* Deprecate random.shuffle() and implement random.permutation() for multi-dimensional matrices.

* instantiate zeros (#2924)

fix dtype

remove TODO

* Update XLA. (#2927)

* Update XLA. (#2929)

Includes a fix that may help with issue #2906.

* jax.random.poisson (#2805)

* jax.random.poisson

The implementation for lam < 10 was directly copied from TensorFlow probability:
https://github.com/tensorflow/probability/blob/v0.10.0-rc0/tensorflow_probability/python/internal/backend/numpy/random_generators.py#L155

I adapted the implementation for lam > 10 from TensorFlow:
https://github.com/tensorflow/tensorflow/blob/v2.2.0-rc3/tensorflow/core/kernels/random_poisson_op.cc

The methods themselves match both TensorFlow and NumPy:
https://github.com/numpy/numpy/blob/v1.18.3/numpy/random/src/distributions/distributions.c#L574

* add a check for even larger lambda

* increment iter count

* remove comment that makes no sense

* Fix chi-squared tests in random_test.py

As far as I can tell, the previous implementation of the chi-squared test
for samples from discrete probability distributions was broken. It should have
been asserting that the p-value was greater 0.01, e.g., as illustrated here:
http://hamelg.blogspot.com/2015/11/python-for-data-analysis-part-25-chi.html

This hid a few other bugs, such a miscalculation of expected frequencies.

Fortunately, the existing random tests for Bernoulli and Categorical *mostly*
still pass, which the exception of multi-dimensional logits for Categorical.
Those tests are disabled by this PR.

* Fix accept condition (based on correct chi-squared test)

* Add moment checks for Poisson

* Add batching test, more Poisson rates

* Update XLA. (#2932)

Mention illegal instruction fix in changelog.

* improve docs and error message for odeint *args (#2931)

cf. #2920

* reduce use of lax on static data (e.g. shapes) (#2933)

* reduce use of lax on static data (e.g. shapes)

* use f-string for error message

* support axis argument in nn.glu (#2879)

* support axis argument in nn.glu

* also add basic correctness test

* Update nn_test.py

* Fixed a few more places where device commitment was lost. (#2913)

* trivial jit computations were forcing commitment to the default device
* a device_put with a device specification would not set the commitment
  if the data was already (uncommitted) on the specified device.
* added tests for the above
* once the above were fixed the LaztTest.test_zeros_ones_compilation
  stated to fail because the `sticky` parameter to lazy_force_computation
  was changing. Fixed this by removing stickyness from the compilation key.
* Expanded docstring for jax.device_put; expanded the
  device placement FAQ entry.

* Fix a codeblock in the "understanding jaxpr" doc. (#2942)

This fixes an issue where the codeblock didn't render properly on the website.

* Update XLA to fix build failures. (#2950)

* Allow ConvDimensionNumbers to be passed into conv_transpose (#2915)

* Fix a number of flaky tests. (#2953)

* relax some test tolerances.
* disable 'random' preconditioner in CG test (#2951).
* ensure that scatter and top-k tests don't create ties.

* Fix spurious rank promotion warning. (#2954)

* DOC: add a table of contents for top level API docs (#2946)

This makes them easier to scan.

* DOC: write a new dosctring for jax.numpy.vectorize (#2944)

* DOC: write a new dosctring for jax.numpy.vectorize

This version is customized entirely for JAX.

* review and typo fixes

* Implementation numpy.ediff1d (#2729)

* Implementation of numpy.ediff1d

* Added testing for numpy.ediff1d implementation

* Made ediff1d jit-compatible

* Implemented corrections: style and more testing

* Adapted tests

* changed tests

* modified tests

* Incorporated changes

* Style changes

* Added line between tests

* Changed op_record test

* Add a note about jax.pmap when leading dim is smaller than num devices. (#2949)

* Cache test_utils.format_shape_and_dtype_string. (#2959)

A significant fraction of time when collecting test cases is spent building shape and dtype strings (which are usually similar and usually thrown away.)

* Raise an error in np.var when array is complex and dtype is not (#2288)

Co-authored-by: vlad <veryfakemail@ya.ru>

* add optional 'forward' argument to lax.scan (#2921)

* add optional 'forward' argument to lax.scan

* switch to reverse; revise disable-jit case

* fix jaxpr.rst

* fix loops.py

Co-authored-by: James Bradbury <jekbradbury@gmail.com>

* Improve JAX test PRNG APIs to fix correlations between test cases. (#2957)

* Improve JAX test PRNG APIs to fix correlations between test cases.

In #2863, we observed that we were missing gradient problems because the random test cases being generated were too similar because they were formed with identically seeded PRNGs. This change updates the test_util.rand_...() functions to take an explicit numpy.random.RandomState, and adds a rng() method to JaxTestCase to form a RandomState seeded on the test case name.

This gives the following properties:
* different test cases receive different seeds
* PRNG seeding is deterministic and independent of execution order and sharding.
* PRNG seeding is deterministic across runs.

* Fix some failing tests.

* Fix more test failures.

Simplify ediff1d implementation and make it more permissive when casting.

* Relax test tolerance of laplace CDF test.

* Fix tests for random.categorical with multi-dimensional logits (#2955)

* Expose functools.reduce initializer argument to tree_util.tree_reduce (#2935)

* Expose `functools.reduce` initializer argument to `tree_util.tree_reduce`.

`functools.reduce` takes an optional `initializer` argument (default=None) which is currently not exposed by `tree_reduce'. This can be useful e.g. for computing an L2 penalty, where you would initialize with 0., and then sum the L2 for each parameter.

Example:
```
def l2_sum(total, param):
  return total + jnp.sum(param**2)

tree_reduce(l2_sum, params, 0.)
```

* Only call functools.reduce with initializer when it is not None.

* Change logic to check for number of args to allow None value as initializer

* Rename seq to tree, and add tree_leaves

* Change reduce to functools.reduce.

* Make tree_reduce self-documenting

* Replace jax.tree_leaves with tree_leaves

* Update to use custom sentinel instead of optional position argument

* jax.tree_leaves -> tree_leaves

* Update README for jaxlib 0.1.46 release. (#2968)

* New and improved _shard_device_array function. (#2958)

This gets the performance of sharding DeviceArray arguments to pmap roughly back to what it was prior to 07571ae. It does so by re-introducing a _shard_device_array function that can handle arbitrary array slices.

Benchmark results compared to 87d9590 (i.e. just prior to the regression):
```
---------Benchmark summary for pmap_shard_device_array---------
  nargs    nshards       mean      %std    relative    mean/baseline
-------  ---------  ---------  --------  ----------  ---------------
     10          8  0.0479975  12.0865      1                1.09631
    100          8  0.32916     5.7446      6.85786          1.10263
    500          8  1.5563      2.68041    32.4246           1.10066
    100          2  0.136431    8.33826     2.84245          1.15886
    100          4  0.198815    5.91716     4.1422           1.11409
    100          8  0.31788     4.80559     6.62285          1.06637
```

This still seems a bit slower than it was before, but gets most of the performance back. We can further optimize in future changes if needed.

Fixes #2958 (hopefully)

* Relax test tolerances, suppress warning messages. (#2967)

* Update jax version to 0.1.66 (#2970)

* Replace np -> jnp, onp -> np in tests. (#2969)

* Replace np -> jnp, onp -> np in examples/ (#2971)

For context, see #2370

* Replace np -> jnp, onp -> np in more places. (#2973)

* Replace np -> jnp, onp -> np in more places.

Context: #2370

* Fix typo in random_test.py

* Suppress pytype error (#2974)

pytype gets confused otherwise:
```
File ".../pxla.py", line 244, in _as_slice_indices: bad option in return type [bad-return-type]
           Expected: Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]
  Actually returned: Tuple[Tuple[Union[Tuple[Union[int, slice], ...], slice], ...], tuple, Tuple[int, ...]]
```

* Add IgammaGradA (#2504)

* Update internal aliases to lax_numpy to jnp instead of np. (#2975)

* Update XLA. (#2977)

* Update grad of while_loop message. (#2976)

The previous error message was misleading as of
ed8dbd2
(see #2414 (comment)
for context).

* Add decorator for performing broadcasting inside translation rules (#2468)

* Add decorator for broadcasting at the translation rule layer.

* Fix broadcasting in igamma gradients.

Co-authored-by: Peter Hawkins <phawkins@google.com>

* Fix some bugs in _shards_device_array path. (#2983)

Also adds more comprehensive unit tests.

* Adjust test tolerances for TPU. (#2984)

Ideally this is temporary, as the tolerances are getting high.

* Adjust test tolerances take 2 (#2985)

* Add jnp.unravel_index (#2966)

* Cleanup: move _wraps into jax.numpy._utils. (#2987)

Why? This prevents circular imports within the numpy submodule.

* DOC: add unravel_index to docs/jax.numpy.rst (forgotten in #2966) (#2989)

* Add implementation of np.searchsorted (#2938)

* Add copyright to new file (#2992)

* Fix IntEnum test when checking is enabled. (#2981)

* Added argument check to all primitives. (#2948)

* Added argument check to all primitives.

The issue that inspired this is that `lax.tie_in` is
easy to misuse if the first argument is not a JAX type, then
it silently disappears. This means that `lax.tie_in((x, x), const)`
is the same as `const` even though `x` is a tracer.

This error would be caught previosuly if core.skip_checks == False
because then `bind` checks its arguments. I have essentially
added an unconditional argument check to `bind`.

In case this is considered too inefficient, we can add argument
checking to individual primivites, e.g., tie_in. For most primitives
if a non-JAX array is passed, the `impl` rule would fire and
`numpy` would report the error somehow, perhaps.

* Merged find_top_trace with check_args

* Fix pytype for copybara import (#2995)

* Undo strict checking of LAX primitives (#2996)

This undoes d08dec5d20

* An experiment for id_print implemented with outfeed

* Added print descriptors, support multiple types
* Added a state-passing mechanism to XLA interpreter

* Fixed scan, and grad. Added multiplexing protocol.

* Ensure that we carry state only for control-flow conditionals that use print

* Added masking transformation, added batch_dims to vmap

* Added support for multiple backends to outfeed receiver

Changed the encoding of the header to be uin32

* Added support for sending all arrays in a single message

* Added error handling for tap function errors

* Implemented pytree support for arg and result.

Enabled outfeed for all arrays as a tuple

* Added error checking when starting compiled computations without starting
the outfeed receiver.

* Improved documentation

* Reimplemented the passing of tokens with a Jaxpr transform

* Unified the eager and jit paths

Added error checking for outfeed_receiver not started to primitive computations

* Use a whitelist to limit visibility of exported names in jax.numpy. (#2978)

* Use a whitelist to limit visibility of exported names in jax.numpy.

Prevents unintentional exports of non-public names in the API.

* Undo the id_print/id_tap feature (PR #2791)

Crashes on Travis with the latest 0.1.46. Need to figure out what is going on

* Implement np.bincount (#2986)

* Use a whitelist to restrict visibility in top-level jax namespace. (#2982)

* Use a whitelist to restrict visibility in top-level jax namespace.

The goal of this change is to capture the way the world is (i.e., not break users), and separately we will work on fixing users to avoid accidentally-exported APIs.

* Prepare version 0.1.47 for jaxlib (#3008)

* Remove `jax.np` from the jax namespace (use `jax.numpy` instead). (#3010)

* Implementation of id_tap/id_print using outfeed. (#3006)

This was already merged as #2791 but reverted due to XLA crashes.

This reverts commit 769d703.

* Use a whitelist to clean up exported symbols in the jax.lax namespace. (#3012)

* Add Colab test notebooks for CPU, GPU, and TPU (#3000)

* Support axis_index_groups in allreduce collectives (#2382)

* support replica groups in allreduce collectives

* add test and fix jaxpr in docs

* switch from XLA replica IDs to JAX axis indices

* fix psum transpose rule

* test other nesting order + imperfect nesting

* update jaxpr.rst

* handle None case

* add note+check that groups  cover the index space

* switch split_axis assert to NotImplementedError

* update CHANGELOG

* improve pmap static broadcasted kwarg error msg (#3018)

fixes #3007

* sort supported dtypes in host_callback_test.py (#3020)

* sort supported dtypes in host_callback_test.py

This fixes issues I ran into with running `pytest -n auto
tests/host_callback_test.py` or similar.

* remove unused import

* Fix links in our developer docs (#3019)

The previous versions weren't valid RST.

Ironically, this was in the section with instructions on how to preview
changes to our documentation!

* Modify linspace so that endpoints equal the inputs. (#3016)

* Implement np.digitize (#3003)

* Implement np.nanvar and np.nanstd (#2310)

* Implement nanvar & nanstd

Add tests for nanvar & nanstd

* Clean up bfloat16 tests for np.nanvar and np.nanstd

* add nanvar & nanstd to the whitelist

ignore numpy ddof warnings

* Use a new variable for static_broadcasted_argnums as a tuple. (#3027)

* Use a new variable for static_broadcasted_argnums as a tuple.

This works around a bug in pytype (b/156151503).

* Fix uses of deprecated onp. in pmap_test (#3028)

* Update numpy references to use np. Added to Changelog (#3029)

Co-authored-by: Peter Hawkins <phawkins@google.com>
Co-authored-by: Jake VanderPlas <jakevdp@google.com>
Co-authored-by: Matthew Johnson <mattjj@csail.mit.edu>
Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
Co-authored-by: Lauro Langosco di Langosco <langosco.lauro@gmail.com>
Co-authored-by: John Aslanides <aslanides@users.noreply.github.com>
Co-authored-by: John Aslanides <jaslanides@google.com>
Co-authored-by: Stephan Hoyer <shoyer@google.com>
Co-authored-by: Daniel Johnson <ddjohnson@google.com>
Co-authored-by: Chris Jones <cjfj@google.com>
Co-authored-by: Jamie Townsend <jamiehntownsend@gmail.com>
Co-authored-by: Roy Frostig <frostig@google.com>
Co-authored-by: Adam Paszke <apaszke@google.com>
Co-authored-by: Jacob Kelly <jacob.jin.kelly@gmail.com>
Co-authored-by: Adam Paszke <adam.paszke@gmail.com>
Co-authored-by: Lucas Beyer <lucasb.eyer.be@gmail.com>
Co-authored-by: Oliver Åstrand <oliver.astrand@gmail.com>
Co-authored-by: James Bradbury <jekbradbury@google.com>
Co-authored-by: William C Grisaitis <wgrisaitis@gmail.com>
Co-authored-by: Matthew Johnson <mattjj@google.com>
Co-authored-by: Yufeng <yufengg@users.noreply.github.com>
Co-authored-by: Trevor Cai <tycai@google.com>
Co-authored-by: MichaelMarien <marien.mich@gmail.com>
Co-authored-by: samuela <skainsworth@gmail.com>
Co-authored-by: Jon Malmaud <malmaud@google.com>
Co-authored-by: David Pfau <pfau@google.com>
Co-authored-by: Abhishek Sharma <abhishekshrm53@gmail.com>
Co-authored-by: Jamie Townsend <jamestownsend@google.com>
Co-authored-by: Anselm Levskaya <levskaya@google.com>
Co-authored-by: Anselm Levskaya <levskaya@gmail.com>
Co-authored-by: Paige Bailey <webpaige@google.com>
Co-authored-by: Eduardo Pignatelli <eduardo.pignatelli@burohappold.com>
Co-authored-by: yurodiviy <44850998+yurodiviy@users.noreply.github.com>
Co-authored-by: vlad <veryfakemail@ya.ru>
Co-authored-by: Vaibhav Balloli <balloli.vb@gmail.com>
Co-authored-by: Martin Sotir <martinsotir@gmail.com>
Co-authored-by: Tom Hennigan <tomhennigan@google.com>
Co-authored-by: Julius Kunze <juliuskunze@gmail.com>
Co-authored-by: Roman Ring <inoryy@gmail.com>
Co-authored-by: tamaranorman <tamaranorman@google.com>
Co-authored-by: joschkabraun <47435119+joschkabraun@users.noreply.github.com>
Co-authored-by: James Bradbury <jekbradbury@gmail.com>
Co-authored-by: Joost Bastings <bastings@users.noreply.github.com>
Co-authored-by: Srinivas Vasudevan <srvasude@google.com>
Co-authored-by: notEvil <a_rappold@gmx.at>
Co-authored-by: Matt Wescott <mattwescott@protonmail.com>
@seberg
Copy link
Member

seberg commented Mar 18, 2022

An additional interesting thing is that complex numbers currently compare with the NaN aware "sort order", but complex scalars do not (see gh-21188, where I will add a test for this).

The scalar version will always return False if NaN is involved (which is probably correct even if we define an order). The array/ufunc version anticipates NaNs.

peterjc added a commit to peterjc/ccn that referenced this issue Sep 5, 2022
Python 3 does not allow sorting of complex numbers.
NumPy currently uses lexical sorting, but plans to stop
doing this by default in favour of having to be explicit.
See: numpy/numpy#15981
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

No branches or pull requests

6 participants