tree ca90e74eae4dd2db2b1b2b3e56f0210e67dcb7fb
parent 458a429d35776b74cf34411cb3fba1697be38297
author Peter Hawkins <phawkins@google.com> 1634667007 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1634667562 -0700

Improve pytype checking of XLA types inside JAX.

Add an explicit `.pyi` file for jax/_src/lib/__init__.pyi, which works around a bug in pytype where the types of modules that are re-exported becomes `Any`.

[XLA:Python] Fix type declaration for sharding specs.

PiperOrigin-RevId: 404313123
Change-Id: Iadeef1f4a996566f4b6e02f850f4932d4f4265d9