jax
CI: 03/17/25 upstream sync
#292
Merged
Overview
Commits
23
Changes
View On GitHub
Commits
  • Use `lax.top_k` instead of `jnp.argsort` in Gumbel top-k trick for weighted sampling without replacement in `jax.random.choice`
    mar-muel committed 1 year ago
  • add experimental lax.optimization_barrier autodiff rules
    mattjj committed 1 year ago
  • Add a C++ implementation of a toplogical sort.
    Google-ML-Automation committed 1 year ago
  • Merge pull request #27174 from mattjj:opt-barrier-ad-rules
    Google-ML-Automation committed 1 year ago
  • mixing modes
    yashk2810 committed 1 year ago
  • Merge pull request #27177 from jax-ml:mixing_modes
    Google-ML-Automation committed 1 year ago
  • Support error checking in explicit mode
    Google-ML-Automation committed 1 year ago
  • Update XLA dependency to use revision
    Google-ML-Automation committed 1 year ago
  • Better docs for jax.lax add/sub/mul/div
    jakevdp committed 1 year ago
  • Change the way that batching.spec_types is updated.
    Google-ML-Automation committed 1 year ago
  • Update XLA dependency to use revision
    Google-ML-Automation committed 1 year ago
  • Merge pull request #27176 from jakevdp:lax-docs
    Google-ML-Automation committed 1 year ago
  • [Mosaic GPU] Add support for fast WGMMA layout changes after 8- to 16-bit upcast
    Google-ML-Automation committed 1 year ago
  • [Mosaic GPU] Add support for changing the layout before the upcast
    Google-ML-Automation committed 1 year ago
  • [pallas:mosaic_gpu] `jnp.reduce_sum` now works for >1D arrays
    Google-ML-Automation committed 1 year ago
  • Update XLA dependency to use revision
    Google-ML-Automation committed 1 year ago
  • Removed trivial docstrings from JAX tests
    Google-ML-Automation committed 1 year ago
  • [Mosaic GPU] Make the s4 -> bf16 upcast more flexible when it comes to vector length
    Google-ML-Automation committed 1 year ago
  • Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt
    Google-ML-Automation committed 1 year ago
  • Merge pull request #27157 from mar-muel:improve-random-choice-performance
    Google-ML-Automation committed 1 year ago
  • [Mosaic GPU] Add initial transform inference rules for `vector.{load,store}`.
    Google-ML-Automation committed 1 year ago
  • Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-149_1
    invalid-email-address committed 1 year ago
  • Trivial change for CI
    charleshofer committed 1 year ago
FAQ Terms Privacy Refunds Impressum
Loading