Commits

Breadcrumbs

History for

onmain

User selector

All users

All time

Commit History

Commits on Aug 11, 2024

• update README and several docs to typed RNG keys

  froystig

  committed
  371935c

  Copy full SHA for 371935c

Commits on Jun 24, 2024

• Adds make_array_from_callback_sharded benchmark

  Show description for 4a91396
  

  marksandler2

  authored and

  jax authors

  committed
  4a91396

  Copy full SHA for 4a91396

Commits on Jun 21, 2024

• [JAX] Add caching to pjit._infer_params.

  Show description for 9e30079
  

  hawkinsp

  authored and

  jax authors

  committed
  9e30079

  Copy full SHA for 9e30079

Commits on Jun 10, 2024

• Merge pull request #21763 from gnecula:export_api

  Show description for 3fe7377
  

  jax authors

  committed
  3fe7377

  Copy full SHA for 3fe7377

• [export] Create the jax.export module APIs.

  Show description for b33aca6
  gnecula

  committed
  b33aca6

  Copy full SHA for b33aca6

Commits on Jun 9, 2024

• Expand `device_put` benchmarks to run with different numbers of arrays and input types

  Show description for 6617a0d
  

  junwhanahn

  authored and

  jax authors

  committed
  6617a0d

  Copy full SHA for 6617a0d

Commits on Jun 6, 2024

• [Mosaic GPU] Add f32 benchmarks for matmul.

  Show description for 24e4bf2
  

  cperivol

  authored and

  jax authors

  committed
  24e4bf2

  Copy full SHA for 24e4bf2

Commits on May 31, 2024

• [Mosaic GPU] Add a simple benchmark.

  Show description for 8eaea2b
  

  cperivol

  authored and

  jax authors

  committed
  8eaea2b

  Copy full SHA for 8eaea2b

Commits on Apr 11, 2024

• Avoid 'from jax import config' imports

  Show description for f090074
  jakevdp

  committed
  f090074

  Copy full SHA for f090074

Commits on Feb 22, 2024

• Deprecate jax.tree_map in favor of jax.tree.map

  jakevdp

  committed
  e59a050

  Copy full SHA for e59a050

Commits on Feb 16, 2024

• Add a benchmark measuring device_put's speed for a 4GB input array

  Show description for 0b542ff
  

  yashk2810

  authored and

  jax authors

  committed
  0b542ff

  Copy full SHA for 0b542ff

Commits on Feb 9, 2024

• [shape_poly] Add another micro-benchmark

  Show description for 1e76422
  

  gnecula

  authored and

  jax authors

  committed
  1e76422

  Copy full SHA for 1e76422

Commits on Feb 5, 2024

• [shape_poly] Add micro-benchmarks for symbolic shape manipulation.

  Show description for 313caba
  

  gnecula

  authored and

  jax authors

  committed
  313caba

  Copy full SHA for 313caba

Commits on Dec 21, 2023

• in partial_eval_custom rule for pjit, cache ClosedJaxpr creation

  Show description for e31018b
  

  mattjj

  and

  yashk2810

  committed
  e31018b

  Copy full SHA for e31018b

Commits on Nov 13, 2023

• Use a Jacobi SVD solver for unbatched SVDs up to 1024x1024 on NVIDIA GPUs.

  Show description for cb182b8
  

  hawkinsp

  authored and

  jax authors

  committed
  cb182b8

  Copy full SHA for cb182b8

Commits on Oct 11, 2023

• Optimize make_array_from_callback for fully replicated shardings by going via batched_device_put

  Show description for fd09b35
  

  yashk2810

  authored and

  jax authors

  committed
  fd09b35

  Copy full SHA for fd09b35

Commits on Sep 21, 2023

• Added serial_dot_products benchmark

  patrick-kidger

  committed
  9d73441

  Copy full SHA for 9d73441

Commits on Sep 13, 2023

• fast dispatch for functions over typed PRNG key arrays

  Show description for 6abefa1
  

  2 peopleauthored and

  jax authors

  committed
  6abefa1

  Copy full SHA for 6abefa1

Commits on Aug 24, 2023

• Add random benchmarks

  Show description for 368d343
  

  Jake VanderPlas

  authored and

  jax authors

  committed
  368d343

  Copy full SHA for 368d343

Commits on Aug 2, 2023

• [jax][benchmark] Added clearing caches for benchmarking compilation time in sparse JAX benchmarks

  Show description for f498442
  

  jax authors

  committed
  f498442

  Copy full SHA for f498442

Commits on Jul 27, 2023

• Update flags to use the ABSL typed flag API.

  Show description for 76cda0a
  

  hawkinsp

  authored and

  jax authors

  committed
  76cda0a

  Copy full SHA for 76cda0a

Commits on May 5, 2023

• Improve the shape incompatible error message by adding the argument/result name path to it.

  Show description for a6254c7
  

  yashk2810

  authored and

  jax authors

  committed
  a6254c7

  Copy full SHA for a6254c7

Commits on Apr 21, 2023

• Change to simpler import for jax.config

  jakevdp

  committed
  fbe4f10

  Copy full SHA for fbe4f10

Commits on Apr 13, 2023

• Change np.prod->math.prod

  Show description for 5521423
  jakevdp

  committed
  5521423

  Copy full SHA for 5521423

Commits on Apr 11, 2023

• Add a C++ safe_zip implementation.

  Show description for 74384e6
  

  hawkinsp

  authored and

  jax authors

  committed
  74384e6

  Copy full SHA for 74384e6

• Add a C++ implementation of safe map.

  Show description for 0dbd467
  

  hawkinsp

  authored and

  jax authors

  committed
  0dbd467

  Copy full SHA for 0dbd467

Commits on Apr 7, 2023

• Remove `experimental_cpp_jit` since that flag is unused and also remove `experimental_cpp_pjit`.

  Show description for 694e43a
  

  yashk2810

  authored and

  jax authors

  committed
  694e43a

  Copy full SHA for 694e43a

Commits on Apr 6, 2023

• Rename jax._src.sharding_utils to jax._src.op_shardings.

  Show description for 452f3c5
  

  hawkinsp

  authored and

  jax authors

  committed
  452f3c5

  Copy full SHA for 452f3c5

Commits on Mar 27, 2023

• 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

  Copy full SHA for cf8c2b8

Commits on Mar 20, 2023

• Add benchmarks for accessing index and replica id in addressable_shards

  Show description for 1faa7a8
  

  yashk2810

  authored and

  jax authors

  committed
  1faa7a8

  Copy full SHA for 1faa7a8

Commits on Mar 15, 2023

• Add benchmarks for np.array, device_put, and _arrays.

  Show description for 4870217
  

  pschuh

  authored and

  jax authors

  committed
  4870217

  Copy full SHA for 4870217

Commits on Mar 13, 2023

• [Fix forward] Rollback the device_put_sharded and device_put_replicated change of using batched_device_put

  Show description for 233911c
  

  yashk2810

  authored and

  jax authors

  committed
  233911c

  Copy full SHA for 233911c

• Split Sharding subclasses out of _src/sharding.py into _src/sharding_impls.py

  Show description for 1925aa1
  

  hawkinsp

  authored and

  jax authors

  committed
  1925aa1

  Copy full SHA for 1925aa1

Commits on Mar 9, 2023

• math_benchmark: add --set_env flag

  Show description for 6f1d829
  

  cota

  authored and

  jax authors

  committed
  6f1d829

  Copy full SHA for 6f1d829

• math_benchmark: add dot op

  Show description for 845d68b
  

  cota

  authored and

  jax authors

  committed
  845d68b

  Copy full SHA for 845d68b