Commits

Branch selector

 main

User selector

yashk2810

All time

Commit History

Commits on Dec 14, 2023

• [roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md.

  Show description for 8bf3a86
  

  yashk2810

  authored and

  jax authors

  committed
  8bf3a86

  Copy full SHA for 8bf3a86

• Finish release of jax and jaxlib 0.4.23

  Show description for 6e1ab7c
  

  yashk2810

  authored and

  jax authors

  committed
  6e1ab7c

  Copy full SHA for 6e1ab7c

• Prepare for emergency jax and jaxlib 0.4.23 release

  Show description for 1c68577
  

  yashk2810

  authored and

  jax authors

  committed
  1c68577

  Copy full SHA for 1c68577

• Finish jax and jaxlib 0.4.22 release

  Show description for 25c16c0
  

  yashk2810

  authored and

  jax authors

  committed
  25c16c0

  Copy full SHA for 25c16c0

Commits on Dec 13, 2023

• Reverts 3c07c10a9a55f9a32390dd10cf3f420bdf3f1ed8

  Show description for b52bcc1
  

  yashk2810

  authored and

  jax authors

  committed
  b52bcc1

  Copy full SHA for b52bcc1

• Fix the pjit flakey test. The test was weirdly written in the first place. The current suspicion is that `x.copy()` inside jit made it flakey. Also delete some duplicate tests from the time of migr…

  Show description for e888806
  

  yashk2810

  authored and

  jax authors

  committed
  e888806

  Copy full SHA for e888806

• Start jax and jaxlib 0.4.22 release

  Show description for 2e63352
  

  yashk2810

  authored and

  jax authors

  committed
  2e63352

  Copy full SHA for 2e63352

• Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md.

  Show description for 3c07c10
  

  yashk2810

  authored and

  jax authors

  committed
  3c07c10

  Copy full SHA for 3c07c10

Commits on Dec 12, 2023

• Add a flag `jax_require_devices_during_lowering` to control if physical devices are passed during lowering to stablehlo. This is temporary to unblock nvidia.

  Show description for f210b0f
  

  yashk2810

  authored and

  jax authors

  committed
  f210b0f

  Copy full SHA for f210b0f

Commits on Dec 9, 2023

• Add a registry for primitives that require device_assignment during lowering

  Show description for 10f6a35
  

  yashk2810

  authored and

  jax authors

  committed
  10f6a35

  Copy full SHA for 10f6a35

Commits on Dec 8, 2023

• Make lowering oblivious to real physical devices. Instead cache lowering on HloSharding only (which is based on logical device numbers)

  Show description for 5fb8cec
  

  yashk2810

  authored and

  jax authors

  committed
  5fb8cec

  Copy full SHA for 5fb8cec

Commits on Dec 7, 2023

• Internal change

  Show description for 78e0e6b
  

  yashk2810

  authored and

  jax authors

  committed
  78e0e6b

  Copy full SHA for 78e0e6b

• Make `GetParameterLayouts` return the correct thing when both CPU and TPU executables exist.

  Show description for 3e68ad1
  

  yashk2810

  authored and

  jax authors

  committed
  3e68ad1

  Copy full SHA for 3e68ad1

Commits on Dec 4, 2023

• Finish jax and jaxlib 0.4.21 release

  Show description for a9bfbd3
  

  yashk2810

  authored and

  jax authors

  committed
  a9bfbd3

  Copy full SHA for a9bfbd3

• Replace device_buffers with addressable_shards in test because device_buffers is deprecated

  Show description for c4d2fc7
  

  yashk2810

  authored and

  jax authors

  committed
  c4d2fc7

  Copy full SHA for c4d2fc7

• Prepare for 0.4.21 release

  Show description for 5942e15
  

  yashk2810

  authored and

  jax authors

  committed
  5942e15

  Copy full SHA for 5942e15

Commits on Dec 2, 2023

• Put the input creating via `jnp` outside the pjit cpp cache miss counting code.

  Show description for 0418370
  

  yashk2810

  authored and

  jax authors

  committed
  0418370

  Copy full SHA for 0418370

• Reverts f0382a5838f4526d21631e804f6fe576bfc3f97e

  Show description for f0bc7e0
  

  yashk2810

  authored and

  jax authors

  committed
  f0bc7e0

  Copy full SHA for f0bc7e0

Commits on Dec 1, 2023

• Add a test to check if arr.delete() is idempotent.

  Show description for 595117b
  

  yashk2810

  authored and

  jax authors

  committed
  595117b

  Copy full SHA for 595117b

Commits on Nov 30, 2023

• Replace apply_primitive internals with `jax.jit`.

  Show description for e624610
  

  yashk2810

  authored and

  jax authors

  committed
  e624610

  Copy full SHA for e624610

Commits on Nov 29, 2023

• Disable test_memory_cosumption test

  Show description for d6637da
  

  yashk2810

  authored and

  jax authors

  committed
  d6637da

  Copy full SHA for d6637da

Commits on Nov 28, 2023

• Use `trace_to_jaxpr_dynamic` for the apply_primitive path. `trace_to_jaxpr_final` is only for final style primitives. Also do some cleanup.

  Show description for cb7c2ed
  

  yashk2810

  authored and

  jax authors

  committed
  cb7c2ed

  Copy full SHA for cb7c2ed

• Remove jaxpr_debug_info from MeshComputation since that information is available via AllArgsInfo

  Show description for f8e22ae
  

  yashk2810

  authored and

  jax authors

  committed
  f8e22ae

  Copy full SHA for f8e22ae

• Typecheck avals and sharding for arguments that were DCE'd.

  Show description for 88d980f
  

  yashk2810

  authored and

  jax authors

  committed
  88d980f

  Copy full SHA for 88d980f

• Fix indexing bug when querying _input_layouts

  Show description for 37f1142
  

  yashk2810

  authored and

  jax authors

  committed
  37f1142

  Copy full SHA for 37f1142

• Simply lower_sharding_computation signature by always taking a closed jaxpr as input. For apply_primitive do the tracing to jaxpr in dispatch.py

  Show description for 81aee23
  

  yashk2810

  authored and

  jax authors

  committed
  81aee23

  Copy full SHA for 81aee23

• `compiled.input_layouts()` should preserve the structure of the original in_tree.

  Show description for 2ed0fc4
  

  2 peopleauthored and

  jax authors

  committed
  2ed0fc4

  Copy full SHA for 2ed0fc4

Commits on Nov 20, 2023

• Only call executale.get_output_memory_kinds() if jax_enable_memories is True

  Show description for d6a9352
  

  yashk2810

  authored and

  jax authors

  committed
  d6a9352

  Copy full SHA for d6a9352

Commits on Nov 18, 2023

• Make the SpecifiedLayout class opaque.

  Show description for c8ef375
  

  yashk2810

  authored and

  jax authors

  committed
  c8ef375

  Copy full SHA for c8ef375

• If a function returns no output, xla_executable.get_output_shardings() returns 1 sharding because for XLA the output is an empty tuple which has a tuple sharding.

  Show description for 493e2f8
  

  yashk2810

  authored and

  jax authors

  committed
  493e2f8

  Copy full SHA for 493e2f8

Commits on Nov 17, 2023

• Make looking up shardings from executable consistent. If `out_shardings` are specified on `jit`, always check it against the `get_output_shardings` from the executable.

  Show description for 3872955
  

  yashk2810

  authored and

  jax authors

  committed
  3872955

  Copy full SHA for 3872955

• Remove DefaultLayout and make `None` same as `DefaultLayout`

  Show description for 439b89e
  

  yashk2810

  authored and

  jax authors

  committed
  439b89e

  Copy full SHA for 439b89e

Commits on Nov 15, 2023

• Make the regex checking of export_tests less strict

  Show description for 118d85c
  

  yashk2810

  authored and

  jax authors

  committed
  118d85c

  Copy full SHA for 118d85c

• Add a private API to allow setting layouts on jitted computations.

  Show description for 5c3da21
  

  2 peopleauthored and

  jax authors

  committed
  5c3da21

  Copy full SHA for 5c3da21

Commits on Nov 9, 2023

• Disable jax memories flag.

  Show description for cf3c041
  

  yashk2810

  authored and

  jax authors

  committed
  cf3c041

  Copy full SHA for cf3c041

Pagination

Previous

Next