trax (pypi): Trax
Keywords: tensorflow|machine|learning|jax|deep-learning|deep-reinforcement-learning|machine-learning|numpy|reinforcement-learning|transformer