Commits

Branch selector

 main

User selector

gnecula

Datepicker

All time

Commit History

Commits on Jul 4, 2024

• [pallas] Improve some error messages and add API tests.

  Show description for a4a9499
  gnecula

  committed
  a4a9499

  Copy full SHA for a4a9499

Commits on Jul 2, 2024

• [pallas] Move tests to the GPU- and TPU- specific directories

  Show description for 242c993
  

  gnecula

  authored and

  jax authors

  committed
  242c993

  Copy full SHA for 242c993

• [export] Disable serialization in export_test if flatbuffers is not installed

  Show description for cfa3c91
  gnecula

  committed
  cfa3c91

  Copy full SHA for cfa3c91

• [pallas] Make pallas_test run on CPU and TPU also.

  Show description for de0fd72
  

  gnecula

  authored and

  jax authors

  committed
  de0fd72

  Copy full SHA for de0fd72

Commits on Jul 1, 2024

• [pallas] Add limited support for shape polymorphism for TPU

  Show description for 9def0f1
  

  gnecula

  authored and

  jax authors

  committed
  9def0f1

  Copy full SHA for 9def0f1

Commits on Jun 29, 2024

• Fix error in custom call registration for some FFI functions

  Show description for 2f808e9
  

  gnecula

  authored and

  jax authors

  committed
  2f808e9

  Copy full SHA for 2f808e9

• [pallas] Added more documentation for grid and BlockSpec.

  Show description for bfdf8f4
  gnecula

  committed
  bfdf8f4

  Copy full SHA for bfdf8f4

Commits on Jun 28, 2024

• Ported threefry2x32 for GPU to the typed XLA FFI

  Show description for cbe5242
  

  gnecula

  authored and

  jax authors

  committed
  cbe5242

  Copy full SHA for cbe5242

• [export] Clean up BUILD targets for jax.experimental.export

  Show description for 24b42ee
  

  gnecula

  authored and

  jax authors

  committed
  24b42ee

  Copy full SHA for 24b42ee

• [export] Add documentation for debugging and for ensuring compatibility.

  Show description for 47f1b3d
  gnecula

  committed
  47f1b3d

  Copy full SHA for 47f1b3d

Commits on Jun 25, 2024

• [pallas] Simplify the generation of the hlo.CustomCall

  Show description for 0b0d5a0
  

  gnecula

  authored and

  jax authors

  committed
  0b0d5a0

  Copy full SHA for 0b0d5a0

• [pallas] Break long lines in the Pallas docs

  Show description for 8528f51
  gnecula

  committed
  8528f51

  Copy full SHA for 8528f51

Commits on Jun 23, 2024

• [export] Fix multi-platform lowering for unknown platform, with donated_argnums

  Show description for d737abd
  gnecula

  committed
  d737abd

  Copy full SHA for d737abd

Commits on Jun 21, 2024

• Fix the eager mode execution for lax.platform_dependent

  Show description for 6e3fc9a
  gnecula

  committed
  6e3fc9a

  Copy full SHA for 6e3fc9a

Commits on Jun 17, 2024

• [shape_poly] Add documentation for workaround with dimension parameters.

  gnecula

  committed
  b1a8c65

  Copy full SHA for b1a8c65

Commits on Jun 15, 2024

• [shape_poly] Add documentation for shape polymorphism

  Show description for b58ff2b
  gnecula

  committed
  b58ff2b

  Copy full SHA for b58ff2b

Commits on Jun 14, 2024

• [export] Fix a couple of missed deprecated fields.

  Show description for 3b5b20c
  gnecula

  committed
  3b5b20c

  Copy full SHA for 3b5b20c

Commits on Jun 13, 2024

• [export] Deprecate jax.experimental.export

  Show description for 7af03a8
  gnecula

  committed
  7af03a8

  Copy full SHA for 7af03a8

• [export] Rename some API entry points

  Show description for 7c3a4db
  gnecula

  committed
  7c3a4db

  Copy full SHA for 7c3a4db

Commits on Jun 12, 2024

• [export] Add documentation for jax.export

  gnecula

  committed
  105cc9a

  Copy full SHA for 105cc9a

• [pallas] Add support for cross-platform lowering

  Show description for 97db0e7
  gnecula

  committed
  97db0e7

  Copy full SHA for 97db0e7

Commits on Jun 11, 2024

• [export] Cleaned up types of [in|out]_shardings

  Show description for e3faf85
  gnecula

  committed
  e3faf85

  Copy full SHA for e3faf85

Commits on Jun 10, 2024

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

  Show description for b33aca6
  gnecula

  committed
  b33aca6

  Copy full SHA for b33aca6

• [pallas] Move the hardware_generation query in the code path that needs it

  Show description for 2ade7e7
  

  gnecula

  authored and

  jax authors

  committed
  2ade7e7

  Copy full SHA for 2ade7e7

Commits on Jun 9, 2024

• [export] Move the export implementation to jax._src.export.

  Show description for 14d87d3
  

  gnecula

  authored and

  jax authors

  committed
  14d87d3

  Copy full SHA for 14d87d3

Commits on Jun 7, 2024

• [export] Remove old deprecated APIs for jax.experimental.export.

  Show description for 3914cb4
  

  gnecula

  authored and

  jax authors

  committed
  3914cb4

  Copy full SHA for 3914cb4

Commits on Jun 6, 2024

• [export] Rename in_shardings and out_shardings fields.

  Show description for 01ee768
  gnecula

  committed
  01ee768

  Copy full SHA for 01ee768

• [export] Add a LoweringParameters.for_export boolean context for exporting

  Show description for 079eea5
  gnecula

  committed
  079eea5

  Copy full SHA for 079eea5

Commits on Jun 5, 2024

• [shape_poly] Add limited support for lax.approx_top_k.

  Show description for dbad518
  

  gnecula

  authored and

  jax authors

  committed
  dbad518

  Copy full SHA for dbad518

• [shape_poly] Move to jax._src in preparation for adding to AOT APIs.

  Show description for 39ac584
  

  gnecula

  authored and

  jax authors

  committed
  39ac584

  Copy full SHA for 39ac584

Commits on Jun 4, 2024

• [export] Add helped methods to create XLACompatibleShardings for in_shardings and out_shardings

  gnecula

  committed
  2f3e02a

  Copy full SHA for 2f3e02a

Commits on Jun 3, 2024

• [docs] Enable doctest for .md documentation files

  Show description for 0223e83
  gnecula

  committed
  0223e83

  Copy full SHA for 0223e83

Commits on Jun 1, 2024

• Copybara import of the project:

  Show description for be1e40d
  

  gnecula

  authored and

  jax authors

  committed
  be1e40d

  Copy full SHA for be1e40d

Commits on May 31, 2024

• [jax2tf] Bump tolerance for FFT tests on CPU.

  Show description for d26819d
  

  gnecula

  authored and

  jax authors

  committed
  d26819d

  Copy full SHA for d26819d

Commits on May 30, 2024

• [shape_poly] Fixes for approx_top_k when aggregated_to_topk=True

  Show description for c6a4731
  gnecula

  committed
  c6a4731

  Copy full SHA for c6a4731

Pagination

Previous

Next