Sergei Lebedev
f936613b06
Upgrade remaining sources to Python 3.9
...
This PR is a follow up to #18881 .
The changes were generated by adding
from __future__ import annotations
to the files which did not already have them and running
pyupgrade --py39-plus --keep-percent-format {jax,tests,jaxlib,examples,benchmarks}/**/*.py
2023-12-13 10:29:45 +00:00
Rebecca Chen
e144f71c33
Fix or ignore some pytype errors.
...
PiperOrigin-RevId: 574633158
2023-10-18 16:21:59 -07:00
Jake VanderPlas
2c64a0ac2c
typing: add some type assertions to typing_test
2023-10-13 08:30:08 -07:00
Peter Hawkins
dea7450e4e
Remove references to jax.config.jax_array, which is always True at head.
...
PiperOrigin-RevId: 516970232
2023-03-15 17:09:11 -07:00
Roy Frostig
cb8dcce2fe
migrate more internal dependencies from jax.core
to jax._src.core
...
PiperOrigin-RevId: 509736368
2023-02-14 23:01:11 -08:00
Jake VanderPlas
f09fd8a4e9
[x64] minor test-only updates for better type safety
2022-11-30 15:18:40 -08:00
Yash Katariya
9e4114f0f1
Move array.py
and sharding.py
from experimental/
to _src/
.
...
PiperOrigin-RevId: 477201711
2022-09-27 10:06:52 -07:00
Yash Katariya
cbf34cb609
Rename the concrete class Array
to ArrayImpl
...
PiperOrigin-RevId: 477017236
2022-09-26 16:18:30 -07:00
Jake VanderPlas
0cb233eec9
Add initial jax.Array base class for instance checks & annotation
2022-09-26 07:48:43 -07:00
Peter Hawkins
ba557d5e1b
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
...
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.
PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Jake VanderPlas
5829c6ae9d
Change case of typing.Dtype -> typing.DType
...
This follows the convention used in numpy.typing.DType.
2022-09-14 15:03:55 -07:00
Jake VanderPlas
13a7034e6a
Internal change
...
PiperOrigin-RevId: 474331907
2022-09-14 10:39:38 -07:00
Jake VanderPlas
b3c31ebe7d
Add typing_test.py
2022-09-13 12:43:51 -07:00