Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Apr 4, 2023

• DCE as early as possible so that committed is not dependent on DCE's vars

  Show description for ffa9d01
  

  yashk2810

  authored and

  jax authors

  committed
  ffa9d01

• Remove _compile_replicated option from compile since it is not needed anymore and some other cosmetic fixes.

  Show description for 14b572f
  

  yashk2810

  authored and

  jax authors

  committed
  14b572f

Commits on Apr 3, 2023

• Rename count_pjit_cache_miss with count_pjit_cpp_cache_miss because it is confusing which cache the first function is taking about as pjit has many caches

  Show description for 78678ee
  

  yashk2810

  authored and

  jax authors

  committed
  78678ee

• Improve the error message of device_indices_map when the sharding is not divisible by the shape rather than raising an opaque assertion error

  Show description for 6f2256a
  

  yashk2810

  authored and

  jax authors

  committed
  6f2256a

Commits on Apr 1, 2023

• Add Deprecation warning if gda_serialization is imported

  Show description for 2432ade
  

  yashk2810

  authored and

  jax authors

  committed
  2432ade

• Rename gda_serialization to array_serialization but keep gda_serialization around until it is included in a jax release so that OSS projects can be moved to array_serialization

  Show description for d27a80d
  

  yashk2810

  authored and

  jax authors

  committed
  d27a80d

Commits on Mar 31, 2023

• Remove dead code from pxla.py

  Show description for 0b31e8b
  

  yashk2810

  authored and

  jax authors

  committed
  0b31e8b

Commits on Mar 30, 2023

• Raise deprecation warnings for {in|out}_axis_resources for pjit and axis_resources for with_sharding_constraint

  Show description for 69c9660
  

  yashk2810

  authored and

  jax authors

  committed
  69c9660

• Relax the tolerance of testCauchyLogCdf

  Show description for 16ca0ca
  

  yashk2810

  authored and

  jax authors

  committed
  16ca0ca

• Delete unused functions from dispatch.py and pjit.py

  Show description for c978df5
  

  yashk2810

  authored and

  jax authors

  committed
  c978df5

Commits on Mar 29, 2023

• Delete _single_device_array_from_buf since everything from JAX is an Array

  Show description for 830cd9f
  

  yashk2810

  authored and

  jax authors

  committed
  830cd9f

• Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago

  Show description for fbc05ee
  

  yashk2810

  authored and

  jax authors

  committed
  fbc05ee

Commits on Mar 28, 2023

• Jax 0.4.7 has been released so assert that length of warnings is 1 in test_cache_read_warning

  Show description for 2f105bd
  

  yashk2810

  authored and

  jax authors

  committed
  2f105bd

• Remove MeshPspecSharding since it has been more than 3 months since it was deprecated (Nov 2, 2022). The replacement name is NamedSharding.

  Show description for 7442faa
  

  yashk2810

  authored and

  jax authors

  committed
  7442faa

• Deprecate FROM_GDA and remove its support from pjit's code since jax.Array inside pjit has sharding inference capabilities by default.

  Show description for 97c8ce3
  

  yashk2810

  authored and

  jax authors

  committed
  97c8ce3

• Remove Cuda 11.4 support. JAX from 0.4.8 release will support cuda 11.8 and cuda 12

  Show description for 86c0b36
  

  yashk2810

  authored and

  jax authors

  committed
  86c0b36

Commits on Mar 27, 2023

• Finish jax and jaxlib 0.4.7 release

  Show description for 670fba3
  

  yashk2810

  authored and

  jax authors

  committed
  670fba3

• Update the commit in workspace too

  Show description for ae4f1fc
  

  yashk2810

  authored and

  jax authors

  committed
  ae4f1fc

• Prepare for jax and jaxlib 0.4.7 release

  Show description for e9cac5e
  

  yashk2810

  authored and

  jax authors

  committed
  e9cac5e

• Add deprecation warning for FROM_GDA usage since that argument is not required anymore.

  Show description for e21aee1
  

  yashk2810

  authored and

  jax authors

  committed
  e21aee1

• Temporarily fix the compilation cache test which is failing on latest jaxlib release

  Show description for 41695cc
  

  yashk2810

  authored and

  jax authors

  committed
  41695cc

• Delete benchmark and pmap_benchmark files as they are legacy and replaced with api_benchmark.py

  Show description for cf8c2b8
  

  yashk2810

  authored and

  jax authors

  committed
  cf8c2b8

Commits on Mar 25, 2023

• Add src argument to device_put as an experimental arg

  Show description for a5d3085
  

  yashk2810

  authored and

  jax authors

  committed
  a5d3085

• Add SDA deprecation warning to pytest.ini

  Show description for 5da4a5a
  

  yashk2810

  authored and

  jax authors

  committed
  5da4a5a

Commits on Mar 24, 2023

• If each host has the full value of the Array, allow fetching it to host. Fixes jax-ml#15162

  Show description for 257ac6a
  

  yashk2810

  authored and

  jax authors

  committed
  257ac6a

• After the SPMD bug fix, always take the _rewriting_take route for getitem instead of bouncing to host.

  Show description for bc231ee
  

  yashk2810

  authored and

  jax authors

  committed
  bc231ee

Commits on Mar 23, 2023

• Deprecated xla_call_p since it has been replaced with pjit.pjit_p

  Show description for a9e48af
  

  yashk2810

  authored and

  jax authors

  committed
  a9e48af

Commits on Mar 21, 2023

• Remove the config.jax_array and jax_jit_pjit_api_merge flag usage since those are always True

  Show description for 9a0de29
  

  yashk2810

  authored and

  jax authors

  committed
  9a0de29

• Raise a better error message when there is a device assignment mismatch via the apply_primitive route.

  Show description for b5c9c0f
  

  yashk2810

  authored and

  jax authors

  committed
  b5c9c0f

• Remove C++ jit support since it has been replaced with Pjit. Keep CompiledFunction alive as a shim which cannot be instantiated but will work for isinstance checks.

  Show description for d05cf13
  

  yashk2810

  authored and

  jax authors

  committed
  d05cf13

Commits on Mar 20, 2023

• Remove pxla.OutputType enum class now that the only output can be jax.Array

  Show description for 58fed70
  

  yashk2810

  authored and

  jax authors

  committed
  58fed70

• Optimize accessing index and replica_id of

  Show description for 021fadf
  

  yashk2810

  authored and

  jax authors

  committed
  021fadf

• Add benchmarks for accessing index and replica id in addressable_shards

  Show description for 1faa7a8
  

  yashk2810

  authored and

  jax authors

  committed
  1faa7a8

Commits on Mar 17, 2023

• Improve the empty mesh error message raised in pjit if mesh is not used and Pspec is passed to in|out_shardings

  Show description for c58e2f6
  

  yashk2810

  authored and

  jax authors

  committed
  c58e2f6

• Remove _PositionalSemantics class since it is not used anymore because jax.Array always has GLOBAL semantics

  Show description for 23d3dfd
  

  yashk2810

  authored and

  jax authors

  committed
  23d3dfd

Pagination

Previous

Next