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.61k
stars
2.82k
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
[ROCm] [ROCm] fix jax and wheelhouse relative paths
#25128
Ruturaj4
closed
12 hours ago
1
Add a check for return codes of `executor.run` so that we propagate error codes correctly
#25127
copybara-service[bot]
closed
2 days ago
0
Test new approach.
#25126
copybara-service[bot]
opened
2 days ago
0
[Mosaic GPU] Restrict WGMMA tests to sm90a
#25125
andportnoy
opened
2 days ago
0
[FFI Lowering] Compiler-only Attrs
#25124
yliu120
opened
2 days ago
3
[Lowering] Stable IR
#25123
yliu120
opened
2 days ago
1
Set option to not eagerly load dialects.
#25122
copybara-service[bot]
opened
2 days ago
0
Fix a weird interaction with `set_local` and empty tuples passed to it.
#25121
copybara-service[bot]
closed
2 days ago
0
[mgpu] Debug print for mlir vectors.
#25120
copybara-service[bot]
closed
1 day ago
0
[mgpu pallas] Layout iota operation.
#25119
copybara-service[bot]
closed
1 day ago
0
Update XLA dependency to use revision
#25118
copybara-service[bot]
closed
2 days ago
0
Use next after tiny as smallest floating point value in complex functions accuracy tests on Mac ARM.
#25117
pearu
opened
2 days ago
1
Add float8_e8m0fnu type support
#25116
wenscarl
opened
2 days ago
0
Install git before actions/checkout
#25115
copybara-service[bot]
closed
2 days ago
0
Add exec_time_optimization_effort and memory_fitting_effort flags - fixes #24715
#25114
jedborovik
opened
3 days ago
1
[AutoPGLE] Fix pgle test after removing pjit cache.
#25113
copybara-service[bot]
closed
2 days ago
0
Skip testAutoPgle tests.
#25112
copybara-service[bot]
opened
3 days ago
0
NaN's produced by jax.numpy.linalg.pinv not produced by numpy
#25111
samskiter
opened
3 days ago
1
[Pallas][Mosaic GPU] Add support for compressing squeezed dims in async_copy + grid fixes
#25110
copybara-service[bot]
closed
22 hours ago
0
Difference between numpy and jax.numpy in advanced indexing axes order
#25109
fzimmermann89
opened
3 days ago
0
[Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction
#25108
copybara-service[bot]
opened
3 days ago
0
Track mapping of platform aliases to compile-only backends
#25107
jaro-sevcik
opened
3 days ago
1
[Mosaic:TPU] Enable broadcast from 1-D vectors
#25106
copybara-service[bot]
closed
2 days ago
0
[Mosaic:TPU] Add relayout for adding minor implicit dim and relax some offset restrictions on similar shape cast
#25105
copybara-service[bot]
opened
3 days ago
0
Add lax.composite primitive
#25104
copybara-service[bot]
opened
3 days ago
0
Add sharding argument to reshape since it also takes a `shape` argument for the output shape
#25103
copybara-service[bot]
closed
3 days ago
0
[JAX] Add Python binding for building a colocated Python program
#25102
copybara-service[bot]
closed
2 days ago
0
Weird defjvp behavior when finding grad of a scalar that depends on the primal
#25101
JadM133
opened
3 days ago
0
Create a `null_mesh_context` internal context manager to handle null contexts properly.
#25100
copybara-service[bot]
closed
3 days ago
0
[Pallas TPU] Add vector support to `pl.debug_print`
#25099
copybara-service[bot]
opened
3 days ago
0
[mosaic_gpu] Fixed unbounded recursion in `FragmentedArray._pointwise`
#25098
copybara-service[bot]
closed
3 days ago
0
Start a new TPU interpret mode for Pallas
#25097
jburnim
opened
3 days ago
0
Update ROCm scripts to match new build.py usage
#25096
nitins17
closed
1 day ago
0
[Mosaic GPU] Support batch dimensions in FA3 MGPU kernel.
#25095
copybara-service[bot]
closed
3 days ago
0
Remove _pjit_lower_cached cache. We can simplify the caching of jit as we have downstream caches and a cpp cache too.
#25094
copybara-service[bot]
closed
3 days ago
0
[SDY] Enable `test_pjit_array_multi_input_multi_output` since Shardy conflict resolution is now complete.
#25093
copybara-service[bot]
closed
3 days ago
0
Reverts e72b449089f6af4ceb18288e36215b3c76e69245
#25092
copybara-service[bot]
opened
3 days ago
0
[cuda] Bump nvidia-cuda-nvcc-cu12 dependency to 12.6.85
#25091
gspschmid
closed
12 hours ago
2
[SDY] remove TODO for enabling Layouts for Shardy post cl/697715276.
#25090
copybara-service[bot]
closed
3 days ago
0
Bump hypothesis from 6.102.4 to 6.119.4
#25089
dependabot[bot]
opened
3 days ago
0
[SDY] enable `cpu_shardy` for JAX shard_alike test.
#25088
copybara-service[bot]
closed
3 days ago
0
Update the render documentation job to use the new self-hosted runners
#25087
copybara-service[bot]
closed
2 days ago
3
[mgpu] FragentedArray.foreach() can now optionally return a new array
#25086
copybara-service[bot]
closed
1 day ago
0
Add abstract mesh context manager to trace_context in the fallback path too (which will be deleted after jax 0.4.36 release)
#25085
copybara-service[bot]
closed
3 days ago
0
[ROCm] Add ROCm version metadata to pip show output for `jax-rocm60-plugin` and `jax-rocm60-pjrt`
#25084
Ruturaj4
closed
3 days ago
2
Automated Code Change
#25083
copybara-service[bot]
closed
3 days ago
0
Document cudaMallocAsync
#25082
nouiz
opened
4 days ago
0
[Mosaic GPU] Simplify logic for pointwise splat operands
#25081
copybara-service[bot]
closed
3 days ago
0
Use xla_extension_version instead of jaxlib_version
#25079
copybara-service[bot]
closed
3 days ago
0
[pallas_mgpu] Map indices primitive
#25078
copybara-service[bot]
opened
4 days ago
0
Previous
Next