flax

ad8a20ff - add tf-text to doc requirements