JAX: Accelerated machine learning research via composable function transformations in Python 10-15-2019