cmarlin

Changing the network parent from tfagents.network to keras.layer (+ __init__ change) helps me to find where the bug was (error message was clearer)

Hi. I wrote a MCTS using tensorflow's tensory only+batching (and XLA). https://gitlab.com/dia.group/tf-muzero So it's possible to the GPU doing the MCTS's stuff (whole MCTS in one gpu function call) !...

> @cmarlin There's already a [JAX implementation here](https://github.com/deepmind/mctx). They addressed the issues you mentioned above using a couple of tricks. It would be very interesting to make full JAX compatible...