Import trax takes too long to load

Question

I was stumped the first time I loaded this library. In my local computer it tooks me at least 40s to load trax on a local Jupyter Notebook and more than 1 minute to load it on a shared Colab environment.

import trax

I'm not sure if it's an issue with my installation or a BUG in the version of trax I'm using?

I'm new in trax, and in fact my experience is with Keras and TensorFlow so I'd like to get an opinion from someone in the trax community, if this is normal or not.

Thanks a lot in advance!

BTW: I'm using trax 1.4.1 with Python 3.9.6 and my local computer has the following specs:

 Intel(R) Core(TM) i5-6200U CPU @ 2.30GHz, 4 cores and 16GB RAM. 

Answer 1

After a while searching it seems that this is normal behavior. In fact there is an issue raised in the official Trax repository about this: import trax takes 17 seconds #1368.

Apparently the fastmath module that contains the trax re-implements for maths operations has plenty of dependencies.

From the issue thread:

This is a well known problem occurring on basically all setups (local, colab, gpu cluster) and it is not a big issue for running long experiments, however it does make local debugging hard. I have tried debugging the import graph with profiler, but without success yet. It looks like even from trax import fastmath has plenty of dependencies - here is the tree generated by importlab library for trax.fastmath.init module: importlab --tree init.py out: https://gist.github.com/syzymon/3bb6f59063f918b4b62b77cdb223da72

So in conclusion, whether it takes you 17 seconds or 40 seconds like me, is a known behavior in Trax.