add cutoff for materialize-and-xfer vs build-on-device
add cutoff for materialize-and-xfer vs build-on-device

see https://github.com/google/jax/pull/140#issuecomment-448433620

google/jax,google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability