Peter Hawkins
fb2824bdbb
[JAX] Add static_argnames support to jax.jit.
...
Requires a new jaxlib build.
Add support for static_argnames in C++ JIT implementation.
PiperOrigin-RevId: 367627359
2021-04-09 07:11:04 -07:00
Peter Hawkins
9fad2441a2
Mark arguments to jax.jit() other than the function as keyword-only.
...
This change is to prevent breakage when options are added or removed.
2021-04-08 10:32:35 -04:00
Skye Wanderman-Milne
f8f373466c
Update README, etc. for jaxlib 0.1.65 release
2021-04-07 17:51:20 -07:00
George Necula
dce31e9631
[jax2tf] Fix handling of float0
2021-04-07 13:48:39 +03:00
Skye Wanderman-Milne
7b42011f7c
Update jax version and changelog
2021-04-01 10:11:52 -07:00
Jake VanderPlas
2a091d2629
Update changelog for #5868
2021-04-01 09:29:22 -07:00
Peter Hawkins
3fc1fdb148
Add a JVP rule for the general case of lax.reduce.
2021-03-30 17:31:47 -04:00
Jake VanderPlas
c11e725ecb
X32 mode: raise OverflowError for large integers
2021-03-30 10:05:03 -07:00
George Necula
d323ad0f2b
[host_callback] Add support for tapping empty arrays
...
We make sure that both the inputs and the outputs of
callbacks can contain empty arrays.
Most platforms do not support empty infeed, so we ensure
we do not send those.
2021-03-30 10:48:58 +03:00
Jake VanderPlas
9790232556
Python integer conversion: always return int64 or OverflowError
2021-03-29 09:26:19 -07:00
Jake VanderPlas
40dac9425c
pre-release omnistaging cleanup
2021-03-25 16:44:58 -07:00
Skye Wanderman-Milne
b68a08adf1
Add programmatic profiling APIs, and rename some existing APIs.
...
This change provides aliases for the renamed APIs so existing code
won't break. We should remove these aliases after the next release.
2021-03-25 11:50:55 -07:00
Peter Hawkins
7052a87ab6
Increase minimum jaxlib version to 0.1.64.
2021-03-24 13:49:48 -04:00
Matthew Johnson
c4a099093c
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
Skye Wanderman-Milne
0cbe2c1c05
Update README, etc. for jaxlib 0.1.64 release
2021-03-18 16:11:40 -07:00
Skye Wanderman-Milne
757247b791
Update README, etc. for jaxlib 0.1.63 release
2021-03-17 10:14:52 -07:00
Peter Hawkins
328930b917
Increase minimum jaxlib version to 0.1.62.
2021-03-16 15:11:36 -04:00
Skye Wanderman-Milne
f06bb9a7f4
Update jaxlib version etc.
2021-03-09 17:55:40 -08:00
Roy Frostig
9c420653c3
move changelog to top level
2021-03-08 10:44:52 -08:00
Jake VanderPlas
94484d85aa
Migrate CHANGELOG.rst -> CHANGELOG.md
2021-02-12 17:03:53 -08:00
George Necula
89514f9278
Moved CHANGELOG to docs ( #2252 )
...
* Moved CHANGELOG to docs
This puts the documentation also on RTD, with TOC.
Also changed its format to .rst, for consistency.
Added GitHub links to the change log.
* Actually add the CHANGELOG.rst
* Added reminder comments to the CHANGELOG.rst
2020-02-23 19:18:06 +01:00
Alexander Botev
43ee917511
Adding broadcast_argnums to pmap for allowing similar behaviour t… ( #1786 )
...
* Adding `static_argnums` to `pmap` for similar behaviour to `static_argnums` of `jit`.
* Removed check for ShardedDeviceArray
* Final clean up and rename.
2020-02-14 07:45:26 -08:00
Matthew Johnson
9e6fe64a66
bump version and update changelog for pypi
2020-02-11 07:22:17 -08:00
George Necula
20f9230f6e
Simplify Jaxpr: remove the bound_subjaxpr field, all subjaxprs are in params.
...
The goal is to make the Jaxpr language more uniform: all higher-order
primitives carry sub-Jaxprs that are part of the parameters, and they
are all called xxx_jaxpr. As a side-effect, some code is simplified
(e.g., the code that searches for sub-jaxprs).
For now the code assumes that all the `call` (final-style) primitives
carry exactly one subjaxpr with the parameter name `call_jaxpr`. These
primitives are still processed differently in the internal code, but
there is no reason any external consumer of a Jaxpr needs to know this.
2020-02-11 10:06:08 +01:00
George Necula
b18a4d8583
Disabled tests known to fail on Mac, and optionally slow tests.
...
Issue: #2166
Added JAX_SKIP_SLOW_TESTS environment variable to skip tests known
to be slow.
2020-02-05 18:02:56 +01:00
George Necula
272620e66c
Added note to CHANGELOG.md
2020-02-04 10:22:54 +01:00
Peter Hawkins
91cd20b173
Update documentation and changelog to mention DLPack and array interface support. ( #2134 )
2020-01-31 11:15:04 -05:00
Roy Frostig
664a4e123d
VJP of cond, via partial eval + transpose ( #2091 )
...
VJP (grad) of lax.cond, via partial eval + transpose
Co-authored-by: Matthew Johnson <mattjj@google.com>
2020-01-30 15:03:00 -08:00
Peter Hawkins
991324f8df
Increase minimum jaxlib version to 0.1.38. ( #2120 )
2020-01-29 14:16:58 -05:00
Peter Hawkins
0fe601227f
Update README.md and CHANGELOG.md.
2020-01-28 10:00:32 -05:00
Skye Wanderman-Milne
19fb494adb
Add jax changelog ( #2022 )
2020-01-17 11:34:59 -08:00
