Commits

Breadcrumbs

History for

onmain

User selector

All users

All time

Commit History

Commits on Aug 20, 2024

• Added PyUnstable_Module_SetGIL to PyInit_cpu_feature_guard

  vfdev-5

  committed
  6546c48

  Copy full SHA for 6546c48

Commits on Aug 19, 2024

• Added py::mod_gil_not_used() to PYBIND11_MODULE register_jax_dialects

  vfdev-5

  committed
  b1b3ea2

  Copy full SHA for b1b3ea2

• Refactor FFI shape inference functions to include dimension check.

  Show description for 30d54ec
  

  dfm

  authored and

  jax authors

  committed
  30d54ec

  Copy full SHA for 30d54ec

• Add support for shape polymorphism in ffi_lowering and move lu_pivots_to_permutation lowering out of jaxlib.

  Show description for dad2f57
  

  dfm

  authored and

  jax authors

  committed
  dad2f57

  Copy full SHA for dad2f57

Commits on Aug 16, 2024

• Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.

  Show description for b6306e3
  

  dfm

  authored and

  jax authors

  committed
  b6306e3

  Copy full SHA for b6306e3

• Determine LAPACK workspace during QR Factorization Kernel runtime

  Show description for acacf88
  

  Paweł Paruzel

  authored and

  jax authors

  committed
  acacf88

  Copy full SHA for acacf88

Commits on Aug 14, 2024

• [Mosaic] Update serde to handle upstream MLIR changes

  Show description for 020513f
  

  tlongeri

  authored and

  jax authors

  committed
  020513f

  Copy full SHA for 020513f

• Move logic about when to dispatch to batched LU decomposition algorithm on GPU into the kernel.

  Show description for ad1bd38
  

  dfm

  authored and

  jax authors

  committed
  ad1bd38

  Copy full SHA for ad1bd38

• Integrate LLVM at llvm/llvm-project@c8b5d30f7077

  Show description for 807dcb5
  

  Google-ML-Automation

  authored and

  jax authors

  committed
  807dcb5

  Copy full SHA for 807dcb5

Commits on Aug 13, 2024

• [Mosaic:TPU] Refactor relayout helper functions to take ctx instead of only target shape.

  Show description for 8f23392
  

  bythew3i

  authored and

  jax authors

  committed
  8f23392

  Copy full SHA for 8f23392

• [Mosaic:TPU] Add shuffled load and store.

  Show description for 2dea3d6
  

  bythew3i

  authored and

  jax authors

  committed
  2dea3d6

  Copy full SHA for 2dea3d6

• Activate Singular Value Decomposition to XLA's FFI

  Show description for 354293d
  

  Paweł Paruzel

  authored and

  jax authors

  committed
  354293d

  Copy full SHA for 354293d

• Determine LAPACK workspaces during SVD kernel runtime

  Show description for 5fc992e
  

  Paweł Paruzel

  authored and

  jax authors

  committed
  5fc992e

  Copy full SHA for 5fc992e

Commits on Aug 12, 2024

• Add support for shape polymorphism with lu_pivots_to_permutation.

  Show description for 3c014a4
  

  dfm

  authored and

  jax authors

  committed
  3c014a4

  Copy full SHA for 3c014a4

• Merge pull request #22667 from ROCm:rocm-jax-triton-add-get_arch_detail

  Show description for be4d52b
  

  jax authors

  committed
  be4d52b

  Copy full SHA for be4d52b

• [ROCm]: Add get_arch_details for triton kernel call

  

  Rahul Batra

  committed
  4b7c198

  Copy full SHA for 4b7c198

Commits on Aug 9, 2024

• [Mosaic TPU][Python] Check validity of VectorLayout on init

  Show description for 77afe25
  

  tlongeri

  authored and

  jax authors

  committed
  77afe25

  Copy full SHA for 77afe25

• [Mosaic] Column shift relayouts for non-native tilings and packed types, except for (1, n) and packed

  Show description for e57a7e3
  

  tlongeri

  authored and

  jax authors

  committed
  e57a7e3

  Copy full SHA for e57a7e3

Commits on Aug 8, 2024

• Deprecate using build_cuda_plugin_from_source flag and rely on jaxlib_build config.

  Show description for 751b574
  

  Jieying Luo

  authored and

  jax authors

  committed
  751b574

  Copy full SHA for 751b574

• Update GPU implementation of `lu_pivots_to_permutation` to infer the permutation size directly from the input dimensions, instead of using an input parameter.

  Show description for 11d9c2d
  

  dfm

  authored and

  jax authors

  committed
  11d9c2d

  Copy full SHA for 11d9c2d

• [Mosaic TPU] Improve an error message in case someone tries to extract a non-32-bit scalar.

  Show description for 04a753a
  

  apaszke

  authored and

  jax authors

  committed
  04a753a

  Copy full SHA for 04a753a

• [Mosaic TPU] Add support for removal of implicit 2nd minor for all 32-bit tilings

  Show description for 42fe45f
  

  apaszke

  authored and

  jax authors

  committed
  42fe45f

  Copy full SHA for 42fe45f

Commits on Aug 7, 2024

• Merge pull request #22909 from ROCm:ci_fix_solver_paths

  Show description for de02988
  

  jax authors

  committed
  de02988

  Copy full SHA for de02988

• [ROCM] Fix BUILD.bazel library source paths

  Ruturaj4

  committed
  a2d7993

  Copy full SHA for a2d7993

• Added another boxDim check to mosaic_gpu_init_tma_desc

  Show description for 28ca734
  

  superbobry

  authored and

  jax authors

  committed
  28ca734

  Copy full SHA for 28ca734

Commits on Aug 6, 2024

• Merge pull request #22901 from ROCm:ci_test_harness_vmap

  Show description for cc96657
  

  jax authors

  committed
  cc96657

  Copy full SHA for cc96657

• [ROCM] Fix hipsolverSsyevd tests due to align with the rocm behavior.

  Ruturaj4

  committed
  707cdd4

  Copy full SHA for 707cdd4

• Re-land FFI port of GPU LU decomposition after fixing XLA FFI memory leak.

  Show description for 23da11b
  

  dfm

  authored and

  jax authors

  committed
  23da11b

  Copy full SHA for 23da11b

Commits on Aug 5, 2024

• Port Eigenvalue Decompositions to XLA's FFI

  Show description for b2a469b
  

  Paweł Paruzel

  authored and

  jax authors

  committed
  b2a469b

  Copy full SHA for b2a469b

Commits on Aug 4, 2024

• Reverts 80560663d3fab4c0c3f87d7c8e52fb9931526dbb

  Show description for 56ff247
  

  John Ryan

  authored and

  jax authors

  committed
  56ff247

  Copy full SHA for 56ff247

Commits on Aug 2, 2024

• [Mosaic TPU] Add support for bf16 reductions

  Show description for f85b8e6
  

  apaszke

  authored and

  jax authors

  committed
  f85b8e6

  Copy full SHA for f85b8e6

• [Mosaic TPU] Add a missing reshape in relayout

  Show description for e88887e
  

  apaszke

  authored and

  jax authors

  committed
  e88887e

  Copy full SHA for e88887e

• [Mosaic TPU] Remove special handling of implicit dim in relayout

  Show description for 959657a
  

  apaszke

  authored and

  jax authors

  committed
  959657a

  Copy full SHA for 959657a

• Enable FFI implementation of GPU Getrf FFI handler.

  Show description for 8056066
  

  dfm

  authored and

  jax authors

  committed
  8056066

  Copy full SHA for 8056066

• [Mosaic TPU] Break out implicit dim changes from relayout

  Show description for 99625ff
  

  apaszke

  authored and

  jax authors

  committed
  99625ff

  Copy full SHA for 99625ff