Google  JAX/Flax

JAX의 구조와 핵심 기술 및 하드웨어 최적화에 대한 심층 분석 + 배포와 사용 시 유의사항을 담은 글입니다.

Running a PyTorch model in the Flax library in JAX requires a bit of work because Flax is designed to work with JAX

JAX is a machine learning library that has gained popularity in the research community. Here's a breakdown of what it is and how it compares to other

JAX 기능: Tensorflow에 있던 XLA 기능을 빼와서 독립적인 모듈로 만들어낸 것. @tf.function 데코레이터만 함수에 붙여주면 JIT 연산이 가능해집니다. fn, fnjit, fntf2 비교 https://brunch.co.kr/@chris-son

Torch 환경에서 작성된 코드를 Jax로 run 하기 위하여Torch2Jax 관련 소스를 살펴보았습니다.나중엔 Torch 환경 코드를 Jax환경에 맞춰port를 할 예정이라 참고하기 좋을 것 같아 정리하기로 했습니다.https://github.com/rdyr

jax packages

.