Code

Macenko JAX backend support#36

Code

Open

andreped wants to merge 16 commits intoEIDOSLAB:developmentEIDOSLAB/torchstain:developmentfrom andreped:jax-backendandreped/torchstain:jax-backendCopy head branch name to clipboard

+219-13

Commits

Commits on Jan 2, 2023

• added jax backend support for macenko

  andreped

  committed
  b297a40

  Copy full SHA for b297a40

Commits on Jan 3, 2023

• fixed bug in jax macenko

  andreped

  committed
  07c2c58

  Copy full SHA for 07c2c58

• fixed masking bug in jax macenko

  andreped

  committed
  a20c69d

  Copy full SHA for a20c69d

Commits on Jan 22, 2023

• made Jax Macenko class a PyTree -> Jax compatible [no ci]

  andreped

  committed
  2366cc5

  Copy full SHA for 2366cc5

• minor bug in _tree_flatten [no ci]

  andreped

  committed
  4eddfc9

  Copy full SHA for 4eddfc9

Commits on Jan 23, 2023

• added macenko jax unit test

  andreped

  committed
  4890bdb

  Copy full SHA for 4890bdb

• run CI on all branches + added jax tests to CI

  andreped

  committed
  d84110f

  Copy full SHA for d84110f

• build full only on main + build quick on all branches

  andreped

  committed
  8013108

  Copy full SHA for 8013108

• renamed name in workflows

  andreped

  committed
  cae7f95

  Copy full SHA for cae7f95

• fixed np.float issue + minor array refactor

  andreped

  committed
  1877664

  Copy full SHA for 1877664

Commits on Jan 29, 2023

• fix: setting jax.jit decorators differently yields correct output

  andreped

  committed
  79b91aa

  Copy full SHA for 79b91aa

• fix: removed unwanted preprocessing in jax unit test

  andreped

  committed
  dd3e8e6

  Copy full SHA for dd3e8e6

• refactor: updated README regarding JAX - match development branch

  andreped

  committed
  9323c03

  Copy full SHA for 9323c03

• fix: updated setup.py to properly support jax during install

  andreped

  committed
  b54c228

  Copy full SHA for b54c228

• fix: run jax quick test CI on latest ubuntu + python 3.7

  andreped

  committed
  50f0fc1

  Copy full SHA for 50f0fc1

• Merge branch 'development' into jax-backend

  andreped

  authored
  dde3db4

  Copy full SHA for dde3db4