jax
Rocm jaxlib v0.9.0
#671
Open
Overview
Commits
6587
Changes
View On GitHub
Commits
  • Add jax.experimental.random to the wheel build
    jakevdp committed 167 days ago
  • fix shard_map transpose explicit sharding zero unsharding bug
    mattjj committed 167 days ago
  • add optional `explain` callback for weakref_lru_cache misses
    Google-ML-Automation committed 167 days ago
  • Merge pull request #34209 from jakevdp:fix-wheel
    Google-ML-Automation committed 167 days ago
  • Merge pull request #34211 from mattjj:andy-customvjp-none
    Google-ML-Automation committed 167 days ago
  • Handle ad.Zero cotangents in _reshard_transpose_fancy.
    Google-ML-Automation committed 167 days ago
  • Update XLA dependency to use revision http://github.com/openxla/xla/commit/76db112da7c2c66afeb550fc1089e6bec297bd4d
    Google-ML-Automation committed 167 days ago
  • [Pallas/Mosaic GPU] Enable more `WarpSpecializedPipelineWGTest`s.
    Google-ML-Automation committed 166 days ago
  • [Mosaic GPU] Add a `bitwidth` field to `Relayout` constraints in layout inference.
    Google-ML-Automation committed 166 days ago
  • [pallas:sc] Allowed specifying tiling in `pltpu.emit_pipeline`
    Google-ML-Automation committed 166 days ago
  • [XLA:MGPU] Port Tiling to C++.
    Google-ML-Automation committed 166 days ago
  • deviceless aot test
    keshavb96 committed 166 days ago
  • Add a thread guard config option.
    Google-ML-Automation committed 166 days ago
  • [Mosaic] Move Float8EXMYType to tpu.td.
    Google-ML-Automation committed 166 days ago
  • Merge pull request #33542 from keshavb96:deviceless_aot_test
    Google-ML-Automation committed 166 days ago
  • Remove redundant test targets that are already executed as a part of `//tests/pallas:tpu_tests`.
    Google-ML-Automation committed 166 days ago
  • respect self.statics in FlatTree.__eq__
    mattjj committed 166 days ago
  • Fix precommit breakage
    Google-ML-Automation committed 166 days ago
  • sick
    mattjj committed 166 days ago
  • revive as many cache miss explanations as reasonably possible
    mattjj committed 166 days ago
  • [pallas:sc] Skip a few tests failing when the compiler uses tiled memrefs
    Google-ML-Automation committed 166 days ago
  • Colocated python: Use a wrapper when storing remote objects at the backend.
    Google-ML-Automation committed 166 days ago
  • skip on jaxlib version
    mattjj committed 166 days ago
  • Add `halt-for-connection` to `build_artifacts.yml` workflow call.
    Google-ML-Automation committed 166 days ago
  • Merge pull request #33839 from jax-ml:pjit-without-linear-util
    Google-ML-Automation committed 166 days ago
  • Update XLA dependency to use revision http://github.com/openxla/xla/commit/9ae3d6dab2c10c8195c8d9862f475904c7cdca91
    Google-ML-Automation committed 166 days ago
  • [Mosaic GPU] Use `isinstance(x, mlir_ty)` instead of the deprecated `mlir_ty.isinstance(x)`.
    Google-ML-Automation committed 165 days ago
  • Remove unnecessary return from placement new in Mosaic GPU extension.
    Google-ML-Automation committed 165 days ago
  • [mosaic] Added a canonicalization rule for memref.dim of tpu.memref_squeeze
    Google-ML-Automation committed 165 days ago
  • [Mosaic GPU] Add basic support for DSMEM
    Google-ML-Automation committed 165 days ago
  • [Mosaic GPU] Add support for using the redux instructions to speed up warp reductions
    Google-ML-Automation committed 165 days ago
  • [Pallas/Mosaic GPU] Disable `test_ragged_dot_transposed` temporarily.
    Google-ML-Automation committed 165 days ago
  • Remove some more stale version guards.
    Google-ML-Automation committed 165 days ago
  • [bug] fix grad-of-vmap-of-dynamic_slice with out-of-bound indices
    jakevdp committed 165 days ago
  • Allow max_size to be None (infinite).
    Google-ML-Automation committed 165 days ago
  • Don't pass axis_index_groups to prim.bind in psum batching rule if primitive is psum_invariant_p. Fixes https://github.com/jax-ml/jax/issues/34249
    Google-ML-Automation committed 165 days ago
  • Migrates `builder.create<Op>()` => `Op::create()`
    Google-ML-Automation committed 165 days ago
  • Merge pull request #34229 from jakevdp:dynamic-slice-grad-fix
    Google-ML-Automation committed 165 days ago
  • Fix Array `__format__` and `__str__` to handle McJAX arrays that are non-addressable by some processes.
    Google-ML-Automation committed 165 days ago
  • [hijax] start implementing custom_vjp on top of hijax primitives
    mattjj committed 165 days ago
  • Merge pull request #34161 from mattjj:custom-vjp3-youandme
    Google-ML-Automation committed 165 days ago
  • Delete the jax_collectives_common_channel_id flag.
    Google-ML-Automation committed 165 days ago
  • [indexing] implement strategy='dynamic_slice'
    jakevdp committed 165 days ago
  • Merge pull request #34225 from jakevdp:index-dynamic-slice
    Google-ML-Automation committed 165 days ago
  • [TPU] Reenable a disabled test that now passes.
    Google-ML-Automation committed 165 days ago
  • [Pallas MGPU] Use `the cp.async.bulk` instruction for large contiguous copies. Currently we’re limited to 256 elements per dimension when using the tensormap in `cp.async.bulk.tensor`.
    Google-ML-Automation committed 165 days ago
  • [Mosaic GPU] Expose the `mbarrier.complete_tx` instruction to manually decrement the transaction count of a barrier.
    Google-ML-Automation committed 165 days ago
  • Shorten splash attention kernel name to address https://github.com/jax-ml/jax/issues/34117
    Google-ML-Automation committed 165 days ago
  • [shmap] fix shmap error logic when subclasses of pspec are used
    mattjj committed 165 days ago
  • Run the GetKeys inside GetOrCreate.
    Google-ML-Automation committed 165 days ago
  • Merge pull request #34262 from mattjj:shmap-error-isleaf
    Google-ML-Automation committed 165 days ago
  • Run g4 fix on the weakref_lru_cache code.
    Google-ML-Automation committed 165 days ago
  • Update rules_ml_toolchain version to accommodate custom redistribution URLs and templates.
    Google-ML-Automation committed 165 days ago
  • Fix deserialization with specified layouts via a ShapeDtypeStruct being ignored.
    Google-ML-Automation committed 165 days ago
  • Remove obsolete filegroup
    Google-ML-Automation committed 165 days ago
  • Update XLA dependency to use revision http://github.com/openxla/xla/commit/36b00de8b53773e7cc5bdd44fe3c140e3ddafa43
    Google-ML-Automation committed 165 days ago
  • don't warn on complex->real cast in dot transpose
    mattjj committed 164 days ago
  • Merge pull request #33708 from mattjj:issue33521
    Google-ML-Automation committed 164 days ago
  • [hijax] add vmap suport to CustomVJP hijax primitive
    mattjj committed 164 days ago
  • Merge pull request #34268 from mattjj:custom-vjp3-youandme-2
    Google-ML-Automation committed 164 days ago
  • Update XLA dependency to use revision http://github.com/openxla/xla/commit/01da52b9afe3a2e694bd926323a649a6e63a3785
    Google-ML-Automation committed 164 days ago
  • [Mosaic] Add abs, sign, erf, atan2, reduce_min, reduce_prod support to mgpu
    oulgen committed 163 days ago
  • [JAX] Add IFRT SerDes to jaxlib deps
    Google-ML-Automation committed 163 days ago
  • Update XLA dependency to use revision http://github.com/openxla/xla/commit/93f1bc2b503855630ef8bb07981850d84ffdce8e
    Google-ML-Automation committed 163 days ago
  • Merge pull request #34120 from oulgen:mgpu-ops
    Google-ML-Automation committed 162 days ago
  • [Mosaic GPU] Test approximate math functions properly + eta reduce args
    Google-ML-Automation committed 162 days ago
  • [Mosaic GPU] Enable redux.sync.f32 on Blackwell
    Google-ML-Automation committed 162 days ago
  • [XLA:MGPU] Port Replicated wrapper to C++.
    Google-ML-Automation committed 162 days ago
  • [mosaic] infer-memref-layout now accepts target shape as a span
    Google-ML-Automation committed 162 days ago
  • [Pallas:MGPU] Lower `lax.sign` consistently for LANE and WG semantics.
    Google-ML-Automation committed 162 days ago
  • [Pallas:MGPU] Add Pallas lowering for `pl.debug_check` under WG semantics.
    Google-ML-Automation committed 162 days ago
  • Add `thread_guard` to the public API.
    Google-ML-Automation committed 162 days ago
  • [Mosaic GPU] Add support for f8 types for WGMMA with lhs in registers under WG semantics.
    Google-ML-Automation committed 162 days ago
  • [Pallas MGPU] Clip the size of contiguous TMA transfers to the size of the ref. This avoids OOB accesses and is more in-line with the tensor map TMAs.
    Google-ML-Automation committed 162 days ago
  • lax: ensure padtype_to_pads returns Python ints
    Prakharprasun committed 162 days ago
  • [jax.collect_profile] Allow arbitrary options to be passed to XProf
    Google-ML-Automation committed 162 days ago
  • [indexing] support newaxis in static/dynamic slice strategies
    jakevdp committed 162 days ago
  • Plumb prim params for call discharge rule (to handle named_computation_p
    Google-ML-Automation committed 162 days ago
  • [sc] Generalized infer-memref-layout to support SC tiling
    Google-ML-Automation committed 162 days ago
  • [Pallas/TPU] Don't lower eqns that have all dropvar outputs (aka DCE at
    Google-ML-Automation committed 162 days ago
  • [sc] Removed `infer_kernel_arguments` from infer-memref-layout
    Google-ML-Automation committed 162 days ago
  • Merge pull request #33974 from Prakharprasun:fix-padtype-to-pads
    Google-ML-Automation committed 162 days ago
  • [Pallas] Allowlist semaphore/prng effects under remat and custom
    Google-ML-Automation committed 162 days ago
  • Add lax.tile_p
    Google-ML-Automation committed 162 days ago
  • Automated Code Change
    Google-ML-Automation committed 162 days ago
  • [Mosaic:TPU] Clean up tpu.memref_slice verifier
    Google-ML-Automation committed 162 days ago
  • Update XLA dependency to use revision http://github.com/openxla/xla/commit/9258aaa31683b9f5e5a1cc21c6c2a3c769df4bfb
    Google-ML-Automation committed 162 days ago
  • Add pallas lowering for jnp.tile using tpu.repeat.
    Google-ML-Automation committed 161 days ago
  • [export] Add support for explicit sharding.
    gnecula committed 161 days ago
  • Handle the case when shardings are GSPMDSharding.
    gnecula committed 161 days ago
  • [Mosaic GPU] Support more reduction kinds and layouts in `MultiDimReductionOp` lowering.
    Google-ML-Automation committed 161 days ago
  • [Pallas/interpreter] Add a prototype for a GPU kernel interpreter.
    Google-ML-Automation committed 161 days ago
  • [Pallas:MGPU][NFC] Update outdated docstring for `scratch_view`.
    Google-ML-Automation committed 161 days ago
  • [Mosaic GPU] Fix issue in `vector_dim` reduction with `vec_len > 2`.
    Google-ML-Automation committed 161 days ago
  • [Mosaic GPU][NFC] Replace usages of `_gpu_ops_gen` with `gpu` dialect whenever possible.
    Google-ML-Automation committed 161 days ago
  • [Mosaic GPU] Fix the return type of `_lift_fast_packed_instr` to match its argument types.
    Google-ML-Automation committed 161 days ago
  • Set the proper memory kinds for output in call_exported
    gnecula committed 161 days ago
  • Merge pull request #33597 from gnecula:export_shit
    Google-ML-Automation committed 161 days ago
  • Merge pull request #34290 from jakevdp:indexing-none
    Google-ML-Automation committed 161 days ago
  • [Pallas:MGPU] Remove broadcast in scalar reduce test.
    Google-ML-Automation committed 161 days ago
  • + more commits ...
FAQ Terms Privacy Refunds Impressum
Loading