mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
25,661 Commits 500 Branches 263 Tags
Commit Graph

33 Commits

Author SHA1 Message Date
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
e6fc7f3e87 refactor: move lax_numpy tensor contractions into their own file 2025-02-10 18:56:18 -08:00
Jake VanderPlas
23c1d62910 internal: move more NumPy APIs to ensure_arraylike 2025-01-23 08:48:13 -08:00
Jake VanderPlas
f4f4bf6a19 Fix type annotations for NumPy 2.2 2024-12-11 14:24:58 -08:00
Jake VanderPlas
f652b6ad6a Set __module__ attribute for objects in jax.numpy 2024-11-15 06:03:54 -08:00
Jake VanderPlas
e1f280c843 CI: enable additional ruff formatting checks 2024-10-16 16:09:54 -07:00
Jake VanderPlas
aa551e66c5 Test that jax.numpy docstrings include examples 2024-09-21 07:39:17 -07:00
rajasekharporeddy
61e1d560d8 Improved docs for polynomial arithmetic functions of jax.numpy 2024-07-02 20:04:22 +05:30
Jake VanderPlas
a43994d464 Fix type annotations for jnp.poly* functions 2024-06-24 09:44:50 -07:00
rajasekharporeddy
eba891e3fa Improve docs for jnp.roots and jnp.polyfit 2024-06-24 16:39:55 +05:30
jax authors
348cbba6b2 Merge pull request #21991 from rajasekharporeddy:testbranch4 
PiperOrigin-RevId: 645770273
 2024-06-22 21:28:10 -07:00
rajasekharporeddy
bad1610ac4 Improved docs for jnp.polyint and jnp.polyder 2024-06-22 23:01:23 +05:30
rajasekharporeddy
c5de7bb92e Improve docs for jnp.poly and polyval 2024-06-22 02:49:43 +05:30
Neil Girdhar
1e580457ba Repair various type errors 2024-03-13 15:13:56 -04: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
Peter Hawkins
816ba91263 Use lower-case PEP 585 names for types. 
Issue https://github.com/google/jax/issues/16537

PiperOrigin-RevId: 542969282
 2023-06-23 15:12:14 -07:00
Jake VanderPlas
760deb310e Remove leading underscores in jax._src.numpy.util 2023-03-13 12:18:36 -07:00
Jake VanderPlas
c8c269f5f5 internal: avoid unused imports in lax_numpy 2023-03-08 10:29:04 -08: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
Yash Katariya
a419e1917a Use jax.Array by default for doctests 
PiperOrigin-RevId: 488719467
 2022-11-15 11:52:22 -08:00
Jake VanderPlas
6a348f9666 [typing] add types for jax.numpy.polynomial 2022-10-05 11:23:45 -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
Peter Hawkins
29d03160e3 Remove _ prefix from functions in jax._src.dtypes. 
to_inexact_dtype and to_complex_dtype are used across the JAX code base,
so they shouldn't have _ prefixes.
 2022-08-12 12:51:09 +00:00
Jake VanderPlas
f6476f7a03 jnp.roots: better support for computation under JIT 2022-06-23 14:48:53 -07:00
Jake VanderPlas
15a19969de [x64] make polynomial_test compatible with strict dtype promotion 2022-06-17 16:52:22 -07:00
Jake VanderPlas
30b687c486 [x64] make jnp.poly* functions work under strict dtype promotion 2022-05-31 14:22:49 -07:00
Jiajie Li
128e51c638 Add polydiv to jax.numpy 
Fix code style, fix tests

Add warning when use polydiv with trim_leading_zeros

Update warning for polydiv

Co-authored-by: Jake Vanderplas <jakevdp@gmail.com>

Enable type check in _CompileAndCheck

Fix cutoff

Fix cut-off in polydiv

Add trim_zeros_tol, remove redundant code in polydiv

Remove unused import

Fix trim_zero_tol usage in polydiv
 2022-04-13 18:31:27 +00:00
Jake VanderPlas
603bb3c5ca lax_numpy: move poly functions into numpy.polynomial 2022-03-17 13:28:54 -07:00
Peter Hawkins
4e21922055 Use imports relative to the jax package consistently, rather than .-relative imports. 
This is more consistent, since currently we use a mix of both styles. It may also help pytype yield more accurate types.

PiperOrigin-RevId: 412057514
 2021-11-24 07:48:29 -08:00
Peter Hawkins
a84426cb8f Switch internal users of jax.ops.index_... to use x.at[x].set() APIs. 2021-09-13 19:48:29 -04:00
Jake VanderPlas
8e789c7380 Run doctest on all source files except jax2tf 2021-04-05 10:39:59 -07:00
Peter Hawkins
aa107cf1f4 Move jax.numpy internals into jax._src.numpy. 2020-10-16 20:35:19 -04:00