PR #187 CI: 12/19/24 upstream sync

Limit self-hosted jobs to JAX main repo

MichaelHudgins committed 1 year ago

Merge branch 'jax-ml:main' into main

MichaelHudgins committed 1 year ago

temporarily un-deprecate several jax.core APIs.

jakevdp committed 1 year ago

Remove deprecated XLA GPU flags.

justinjfu committed 1 year ago

Add flag desc to gpu_performance_tips.md

sfvaroglu committed 1 year ago

Allow tracing and lowering (with lowering_platforms specified) to work with an AbstractMesh. Such a computation cannot be compiled.

Google-ML-Automation committed 1 year ago

Remove autotune sharing.

Google-ML-Automation committed 1 year ago

Use a broadcasted gather in the sort JVP, rather than forming explicit iotas.

hawkinsp committed 1 year ago

Merge pull request #25454 from sfvaroglu:sevin/slop_factor_doc

Google-ML-Automation committed 1 year ago

Merge pull request #25459 from hawkinsp:sort

Google-ML-Automation committed 1 year ago

get_githash: fix support for missing git

adamjstewart committed 1 year ago

Migrate JAX MLIR Python dialect extensions to nanobind.

Google-ML-Automation committed 1 year ago

Set device_assignment to None when only AbstractMesh exist in the computation

Google-ML-Automation committed 1 year ago

Merge pull request #25451 from jakevdp:undep-core

Google-ML-Automation committed 1 year ago

Support FP8 for dot_product_attention

wenscarl committed 1 year ago

Merge pull request #25444 from justinjfu:gpu_docs_update

Google-ML-Automation committed 1 year ago

Merge branch 'jax-ml:main' into main

MichaelHudgins committed 1 year ago

Merge pull request #24536 from MichaelHudgins:main

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Implement `process_call` and `process_custom_jvp` for LinearizeTrace

dougalm committed 1 year ago

Finalize deprecation of jnp.round_

Google-ML-Automation committed 1 year ago

fix `gamma_p` in vmap-based impl rule mode

froystig committed 1 year ago

Merge pull request #25481 from jax-ml:custom-linearize-process-call

Google-ML-Automation committed 1 year ago

Add an experimental Cloud TPU presubmit job

Google-ML-Automation committed 1 year ago

Merge pull request #25487 from froystig:gamma-impl-vmap-fix

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Fix layout inference traversal to traverse ops recursively.

Google-ML-Automation committed 1 year ago

Log DeprecationWarnings once per method/class

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

[Mosaic:TPU] Replicating retilings with increasing tile size for (a) replicated 2nd minor or (b) 32-bit single-row

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

[export] Expand exporting to work with AbstractMesh.

gnecula committed 1 year ago

[mosaic_gpu] Allow calling `reduce_sum` on a fragmented array in splat layout

Google-ML-Automation committed 1 year ago

Drop the frequency of Cloud TPU tests

Google-ML-Automation committed 1 year ago

Ignore a false-positive Hypothesis health check

Google-ML-Automation committed 1 year ago

Bump the oldest supported libtpu to conform to the 12 week window

Google-ML-Automation committed 1 year ago

Merge pull request #22670 from wenscarl:fp8_sdpa

Google-ML-Automation committed 1 year ago

Add Python.h include to fix Windows build.

Google-ML-Automation committed 1 year ago

Avoid assuming that jnp.sin will be traced in abstract mesh tests

Google-ML-Automation committed 1 year ago

Merge pull request #25435 from gnecula:export_abs

Google-ML-Automation committed 1 year ago

Fix some flaky LAX autodiff tests.

Google-ML-Automation committed 1 year ago

Raise the timeout for Cloud TPU nightly CI

Google-ML-Automation committed 1 year ago

Relax tolerance for LAX reduction test in float16.

Google-ML-Automation committed 1 year ago

Minor typo fixin doc.

Google-ML-Automation committed 1 year ago

Fix incorrect capitalization in scan error message

jakevdp committed 1 year ago

Merge pull request #25511 from jakevdp:fix-scan-err

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Add utility script and env for running the CI scripts under Docker

Google-ML-Automation committed 1 year ago

Re-deprecate a number of symbols from jax.core

jakevdp committed 1 year ago

Merge pull request #25508 from jakevdp:dep-core

Google-ML-Automation committed 1 year ago

Fix `save_from_both_policies` in presence of `save_and_offload_only_these_names` by comparing the enum

Google-ML-Automation committed 1 year ago

[Pallas TPU] Use vector.broadcast instead of vector.BroadcastOp to fix type check failure

Google-ML-Automation committed 1 year ago

Allow `lax.ragged_all_to_all` input and output operands to have different ragged dimension sizes.

Google-ML-Automation committed 1 year ago

[Mosaic:TPU][NFC] Clean up local variable

Google-ML-Automation committed 1 year ago

[Mosaic:TPU][NFC] use_implicit_shape instead of Reshape + don't create unused constant

Google-ML-Automation committed 1 year ago

[Mosaic GPU] Add layout inference and initial lowering for `vector.{load,store}`.

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Add support for bf16 abs

Google-ML-Automation committed 1 year ago

Make gmm TPU kernel tests significantly cheaper

Google-ML-Automation committed 1 year ago

Reduce the max_examples for splash attention tests

Google-ML-Automation committed 1 year ago

Put abstract_mesh on every eqn so that we can preserve it during `eval_jaxpr` and `check_jaxpr` roundtrip.

Google-ML-Automation committed 1 year ago

Cleanup: toward merging core.concrete_aval & xla.abstractify

jakevdp committed 1 year ago

[Mosaic GPU] Add end-to-end lowering example for a pointwise kernel using the dialect and layout inference.

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Add support for the interleaved pack format to tpu.unpack_subelements

Google-ML-Automation committed 1 year ago

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

Google-ML-Automation committed 1 year ago

Merge pull request #25473 from adamjstewart:build/no-git

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Merge pull request #25456 from jakevdp:xla-abstractify

Google-ML-Automation committed 1 year ago

Streamline some core.concrete_aval compute paths

jakevdp committed 1 year ago

Raise rather than return error

jakevdp committed 1 year ago

Release v0.4.38

hawkinsp committed 1 year ago

Merge pull request #25538 from jakevdp:raise-runtimeerror

Google-ML-Automation committed 1 year ago

Enable PJRT compatibility in cloud TPU CI.

Google-ML-Automation committed 1 year ago

Merge pull request #25534 from jakevdp:faster-avals

Google-ML-Automation committed 1 year ago

Always use the same code for array avals

jakevdp committed 1 year ago

Merge pull request #25544 from jakevdp:array-aval

Google-ML-Automation committed 1 year ago

Merge branch 'release/0.4.38'

hawkinsp committed 1 year ago

Update version numbers after release.

hawkinsp committed 1 year ago

[Mosaic GPU] Skip testing uint64 unless 64-bit types are enabled

andportnoy committed 1 year ago

Add Pallas Philox implementation.

Google-ML-Automation committed 1 year ago

Merge pull request #25550 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

Don't create a mesh inside jit in sharded-computation tutorial

Google-ML-Automation committed 1 year ago

Merge pull request #25548 from andportnoy:aportnoy/mosaic-gpu-scalar-argument-64-bit-fix

Google-ML-Automation committed 1 year ago

fix squeeze lowering required sharding argument

shangz-ai committed 1 year ago

Delete enable_memories code in C++ since that flag is always True and cannot be turned off now.

Google-ML-Automation committed 1 year ago

Fix the breakage caused by deleted enable_memories config

Google-ML-Automation committed 1 year ago

Remove core.concrete_aval and replace with abstractify

jakevdp committed 1 year ago

Merge pull request #25555 from jakevdp:combine-avals

Google-ML-Automation committed 1 year ago

Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is required for sharding in types to work properly when out_avals contain UNCONSTRAINED specs.

Google-ML-Automation committed 1 year ago

[sharding_in_types] Add tests allowing inputs and outputs of jit to have different axis_types on their mesh than the axis_types on the surrounding mesh context

Google-ML-Automation committed 1 year ago

improve checkpoint / remat concreteness error with static_argnums

mattjj committed 1 year ago

[Mosaic:TPU] Fix bug after cl/707025084

Google-ML-Automation committed 1 year ago

Merge pull request #24516 from mattjj:improve-concreteness-error-in-remat-3

Google-ML-Automation committed 1 year ago

ref errors

mattjj committed 1 year ago

[Mosaic:TPU] Allow null parts for tpu.pack_subelements, meaning "don't care"

Google-ML-Automation committed 1 year ago

Merge pull request #25449 from mattjj:ref-errors-3

Google-ML-Automation committed 1 year ago

Reverts b56dc63160eaccd7df05d03b1c38f804ff85f564

Google-ML-Automation committed 1 year ago

Bump xla_extension_version after jaxlib release.

hawkinsp committed 1 year ago

Increase the minimum NumPy version to v1.25.

hawkinsp committed 1 year ago

Use StableHLO acos and update complex acos accuracy tests.

pearu committed 1 year ago

Merge pull request #25559 from hawkinsp:minver

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Add support for exp, exp2 and log in bf16 on TPUv6

Google-ML-Automation committed 1 year ago

jax CI: 12/19/24 upstream sync #187 Merged

jax
CI: 12/19/24 upstream sync
#187

Merged