~ which gcc

  /usr/bin/gcc

  Apple clang version 13.1.6 (clang-1316.0.21.2.5)
  Target: arm64-apple-darwin21.5.0
  Thread model: posix
  InstalledDir: /Library/Developer/CommandLineTools/usr/bin

I couldn't get the "mega" option to work, I got an error "TypeError: lax.dynamic_update_slice requires arguments to have the same dtypes, got float32, float16" (looks like a known issue https://github.com/kuprel/min-dalle/issues/2)

Edit: installing flax 0.4.2 fixes this issue, thank all!