Commits

Breadcrumbs

History for

onjax-v0.3.8

User selector

All users

Datepicker

All time

Commit History

Commits on Apr 6, 2022

• Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

  Show description for 6a7a346
  

  yashk2810

  authored and

  jax authors

  committed
  6a7a346

  Copy full SHA for 6a7a346

Commits on Mar 29, 2022

• Speed up GDA initialization by making local_shards lazy and hiding checks behind config.jax_enable_checks flag.

  Show description for e08bc27
  

  yashk2810

  authored and

  jax authors

  committed
  e08bc27

  Copy full SHA for e08bc27

• Fix the mesh_axes of benchmarks to be Pspecs

  Show description for a68b0f3
  

  yashk2810

  authored and

  jax authors

  committed
  a68b0f3

  Copy full SHA for a68b0f3

Commits on Mar 17, 2022

• Add small and big matmul to api_benchmarks.

  Show description for 5d7f639
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  5d7f639

  Copy full SHA for 5d7f639

Commits on Feb 1, 2022

• Add the cache back now that Mesh's __hash__ is also being hashed on self.devices.shape.

  Show description for 4e47de6
  

  yashk2810

  authored and

  jax authors

  committed
  4e47de6

  Copy full SHA for 4e47de6

• Couple of changes because of the serialization inconsistencies being observed.

  Show description for f80887e
  

  yashk2810

  authored and

  jax authors

  committed
  f80887e

  Copy full SHA for f80887e

Commits on Jan 18, 2022

• Cache the expensive computations in GDA. For example get_shard_indices_replica_ids can be the same for multiple variables in a neural network (global_shape, mesh_axes and global_mesh) can be the …

  Show description for 7f192c1
  

  yashk2810

  authored and

  jax authors

  committed
  7f192c1

  Copy full SHA for 7f192c1

Commits on Jan 13, 2022

• Optimizations for GDA to make creating GDA faster.

  Show description for 0532a63
  

  yashk2810

  authored and

  jax authors

  committed
  0532a63

  Copy full SHA for 0532a63

Commits on Jan 11, 2022

• Benchmarks for GDA. Also move create_global_mesh to test_utils since it was replicated in a lot of places.

  Show description for fbb8b9f
  

  yashk2810

  authored and

  jax authors

  committed
  fbb8b9f

  Copy full SHA for fbb8b9f

Commits on Dec 21, 2021

• [JAX] Requires indices to be sorted and of int32 in _sparse_bcoo_matvec test.

  Show description for 6a0bfdd
  

  tlu7

  authored and

  jax authors

  committed
  6a0bfdd

  Copy full SHA for 6a0bfdd

Commits on Oct 26, 2021

• benchmarks: add compilation of sparse operations

  Show description for b8296bb
  

  Jake VanderPlas

  authored and

  jax authors

  committed
  b8296bb

  Copy full SHA for b8296bb

• benchmarks: add JIT versions of sparse.BCOO benchmarks

  Show description for c62452f
  

  Jake VanderPlas

  authored and

  jax authors

  committed
  c62452f

  Copy full SHA for c62452f

Commits on Oct 25, 2021

• benchmarks: add sparse.BCOO todense/fromdense/matvec benchmarks

  Show description for 5a37337
  

  Jake VanderPlas

  authored and

  jax authors

  committed
  5a37337

  Copy full SHA for 5a37337

Commits on Oct 7, 2021

• Add a warmup loop to pmap_simple_8_devices_100_args benchmark so as to not measure the compile time.

  Show description for bfbdfa8
  

  yashk2810

  authored and

  jax authors

  committed
  bfbdfa8

  Copy full SHA for bfbdfa8

Commits on Oct 5, 2021

• [JAX] Fix pylint errors.

  Show description for 256e722
  

  hawkinsp

  authored and

  jax authors

  committed
  256e722

  Copy full SHA for 256e722

Commits on Sep 16, 2021

• Remove jax.api.

  Show description for 6a1b626
  hawkinsp

  committed
  6a1b626

  Copy full SHA for 6a1b626

Commits on Aug 26, 2021

• Add a benchmark with many arguments.

  Show description for 6cb8737
  

  jblespiau

  authored and

  jax authors

  committed
  6cb8737

  Copy full SHA for 6cb8737

Commits on Aug 3, 2021

• Internal only.

  Show description for afdd195
  

  jblespiau

  authored and

  jax authors

  committed
  afdd195

  Copy full SHA for afdd195

Commits on May 3, 2021

• [JAX] Prune unused inputs in jit.

  Show description for 850bd66
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  850bd66

  Copy full SHA for 850bd66

Commits on Apr 14, 2021

• Remove special casing on npy_value when indexing sharded arrays.

  Show description for 9d56552
  

  tomhennigan

  authored and

  jax authors

  committed
  9d56552

  Copy full SHA for 9d56552

Commits on Mar 19, 2021

• add simple single-primitive eager benchmarks

  mattjj

  committed
  9802f33

  Copy full SHA for 9802f33

Commits on Mar 4, 2021

• Improve API benchmarks.

  Show description for cdd36b1
  hawkinsp

  committed
  cdd36b1

  Copy full SHA for cdd36b1

Commits on Feb 25, 2021

• Sort imports.

  Show description for 7b94c44
  

  tomhennigan

  authored and

  jax authors

  committed
  7b94c44

  Copy full SHA for 7b94c44

Commits on Jan 21, 2021

• Revert import path changes to examples/ and benchmarks/

  Show description for 160dfd3
  

  hawkinsp

  authored and

  jax authors

  committed
  160dfd3

  Copy full SHA for 160dfd3

Commits on Jan 20, 2021

• Small cleanups to dependency structure.

  Show description for 929a684
  

  hawkinsp

  authored and

  jax authors

  committed
  929a684

  Copy full SHA for 929a684

Commits on Jan 11, 2021

• [JAX] Move jax.util to jax._src_util.

  Show description for 3ac809e
  

  hawkinsp

  authored and

  jax authors

  committed
  3ac809e

  Copy full SHA for 3ac809e

Commits on Sep 18, 2020

• fix benchmark sums (jax-ml#4329)

  jakevdp

  authored
  6a89f60

  Copy full SHA for 6a89f60

Commits on Aug 27, 2020

• Add benchmarks for specifically the dispatch time. (jax-ml#4128)

  Show description for e95d570
  jblespiau

  authored
  e95d570

  Copy full SHA for e95d570

Commits on Aug 18, 2020

• Cleanup: avoid jnp.prod & np.prod on array shapes (jax-ml#4086)

  jakevdp

  authored
  29aa9bf

  Copy full SHA for 29aa9bf

Commits on Jul 14, 2020

• Cleanup: convert uses of import numpy as onp in library code (jax-ml#3754)

  jakevdp

  authored
  a7c2cde

  Copy full SHA for a7c2cde

Commits on Jul 9, 2020

• Add simple JAX API microbenchmarks. (jax-ml#3679)

  chr1sj0nes

  authored
  c1aeb8b

  Copy full SHA for c1aeb8b

Commits on Jun 9, 2020

• Fix remaining flakes and use exclude within setup.cfg (jax-ml#3371)

  jakevdp

  authored
  6aa8f24

  Copy full SHA for 6aa8f24

Commits on May 5, 2020

• Replace np -> jnp, onp -> np in more places. (jax-ml#2973)

  Show description for b1bc841
  hawkinsp

  authored
  b1bc841

  Copy full SHA for b1bc841

Commits on Apr 28, 2020

• Add pmap_shard_device_array_benchmark. (jax-ml#2864)

  Show description for 4b03343
  skye

  authored
  4b03343

  Copy full SHA for 4b03343

Commits on Apr 13, 2020

• Add --export_dir and --baseline_dir flags to benchmark.py. (jax-ml#2677)

  Show description for 8c2901c
  skye

  authored
  8c2901c

  Copy full SHA for 8c2901c

Pagination

Previous

Next