Jake Vanderplas

Yeah, #14688 was eventually blocked by the fact that pytype doesn't properly suppot `ParamSpec`. I'm not sure whether that's changed in the meantime.

Some of the errors reflect that the `jit` annotation in this PR is not correct. For example this one: ``` jax/_src/lax/linalg.py:1665: error: Argument 1 to "__call__" of "Wrapped" has incompatible...

I think, setting aside caveats about how hard this might be to land, this is generally a change we want, and one we've been hoping to add for a long...

Just to be clear, the `pytype` blocker is not about revealing existing errors, it's about the fact that pytype fails loudly and completely when it sees `covariant=True`. If that hasn't...

Yes, if this didn't use covariant typevars, it would be easier to merge. But my understanding from #14688 was that covariant typevars are required in order to correctly annotate `jit`.

Pulling in to run internal pytype tests

> I'm curious: If this ends up getting merged, will the Google-ML-Automation bot include my github username in the final commit? Or would someone on the inside need to add...

Thanks for the report! I believe this is working as intended: `astype` in general forces a copy when converting between different dtypes. It sounds like what you're looking for is...

Thanks for the report! Here's a ref to the T4 architecture spec: https://images.nvidia.com/aem-dam/Solutions/design-visualization/technologies/turing-architecture/NVIDIA-Turing-Architecture-Whitepaper.pdf T4 doesn't support bfloat16, but JAX (via the XLA GPU compiler) should be falling back to float32....

The window transform documentation you linked to is auto-generated from the Vega-Lite schema, so there's no way to change it within Altair itself. You might submit a PR to Vega-Lite...