Commits

Branch selector

 main

User selector

dfm

Datepicker

All time

Commit History

Commits on Dec 20, 2024

• Accelerate deprecation of legacy JAX FFI calling convention.

  Show description for 4216f8f
  

  dfm

  authored and

  Google-ML-Automation

  committed
  4216f8f

  Copy full SHA for 4216f8f

• Relax some test tolerances in N-D FFT tests.

  Show description for 0b190bb
  

  dfm

  authored and

  Google-ML-Automation

  committed
  0b190bb

  Copy full SHA for 0b190bb

Commits on Dec 19, 2024

• Add support for N-D FFTs with D>3.

  dfm

  committed
  c6131ee

  Copy full SHA for c6131ee

Commits on Dec 16, 2024

• Relax tolerance for LAX reduction test in float16.

  Show description for ed4e982
  

  dfm

  authored and

  Google-ML-Automation

  committed
  ed4e982

  Copy full SHA for ed4e982

• Fix some flaky LAX autodiff tests.

  Show description for a2e9699
  

  dfm

  authored and

  Google-ML-Automation

  committed
  a2e9699

  Copy full SHA for a2e9699

Commits on Dec 10, 2024

• Deduplicate some GPU plugin definition code.

  Show description for 593143e
  

  dfm

  authored and

  Google-ML-Automation

  committed
  593143e

  Copy full SHA for 593143e

Commits on Dec 9, 2024

• Fix expected exception type in pallas grad tests.

  Show description for 978d35f
  

  dfm

  authored and

  Google-ML-Automation

  committed
  978d35f

  Copy full SHA for 978d35f

• Port symmetric tridiagonal reduction GPU kernel to FFI.

  Show description for 32df37e
  

  dfm

  authored and

  Google-ML-Automation

  committed
  32df37e

  Copy full SHA for 32df37e

• Add error message when using custom_vmap with reverse-mode AD, and add docstrings.

  Show description for 092d2a0
  

  dfm

  authored and

  Google-ML-Automation

  committed
  092d2a0

  Copy full SHA for 092d2a0

Commits on Dec 7, 2024

• Temporarily allow bfloat16 dot algorithms on CPU.

  Show description for 1f4d184
  

  dfm

  authored and

  Google-ML-Automation

  committed
  1f4d184

  Copy full SHA for 1f4d184

Commits on Nov 25, 2024

• Refactor FFI examples to consolidate several examples into one submodule.

  dfm

  committed
  84a9cba

  Copy full SHA for 84a9cba

Commits on Nov 22, 2024

• Use optimize='auto' for multi_dot.

  dfm

  committed
  236d4c6

  Copy full SHA for 236d4c6

Commits on Nov 21, 2024

• Fix false positive debug_nans error caused by NaNs that are properly handled in jax.scipy.stats.gamma

  Show description for 96c0129
  

  dfm

  authored and

  Google-ML-Automation

  committed
  96c0129

  Copy full SHA for 96c0129

Commits on Nov 19, 2024

• Add test utility for accessing jaxlib version tuple.

  Show description for a59bbb7
  

  dfm

  authored and

  Google-ML-Automation

  committed
  a59bbb7

  Copy full SHA for a59bbb7

• Add missing version guard in GPU tests for jnp.poly.

  Show description for 3556a83
  

  dfm

  authored and

  Google-ML-Automation

  committed
  3556a83

  Copy full SHA for 3556a83

Commits on Nov 18, 2024

• Add a GPU implementation of lax.linalg.eig.

  Show description for ccb3317
  

  dfm

  authored and

  Google-ML-Automation

  committed
  ccb3317

  Copy full SHA for ccb3317

Commits on Nov 14, 2024

• Add shard map replication rule for ffi_call.

  dfm

  committed
  41a0493

  Copy full SHA for 41a0493

• Fix pgle test handling of directory removal.

  Show description for de07be5
  

  dfm

  authored and

  Google-ML-Automation

  committed
  de07be5

  Copy full SHA for de07be5

Commits on Nov 13, 2024

• Add a shard map replication rule for cond_p.

  dfm

  committed
  dfabcb0

  Copy full SHA for dfabcb0

Commits on Nov 12, 2024

• Add GPU overflow bugfix (jax-ml#24846) to changelog.

  dfm

  committed
  5808170

  Copy full SHA for 5808170

• Add an example FFI call to demonstrate the use of global state.

  dfm

  committed
  f086483

  Copy full SHA for f086483

• Update some outdated syntax in FFI tutorial.

  dfm

  committed
  f757054

  Copy full SHA for f757054

• Remove GPU test with unreasonably large memory footprint.

  Show description for a99ccd9
  

  dfm

  authored and

  Google-ML-Automation

  committed
  a99ccd9

  Copy full SHA for a99ccd9

• Fix overflow error in GPU batched linear algebra kernels.

  Show description for 21e98b5
  

  dfm

  authored and

  Google-ML-Automation

  committed
  21e98b5

  Copy full SHA for 21e98b5

• Allow more output storage types for some dot algorithms.

  Show description for 9bb6366
  

  dfm

  authored and

  Google-ML-Automation

  committed
  9bb6366

  Copy full SHA for 9bb6366

Commits on Nov 11, 2024

• Allow 64-bit output types from ffi_call regardless of enable_x64 flag.

  dfm

  committed
  478ea0d

  Copy full SHA for 478ea0d

Commits on Nov 8, 2024

• Fix pre-commit to run on all files in CI.

  dfm

  committed
  4a36567

  Copy full SHA for 4a36567

Commits on Nov 1, 2024

• Update pre-commit workflow to cache on jax version

  dfm

  committed
  b7bdee9

  Copy full SHA for b7bdee9

Commits on Oct 31, 2024

• Move the CUDA end-to-end example to FFI examples workflow + hosted

  Show description for ce8dba9
  dfm

  committed
  ce8dba9

  Copy full SHA for ce8dba9

• Run dot algorithm tests with PJRT plugin.

  dfm

  committed
  52ad605

  Copy full SHA for 52ad605

Commits on Oct 29, 2024

• Allow dot algorithms in default_matmul_precision config.

  dfm

  committed
  03854cf

  Copy full SHA for 03854cf

• Fix segfault caused by uninitialized LAPACK in FFI test.

  dfm

  committed
  1785479

  Copy full SHA for 1785479

Commits on Oct 26, 2024

• Fix lint at head

  dfm

  committed
  ad1d864

  Copy full SHA for ad1d864

Commits on Oct 25, 2024

• Add support for layouts and other advanced features in ffi_call.

  dfm

  committed
  21f3353

  Copy full SHA for 21f3353

• Re-enable jax2tf test for dot algorithm with stricter TF version check.

  dfm

  committed
  33a46e8

  Copy full SHA for 33a46e8

Pagination

Previous

Next