Reformer does not export to SavedModel format.

[StewartSethA]

Description

...

Environment information

tensorflow-gpu installed through anaconda, trax installed via pip.

OS: Ubuntu 20.04

$ pip freeze | grep trax
trax==1.3.9

$ pip freeze | grep tensor
tensorboard==2.6.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.0
tensorflow==2.5.1
tensorflow-addons==0.13.0
tensorflow-datasets==4.4.0
tensorflow-estimator==2.5.0
tensorflow-gpu==2.6.0
tensorflow-hub==0.12.0
tensorflow-metadata==1.2.0
tensorflow-serving-api==1.15.0
tensorflow-text==2.5.0

$ pip freeze | grep jax
jax==0.2.17
jaxlib==0.1.69

$ python -V
Python 3.6.9 :: Anaconda, Inc.

For bugs: reproduction and error logs

# Steps to reproduce:
...

# Error logs:
...