Code

feat(jax): checkpoint I/O#4236

Code

Merged

iProzd merged 18 commits intodeepmodeling:develdeepmodeling/deepmd-kit:develfrom njzjz:save-jaxnjzjz/deepmd-kit:save-jaxCopy head branch name to clipboardOct 24, 2024

+156-10

Commits

Commits on Oct 16, 2024

• feat(jax): energy model (no grad support)

  Show description for f6b30bf
  njzjz

  committed
  f6b30bf

• address comments

  Show description for 3e13a33
  njzjz

  committed
  3e13a33

Commits on Oct 21, 2024

• export EnergyModel

  Show description for fa9a000
  njzjz

  committed
  fa9a000

• jax2tf

  Show description for d3b5ce3
  njzjz

  committed
  d3b5ce3

• apply flax_module to energy model

  Show description for 91ec084
  njzjz

  committed
  91ec084

• checkpoint

  Show description for 14fab30
  njzjz

  committed
  14fab30

• fix atomic model

  Show description for 390ace5
  njzjz

  committed
  390ace5

• apply flax_module

  Show description for a4cd627
  njzjz

  committed
  a4cd627

• checkpoint

  Show description for 71a4b55
  njzjz

  committed
  71a4b55

• done

  Show description for 5024f70
  njzjz

  committed
  5024f70

• add dependencies

  Show description for f0bc8b8
  njzjz

  committed
  f0bc8b8

Commits on Oct 23, 2024

• Merge branch 'devel' into save-jax

  Show description for 2a177df
  njzjz

  authored
  2a177df

• [pre-commit.ci] auto fixes from pre-commit.com hooks

  Show description for 8c25e87
  pre-commit-ci[bot]

  committed
  8c25e87

• Remove jax2tf

  njzjz

  authored
  7e003b2

• Update pyproject.toml

  Show description for 536bbcd
  njzjz

  authored
  536bbcd

• update the model with the state

  Show description for f90fc52
  njzjz

  authored
  f90fc52

• fix is_available

  Show description for 9b571d1
  njzjz

  committed
  9b571d1

• bugfix

  Show description for fb3df8b
  njzjz

  committed
  fb3df8b