Commits

Branch selector

 main

User selector

zhangqiaorjc

Datepicker

All time

Commit History

Commits on Mar 25, 2024

• Reverts 2a4e1ca

  Show description for 19e6156
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  19e6156

  Copy full SHA for 19e6156

Commits on Feb 15, 2024

• Add Pallas attention kernel for GPU serving.

  Show description for 9fcf9e5
  

  2 peopleauthored and

  jax authors

  committed
  9fcf9e5

  Copy full SHA for 9fcf9e5

Commits on Jan 5, 2024

• Fix a type annotation typo.

  Show description for 1ed6a81
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  1ed6a81

  Copy full SHA for 1ed6a81

Commits on Nov 13, 2023

• Avoids some unnecessary dtype casts in Pallas attention to make fwd pass faster. Use num_warps=8 to make bwd faster. Minor fixes to make code more readable.

  Show description for 24b7350
  

  2 peopleauthored and

  jax authors

  committed
  24b7350

  Copy full SHA for 24b7350

Commits on Nov 7, 2023

• Move code around to make _mha_backward xla implementation avoid unnecessary kernel computation.

  Show description for 5b1edac
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  5b1edac

  Copy full SHA for 5b1edac

Commits on Sep 11, 2023

• Add default jvp and transpose rule for jax.lax.reduce_precision.

  Show description for d4adf00
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  d4adf00

  Copy full SHA for d4adf00

Commits on Apr 2, 2023

• Avoid re-constructing set. Expensive at scale.

  Show description for cf599c7
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  cf599c7

  Copy full SHA for cf599c7

Commits on Jan 30, 2023

• Allow jnp.nan_to_num handle integer types like numpy.

  Show description for 65ef487
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  65ef487

  Copy full SHA for 65ef487

Commits on Jan 20, 2023

• Expose fp8 from jnp. Add the missing import.

  Show description for 3cb5c93
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  3cb5c93

  Copy full SHA for 3cb5c93

• Expose fp8 types from jnp namespace.

  Show description for 650e1ef
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  650e1ef

  Copy full SHA for 650e1ef

Commits on Jan 18, 2023

• Store sorted flattened dict keys in PyTree as a c++ vector instead of py::list to avoid creating new python object on every single dict flatten. For deeply nested dict, this avoids excessive gc pre…

  Show description for d58266e
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  d58266e

  Copy full SHA for d58266e

Commits on Jan 14, 2023

• Expose fp8 in jax dtypes and mlir builder.

  Show description for d203926
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  d203926

  Copy full SHA for d203926

Commits on Dec 23, 2022

• Minor docstring fix for custom_gradient.

  Show description for 0c51ca2
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  0c51ca2

  Copy full SHA for 0c51ca2

• Add examples to lax.psum to illustrate axis_index_groups better.

  Show description for 9fda20f
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  9fda20f

  Copy full SHA for 9fda20f

Commits on Dec 3, 2022

• Skip test_lstm on CPU and TPU for jax OSS build.

  Show description for 55d6daa
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  55d6daa

  Copy full SHA for 55d6daa

Commits on Nov 28, 2022

• Fix cudnn_header OSS BUILD dep.

  Show description for c54bc90
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  c54bc90

  Copy full SHA for c54bc90

• Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.

  Show description for 4d1c4bc
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  4d1c4bc

  Copy full SHA for 4d1c4bc

Commits on Nov 23, 2022

• Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.

  Show description for 78963b6
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  78963b6

  Copy full SHA for 78963b6

Commits on Oct 21, 2022

• Add delete method to GlobalDeviceArray and ShardedBuffer.

  Show description for 4e8fbd0
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  4e8fbd0

  Copy full SHA for 4e8fbd0

Commits on Jun 22, 2022

• Remove broken image link.

  Show description for be71989
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  be71989

  Copy full SHA for be71989

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 Dec 9, 2021

• Upstream mesh utils to JAX core.

  Show description for 64569ab
  

  3 peopleauthored and

  jax authors

  committed
  64569ab

  Copy full SHA for 64569ab

Commits on Dec 8, 2021

• Ensure semantically equivalent PartitionSpecs are considered equivalent when

  Show description for 2143f3d
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  2143f3d

  Copy full SHA for 2143f3d

Commits on Nov 10, 2021

• update version and changelog for pypi

  zhangqiaorjc

  committed
  ad4cb94

  Copy full SHA for ad4cb94

Commits on Oct 26, 2021

• Add jax.distributed.initialize for multi-host GPU.

  zhangqiaorjc

  committed
  0be30fb

  Copy full SHA for 0be30fb

Commits on Oct 13, 2021

• Add checkpoint policy to save dots w/o batch dim.

  zhangqiaorjc

  committed
  d9c6ddc

  Copy full SHA for d9c6ddc

Commits on Sep 15, 2021

• Update docs for nan checker and tfds example.

  zhangqiaorjc

  committed
  aec40b4

  Copy full SHA for aec40b4

Commits on Aug 27, 2021

• Set ShardedArray device_buffer aval.

  Show description for 619f2b9
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  619f2b9

  Copy full SHA for 619f2b9

Commits on Aug 11, 2021

• Internal backend config changes.

  Show description for 8d6ff96
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  8d6ff96

  Copy full SHA for 8d6ff96

Commits on Aug 10, 2021

• Use absl::Status::message() instead of error_message().

  Show description for a93eaf3
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  a93eaf3

  Copy full SHA for a93eaf3

• Fix return value for MakeBatchPointers.

  Show description for 2afba31
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  2afba31

  Copy full SHA for 2afba31

Commits on Aug 9, 2021

• [JAX] Skip deleted buffer and executable for heap profile.

  Show description for 57ce886
  

  zhangqiaorjc

  authored and

  jax authors

  committed
  57ce886

  Copy full SHA for 57ce886

Commits on Jul 12, 2021

• Change test matrix to use latest jaxlib release.

  zhangqiaorjc

  committed
  0216a31

  Copy full SHA for 0216a31

• Add a test to repro bugs in TFRT CPU backend.

  zhangqiaorjc

  committed
  72b436f

  Copy full SHA for 72b436f

• Update changelog for jaxlib-0.1.69.

  zhangqiaorjc

  committed
  82e7495

  Copy full SHA for 82e7495

Pagination

Previous

Next