jax authors
11da3df238
Merge pull request #21096 from gspschmid:gschmid/sourcemaps
...
PiperOrigin-RevId: 631769572
2024-05-08 05:44:08 -07:00
Georg Stefan Schmid
b0b322d486
Add sourcemap module to generate TC39-compliant source maps
2024-05-08 01:54:25 -07:00
jax authors
335f27b0b6
Update XLA dependency to use revision
...
c6df436f9e
2024-05-07 19:58:28 -07:00
jax authors
e9e4e5341e
[jax:mosaic-gpu] FragmentedArray can do tiled load.
...
PiperOrigin-RevId: 631611060
2024-05-07 18:13:55 -07:00
Jevin Jiang
79f11d5495
[Pallas] Fix some typos.
...
PiperOrigin-RevId: 631592201
2024-05-07 16:52:38 -07:00
Yash Katariya
395d3cb79e
Bump minimum jaxlib version to 0.4.27
...
xla_extension_version is 261 and mlir_api_version is 56
PiperOrigin-RevId: 631579739
2024-05-07 16:07:59 -07:00
Jieying Luo
5ba56bb075
Recommend the plugin in the CUDA installation instructions.
...
PiperOrigin-RevId: 631555876
2024-05-07 14:47:39 -07:00
jax authors
78e10eea79
Merge pull request #21115 from jakevdp:multi-dot
...
PiperOrigin-RevId: 631545161
2024-05-07 14:14:02 -07:00
jax authors
2b3251ed46
Merge pull request #21092 from jakevdp:dot-doc
...
PiperOrigin-RevId: 631536980
2024-05-07 13:51:13 -07:00
Jake VanderPlas
09810be0cd
Implement jnp.linalg.multi_dot using opt_einsum
2024-05-07 13:40:25 -07:00
jax authors
5f702674f7
Merge pull request #21103 from superbobry:mosaic-gpu-fix
...
PiperOrigin-RevId: 631521771
2024-05-07 13:11:43 -07:00
jax authors
7153738551
Merge pull request #21104 from superbobry:triton-fixes
...
PiperOrigin-RevId: 631521767
2024-05-07 13:06:33 -07:00
jax authors
174405c953
The Bazel version used in JAX is bumped from 6.1.2 to 6.5.0.
...
The update is needed for Windows/Clang builds and for the future hermetic CUDA implementation.
PiperOrigin-RevId: 631519200
2024-05-07 12:58:37 -07:00
jax authors
b6fea3734a
Merge pull request #21111 from jakevdp:fix-changelog
...
PiperOrigin-RevId: 631493089
2024-05-07 11:34:06 -07:00
jax authors
9524188e45
Merge pull request #21110 from jakevdp:upstream-nightly
...
PiperOrigin-RevId: 631490005
2024-05-07 11:25:19 -07:00
Jake VanderPlas
c18851b65d
CHANGELOG: move change from 0.4.27 to 0.4.28
2024-05-07 11:16:11 -07:00
Jake VanderPlas
496795e95f
CI: fix typo in workflow
2024-05-07 11:14:11 -07:00
Yash Katariya
5031a1ddc4
Finish jax and jaxlib 0.4.27 release
...
PiperOrigin-RevId: 631486157
2024-05-07 11:14:09 -07:00
Jake VanderPlas
034d843e8c
jax.numpy: better docs for matmul-like functions
2024-05-07 11:01:54 -07:00
jax authors
3dd8af93b4
Merge pull request #21090 from jakevdp:extract
...
PiperOrigin-RevId: 631480474
2024-05-07 10:59:10 -07:00
jax authors
274e7355da
Merge pull request #21105 from albanie:patch-3
...
PiperOrigin-RevId: 631472758
2024-05-07 10:38:06 -07:00
jax authors
daab7a0329
Handle ellipsis ... in _attempt_rewriting_take_via_slice.
...
Previously `model['some_array'][:,0,0,:]` would generate a `slice`, while `model['some_array'][...,0,0,:]` would generate a `gather`. Now both of these generate `slice` eqns.
PiperOrigin-RevId: 631469837
2024-05-07 10:30:08 -07:00
Samuel
fdd8b61094
Minor doc fix to jaxpr.rst
...
Update the written description of the example `jaxpr` to match the given snippet of code.
2024-05-07 17:36:19 +01:00
Sergei Lebedev
2ca3264052
Updated the Pallas GPU lowering to work with older jaxlib versions
...
Triton changed the signatures of LoadOp and DotOp upstream, and the lowering
code was not ready to handle both old and new signatures.
2024-05-07 17:30:52 +01:00
Sergei Lebedev
8ccbebae4b
Fixed Mosaic GPU build following #21029
2024-05-07 17:08:00 +01:00
Jake VanderPlas
0c26a34df2
Add optional size argument to jnp.compress & jnp.extract.
2024-05-07 08:47:34 -07:00
Jake VanderPlas
9b79f6520a
Remove deprecated kind argument from jnp.sort and jnp.argsort.
...
PiperOrigin-RevId: 631429900
2024-05-07 08:18:59 -07:00
jax authors
500da57e91
Merge pull request #21077 from merrymercy:patch-1
...
PiperOrigin-RevId: 631409738
2024-05-07 07:07:04 -07:00
Yash Katariya
70b4477296
Start jax and jaxlib 0.4.27 release
...
PiperOrigin-RevId: 631409685
2024-05-07 07:01:24 -07:00
Adam Paszke
326adc01a5
[Mosaic GPU] Adjust memref.expand_shape construction to pass in the new args
...
PiperOrigin-RevId: 631404097
2024-05-07 06:36:36 -07:00
jax authors
3e5a18f1e2
Update XLA dependency to use revision
...
873d09720f
2024-05-06 20:37:43 -07:00
jax authors
cb0c49850c
Merge pull request #21081 from hawkinsp:sourcemap
...
PiperOrigin-RevId: 631236806
2024-05-06 17:33:12 -07:00
jax authors
4de346485d
Fix that the insufficient output HBM buffer init would cause the <unk> token generated for quantized int8 model.
...
PiperOrigin-RevId: 631235764
2024-05-06 17:28:13 -07:00
jax authors
eee2783e85
Merge pull request #21070 from shuhand0:rel0.0.7
...
PiperOrigin-RevId: 631218770
2024-05-06 16:22:15 -07:00
jax authors
f6d88525a8
Merge pull request #20327 from selamw1:add_examples
...
PiperOrigin-RevId: 631186425
2024-05-06 14:30:06 -07:00
Shuhan Ding
aac36799fd
fix jaxlib config name
2024-05-06 13:51:22 -07:00
Selam Waktola
9caf59d68b
improve documentation for ix_
2024-05-06 13:43:55 -07:00
jax authors
3d3cb0bd2c
Merge pull request #20842 from Micky774:array-api-default-promotion
...
PiperOrigin-RevId: 631168892
2024-05-06 13:39:56 -07:00
jax authors
1013b1ad85
Merge pull request #21079 from jakevdp:tensorinv
...
PiperOrigin-RevId: 631168828
2024-05-06 13:35:03 -07:00
jax authors
bb6aa12fc1
Merge pull request #21087 from jakevdp:upstream-print-version
...
PiperOrigin-RevId: 631154080
2024-05-06 12:48:29 -07:00
Peter Hawkins
d014f5dc5f
Compute source maps when pretty-printing jaxprs.
...
This change is in preparation for adding support for emitting https://tc39.es/source-map/ source map information for jaxprs, so the relationship between a jaxpr and its Python code can be visualized using tooling for that purpose.
This change adds a new `source_map()` pretty printer document, which causes the pretty-printer to populate a source_map side output during pretty printing.
The change also teaches the core jaxpr pretty printer to populate source map information on each equation.
2024-05-06 15:45:25 -04:00
Jake VanderPlas
7d96e78a55
CI: print numpy/scipy version in upstream job
2024-05-06 11:14:38 -07:00
jax authors
a265e42192
Add an experimental, Clang version of the Windows CI job.
...
Once proven to work, this job will be deleted, and the MSVC job changed to use Clang.
PiperOrigin-RevId: 631122967
2024-05-06 11:10:09 -07:00
Jake VanderPlas
4a363156b9
jnp.linalg tensorinv & tensorsolve: improve implementation & docs
2024-05-06 11:08:36 -07:00
jax authors
d26bd735f6
Merge pull request #21084 from jakevdp:fix-upstream
...
PiperOrigin-RevId: 631112318
2024-05-06 10:49:34 -07:00
jax authors
7e9ef1e4d2
Merge pull request #21078 from jakevdp:numpy-linalg-doc
...
PiperOrigin-RevId: 631112228
2024-05-06 10:44:42 -07:00
jax authors
fb65ba4adf
Add a config for using Clang on Windows.
...
PiperOrigin-RevId: 631112031
2024-05-06 10:39:28 -07:00
jax authors
8ba5c64794
Pass bazel_options directly to the Bazel command, instead of into .bazelrc.
...
PiperOrigin-RevId: 631099970
2024-05-06 10:05:19 -07:00
Jake VanderPlas
40b2d4852e
jnp.linalg: improve API documentation
2024-05-06 09:22:59 -07:00
Jake VanderPlas
6f7ebff585
random_lax_test: fix kstest for newer NumPy
2024-05-06 09:20:07 -07:00
