Sergei Lebedev
194884d311
Migrated to mypy 1.14.1 with --allow_redefinition
...
I initially wanted to upgrade to 1.15, but it seems to have a bug in how
ternary expressions are type checked. For example,
def f(x: int) -> str: ...
def g(x: int) -> str: ...
callback = f if ... else g # has type object!
2025-02-13 15:38:28 +00:00
Jake VanderPlas
f652b6ad6a
Set __module__ attribute for objects in jax.numpy
2024-11-15 06:03:54 -08:00
Jake VanderPlas
adf1492843
Add some missing jax.numpy documentation
2024-10-25 13:14:44 -07:00
Jake VanderPlas
a3d6cf007e
First pass at ufunc interfaces for several jax.numpy functions
2024-08-30 11:53:02 -07:00
Peter Hawkins
7f4ef63cd8
Run pyupgrade --py310-plus.
...
Also apply manual fixes to import sorting and unused imports.
2024-06-26 16:10:18 -04:00
rajasekharporeddy
b93da3873b
Fix Typos
2024-06-17 13:55:46 +05:30
Jake VanderPlas
1f9a2dddb8
ufunc: fix implements wrapper for at
2024-04-05 09:42:49 -07:00
jax authors
9b27d43e70
Import submodules from jax._src explicitly, instead of relying on import side-effects. It will lead to the missing x-refs in code search according to go/pywald-sawmill-analysis.
...
PiperOrigin-RevId: 604788105
2024-02-06 15:47:16 -08:00
Jake VanderPlas
43a9faa06a
Rename _wraps to implements
2024-01-24 14:14:19 -08:00
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
Lukas Geiger
52d7f4911c
Prefer expand_dims over reshape
2023-11-16 01:15:48 +00:00
Jake VanderPlas
41a7d66686
jnp.frompyfunc: fix .at() edge case
2023-10-09 11:24:25 -07:00
Jake VanderPlas
7b6a88c9aa
[typing] better annotations for jnp.ufunc
2023-09-21 12:48:14 -07:00
Jake VanderPlas
cb7c7ad942
jnp.ufunc: add fast paths for add/prod reductions
2023-08-28 08:30:23 -07:00
Jake VanderPlas
ac1233b453
jnp.ufunc: support where argument in ufunc.reduce
2023-08-24 15:01:44 -07:00
Jake VanderPlas
61f50bd3b6
jnp.ufunc: minor cleanups & test fixes
2023-08-14 15:19:46 -07:00
Jake VanderPlas
b3a02e1b62
jnp.ufunc: add __hash__ method and jit methods by default
...
This allows the JIT cache to work properly with ufunc methods, because bound
methods are created with a new ID each time.
2023-08-14 13:06:18 -07:00
Jake VanderPlas
ad8e719b82
Add jnp.ufunc and jnp.frompyfunc
2023-08-10 14:58:18 -07:00
