43 Commits

Author SHA1 Message Date
Jake VanderPlas
a5737f82af custom prng: remove stackable override for jnp.concatenate 2023-04-24 12:26:58 -07:00
Jake VanderPlas
8f72454bdf Add internal jax.lax.asarray utility 2023-03-30 10:21:55 -07:00
Jake VanderPlas
b308312986 jnp.arange: better validation of inputs 2023-03-14 16:41:58 -07:00
Jake VanderPlas
760deb310e Remove leading underscores in jax._src.numpy.util 2023-03-13 12:18:36 -07:00
Peter Hawkins
a4412e2715 Remove internal ndarray type name. Use Array throughout.
jax.numpy.ndarray remains an exported alias for jax.Array.

PiperOrigin-RevId: 513046188
2023-02-28 14:51:08 -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
Jake VanderPlas
8562e76e94 jax.numpy docstrings: remove empty parameters section 2022-10-28 14:13:29 -07:00
Jake VanderPlas
2f27d516d7 [typing] annotate next part of lax_numpy.py 2022-10-25 12:36:26 -07:00
Jake VanderPlas
48e680c839 CI: avoid raising error when wrapped function is None 2022-10-24 08:57:53 -07:00
Jake VanderPlas
069866e07a Add types to jax/_src/numpy/util.py 2022-10-04 10:07:38 -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
jax authors
d20dcf4b50 Merge pull request #11857 from jakevdp:fix-bool-args
PiperOrigin-RevId: 467259299
2022-08-12 11:45:35 -07:00
Jake VanderPlas
3f06195994 jax.numpy: improve support for boolean inputs 2022-08-12 09:51:25 -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
9769a0accf DOC: ensure that _wraps() generates correct links to wrapped functions 2022-07-21 11:12:35 -07:00
Matthew Johnson
6bb90fde9e [dynamic shapes] revive iree 2022-07-06 15:01:16 -07:00
Jake VanderPlas
e92e23e5f8 Use equality rather than identity when checking for float0
Why? This is required due to changes to dtype canonicalization in numpy v1.23; see #11221
2022-06-23 11:46:20 -07:00
Jake VanderPlas
297a2969a5 [x64] make fft functionality compatible with strict dtype promotion 2022-06-15 10:10:44 -07:00
Jake VanderPlas
e8690f6ba3 [x64] preserve weak types in promote_dtypes_inexact 2022-06-14 09:34:43 -07:00
Jake VanderPlas
5f7cd72130 [x64] use explicit casting rules for promote_dtypes_inexact 2022-05-24 15:51:44 -07:00
Sergei Lebedev
be140981ac ENH _wraps() in jax._src.numpy.util is now returns a generic function
This frees type checkers from the need to explicitly infer the return type
of _wraps() at each call site.
2022-05-23 19:44:11 +01:00
Jake VanderPlas
4f6ea7b296 docs: use intersphinx links for wrapped functions 2022-03-29 14:43:59 -07:00
Jake VanderPlas
ddf23dead3 lax_numpy.py: factor out some common utilities
Re-lands part of #9724

PiperOrigin-RevId: 433838553
2022-03-10 13:35:18 -08:00
jax authors
3766dd2120 Rollback of:
d09d7b8d1363eab1c14051eb2376e605366537f9 by Jake VanderPlas <jakevdp@google.com>:

Factor-out pieces of lax_numpy.py

PiperOrigin-RevId: 431833044
2022-03-01 19:39:31 -08:00
Jake VanderPlas
d09d7b8d13 Factor-out pieces of lax_numpy.py 2022-03-01 09:40:37 -08:00
Jake VanderPlas
22ff25bb8e DOC: add ability to document extra_params within _wraps 2022-02-10 16:54:57 -08:00
Jake VanderPlas
f2222bb1cf CI: error if docstring rewrite fails 2022-02-07 14:43:00 -08:00
Jake VanderPlas
b9b79bab31 maint: update pre-commit package versions & fix new mypy errors 2022-01-31 13:39:11 -08:00
Jake VanderPlas
c4b97b25d2 Fix auto-generated docstrings for JIT-compiled functions 2022-01-24 09:19:51 -08:00
Jake VanderPlas
c473fdee63 docs: tweaks to make docs build on Python 3.8 2021-04-01 15:50:15 -07:00
Jake VanderPlas
d6ccaca62d Make numpy doc parsing more robust 2021-04-01 11:59:10 -07:00
Jake VanderPlas
5c098b11c5 DOC: remove unimplemneted parameters from lax.numpy docstrings 2021-03-25 14:47:18 -07:00
Jake VanderPlas
6ba38774a3 DOC: strip versionadded markup from wrapped docstrings 2021-03-23 16:13:06 -07:00
Jake VanderPlas
1e527d94b7 Don't raise exception for failed docstring parsing 2021-03-22 13:48:00 -07:00
Jake VanderPlas
8a8080c8b8 Rewrite numpy docstring parser & remove extraneous sections 2021-03-05 09:16:41 -08:00
Jake VanderPlas
8cf2f6d27e BUG: fix docstrings for wrapped objects with no qualname 2021-03-04 12:43:33 -08:00
Jake VanderPlas
d8074aca10 DOC: fix update_numpydoc to prevent deleting docstring content 2021-02-22 16:45:40 -08:00
Jake VanderPlas
68af8fddbd Set qualname in jax.numpy.util._wraps 2021-02-16 17:05:29 -08:00
Jake VanderPlas
ba9015a61d Remove signatures from wrapped docstrings 2021-01-28 17:01:20 -08:00
Jake VanderPlas
fe9a3e7024 util._wraps: correctly handle initial line indentation 2021-01-20 13:40:26 -08:00
Jake VanderPlas
e1ff14e56a Dedent wrapped doc sections individually 2021-01-15 15:02:25 -08:00
Jake VanderPlas
3ae68b669e Remove indentation from wrapped docstrings 2021-01-14 16:27:58 -08:00
Peter Hawkins
aa107cf1f4 Move jax.numpy internals into jax._src.numpy. 2020-10-16 20:35:19 -04:00