Jake Vanderplas

Implement Vega-Lite 4.0 transforms

2

Vega-Lite 4.0 adds a number of additional transforms, including ``density``, ``loess``, ``pivot``, ``quantile``, and ``regression``. These can't be fully supported until Altair 4.0 comes out, but we could start implementing...

Currently the altair object transform only supports simple chart specifications. With a bit of work, we could also support layered, faceted, concatenated, and repeated charts as well.

Add jax.tree_util.register_simple

1

This PR proposes a new API, `jax.tree_util.register_simple` that will simplify pytree registration and promote best practices in the most common of cases. Example usage with `dataclass`, although this will work...

Goal: make JAX support https://data-apis.org/array-api/latest/ Related to #19246 ## TODO 1. Initial Implementation - [x] Add initial implementation in `jax.experimental.array_api` #16099 - [x] Add CI test based on https://github.com/data-apis/array-api-tests #16099...

Part of #21461 Preview: - https://jax--22189.org.readthedocs.build/en/22189/_autosummary/jax.numpy.sort.html - https://jax--22189.org.readthedocs.build/en/22189/_autosummary/jax.numpy.argsort.html

Part of #21461

The reason to avoid this is that overloaded signatures can be very confusing. This is the current documentation page for `jax.numpy.where`: This is the same page built with this PR:...

`jnp.unique` along an axis is very slow for large arrays

7

For example: ```python import jax.numpy as jnp x = jnp.zeros((2, 1000)) jnp.unique(x, axis=0) # takes ~2.5 minutes on a Colab CPU ``` The issue is that XLA has no `unique`...

Tracking issue: inline docstrings

2

For many public APIs, JAX currently uses the `jax._src.numpy.util.implements` helper to define docstrings dynamically at runtime. This has several drawbacks: 1. The docstrings extracted from NumPy do not always precisely...