TRAX v1.4.0

PiperOrigin-RevId: 405669883

Author: afrozenator

setup.py

@@ -22,7 +22,7 @@
 
 setup(
  name='trax',
- version='1.3.9',
+ version='1.4.0',
  description='Trax',
  long_description=(
  'Trax helps you understand deep learning. We start with basic maths and'