issues
search
google
/
jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
28.02k
stars
2.57k
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
add error for remat + rbg for #21046
#21049
mattjj
opened
1 hour ago
0
Better documentation for several jax.numpy functions
#21048
jakevdp
opened
1 hour ago
0
Make linkcode_resolve() a bit more robust
#21047
shoyer
opened
4 hours ago
0
buggy interaction: remat, automatic partitioning, and unsafe `rbg`-based RNGs
#21046
froystig
opened
4 hours ago
0
jax.random seems to have unnecessary buffer allocations on stack
#21045
py4
closed
4 hours ago
5
[omnitracing] partially un-regress dispatch time
#21044
mattjj
closed
6 hours ago
0
Remove remaining top-level scipy imports
#21043
jakevdp
closed
6 hours ago
0
DOC: Improve remaining jax.scipy docstrings
#21042
jakevdp
closed
6 hours ago
0
[Mosaic GPU] Prepare for writing warp-specialized kernels
#21040
copybara-service[bot]
opened
14 hours ago
0
DOC: Fix typos and math rendering in `jax.random` docs
#21038
rajasekharporeddy
closed
10 hours ago
0
Made has_side_effect= parameter of mlir.emit_python_callback keyword-only
#21037
copybara-service[bot]
closed
9 hours ago
0
[export] Add backwards compatibility test for Pallas call on GPUs.
#21036
copybara-service[bot]
closed
17 hours ago
0
jet_test: float32 matmul precision
#21035
olupton
closed
11 hours ago
0
Add support for TPU delay in Mosaic
#21034
copybara-service[bot]
opened
22 hours ago
1
fix vmap-grad-shmap bug with spmd_axis_name
#21032
mattjj
opened
1 day ago
0
spsolve exits with error when inverting matrix sum
#21031
buvoli
opened
1 day ago
0
[Pallas TPU] Increase clarity of dot 2D shape enforcement error.
#21030
copybara-service[bot]
closed
11 hours ago
1
Bundle MLIR .pyi files with jaxlib
#21029
superbobry
closed
1 day ago
0
Add sharding on output avals for shard_map.
#21028
copybara-service[bot]
opened
1 day ago
0
Allow multiple indexers when doing discharge or swap in pallas
#21027
copybara-service[bot]
closed
1 day ago
1
Fix cuda array interface with old jaxlib.
#21026
copybara-service[bot]
closed
1 day ago
0
dlpack.callback now supports mutable_results=
#21025
copybara-service[bot]
opened
1 day ago
0
Do not require a capsule to have a specific name in the CUDA plugin
#21024
copybara-service[bot]
closed
1 day ago
0
reintroduce the Threefry GPU kernel lowering, under a flag
#21023
copybara-service[bot]
closed
1 day ago
0
Add `sharding` on `ShapedArray` and in the type system of jaxpr.
#21022
copybara-service[bot]
opened
2 days ago
0
Use self in _ProfileState.reset
#21021
lukebaumann
closed
2 days ago
0
make _ProfileState's reset point to self.
#21020
sadikneipp
opened
2 days ago
0
Finalize deprecation of lax.linalg positional args
#21019
copybara-service[bot]
closed
2 days ago
0
DOC: Improve docstrings for jax.scipy.linalg
#21018
jakevdp
closed
1 day ago
0
[omnitracing] pop frames, see if anything downstream breaks
#21017
mattjj
closed
2 days ago
0
skip pallas/gmm_test.py if we don't have hypothesis
#21016
mattjj
closed
2 days ago
0
psum_scatter does not allow scatter_dimension to be negative
#21015
jewillco
closed
1 day ago
3
[XLA:Mosaic] Support retiling from (8, 128, -2) to (8, 128) for 32-bit data.
#21014
copybara-service[bot]
opened
2 days ago
0
Refactor `array_api` namespace, relying more directly on `jax.numpy`
#21013
Micky774
closed
10 hours ago
1
bump shard count for `random_lax_test`
#21012
copybara-service[bot]
closed
2 days ago
0
Return to the `win-2019` image for the `windows_ci` workflow.
#21009
copybara-service[bot]
closed
2 days ago
1
Don't create temp directory when module is getting imported.
#21008
copybara-service[bot]
opened
2 days ago
1
Adds rewrite patterns for `arith` and `math` operations with `bf16` operands/results that are not supported by the underlying hardware.
#21007
copybara-service[bot]
opened
2 days ago
1
Try fixing the MSVC path for `windows_ci` workflow once more.
#21006
copybara-service[bot]
closed
2 days ago
1
Point the `windows_ci` workflow to the correct VC directory.
#21005
copybara-service[bot]
closed
2 days ago
1
Added a new callback API which uses DLPack tensors to keep the data on the device
#21003
copybara-service[bot]
opened
2 days ago
0
jnp.fft.ifft imprecision for GPU
#21002
chih-kang-huang
opened
2 days ago
0
Fix warnings in CI from compilation_cache_test.
#21001
hawkinsp
closed
2 days ago
0
Use the `windows-2022` image for running the `windows_ci` workflow.
#21000
copybara-service[bot]
closed
2 days ago
1
[Mosaic GPU] Fix the diagnostic dump infrastructure
#20999
copybara-service[bot]
closed
2 days ago
0
Ported threefry2x32 for GPU to the typed XLA FFI
#20997
copybara-service[bot]
opened
2 days ago
0
⚠️ Nightly upstream-dev CI failed ⚠️
#20996
github-actions[bot]
opened
2 days ago
0
[Mosaic GPU] Only call kernel initializer from inside a custom call
#20995
copybara-service[bot]
closed
2 days ago
0
Allow replacing jaxpr `debug_info` with `None`.
#20994
copybara-service[bot]
closed
2 days ago
0
[XLA:CPU] Enable constant host offloading
#20993
copybara-service[bot]
closed
2 days ago
1
Next