Skip to content

Issues: google/jax

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Clear current search query, filters, and sorts
60 Open 263 Closed
60 Open 263 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
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Issues list

Pallas kernel crash at llo::CouldLtS32 when interpret=False on TPU bug Something isn't working
#23181 opened Aug 21, 2024 by nathom
@sharadmv @justinjfu
3
[Pallas/TPU] Correctly populate new_avals when we are using sem indexers in dma_start discharge rule
#23135 opened Aug 20, 2024 by copybara-service bot Loading…
@sharadmv
[pallas] Request to add support for select/argmax primitive in Pallas TPU enhancement New feature or request pallas Issues pertaining to Pallas (GPU or TPU)
#22508 opened Jul 18, 2024 by Lime-Cakes
@sharadmv
1
pallas hello world not implemented for int8 bug Something isn't working pallas Issues pertaining to Pallas (GPU or TPU)
#22204 opened Jun 30, 2024 by vlad17
@sharadmv
JAX profiling tools support enhancement New feature or request
#21338 opened May 21, 2024 by PrajP
@sharadmv
1
Pallas/Mosiac documentation and examples documentation enhancement New feature or request pallas Issues pertaining to Pallas (GPU or TPU)
#21337 opened May 21, 2024 by PrajP
@sharadmv @apaszke
2
simple pallas kernel hangs when input size exceeds some threshold bug Something isn't working pallas Issues pertaining to Pallas (GPU or TPU)
#21282 opened May 17, 2024 by zhixuan-lin
@sharadmv @apaszke
2
Add type information for pallas bitcast/roll pull ready Ready for copybara import and testing
#20576 opened Apr 3, 2024 by kkiningh Loading…
@sharadmv
5
Unimplemented primitive in Pallas GPU lowering: scatter-add bug Something isn't working pallas Issues pertaining to Pallas (GPU or TPU)
#20168 opened Mar 11, 2024 by bsaoptima
@sharadmv
1
Pallas Kernel using Smem/SReg failed to lower better_errors Improve the error reporting pallas Issues pertaining to Pallas (GPU or TPU)
#20046 opened Mar 1, 2024 by tankbattle
@sharadmv
2
More type annotations of Pallas primitives pull ready Ready for copybara import and testing
#19608 opened Feb 1, 2024 by trishume Loading…
@sharadmv
4
Write docs for new Pallas collective APIs. documentation pallas Issues pertaining to Pallas (GPU or TPU)
#19420 opened Jan 19, 2024 by epiqueras Loading…
@sharadmv
17
Pallas: jnp.indices fails to lower to mosaic bug Something isn't working pallas Issues pertaining to Pallas (GPU or TPU)
#19291 opened Jan 10, 2024 by j2kun
@sharadmv
[Pallas] NotImplementedError: Unimplemented primitive in Pallas GPU lowering: slice. enhancement New feature or request pallas Issues pertaining to Pallas (GPU or TPU)
#19214 opened Jan 5, 2024 by vvvm23
@sharadmv
6
Wrong behaviour of pallas.atomic_cas bug Something isn't working NVIDIA GPU Issues specific to NVIDIA GPUs pallas Issues pertaining to Pallas (GPU or TPU)
#18909 opened Dec 10, 2023 by southfreebird
@sharadmv @apaszke
Unimplemented "discharged_consts" when discharging while loop effects bug Something isn't working pallas Issues pertaining to Pallas (GPU or TPU)
#18357 opened Nov 2, 2023 by brianwa84
@sharadmv
1
give users a way to type-query effects, e.g. in eval_shape enhancement New feature or request
#18239 opened Oct 23, 2023 by mattjj
@froystig @sharadmv @mattjj
4
Allow import or other statements in jdb enhancement New feature or request
#18059 opened Oct 11, 2023 by jjyyxx
2 tasks done
@sharadmv
perf(Pallas): Controlling accumulation dtype for dot_general and Pallas dot enhancement New feature or request pallas Issues pertaining to Pallas (GPU or TPU)
#17612 opened Sep 14, 2023 by jon-chuang
@sharadmv
3
Better error messages for shape and dtype mismatch.
#17546 opened Sep 11, 2023 by mehdiataei Loading…
@sharadmv
jax.debug.breakpoint() give numpy array and with bfloat16, it comptue bug Something isn't working
#17402 opened Sep 1, 2023 by nouiz
@sharadmv
1
Change scan partial_eval to use partial_eval_jaxpr_dynamic
#17357 opened Aug 29, 2023 by sharadmv Loading…
@sharadmv
feat(pallas): Optimize Pallas Attention + Benchmark
#17328 opened Aug 28, 2023 by jon-chuang Loading…
@sharadmv
17
jaxlib.xla_extension.XlaRuntimeError due to failed: Shared memory requested exceeds device resources bug Something isn't working NVIDIA GPU Issues specific to NVIDIA GPUs
#17320 opened Aug 28, 2023 by luise1030
@sharadmv
2
vmap of pure_callback(..., vectorized=True) fails on broadcasts bug Something isn't working
#17187 opened Aug 18, 2023 by Gattocrucco
@sharadmv
9
ProTip! Adding no:label will show everything without a label.