actually it is not the original code of alphazero , the original code was tensorflow/TPU https://en.wikipedia.org/wiki/AlphaGo_Zero#Training and this code is in JAX