Neil Girdhar

Feature request: Allow unpacking of type variable bound to typed dict

4

I'd like to have a family of classes where overrides have additional parameters to various methods: ```python from typing import Generic, Unpack from typing_extensions import TypedDict, TypeVar, override class EmptyDict(TypedDict,...

Consider adding a "TopLevel" keyword that stands in for the top level bound of a corresponding type variable

6

Summary --- Please consider a new typing keyword that stands in for the top level type of a corresponding type variable. Motivation --- Consider a test [here](https://github.com/NeilGirdhar/efax/blob/118b3e1f39dd0cdd99f4fe19b13c666b27a4236f/tests/test_flatten.py) and the corresponding...

Add PartialApplication

4

Pitch --- Add [`typing.PartialApplication`](https://en.wikipedia.org/wiki/Partial_application) to facilitate the implementation of: - `__get__`, and - `functools.partial`, both of which are practically impossible to natively (without plugins) annotate. The `__get__` method is currently...

Provide guidance on PEP 698's strict enforcement feature of typing.override

5

PEP 698 (`typing.override`) suggests that type-checkers provide a strict enforcement option that would mark method overrides. However, no guidance is provided about the methods `__init__` and `__new__`. These methods are...

Proposal: Programmatically create types

18

This proposal is aimed at solving two related problem. First, when defining a multi-ary operator on Numpy arrays, e.g., [leaky integral](https://github.com/NeilGirdhar/tjax/blob/0c1b13629a197f38b73f431926cf536d8c682315/tjax/_src/leaky_integral.py), then you ideally want to bake in Numpy type...

[typing] Use ParamSpec in JIT annotation

33

This pull request would be a huge improvement for Jax users who use type checkers like MyPy or Pyright, which now support `ParamSpec`. Consider: ```python from jax import Array, jit...

Fix annotations

8

This pull request probably needs some reviewer guidance. It fixes all (but one) of the MyPy errors after the jit-annotation [PR](https://github.com/google/jax/pull/14688). I will comment on various changes in line. The...

Add .editorconfig file

1

This file helps non-Googler contributors who have their editors set up differently for different projects. Most modern editors will automatically use the Google style when they see this file: *...

Bad interaction between nnx.Rngs and custom derivatives

6

``` File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback return fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^ File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/custom_derivatives.py", line 613, in __call__ in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/custom_derivatives.py", line 613, in...

Remove unused imports

2

Not a big deal, but the unused imports are noisy :)