issues
search
jax-ml
/
jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.44k
stars
2.8k
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
Pattern match dot algorithm spec to preset name
#24820
dfm
opened
13 hours ago
0
[shape_poly] Remove caching for the symbolic shape evaluator
#24818
gnecula
opened
20 hours ago
0
[Mosaic GPU] Ensure that lowering `InitializeBarrierOp` preserves the result's type.
#24817
copybara-service[bot]
opened
1 day ago
0
[Mosaic GPU] Add `base_pointer` argument to `InitializeBarrierOp`.
#24816
copybara-service[bot]
opened
1 day ago
0
Make GPU work with copy=True and device_put since same device pinned_host -> pinned_host copy is possible.
#24815
copybara-service[bot]
closed
1 day ago
0
jnp.bincount: support boolean inputs
#24814
jakevdp
opened
1 day ago
0
bincount rejects bool
#24813
carlosgmartin
opened
1 day ago
0
Skip test_jnp_einsum_grad_y_pallas on gpu due to ooms
#24812
copybara-service[bot]
opened
1 day ago
0
Adds coverage for spmd-axisname-filtering in shard_map transpose.
#24811
copybara-service[bot]
opened
1 day ago
0
Specialize ufunc.reduce for monoidal binary ufuncs.
#24810
jakevdp
opened
1 day ago
0
[Pallas] Add a cost estimator for Pallas/JAX functions.
#24809
copybara-service[bot]
opened
1 day ago
0
Remove unused import
#24808
jakevdp
closed
1 day ago
0
Division by self not always "1.0" on JAX GPU, but consistently gives "1.0" on JAX CPU.
#24807
mattlevine22
opened
1 day ago
1
Internal: create decorators for defining ufuncs
#24806
jakevdp
opened
1 day ago
0
[Mosaic GPU] Add CUPTI profiler alongside events-based implementation
#24805
andportnoy
opened
1 day ago
0
Remove unused import in shard_alike_test.py
#24804
copybara-service[bot]
closed
1 day ago
0
Put the set of current spmd axis names in the axis env instead of spelunking
#24803
copybara-service[bot]
closed
1 day ago
0
Add tests for jnp.einsum in Pallas
#24802
copybara-service[bot]
closed
1 day ago
0
[MOSAIC:GPU] Add `async_load`, `async_store`, and supporting attributes to the MLIR Mosaic GPU Dialect.
#24801
copybara-service[bot]
closed
1 day ago
0
Allow more output storage types for some dot algorithms.
#24800
copybara-service[bot]
opened
1 day ago
0
pre-commit run --all failed for ruff and mypy
#24799
apivovarov
closed
1 day ago
1
Flaky test tests/fft_test.py::FftTest::testFftfreq5
#24798
apivovarov
opened
1 day ago
2
[AutoPGLE] Explicitly ignore host callback pointers
#24797
copybara-service[bot]
opened
1 day ago
0
Check mem_stats key in test_lax_full_like_efficient
#24796
apivovarov
opened
1 day ago
0
Sparse reshape throws error when `n_dense>0` and some target dimension has size 1
#24795
cherrywoods
opened
1 day ago
1
`jax.lax.DotAlgorithm` generates an unnecessary conversion on H100
#24794
andportnoy
opened
1 day ago
3
Fix parallel pgle-tests execution.
#24793
copybara-service[bot]
opened
1 day ago
0
Update some outdated syntax in FFI tutorial
#24792
dfm
opened
1 day ago
0
[Mosaic GPU] Make sure to free the cloned MLIR module when debugging
#24791
copybara-service[bot]
closed
1 day ago
0
Add new CI scripts for running Pytests
#24790
copybara-service[bot]
opened
1 day ago
0
Disable lax_test on ARM in Google's internal CI.
#24789
copybara-service[bot]
closed
1 day ago
0
Add commands to run Bazel GPU (non-RBE) jobs
#24788
copybara-service[bot]
opened
1 day ago
0
Nightly builds on Mac ARM fail complex function numerical tests
#24787
hawkinsp
opened
1 day ago
1
Don't perform size 0 slices into scipy rotations.
#24786
hawkinsp
closed
1 day ago
1
Add new CI scripts for building JAX artifacts
#24785
copybara-service[bot]
opened
1 day ago
0
Disable the paged_attention test on TPU v5p.
#24784
copybara-service[bot]
closed
1 day ago
0
Disable for_loop_test on TPU v5p.
#24783
copybara-service[bot]
opened
1 day ago
0
[pallas] Allow user to pass 64-bit indices to `pl.{load,store,...}`.
#24782
copybara-service[bot]
opened
1 day ago
0
[Mosaic TPU] Support 1D concat: set implicit_dim to kSecondMinor to treat 1D (N,) as (1, N) and then tile it as (1, 128)
#24781
copybara-service[bot]
opened
1 day ago
0
[mgpu] Broadcast the LHS fragmented array if it is splat.
#24780
copybara-service[bot]
opened
1 day ago
0
Added suppressions for pytype --none-is-not-bool.
#24779
copybara-service[bot]
opened
1 day ago
0
(follow-up #23852) add missing `typename` keyword to work with `gcc`
#24778
cainmagi
closed
1 day ago
1
[Pallas] Increase test coverage of pl.dot.
#24777
copybara-service[bot]
opened
1 day ago
0
Automated Code Change
#24776
copybara-service[bot]
closed
1 day ago
0
Add typing overloads for jax.extend.ffi.ffi_call() to aid type checkers
#24775
copybara-service[bot]
closed
1 day ago
0
don't warn on unused `type: ignore`
#24774
mattjj
closed
1 day ago
0
Fix pre-commit to run on all files in CI
#24773
dfm
closed
1 day ago
0
Allow 64-bit output types from ffi_call regardless of enable_x64 flag
#24772
dfm
opened
2 days ago
0
Adding support for copy_p primitive to jet.
#24771
copybara-service[bot]
closed
2 days ago
0
jax.device_get: handle generic extended dtypes
#24770
jakevdp
opened
2 days ago
0
Next