Commits

Branch selector

 main

User selector

mattjj

Datepicker

All time

Commit History

Commits on Dec 20, 2024

• add mutable array ref error checks to cond and custom_vjp

  mattjj

  committed
  b6482f1

  Copy full SHA for b6482f1

Commits on Dec 19, 2024

• add test for partial-auto ppermute

  Show description for 9f42b99
  

  mattjj

  authored and

  Google-ML-Automation

  committed
  9f42b99

  Copy full SHA for 9f42b99

• add mutable array ref error checks to scan

  mattjj

  committed
  e528562

  Copy full SHA for e528562

Commits on Dec 18, 2024

• remove whitespace from distributed arrays doc

  mattjj

  committed
  c657682

  Copy full SHA for c657682

• fix bug with jax.remat static_argnums not supporting int

  Show description for bb16d5a
  

  mattjj

  authored and

  Google-ML-Automation

  committed
  bb16d5a

  Copy full SHA for bb16d5a

• ref errors

  mattjj

  committed
  42ac4ca

  Copy full SHA for 42ac4ca

• improve checkpoint / remat concreteness error with static_argnums

  mattjj

  committed
  9acd4a9

  Copy full SHA for 9acd4a9

Commits on Dec 5, 2024

• remove vestigial ad.reducing_transposes table

  Show description for 6172a1f
  mattjj

  committed
  6172a1f

  Copy full SHA for 6172a1f

Commits on Nov 8, 2024

• don't warn on unused type: ignore

  mattjj

  committed
  6115060

  Copy full SHA for 6115060

Commits on Nov 1, 2024

• support exec_time_optimization_effort and memory_fitting_effort xla compilation

  Show description for 0f3ba42
  

  mattjj

  authored and

  Google-ML-Automation

  committed
  0f3ba42

  Copy full SHA for 0f3ba42

• remove busted example from shmap jep

  mattjj

  committed
  26f70c9

  Copy full SHA for 26f70c9

Commits on Oct 29, 2024

• fix jax.custom_gradient to allow closing over non-autodiff tracers

  mattjj

  committed
  86a47a7

  Copy full SHA for 86a47a7

Commits on Oct 24, 2024

• improve concreteness error message in remat

  mattjj

  committed
  4231128

  Copy full SHA for 4231128

Commits on Oct 2, 2024

• fix

  mattjj

  committed
  1fe3b54

  Copy full SHA for 1fe3b54

• add jax2tf patch to ignore these primitives

  mattjj

  committed
  15f9ac4

  Copy full SHA for 15f9ac4

Commits on Oct 1, 2024

• add checkify rule for remat

  Show description for 11fdda9
  mattjj

  committed
  11fdda9

  Copy full SHA for 11fdda9

Commits on Sep 25, 2024

• simplify conversion logic involving extended dtypes

  Show description for 0a73d74
  mattjj

  committed
  0a73d74

  Copy full SHA for 0a73d74

Commits on Sep 21, 2024

• add jax.experimental.primal_tangent_dtype helper

  Show description for 43cc70b
  mattjj

  committed
  43cc70b

  Copy full SHA for 43cc70b

Commits on Sep 19, 2024

• custom_vjp: don't drop tangents just because they have a different dtype than the primal

  Show description for 7571b9e
  mattjj

  committed
  7571b9e

  Copy full SHA for 7571b9e

Commits on Sep 14, 2024

• tweak error logic to save a comment :)

  mattjj

  committed
  02bb3d1

  Copy full SHA for 02bb3d1

Commits on Sep 12, 2024

• shmap in_spec None shouldn't require hashability

  Show description for 358f00d
  

  mattjj

  and

  froystig

  committed
  358f00d

  Copy full SHA for 358f00d

Commits on Sep 10, 2024

• trying on another readme header

  mattjj

  committed
  95cb347

  Copy full SHA for 95cb347

Commits on Sep 7, 2024

• tweak readme title to be more about what jax can do for you, dear user

  Show description for 8ab5031
  mattjj

  committed
  8ab5031

  Copy full SHA for 8ab5031

Commits on Sep 2, 2024

• tweak shmap implementation to work better with leak checker

  mattjj

  committed
  f2bef6b

  Copy full SHA for f2bef6b

Commits on Aug 23, 2024

• add experimental jax.no_tracing context manager

  mattjj

  committed
  670a648

  Copy full SHA for 670a648

Commits on Aug 19, 2024

• fix basic scan bug with attrs

  mattjj

  committed
  ef82cb2

  Copy full SHA for ef82cb2

Commits on Aug 8, 2024

• fix jax-ml#22944

  mattjj

  committed
  44ae9b3

  Copy full SHA for 44ae9b3

Commits on Aug 3, 2024

• improve while_loop carry pytree/type mismatch errors

  Show description for bdcd358
  mattjj

  committed
  bdcd358

  Copy full SHA for bdcd358

Commits on Jul 31, 2024

• add temporary flag to suppress an error message, to unblock a user

  mattjj

  committed
  abfb1ce

  Copy full SHA for abfb1ce

Commits on Jul 25, 2024

• deflake cond memory leak regression test

  mattjj

  committed
  f72a3f8

  Copy full SHA for f72a3f8

• remove pdot and xeinsum (since xmap is gone)

  mattjj

  committed
  88d1cd7

  Copy full SHA for 88d1cd7

• remove inlined jax.nn.initializers definitions, resolving TODO of levskaya et al

  Show description for c8ea86c
  

  mattjj

  authored and

  jax authors

  committed
  c8ea86c

  Copy full SHA for c8ea86c

• remove named_shapes (since xmap is now gone)

  mattjj

  committed
  3f9eb40

  Copy full SHA for 3f9eb40

Commits on Jul 23, 2024

• fix memory leak in cond jaxpr tracig

  Show description for 8db862c
  mattjj

  committed
  8db862c

  Copy full SHA for 8db862c

• make grad-of-pmap with out_axes=None raise NotImplementedError

  Show description for 405872d
  mattjj

  committed
  405872d

  Copy full SHA for 405872d

Pagination

Previous

Next