Docker used to run trax. Please see our Dockerfile below as a reference. Additionally we identified matplotlib-backend as the root cause of the import-crash. We set the matplotlib-backend to GTK3Agg and by invoking matplotlib.pyplot before importing trax

Dockerfile

ARG UBUNTU_VERSION=18.04
ARG ARCH=
ARG CUDA=10.1
FROM nvidia/cuda${ARCH:+-$ARCH}:${CUDA}-base-ubuntu${UBUNTU_VERSION} as base
ARG ARCH
ARG CUDA
ARG CUDNN=7.6.5.32-1
ARG CUDNN_MAJOR_VERSION=7
ARG LIB_DIR_PREFIX=x86_64
ARG LIBNVINFER=6.0.1-1
ARG LIBNVINFER_MAJOR_VERSION=6

ENV PATH "$PATH:/root/.local/bin"
ENV PYTHONPATH "$PYTHONPATH:/root/dev/git/prinvision"

# Needed for string substitution
SHELL ["/bin/bash", "-c"]
# Pick up some TF dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        cuda-command-line-tools-${CUDA/./-} \
        libcublas10=10.1.0.105-1 \
        cuda-nvrtc-${CUDA/./-} \
        cuda-cufft-${CUDA/./-} \
        cuda-curand-${CUDA/./-} \
        cuda-cusolver-${CUDA/./-} \
        cuda-cusparse-${CUDA/./-} \
        curl \
        libcudnn7=${CUDNN}+cuda${CUDA} \
        libcudnn7-dev=${CUDNN}+cuda${CUDA} \
        libfreetype6-dev \
        libhdf5-serial-dev \
        libzmq3-dev \
        pkg-config \
        software-properties-common \
        unzip \
        mariadb-server-10.1

# Install TensorRT if not building for PowerPC
RUN [[ "${ARCH}" = "ppc64le" ]] || { apt-get update && \
        apt-get install -y --no-install-recommends libnvinfer${LIBNVINFER_MAJOR_VERSION}=${LIBNVINFER}+cuda${CUDA} \
        libnvinfer-plugin${LIBNVINFER_MAJOR_VERSION}=${LIBNVINFER}+cuda${CUDA} \
        && apt-get clean \
        && rm -rf /var/lib/apt/lists/*; }

# For CUDA profiling, TensorFlow requires CUPTI.
ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH

# Link the libcuda stub to the location where tensorflow is searching for it and reconfigure
# dynamic linker run-time bindings
RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 \
    && echo "/usr/local/cuda/lib64/stubs" > /etc/ld.so.conf.d/z-cuda-stubs.conf \
    && ldconfig

ARG USE_PYTHON_3_NOT_2
ARG _PY_SUFFIX=${USE_PYTHON_3_NOT_2:+3}
ARG PYTHON=python${_PY_SUFFIX}
ARG PIP=pip${_PY_SUFFIX}

# See http://bugs.python.org/issue19846
ENV LANG C.UTF-8

# Some TF tools expect a "python" binary
#RUN ln -s $(which ${PYTHON}) /usr/local/bin/python

# Installs the latest version by default.
RUN apt-get update
RUN apt -y install python3 python3-dev curl python3-distutils git ca-certificates libssl-dev libmysqlclient-dev nano build-essential libffi-dev libxml2-dev libxslt1-dev zlib1g-dev
RUN apt -y install wget freeglut3 freeglut3-dev libxi-dev libxmu-dev

#Recreate the original root directory for the project "prinvision"
RUN mkdir -p ~/dev/git
RUN mkdir -p ~/dev/pv_data

#Install necessary dependencies and desired repositories or utilities
RUN curl https://bootstrap.pypa.io/get-pip.py | python3
RUN pip install --upgrade pip
RUN pip install --user tensorflow==2.1.0
RUN pip install --user transformers==2.3.0
RUN pip install --user mysqlclient==1.4.6
RUN pip install --user peewee==3.13.1
RUN pip install --user matplotlib==3.1.3
RUN pip install --user scikit-learn==0.22.1
RUN pip install --user python-Levenshtein==0.12.0
RUN pip install --user sentencepiece==0.1.85

#Install jaxlib
RUN pip install --upgrade https://storage.googleapis.com/jax-releases/cuda101/jaxlib-0.1.39-cp36-none-linux_x86_64.whl

#Install Jax for GPU support
RUN pip install --upgrade --user jax==0.1.59

#Install Jupyter for notebook environment
RUN pip install --upgrade --user jupyter==1.0.0

RUN pip install --user trax==1.2.3