Code

register lax.cond_p#34

Code

Merged

patrick-kidger merged 11 commits intopatrick-kidger:mainpatrick-kidger/quax:mainfrom vadmbertr:cond_primitivevadmbertr/quax:cond_primitiveCopy head branch name to clipboardNov 11, 2024

+200-2

Commits

Commits on Oct 24, 2024

• register lax.cond_p

  vadmbertr

  committed
  7cbbd96

  Copy full SHA for 7cbbd96

Commits on Oct 25, 2024

• update test_cond.py

  vadmbertr

  committed
  b2ec956

  Copy full SHA for b2ec956

• fix union syntax for Python<3.10

  vadmbertr

  committed
  c360646

  Copy full SHA for c360646

• fix union syntax for Python<3.10

  vadmbertr

  committed
  6bd9d15

  Copy full SHA for 6bd9d15

• add select_n and stop_gradient primitives in order to handle vmapping over the pred of lax.cond

  vadmbertr

  committed
  78e8a67

  Copy full SHA for 78e8a67

• gruff format

  vadmbertr

  committed
  cd22172

  Copy full SHA for cd22172

• more concise impl of select and stop_gradient

  vadmbertr

  committed
  b96b1b9

  Copy full SHA for b96b1b9

• remove print left behind

  vadmbertr

  committed
  b68e80e

  Copy full SHA for b68e80e

Commits on Oct 30, 2024

• handles unprescribed number of branches (i.e. switch). enforce exact same output type. add tests for those. remove stop_gradient and select primitives 'defaults'

  vadmbertr

  committed
  0685f83

  Copy full SHA for 0685f83

Commits on Nov 6, 2024

• simplify eval of quaxed jaxpr and out tree of the branches

  vadmbertr

  committed
  e130e9b

  Copy full SHA for e130e9b

Commits on Nov 11, 2024

• compat with older jax versions

  vadmbertr

  committed
  95cdab7

  Copy full SHA for 95cdab7