xla
xla copied to clipboard
openxla
A machine learning compiler for GPUs, CPUs, and ML accelerators
PR https://github.com/openxla/xla/pull/17722 supports limiting the CPU ISA that LLVM will codegen. It currently only supports x86 ISAs from SSE4_2 up to AMX_FP16 and imposes a strict ordering (SSE4_2 < AVX...
Add more elementwise ops with bf16 support on Hopper. There are a few more elementwise ops where we don't need to convert to f32 when running on Hopper.
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
PR #17636: [NVIDIA GPU] Enhance concurrency handling in cross-rank address sharing Imported from GitHub PR https://github.com/openxla/xla/pull/17636 This is a followup PR to https://github.com/openxla/xla/pull/15144. A distributed cache is maintained when device...
RIght now, users can't reuse a jitted method that includes scheduling ids more than once in a JAX program. Here is a very stripped down code JAX example that showcases...
Major deps udpate: -Protobuf 5.28.3 -Grpc 1.68.2 -Abseil: LTS 20240116.3 -Plus some other transitive/related deps, riegeli and pybind11 in particular. -rules_python & rules_cc will be updated in a subsequent CL...
Add int2 dtypes to TensorFlow. These types are used by JAX / XLA. This provides coverage for TensorFlow for JAX / XLA dependency.
[IFRT] Define Layout and common subclasses `xla::ifrt::Layout` is a first-class type in IFRT describing how elements of an array shard are laid out in memory. Similar to existing `xla::ifrt::DType`, `xla::ifrt::Shape`,...
Make error messages from dtype conversion more readable
[OpenXLA] Fix color of JAX logo in OpenXLA diagram
Add sharding devices to XlaCompileOptions and plumb them through from JAX. This is necessary to support MPMD parallelism in McJAX, since the PjRt-IFRT executable's output shardings can no longer be...