jax-ml/jax

Issues

• How do you remat GSPMD inserted all-gathers?

  #25010 opened a month ago by ptoulme-aws
  2

• `CUDA_ERROR_ILLEGAL_ADDRESS`

  #25002 opened a month ago by PhilipVinc
  1

• Numeric stability issue in higher-order tensor einsum

  #25024 opened a month ago by dest1n1s
  4

• An example for `ffi` input/output aliasing

  #24986 opened a month ago by saeedmaleki
  1

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

  #25023 opened a month ago by caixiiaoyang
  1

• time.sleep affects the execution time of JAX

  #24941 opened a month ago by horse6
  8

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

  #24995 opened a month ago by carlosgmartin
  7

• libtpu.so Not Found for JAX TPU 0.4.14

  #24987 opened a month ago by LeoXinhaoLee
  2

• `debug_nans` error always says the de-optimized function did not produce NaNs

  #24955 opened a month ago by emilyfertig
  5

• AOT compilation and serialization

  #24982 opened a month ago by jkup14
  6

• FloatingPointError in jax.scipy.stats

  #24939 opened a month ago by IlayMenahem
  2

• Slow import time: can it be reduced?

  #24967 opened a month ago by jeertmans
  2

• [GPU] FlashAttention performance lags behind PyTorch

  #24934 opened a month ago by neel04
  4

• ⚠️ Nightly upstream-dev CI failed ⚠️

  #24875 opened a month ago by github-actions
  1

• Bug in the latest libtpu nightly release

  #24829 opened a month ago by knyazer
  3

• Small Einsum is hanging

  #24929 opened a month ago by ryan112358
  2

• Rework comparisons in the docs of JAX vs. NumPy PRNG

  #24927 opened a month ago by emilyfertig
  1

• Allow einsum to support naive contraction strategy

  #24915 opened a month ago by ryan112358
  0

• array.at.set is incredibly slow for complex128 dtype

  #24872 opened a month ago by chrisrothUT
  5

• Results do not match the reference. This is likely a bug/unexpected loss of precision

  #24909 opened a month ago by yanboyang97
  0

• CUDA12 plugin segfaults when older version of JAX is installed

  #24901 opened a month ago by dime10
  0

• Nightly builds on Mac ARM fail complex function numerical tests

  #24787 opened 2 months ago by hawkinsp
  2

• pure_callback using a threadpool to accelerate vmap

  #24756 opened 2 months ago by Joshuaalbert
  3

• Missing annotations

  #24888 opened a month ago by PerilousApricot
  4

• CUDA_ERROR_SYSTEM_NOT_READY: system not yet initialized

  #24866 opened a month ago by carlosgmartin
  1

• pjit handling of static_argnames is broken when using "dynamic" python strings.

  #24857 opened a month ago by etarassov
  5

• Unsupported type in metal PJRT plugin with rng_bit_generator

  #24867 opened a month ago by dmarro89
  0

• cond error by using Tracer as a condition

  #24858 opened a month ago by 2300504237
  3

• ⚠️ Nightly upstream-dev CI failed ⚠️

  #24824 opened a month ago by github-actions
  1

• AOT errors when setting compile options

  #24869 opened a month ago by man2machine
  0

• Error with `at[index].set(value)` when using `jit` and `shard_map` in multi-host setting.

  #24768 opened 2 months ago by lollcat
  2

• support q offset w.r.t k/v in flash_attention function

  #24830 opened a month ago by yamingx
  2

• Wrong determinant results for large batch

  #24843 opened a month ago by ChenAo-Phys
  5

• `jax.lax.DotAlgorithm` generates an unnecessary conversion on H100

  #24794 opened a month ago by andportnoy
  3

• InconclusiveDimensionOperation: Symbolic dimension comparison 'b' < '2147483647' is inconclusive.

  #24730 opened a month ago by njzjz
  3

• 7x7 `nnx.Conv` using `float32` parameter dtype overflows(?) to `nan` when sharded

  #24848 opened a month ago by joshhansen
  1

• potential jax 0.4.35 release issue?

  #24826 opened a month ago by haohuanw
  1

• cudaErrorSymbolNotFound : named symbol not found

  #24749 opened 2 months ago by amanjitsk
  3

• `shard_map` doesn't work with `jnp.insert`

  #24762 opened 2 months ago by mrlazy1708
  0

• bincount rejects bool

  #24813 opened a month ago by carlosgmartin
  0

• Deepcopy of pjit functions failing (or not supported)

  #24838 opened a month ago by jlperla
  0

• `scipy.sparse.csgraph.connected_components` implementation

  #24737 opened 2 months ago by vboussange
  2

• TPU with sharding: grpc initialization failure

  #24821 opened a month ago by knyazer
  2

• Sparse reshape throws error when `n_dense>0` and some target dimension has size 1

  #24795 opened 2 months ago by cherrywoods
  1

• Division by self not always "1.0" on JAX GPU, but consistently gives "1.0" on JAX CPU.

  #24807 opened 2 months ago by mattlevine22
  1

• pre-commit run --all failed for ruff and mypy

  #24799 opened 2 months ago by apivovarov
  1

• Flaky test tests/fft_test.py::FftTest::testFftfreq5

  #24798 opened 2 months ago by apivovarov
  2

• TracerBoolConversionError

  #24740 opened 2 months ago by hangita101
  4

• Facing Scaling issue on cpu (arm and x86).

  #24753 opened 2 months ago by choudhary-devang
  0

• median FloatingPointError: invalid value (nan) encountered in jit(convert_element_type)

  #24732 opened 2 months ago by zhangylch
  2