컴맹의 컴퓨터공학과에서 살아남기

JAX에서 모델을 사용할 때, 주로 Flax를 이용하여 모델을 다룬다. Flax official page: https://flax.readthedocs.io/en/latest/index.html Flax Functional API Flax’s functional API radically redefines what Modules can do via lifted transformations like vmap, scan, etc, while also enabling seamless integration with other JAX libraries like Optax and Chex. flax.readthedocs.io Flax에도 Pytorch에서 nn.Module과 동일한 Flax.linen.Module이 ..

JAX 사용할 때, 계속해서 "AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)" 이런 오류가 발생. 고치면 또 까먹고 고치면 또 까먹어서 정리하는 글. 저 에러는 jaxlib랑 jax가 버전이 맞지 않아서 발생하는 오류라고 한다. 각 버전을 먼저 확인해주기 위해 아래를 실행: pip list | grep jax 그리고 각 jaxlib랑 jax를 업데이트 진행하면 해결 가능. pip install -U jax jaxlib 업데이트 후 각 버전 확인해보면 버전이 동일해진걸 확인할 수 있다. 그런데 가끔 colab 환경에서 진행할때 업데이트를 해도 에..