TypeError: external_configurable() got an unexpected keyword argument 'blacklist'

[NirajanBekoju]

Description

I got this error while importing the trax

$ pip freeze | grep trax

trax                         1.3.1

$ pip freeze | grep tensor

mesh-tensorflow==0.1.21
tensor2tensor==1.15.7
tensorboard==2.9.1
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow==2.9.1
tensorflow-addons==0.17.1
tensorflow-datasets==4.6.0
tensorflow-estimator==2.9.0
tensorflow-gan==2.1.0
tensorflow-hub==0.12.0
tensorflow-io-gcs-filesystem==0.26.0
tensorflow-metadata==1.9.0
tensorflow-probability==0.7.0
tensorflow-text==2.9.0

$ pip freeze | grep jax

jax==0.3.14
jaxlib==0.3.14

$ python -V

Python 3.8.10

For bugs: reproduction and error logs

import trax

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [14], in <cell line: 1>()
----> 1 import trax

File ~/.local/lib/python3.8/site-packages/trax/__init__.py:19, in <module>
     16 """Trax top level import."""
     18 from trax import fastmath
---> 19 from trax import layers
     20 from trax import models
     21 from trax import optimizers

File ~/.local/lib/python3.8/site-packages/trax/layers/__init__.py:66, in <module>
     63 FilterResponseNorm = layer_configure(FilterResponseNorm)
     64 ThresholdedLinearUnit = layer_configure(ThresholdedLinearUnit)
---> 66 DotProductCausalAttention = layer_configure(
     67     DotProductCausalAttention, blacklist=['mode'])
     68 SelfAttention = layer_configure(SelfAttention, blacklist=['mode'])
     69 LSHSelfAttention = layer_configure(LSHSelfAttention, blacklist=['mode'])

File ~/.local/lib/python3.8/site-packages/trax/layers/__init__.py:43, in layer_configure(*args, **kwargs)
     41 def layer_configure(*args, **kwargs):
     42   kwargs['module'] = 'trax.layers'
---> 43   return gin.external_configurable(*args, **kwargs)

TypeError: external_configurable() got an unexpected keyword argument 'blacklist'