jax-ml jax issues - Githubissues

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.59k stars 2.82k forks source link

issues

Newest

Newest Most commented Recently updated Oldest Least commented Least recently updated

Add lax.composite primitive

#25104 copybara-service[bot] opened 14 minutes ago
0
Add sharding argument to reshape since it also takes a `shape` argument for the output shape

#25103 copybara-service[bot] opened 21 minutes ago
0
[JAX] Add Python binding for building a colocated Python program

#25102 copybara-service[bot] opened 26 minutes ago
0
Weird defjvp behavior when finding grad of a scalar that depends on the primal

#25101 JadM133 opened 55 minutes ago
0
Create a `null_mesh_context` internal context manager to handle null contexts properly.

#25100 copybara-service[bot] opened 1 hour ago
0
[Pallas TPU] [WIP] Add vector support to `pl.debug_print`

#25099 copybara-service[bot] opened 1 hour ago
0
[mosaic_gpu] Fixed unbounded recursion in `FragmentedArray._pointwise`

#25098 copybara-service[bot] opened 2 hours ago
0
Start a new TPU interpret mode for Pallas

#25097 jburnim opened 4 hours ago
0
Update ROCm scripts to match new build.py usage

#25096 nitins17 opened 4 hours ago
0
[Mosaic GPU] Support batch dimensions in FA3 MGPU kernel.

#25095 copybara-service[bot] closed 4 hours 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 4 hours ago
0
[SDY] Enable `test_pjit_array_multi_input_multi_output` since Shardy conflict resolution is now complete.

#25093 copybara-service[bot] closed 5 hours ago
0
Reverts e72b449089f6af4ceb18288e36215b3c76e69245

#25092 copybara-service[bot] opened 5 hours ago
0
[cuda] Bump nvidia-cuda-nvcc-cu12 dependency to 12.6.85

#25091 gspschmid opened 5 hours ago
0
[SDY] remove TODO for enabling Layouts for Shardy post cl/697715276.

#25090 copybara-service[bot] closed 4 hours ago
0
Bump hypothesis from 6.102.4 to 6.119.4

#25089 dependabot[bot] opened 6 hours ago
0
[SDY] enable `cpu_shardy` for JAX shard_alike test.

#25088 copybara-service[bot] closed 5 hours ago
0
Update the "CI" workflow to use the new self-hosted runners

#25087 copybara-service[bot] opened 6 hours ago
1
[mgpu] FragentedArray.foreach() can now optionally return a new array

#25086 copybara-service[bot] opened 7 hours 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 6 hours ago
0
[ROCm] Add ROCm version metadata to pip show output for `jax-rocm60-plugin` and `jax-rocm60-pjrt`

#25084 Ruturaj4 opened 7 hours ago
2
Automated Code Change

#25083 copybara-service[bot] closed 7 hours ago
0
Document cudaMallocAsync

#25082 nouiz opened 9 hours ago
0
[Mosaic GPU] Simplify logic for pointwise splat operands

#25081 copybara-service[bot] closed 8 hours ago
0
Use xla_extension_version instead of jaxlib_version

#25079 copybara-service[bot] opened 14 hours ago
0
[pallas_mgpu] Map indices primitive

#25078 copybara-service[bot] opened 22 hours ago
0
PRNGKey error

#25076 hangita101 opened 1 day ago
1
cudnn handle (500): cudaErrorSymbolNotFound : named symbol not found

#25075 PhilipVinc opened 2 days ago
0
When host local inputs on all hosts are the same, use `_DeferredShardArg` to do the transfers instead of `jit` to avoid blocking.

#25074 copybara-service[bot] closed 2 days ago
0
Test lax_scipy_spectral_dac error message

#25073 copybara-service[bot] opened 2 days ago
0
Move `tsl/platform/{build_config,build_config_root,rules_cc}.bzl` to `xla/tsl/platform`

#25072 copybara-service[bot] opened 3 days ago
0
Fix member access to xla backend. The correct member is `client` instead of `backend`

#25071 copybara-service[bot] closed 2 days ago
0
Add linearization rule for pjit_p

#25070 dougalm closed 3 days ago
0
Possible leak in random number generation

#25069 RadostW opened 3 days ago
1
[Mosaic TPU] Support packed type matmul with arbitrary shapes.

#25068 copybara-service[bot] opened 3 days ago
0
Test PR for upstream changes

#25067 charleshofer closed 3 days ago
1
`AssertionError: Unexpected XLA layout override` when adding two `from_dlpack` arrays

#25066 samuela opened 3 days ago
1
Use with_spec where possible to clean up the code a bit

#25065 copybara-service[bot] closed 3 days ago
0
[sharding_in_types] Make argmax and argmin work with sharding_in_types. This also requires adding `reduce_p` sharding rule

#25064 copybara-service[bot] closed 3 days ago
0
[pallas_mgpu] Allow trees (eg tuples) to be returned from cond_p expressions.

#25063 copybara-service[bot] opened 3 days ago
0
[mgpu] Pointwise min

#25061 copybara-service[bot] opened 3 days ago
0
[mgpu/pallas_mgpu] Pointwise sign function

#25060 copybara-service[bot] opened 3 days ago
0
[mgpu/pallas_mgpu] Pointwise tanh support

#25059 copybara-service[bot] opened 3 days ago
0
[pallas:mosaic_gpu] Add test for FragmentedArray.bitcast.

#25058 copybara-service[bot] closed 12 hours ago
0
Rank-one updates to eigenvalue decompositions

#25057 mishavanbeek opened 3 days ago
2
Added stream annotation support via @compute_on('stream:#') decorator

#25056 chaserileyroberts opened 3 days ago
1
Use optimize='auto' for `jnp.linalg.multi_dot`

#25055 dfm opened 3 days ago
0
Implement SVD algorithm based on QR for CPU targets

#25053 JanLuca opened 3 days ago
3
Rework custom hermetic python instructions.

#25052 copybara-service[bot] opened 3 days ago
0
jax.numpy.linalg.multi_dot is O(2^N) in the number of matrices being multiplied

#25051 rohan-hitchcock opened 3 days ago
6