jax
CI: 12/23/24 upstream sync
#190
Merged
Overview
Commits
52
Changes
View On GitHub
Commits
  • [ROCm] ci build and dockerfile changes
    Ruturaj4 committed 1 year ago
  • Remove internal uses of api_util.shaped_abstractify
    jakevdp committed 1 year ago
  • Add support for N-D FFTs with D>3.
    dfm committed 1 year ago
  • Update documentation and add setup.py pypi bindings
    Ruturaj4 committed 1 year ago
  • Add experimental support for building JAX CPU and GPU wheels with GCC.
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25148 from ROCm:ci_install_sys_libs-upstream
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25440 from ROCm:jax_pypi-upstream
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25606 from dfm:fft-nd
    Google-ML-Automation committed 1 year ago
  • Update XLA dependency to use revision
    Google-ML-Automation committed 1 year ago
  • DOC: clarify API compatibility discussion & mention EffVer
    jakevdp committed 1 year ago
  • Merge pull request #25614 from jakevdp:dep-shaped-abstractify
    Google-ML-Automation committed 1 year ago
  • Migrate shardy dialect extension to nanobind.
    Google-ML-Automation committed 1 year ago
  • add test for partial-auto ppermute
    Google-ML-Automation committed 1 year ago
  • [Mosaic GPU] Prototype of a warp-specialized pipeline emitter for Mosaic GPU.
    Google-ML-Automation committed 1 year ago
  • [Mosaic:TPU] For trunc, expand supported tilings, offsets and bitwidths
    Google-ML-Automation committed 1 year ago
  • [Mosaic TPU] Add support for sqrt and rsqrt in bf16 on TPUv6
    Google-ML-Automation committed 1 year ago
  • Ensure that the two offsets of a dynamic_slice have the same dtype regardless
    Google-ML-Automation committed 1 year ago
  • Use common set of build options when building jaxlib+plugin artifacts together
    Google-ML-Automation committed 1 year ago
  • [Mosaic] Remove TODOs that are already addressed or obsolete.
    Google-ML-Automation committed 1 year ago
  • [Mosaic] Extend macros to handle parentheses.
    Google-ML-Automation committed 1 year ago
  • Add more input validation to jax.distributed.initialize.
    emilyfertig committed 1 year ago
  • Delete unused internal symbols
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25489 from emilyfertig:distributed-init-input-validation
    Google-ML-Automation committed 1 year ago
  • add mutable array ref error checks to cond and custom_vjp
    mattjj committed 1 year ago
  • [Mosaic:TPU] Fix trunc infer rule after cl/708011538
    Google-ML-Automation committed 1 year ago
  • [Mosaic TPU] Add relayout-insertion pass and support bitwidth change for i1 vector relayout
    Google-ML-Automation committed 1 year ago
  • [Mosaic:TPU] Roll back cl/708011538 and cl/708112341
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25625 from mattjj:ref-errors-5
    Google-ML-Automation committed 1 year ago
  • Relax some test tolerances in N-D FFT tests.
    Google-ML-Automation committed 1 year ago
  • [pallas:mosaic_gpu] Change the fori tests to also take the while_p path and fix the bug.
    Google-ML-Automation committed 1 year ago
  • Unify abstractify & shaped_abstractify rules
    jakevdp committed 1 year ago
  • Make jax.Arrays a necessary part of the cycle in the GC guard test
    Google-ML-Automation committed 1 year ago
  • Clarify documentation for output_offsets operand of ragged_all_to_all.
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25616 from jakevdp:abstractify
    Google-ML-Automation committed 1 year ago
  • [Mosaic GPU] Add a lowering for simple `async_load` and `async_store` ops.
    Google-ML-Automation committed 1 year ago
  • Update XLA dependency to use revision
    Google-ML-Automation committed 1 year ago
  • Migrate _mlir Python binding target to nanobind.
    Google-ML-Automation committed 1 year ago
  • [Mosaic:TPU] Roll forward of cl/708011538 (expanded trunc support), minus changes in infer-vector-layout
    Google-ML-Automation committed 1 year ago
  • Avoid calls to warnings.catch_warnings in JAX core code.
    hawkinsp committed 1 year ago
  • Add int4/uint4 support to bitcast_convert_type
    jakevdp committed 1 year ago
  • Merge pull request #25646 from jakevdp:int4-casting
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25648 from hawkinsp:warnings3
    Google-ML-Automation committed 1 year ago
  • Support int4/uint4 in jnp.ndarray.view
    jakevdp committed 1 year ago
  • Merge pull request #25416 from jakevdp:api-compat
    Google-ML-Automation committed 1 year ago
  • Accelerate deprecation of legacy JAX FFI calling convention.
    Google-ML-Automation committed 1 year ago
  • Change the namespace name to avoid using `export` c++ keyword on namespace.
    Google-ML-Automation committed 1 year ago
  • [JAX] Add a test using inputs with different device orders for a single colocated Python call
    Google-ML-Automation committed 1 year ago
  • Move output path to be inside the wheel build command execution loop
    Google-ML-Automation committed 1 year ago
  • Merge pull request #25650 from jakevdp:view-int4
    Google-ML-Automation committed 1 year ago
  • Update XLA dependency to use revision
    Google-ML-Automation committed 1 year ago
  • [Jax][Pallas][Mosaic] Implement platform dependent diag, with branch selection driven by constant prop in mosaic lowering.
    Google-ML-Automation committed 1 year ago
  • Move ragged tests under a new class.
    Google-ML-Automation committed 1 year ago
FAQ Terms Privacy Refunds Impressum
Loading