Skip to content

Pull requests: jax-ml/jax

New pull request New
373 Open 15,220 Closed
373 Open 15,220 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[Mosaic TPU] Force offset to 0 when inferring input has offset out of the first tile.
#23987 opened Sep 28, 2024 by copybara-service bot Loading…
[Mosaic TPU] Canonicalize arith.select's condition to vector if other types are vector.
#23986 opened Sep 27, 2024 by copybara-service bot Loading…
Fix typo in Pallas TPU matmul doc. I think the logical layout of the input array is non-transposed, rather than transposed?
#23985 opened Sep 27, 2024 by copybara-service bot Loading…
jnp.mask_indices: add docs & tests documentation
#23984 opened Sep 27, 2024 by jakevdp Loading…
@jakevdp
[Pallas/MGPU] Undo transforms on refs before giving them back to the users
#23974 opened Sep 27, 2024 by copybara-service bot Loading…
@apaszke
[jax.distributed] Allow enabling grpc channel compression
#23969 opened Sep 27, 2024 by gspschmid Loading…
[Pallas TPU] Remove the skipTest for OpsExtraTest on TPU
#23967 opened Sep 27, 2024 by copybara-service bot Loading…
@ayaka14732
Update device_memory_profiling.md
#23964 opened Sep 27, 2024 by garymm Loading…
[Pallas] Simplify OpsExtraTest by skipping 64-bit tests on 32-bit environments
#23960 opened Sep 26, 2024 by copybara-service bot Loading…
@ayaka14732
Exposed sequential iteration index via pl.program_id in Pallas Mosaic GPU
#23956 opened Sep 26, 2024 by copybara-service bot Loading…
@superbobry
Add support for passing array attributes via ffi_call and an example demonstrating the different attribute APIs supported by the FFI
#23951 opened Sep 26, 2024 by dfm Loading…
@dfm
3
Activate FFI implementation of the QR decomposition.
#23948 opened Sep 26, 2024 by copybara-service bot Loading…
@dfm
Introducing partial discharge rules and implementations for cond_p and for_p
#23947 opened Sep 26, 2024 by copybara-service bot Loading…
[Mosaic TPU] Support bitcast with non-native tiling and reset offsets when not falling within the first tile.
#23934 opened Sep 26, 2024 by copybara-service bot Loading…
[Pallas] Add checkify support for checking pipeline grid bounds.
#23932 opened Sep 25, 2024 by copybara-service bot Loading…
@justinjfu
Add custom_jvp to jax.numpy.ldexp.
#23923 opened Sep 25, 2024 by carlosgmartin Loading…
@jakevdp
9
[Mosaic TPU] Lock down local DMAs with an SMEM source buffer temporarily
#23903 opened Sep 25, 2024 by copybara-service bot Loading…
@apaszke
[Pallas] Clean up lowering exceptions.
#23897 opened Sep 25, 2024 by copybara-service bot Loading…
@justinjfu
[ROCm] Apply some fixes for issues with ROCm builds
#23891 opened Sep 24, 2024 by mrodden Loading…
@dfm
1
DOC: Improved documentation for jax.numpy.isscalar
#23890 opened Sep 24, 2024 by selamw1 Loading…
7
allow dynamic tensor in multi_broadcast_in_dim
#23886 opened Sep 24, 2024 by erick-xanadu Loading…
@mattjj
1
Deprecate the vectorized argument to pure_callback and ffi_call pull ready Ready for copybara import and testing
#23881 opened Sep 24, 2024 by dfm Loading…
@dfm
4
Do not attempt checking dynamic dimensions values
#23879 opened Sep 24, 2024 by erick-xanadu Loading…
1
Fix clean_up_dead_vars when jax_dynamic_shapes is True
#23877 opened Sep 24, 2024 by erick-xanadu Loading…
1
fixing jnp.broadcast_to throwing an exception when passing an instance of nnx.Variable
#23865 opened Sep 24, 2024 by YushaArif99 Loading…
5
ProTip! Adding no:label will show everything without a label.