{"library":"flax","install":[{"cmd":"pip install flax jax[cpu]","imports":["import flax.linen as nn","from flax.core import FrozenDict","from flax.training import train_state"]},{"cmd":"pip install flax jax[cuda12_pip]","imports":[]}]}