search

LTH14 / mage

A PyTorch implementation of MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis
MIT License
529 stars 25 forks source link

JAX version code #10

Closed magicknight closed 1 year ago

magicknight commented 1 year ago

I am interested in the repository and noticed that a large portion of the code is based on MAE and VQGAN, and the original implementation was in JAX/TPU. I was wondering if it would be possible to obtain the JAX version code that was used to create this repository? Thank you.

LTH14 commented 1 year ago

Hi, thanks for your interest! Currently, we do not have a plan to release the JAX code as it depends on some of Google's internal codebase. From our experiment, this PyTorch version should give very similar performance as the JAX version.