[anton_ai]

I'm as well going from TF2 to JAX, it's like TF3 and I hope that google will just put the keras team to work on a JAX version