jax-ml/jax issues and pull requests

#23924 - Fix make_remote_async_copy -> make_async_remote_copy in async doc

Pull Request - State: closed - Opened by jburnim about 2 months ago
Labels: pull ready

#23924 - Fix make_remote_async_copy -> make_async_remote_copy in async doc

Pull Request - State: closed - Opened by jburnim about 2 months ago
Labels: pull ready

#23923 - Add custom_jvp to jax.numpy.ldexp.

Pull Request - State: open - Opened by carlosgmartin about 2 months ago - 7 comments

#23921 - Raise `ValueError` when `axis1==axis2` for `jnp.trace`

Pull Request - State: closed - Opened by rajasekharporeddy about 2 months ago - 1 comment
Labels: pull ready

#23920 - Added support for `lax.cond_p` to Pallas Mosaic GPU lowering

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23920 - Added support for `lax.cond_p` to Pallas Mosaic GPU lowering

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23919 - Add a fused multiply-add function

Issue - State: open - Opened by carlosgmartin about 2 months ago - 7 comments
Labels: enhancement

#23901 - #sdy Add CPU targets in JAX.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23899 - DOC: Better doc for `jnp.trace`

Pull Request - State: closed - Opened by rajasekharporeddy about 2 months ago - 2 comments
Labels: documentation, pull ready

#23899 - DOC: Better doc for `jnp.trace`

Pull Request - State: closed - Opened by rajasekharporeddy about 2 months ago - 2 comments
Labels: documentation, pull ready

#23897 - [Pallas] Clean up lowering exceptions.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23896 - Fix printing of saved_residual for `jit` by looking for `pjit` as the primitive instead of `xla_call` which was removed 2 years ago

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23895 - JAX dct is not accurate compared to scipy

Issue - State: open - Opened by unalmis about 2 months ago - 2 comments

#23894 - Enable some more H100 tests in CI.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23893 - Fix some warnings causing CI failures on ARM.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23892 - [mosaic] Fix a warning causing CI failures.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23891 - [ROCm] Apply some fixes for issues with ROCm builds

Pull Request - State: closed - Opened by mrodden about 2 months ago - 2 comments
Labels: pull ready

#23891 - [ROCm] Apply some fixes for issues with ROCm builds

Pull Request - State: open - Opened by mrodden about 2 months ago

#23890 - DOC: Improved documentation for jax.numpy.isscalar

Pull Request - State: closed - Opened by selamw1 about 2 months ago - 2 comments

#23890 - DOC: Improved documentation for jax.numpy.isscalar

Pull Request - State: closed - Opened by selamw1 about 2 months ago - 2 comments

#23889 - [Pallas TPU] Add a note in the Pallas Quickstart documentation about the instructions of running the existing example on TPU

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23888 - Update tf2jax to use new gather/scatter batching dims.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23887 - `lax.cond` is broken within `compute_on('device_host')`

Issue - State: open - Opened by hanzhi713 about 2 months ago
Labels: bug

#23887 - `lax.cond` is broken within `compute_on('device_host')`

Issue - State: closed - Opened by hanzhi713 about 2 months ago - 1 comment
Labels: bug

#23886 - allow dynamic tensor in multi_broadcast_in_dim

Pull Request - State: open - Opened by erick-xanadu about 2 months ago - 1 comment

#23886 - allow dynamic tensor in multi_broadcast_in_dim

Pull Request - State: open - Opened by erick-xanadu about 2 months ago - 1 comment

#23885 - [Pallas TPU] Improve error message when trying to store a scalar to VMEM

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23884 - [Pallas TPU] Confusing error message when storing scalars to VMEM

Issue - State: closed - Opened by ayaka14732 about 2 months ago
Labels: bug, pallas

#23883 - Lower the shard count for sparse_bcoo_bcsr_test on TPU as well.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23883 - Lower the shard count for sparse_bcoo_bcsr_test on TPU as well.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23882 - Jax[cuda] always allocates memory

Issue - State: open - Opened by GJBoth about 2 months ago - 1 comment
Labels: bug

#23882 - Jax[cuda] always allocates memory

Issue - State: open - Opened by GJBoth about 2 months ago - 1 comment
Labels: bug

#23881 - Deprecate the `vectorized` argument to pure_callback and ffi_call

Pull Request - State: closed - Opened by dfm about 2 months ago - 2 comments
Labels: pull ready

#23881 - Deprecate the `vectorized` argument to pure_callback and ffi_call

Pull Request - State: open - Opened by dfm about 2 months ago
Labels: pull ready

#23880 - Set PYTHONWARNINGS=error in bazel tests.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23880 - Set PYTHONWARNINGS=error in bazel tests.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23879 - Do not attempt checking dynamic dimensions values

Pull Request - State: open - Opened by erick-xanadu about 2 months ago - 1 comment

#23879 - Do not attempt checking dynamic dimensions values

Pull Request - State: open - Opened by erick-xanadu about 2 months ago - 1 comment

#23878 - [pallas:triton] Elide `program_id` calls where launch grid dimension is 1.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23877 - Fix `clean_up_dead_vars` when `jax_dynamic_shapes` is `True`

Pull Request - State: open - Opened by erick-xanadu about 2 months ago

#23877 - Fix `clean_up_dead_vars` when `jax_dynamic_shapes` is `True`

Pull Request - State: open - Opened by erick-xanadu about 2 months ago

#23876 - [Pallas] Update export compatibility tests

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23875 - Expose some APIs for querying trace state. This will let us move users away from

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23874 - DOC: Improve docs for `jax.numpy`: `arcsin`, `arccos` and `arctan`

Pull Request - State: open - Opened by rajasekharporeddy about 2 months ago

#23873 - Fix an incorrect output for jnp.cumsum.

Pull Request - State: closed - Opened by hawkinsp about 2 months ago
Labels: pull ready

#23873 - Fix an incorrect output for jnp.cumsum.

Pull Request - State: closed - Opened by hawkinsp about 2 months ago
Labels: pull ready

#23872 - [pallas] Allow `TransformedRef` to be passed to `pl.load` / `pl.store`, when `idx = ()`.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23872 - [pallas] Allow `TransformedRef` to be passed to `pl.load` / `pl.store`, when `idx = ()`.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23871 - [Mosaic GPU] Add support for input_output_aliases

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23870 - Fix stdout buffer error in JAX UT api_test.py

Pull Request - State: closed - Opened by Zantares about 2 months ago - 2 comments
Labels: pull ready

#23869 - Unable to use residual offloading with scan and remat

Issue - State: open - Opened by qGentry about 2 months ago - 5 comments
Labels: bug

#23868 - Added support for % and `select` to `mgpu.FragmentedArray`

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23868 - Added support for % and `select` to `mgpu.FragmentedArray`

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23867 - Checkify doesn't work with remat.

Issue - State: open - Opened by giorgio-arena about 2 months ago
Labels: bug

#23867 - Checkify doesn't work with remat.

Issue - State: open - Opened by giorgio-arena about 2 months ago
Labels: bug

#23866 - Wrong result when Jitting a function

Issue - State: open - Opened by mtagliazucchi about 2 months ago - 1 comment
Labels: bug

#23866 - Wrong result when Jitting a function

Issue - State: open - Opened by mtagliazucchi about 2 months ago - 1 comment
Labels: bug

#23865 - fixing `jnp.broadcast_to` throwing an exception when passing an instance of `nnx.Variable`

Pull Request - State: open - Opened by YushaArif99 about 2 months ago - 1 comment

#23863 - Support pmin and pmax in check_rep.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23863 - Support pmin and pmax in check_rep.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23862 - trigger array API tests for all PRs.

Pull Request - State: closed - Opened by jakevdp about 2 months ago
Labels: pull ready

#23862 - trigger array API tests for all PRs.

Pull Request - State: closed - Opened by jakevdp about 2 months ago
Labels: pull ready

#23861 - Remove warning filters from pyproject.toml, add local warning

Pull Request - State: closed - Opened by hawkinsp about 2 months ago - 3 comments
Labels: pull ready

#23861 - Remove warning filters from pyproject.toml, add local warning

Pull Request - State: closed - Opened by hawkinsp about 2 months ago - 3 comments
Labels: pull ready

#23860 - Add more informative error when static argument is passed to non-static JIT parameter

Pull Request - State: closed - Opened by jakevdp about 2 months ago
Labels: pull ready

#23860 - Add more informative error when static argument is passed to non-static JIT parameter

Pull Request - State: closed - Opened by jakevdp about 2 months ago
Labels: pull ready

#23858 - [Pallas] Add scalar f32 downcast test cases.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23858 - [Pallas] Add scalar f32 downcast test cases.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23857 - Pallas pipeline API tweaks for more advanced pipelining patterns.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23857 - Pallas pipeline API tweaks for more advanced pipelining patterns.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23856 - [WIP] jax pytree data serialization

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23855 - `shaped_abstractify` needs a more informative error message

Issue - State: closed - Opened by GallagherCommaJack about 2 months ago
Labels: enhancement, better_errors

#23855 - `shaped_abstractify` needs a more informative error message

Issue - State: closed - Opened by GallagherCommaJack about 2 months ago
Labels: enhancement, better_errors

#23854 - segfault with `map` and `pure_callback`

Issue - State: open - Opened by Michael-T-McCann about 2 months ago
Labels: bug

#23854 - segfault with `map` and `pure_callback`

Issue - State: open - Opened by Michael-T-McCann about 2 months ago
Labels: bug

#23853 - Enable the activation offloading test

Pull Request - State: open - Opened by zhenying-liu about 2 months ago

#23853 - Enable the activation offloading test

Pull Request - State: closed - Opened by zhenying-liu about 2 months ago
Labels: pull ready

#23852 - [ROCM] add missing typename keyword to work with gcc

Pull Request - State: closed - Opened by Ruturaj4 about 2 months ago
Labels: pull ready

#23852 - [ROCM] add missing typename keyword to work with gcc

Pull Request - State: closed - Opened by Ruturaj4 about 2 months ago
Labels: pull ready

#23851 - Fix another errata in block-sparse kernel tutorial.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23851 - Fix another errata in block-sparse kernel tutorial.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23850 - Re-enable tsan tests after fix.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23850 - Re-enable tsan tests after fix.

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23849 - Use `jax.make_array_from_process_local_data` API in distributed data loading doc

Pull Request - State: closed - Opened by copybara-service[bot] about 2 months ago

#23843 - Hide JAX's internal tracing state and update libraries to use limited trace-state-querying APIs as needed. This is prep work for stackless which will change those internals while preserving the API.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

#23843 - Hide JAX's internal tracing state and update libraries to use limited trace-state-querying APIs as needed. This is prep work for stackless which will change those internals while preserving the API.

Pull Request - State: open - Opened by copybara-service[bot] about 2 months ago

GitHub / jax-ml/jax issues and pull requests