Code

Add Export for JAX Backend#819

Code

Merged

fchollet merged 27 commits intokeras-team:mainkeras-team/keras-core:mainfrom nkovela1:exportnkovela1/keras-core:exportCopy head branch name to clipboardSep 8, 2023

+209-82

Commits

Commits on Aug 16, 2023

• Add JAX export workflow example

  nkovela1

  committed
  4b29990

  Copy full SHA for 4b29990

Commits on Aug 23, 2023

• Merge branch 'keras-team:main' into export

  nkovela1

  authored
  d3fb99c

  Copy full SHA for d3fb99c

• Merge branch 'keras-team:main' into export

  nkovela1

  authored
  860edfd

  Copy full SHA for 860edfd

Commits on Aug 25, 2023

• Remove numpy conversion and add JAX to export tests

  nkovela1

  committed
  91820cf

  Copy full SHA for 91820cf

• Initial changes to export library for JAX

  nkovela1

  committed
  fc9ecb2

  Copy full SHA for fc9ecb2

Commits on Aug 28, 2023

• Add JAX tracing compatibility, edit variable collection flow

  nkovela1

  committed
  09af61c

  Copy full SHA for 09af61c

• Remove pdb trace

  nkovela1

  committed
  08fe8b0

  Copy full SHA for 08fe8b0

Commits on Aug 29, 2023

• Sync and fix merge conflict with TF dependency change

  nkovela1

  committed
  da76d81

  Copy full SHA for da76d81

• Fix test

  nkovela1

  committed
  7d662cd

  Copy full SHA for 7d662cd

• Add JAX instance checking

  nkovela1

  committed
  91e0234

  Copy full SHA for 91e0234

• Split tf.function test, mark skipped for JAX backend

  nkovela1

  committed
  f9b67ca

  Copy full SHA for f9b67ca

• Convert added variables to tf.Variables

  nkovela1

  committed
  8c83b49

  Copy full SHA for 8c83b49

• Add dict and list input pytree support

  nkovela1

  committed
  1916657

  Copy full SHA for 1916657

• Add subclassed model export test

  nkovela1

  committed
  4fa4bb8

  Copy full SHA for 4fa4bb8

• Add save spec and setter method

  nkovela1

  committed
  3929b9a

  Copy full SHA for 3929b9a

Commits on Aug 30, 2023

• Add save spec functionality to export

  nkovela1

  committed
  a6cea92

  Copy full SHA for a6cea92

• Add save spec functionality to export

  nkovela1

  committed
  13f3c44

  Copy full SHA for 13f3c44

• Split and temporarily disable test, awaiting JAX2TF fix or workaround

  nkovela1

  committed
  b2f3cde

  Copy full SHA for b2f3cde

• Fix formatting and nits

  nkovela1

  committed
  d28586f

  Copy full SHA for d28586f

• Delete JAX workflow testing sample

  nkovela1

  committed
  7e683cd

  Copy full SHA for 7e683cd

Commits on Aug 31, 2023

• Bring save spec creation to export

  nkovela1

  committed
  4aedb00

  Copy full SHA for 4aedb00

• Assert called separate from built

  nkovela1

  committed
  e7a2214

  Copy full SHA for e7a2214

Commits on Sep 7, 2023

• Make module a private attribute

  nkovela1

  committed
  42412b6

  Copy full SHA for 42412b6

• Simplify JAX flow to use Autotrackable

  nkovela1

  committed
  50b2be5

  Copy full SHA for 50b2be5

• Fix formatting

  nkovela1

  committed
  a4d9505

  Copy full SHA for a4d9505

• Fix nits

  nkovela1

  committed
  7de39f9

  Copy full SHA for 7de39f9

• Small grammar fix

  nkovela1

  committed
  dcd4142

  Copy full SHA for dcd4142