Commits

Branch selector

 main

User selector

tomhennigan

Datepicker

All time

Commit History

Commits on Nov 5, 2022

• Make Sharding instances picklable.

  Show description for 5b45357
  

  tomhennigan

  authored and

  jax authors

  committed
  5b45357

  Copy full SHA for 5b45357

Commits on Jul 15, 2022

• Reduce the verbosity of treedef printing for custom nodes.

  Show description for 1072025
  

  tomhennigan

  authored and

  jax authors

  committed
  1072025

  Copy full SHA for 1072025

Commits on May 20, 2022

• Only register backend factories once when jax_platforms config value is set.

  Show description for b308874
  

  tomhennigan

  authored and

  jax authors

  committed
  b308874

  Copy full SHA for b308874

Commits on May 6, 2022

• Return wrapper type in pytree rule for TF's DictWrapper.

  Show description for cec717d
  

  tomhennigan

  authored and

  jax authors

  committed
  cec717d

  Copy full SHA for cec717d

Commits on Jan 28, 2022

• Strip named_shape and weak_type from aval when donating buffers.

  Show description for ace8c0a
  

  tomhennigan

  authored and

  jax authors

  committed
  ace8c0a

  Copy full SHA for ace8c0a

Commits on Dec 22, 2021

• Add is_leaf to tree_{leaves,structure}.

  Show description for 2f62574
  

  tomhennigan

  authored and

  jax authors

  committed
  2f62574

  Copy full SHA for 2f62574

Commits on Dec 20, 2021

• Add build time support for AOT compilation to TF graphs.

  Show description for 2a6147a
  

  tomhennigan

  authored and

  jax authors

  committed
  2a6147a

  Copy full SHA for 2a6147a

Commits on Nov 17, 2021

• Ensure that size property of large ShardedDeviceArrays does not overflow.

  Show description for bb3f198
  

  tomhennigan

  authored and

  jax authors

  committed
  bb3f198

  Copy full SHA for bb3f198

Commits on Aug 6, 2021

• Update JAX flag parsing to work when some elements of argv aren't absl friendly.

  Show description for c94f412
  

  tomhennigan

  authored and

  jax authors

  committed
  c94f412

  Copy full SHA for c94f412

Commits on Jul 23, 2021

• Add pure JAX implementation of tridiagonal_solve for non-GPU platforms.

  tomhennigan

  committed
  2994d2f

  Copy full SHA for 2994d2f

Commits on Jul 16, 2021

• Avoid sharing handles across streams.

  Show description for afbd831
  tomhennigan

  committed
  afbd831

  Copy full SHA for afbd831

Commits on Jul 15, 2021

• Compute gtsv2 buffer size ahead of time and pass in to kernel.

  Show description for afa0d57
  tomhennigan

  committed
  afa0d57

  Copy full SHA for afa0d57

• Add source location and expression to error messages for CUDA API calls.

  Show description for d6e56f2
  tomhennigan

  committed
  d6e56f2

  Copy full SHA for d6e56f2

Commits on Jun 2, 2021

• Add lax.linalg.tridiagonal_solve(..), lowering to cusparse_gtsv2<T>() on GPU.

  Show description for ffac40a
  tomhennigan

  committed
  ffac40a

  Copy full SHA for ffac40a

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 Apr 7, 2021

• Remove jax.argnums_partial.

  tomhennigan

  committed
  441ded0

  Copy full SHA for 441ded0

Commits on Mar 12, 2021

• Update README.md

  Show description for 4f74b33
  

  tomhennigan

  and

  8bitmp3

  authored
  4f74b33

  Copy full SHA for 4f74b33

• Link to DeepMind JAX blog post.

  tomhennigan

  committed
  0fd83a3

  Copy full SHA for 0fd83a3

Commits on Feb 25, 2021

• Sort imports.

  Show description for 7b94c44
  

  tomhennigan

  authored and

  jax authors

  committed
  7b94c44

  Copy full SHA for 7b94c44

• Add autodidax to dev docs toctree and ignore md file to silence warning.

  Show description for 090dd21
  tomhennigan

  committed
  090dd21

  Copy full SHA for 090dd21

Commits on Feb 4, 2021

• Add jax.default_backend() which returns the default platform name.

  Show description for 7adb1e3
  tomhennigan

  committed
  7adb1e3

  Copy full SHA for 7adb1e3

Commits on Oct 1, 2020

• Remove scalar float0 XLA representation.

  Show description for f18b914
  tomhennigan

  committed
  f18b914

  Copy full SHA for f18b914

Commits on Aug 26, 2020

• Use omnistaging env var even when not using absl flags for config. (jax-ml#4152)

  tomhennigan

  authored
  f0fb7d0

  Copy full SHA for f0fb7d0

Commits on Jul 30, 2020

• Compare treedefs by num_leaves not traversal_ in tree_transpose. (jax-ml#3659)

  Show description for bf041fb
  tomhennigan

  authored
  bf041fb

  Copy full SHA for bf041fb

Commits on Jul 12, 2020

• Import profiler in jax/__init__.py (jax-ml#3719)

  tomhennigan

  authored
  aa75209

  Copy full SHA for aa75209

Commits on Jul 4, 2020

• Correctly handle bfloat16 input and output from jax2tf functions. (jax-ml#3652)

  Show description for 904f34a
  tomhennigan

  authored
  904f34a

  Copy full SHA for 904f34a

Commits on Jun 30, 2020

• Retain original docstring when vmap'ing functions. (jax-ml#3592)

  tomhennigan

  authored
  2ef39ab

  Copy full SHA for 2ef39ab

Commits on Jun 23, 2020

• Cast int8 to bool for lax.not in jax2tf. (jax-ml#3519)

  tomhennigan

  authored
  ca5b0b1

  Copy full SHA for ca5b0b1

Commits on Jun 11, 2020

• Improve error message when passing an invalid dtype. (jax-ml#3405)

  Show description for c7a021e
  tomhennigan

  authored
  c7a021e

  Copy full SHA for c7a021e

Commits on May 31, 2020

• Add support for buffer donation in jit and pmap. (jax-ml#2936)

  Show description for 6124f70
  tomhennigan

  authored
  6124f70

  Copy full SHA for 6124f70

Commits on May 12, 2020

• Avoid recompilation of rolled loops in threefry2x32. (jax-ml#3069)

  tomhennigan

  authored
  abdf504

  Copy full SHA for abdf504

Commits on May 4, 2020

• Add a note about jax.pmap when leading dim is smaller than num devices. (jax-ml#2949)

  tomhennigan

  authored
  4c2c5ad

  Copy full SHA for 4c2c5ad

Commits on May 1, 2020

• Explicitly broadcast values in nn.one_hot and nn.initializers.orthogonal. (jax-ml#2901)

  Show description for 0736679
  tomhennigan

  authored
  0736679

  Copy full SHA for 0736679

Commits on Apr 1, 2020

• Replace CHECK with assert. (jax-ml#2569)

  tomhennigan

  authored
  1d09b6b

  Copy full SHA for 1d09b6b

• Add jax.tree_util.all_leaves(iterable).

  Show description for ca23be6
  

  tomhennigan

  authored and

  gnecula

  committed
  ca23be6

  Copy full SHA for ca23be6

Pagination

Previous

Next