jax-ml/jax

Issues

• unset JAX_PLATFORMS finds cuda, but JAX_PLATFORMS=gpu tries to use rocm (and fails)

  #25315 opened a month ago by Joshuaalbert
  2

• [Pallas] Add lowering support for tl.join

  #25321 opened a month ago by ggengnv
  0

• Slowdown of jax::lax::while_loop on CPU between 0.4.31 and 0.4.33

  #25190 opened a month ago by bondquant
  9

• Python crashes when trying to compute the gradient of jnp.take

  #25284 opened a month ago by bitsandscraps
  4

• Missing `jax.core.ConcreteArray` in v0.4.36: Migration Guide and API Stability Concerns

  #25314 opened a month ago by daskol
  4

• About FFI failure: Failed to destroy GPU Graph

  #25141 opened 2 months ago by victoryang00
  3

• unexpected `vmap` error due to commit `c36e1f7`

  #25289 opened a month ago by marcocuturi
  5

• LLVM ERROR: Failed to infer result type(s). on Jax Metal 0.1.1

  #25302 opened a month ago by dlwh
  0

• Memory leaked over time when calling the same computation

  #25184 opened a month ago by jonatanklosko
  2

• [//tests/mosaic:gpu_test_gpu] Wrong results on H100

  #25218 opened a month ago by gpupuck
  0

• Regressions in `debug_nans` (and `debug_infs`)

  #25299 opened a month ago by emilyfertig
  0

• Performance Slowdown in Multi-Node JAX with mpi4py Initialization

  #25293 opened a month ago by shyams2
  0

• Schur decomposition test assumes the decomposition is unique, but it isn't

  #25287 opened a month ago by hawkinsp
  0

• Value becoming a tracer error

  #25186 opened a month ago by alexandermm
  9

• Difference between numpy and jax.numpy in advanced indexing axes order

  #25109 opened a month ago by fzimmermann89
  2

• Possible leak in random number generation

  #25069 opened 2 months ago by RadostW
  8

• Request for jax.scipy.special.airy to be implemented

  #25244 opened a month ago by elliottperryman
  0

• ⚠️ Nightly upstream-dev CI failed ⚠️

  #25225 opened a month ago by github-actions
  0

• [Pallas TPU] `pl.load` does not work inside TPU kernels

  #25230 opened a month ago by ayaka14732
  0

• Allow Pallas Triton to be serialized as PTX

  #25196 opened a month ago by sbodenstein
  0

• Unable to run FFI example with stateful function

  #25185 opened a month ago by Luthaf
  2

• ⚠️ Nightly upstream-dev CI failed ⚠️

  #25146 opened a month ago by github-actions
  1

• [Pallas] Uable to convert negative values from float16/float32 to int8/int32 in pallas

  #25047 opened a month ago by shangz-ai
  4

• NaN's produced by jax.numpy.linalg.pinv not produced by numpy

  #25111 opened 2 months ago by samskiter
  2

• [//tests/pallas:export_back_compat_pallas_test_gpu] INTERNAL: RET_CHECK failure (external/xla/xla/service/gpu/ir_emitter_unnested.cc:1388) triton_module Failed to parse Triton module: ML�R

  #25212 opened a month ago by gpupuck
  0

• jax.numpy.linalg.multi_dot is O(2^N) in the number of matrices being multiplied

  #25051 opened a month ago by rohan-hitchcock
  7

• `AssertionError: Unexpected XLA layout override` when adding two `from_dlpack` arrays

  #25066 opened 2 months ago by samuela
  1

• Failure to build jaxlib, AMD GPU

  #25204 opened a month ago by Wintoplay
  2

• [Lowering] Stable IR

  #25123 opened 2 months ago by yliu120
  4

• Allow using jax.jit as decorator with arguments

  #25194 opened a month ago by cool-RR
  3

• in jax.lax.switch docstring, explain that all branch output shapes must match

  #25140 opened 2 months ago by mattjj
  0

• Weird defjvp behavior when finding grad of a scalar that depends on the primal

  #25101 opened a month ago by JadM133
  2

• Array.view(jnp.complex64) doesn't work on TPU with jax_enable_x64

  #25189 opened a month ago by houeland
  2

• Support serializing Mosaic GPU to PTX

  #25197 opened a month ago by sbodenstein
  0

• Installation fails on a cluster

  #25195 opened a month ago by PhilipVinc
  1

• `jax.experimental.pallas.debug_print` does not print output on cloud TPU

  #25192 opened a month ago by byi8220
  1

• How do you remat GSPMD inserted all-gathers?

  #25010 opened 2 months ago by ptoulme-aws
  2

• Generating random numbers in pallas on gpu

  #25188 opened a month ago by lengstrom
  0

• Sum along first axis is slow

  #25187 opened a month ago by ricardoV94
  0

• Cannot generate from `jax.experimental.sparse.random_bcoo` for arrays larger than the int32 max

  #25182 opened a month ago by lengstrom
  1

• Rank-one updates to eigenvalue decompositions

  #25057 opened 2 months ago by mishavanbeek
  4

• PRNGKey error

  #25076 opened 2 months ago by itahang
  2

• `CUDA_ERROR_ILLEGAL_ADDRESS`

  #25002 opened 2 months ago by PhilipVinc
  1

• [FFI Lowering] Compiler-only Attrs

  #25124 opened 2 months ago by yliu120
  3

• cudnn handle (500): cudaErrorSymbolNotFound : named symbol not found

  #25075 opened 2 months ago by PhilipVinc
  0

• Numeric stability issue in higher-order tensor einsum

  #25024 opened 2 months ago by dest1n1s
  4

• An example for `ffi` input/output aliasing

  #24986 opened 2 months ago by saeedmaleki
  1

• Running the executable compiled directly from jax.jit is more than three times slower than jax.jit itself.

  #25023 opened 2 months ago by caixiiaoyang
  1

• Add a way to get the tree size of a jaxpr

  #24995 opened 2 months ago by carlosgmartin
  7

• libtpu.so Not Found for JAX TPU 0.4.14

  #24987 opened 2 months ago by LeoXinhaoLee
  2