Importing substrates is slow

[carlosgmartin]

It takes a noticeably long time to just import TFP substrates:

$ time python3 -c "import tensorflow_probability.substrates.jax"
python3 -c "import tensorflow_probability.substrates.jax"  4.72s user 0.77s system 116% cpu 4.724 total

$ time python3 -c "import tensorflow_probability.substrates.numpy"
python3 -c "import tensorflow_probability.substrates.numpy"  4.75s user 0.86s system 115% cpu 4.843 total

Might it be possible to speed up these imports somehow?

[brianwa84]

Do you have tensorflow installed? Uninstalling that might speed up the import, if you're only going to use JAX or numpy. It would be nice to disentangle how much of the import time comes from importing the respective lib (jax | numpy). I would think that would account for a decent chunk. We could push lazy module loading deeper into the library, if it's actually beneficial, but we'd want to see some more precise measurements. Brian Patton | Software Engineer | ***@***.***

…

On Mon, Apr 17, 2023 at 4:17 PM Carlos Martin ***@***.***> wrote: It takes a noticeably long time to just import TFP substrates: $ time python3 -c "import tensorflow_probability.substrates.jax" python3 -c "import tensorflow_probability.substrates.jax" 4.72s user 0.77s system 116% cpu 4.724 total $ time python3 -c "import tensorflow_probability.substrates.numpy" python3 -c "import tensorflow_probability.substrates.numpy" 4.75s user 0.86s system 115% cpu 4.843 total Might it be possible to speed up these imports somehow? — Reply to this email directly, view it on GitHub <#1714>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AFJFSIYSG2WPZ5ECJ5QXIETXBWQNZANCNFSM6AAAAAAXBVVAXI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[carlosgmartin]

@brianwa84 Uninstalling tensorflow indeed reduced the time. With tensorflow:

python3 -c "import tensorflow_probability.substrates.jax"  5.85s user 1.28s system 90% cpu 7.858 total
python3 -c "import tensorflow_probability.substrates.numpy"  5.16s user 0.94s system 109% cpu 5.590 total

Without tensorflow:

python3 -c "import tensorflow_probability.substrates.jax"  1.57s user 0.35s system 153% cpu 1.254 total
python3 -c "import tensorflow_probability.substrates.numpy"  1.12s user 0.27s system 241% cpu 0.575 total