Code

Fix the JAX Scan dispatcher#1202

Code

Draft

rlouf wants to merge 6 commits intoaesara-devs:mainaesara-devs/aesara:mainfrom rlouf:rewrite-jax-scanrlouf/aesara:rewrite-jax-scanCopy head branch name to clipboard

+966-160

Commits

Commits on Dec 15, 2022

• Pass updates to compare_jax_and_py

  rlouf

  committed
  2657627

  Copy full SHA for 2657627

• Refactor the JAX Scan dispatcher

  rlouf

  committed
  3918d9f

  Copy full SHA for 3918d9f

• Support while loops in the JAX Scan dispatcher

  rlouf

  committed
  2fd8122

  Copy full SHA for 2fd8122

• Support mit-mots in the JAX backend

  rlouf

  committed
  0932c8e

  Copy full SHA for 0932c8e

Commits on Dec 16, 2022

• Pass through ScalarFromTensor when input is scalar

  rlouf

  committed
  1d2cc07

  Copy full SHA for 1d2cc07

Commits on Feb 21, 2023

• Preprend the initial state to the result of the scan loop

  rlouf

  committed
  b09a40e

  Copy full SHA for b09a40e