Usage of certain metrics from trax fails

[AlyHdr]

Description

I'm trying to build a simple sentiment analysis model using the code from the tutorial you provided here: https://trax-ml.readthedocs.io/en/latest/notebooks/trax_intro.html However I'm not able to use certain metrics, although they exist in the code base (like CategoryCrossEntropy,WeightedCategoryCrossEntropy...etc)

Environment information

OS: Kaggle and Colab notebooks

$ pip freeze | grep trax
trax==1.3.6

$ pip freeze | grep tensor

mesh-tensorflow==0.1.17
tensorboard==2.4.0
tensorboard-plugin-wit==1.7.0
tensorboardX==2.1
tensorflow @ file:///tmp/tensorflow_cpu/tensorflow-2.3.1-cp37-cp37m-linux_x86_64.whl
tensorflow-addons @ file:///tmp/tfa_cpu/tensorflow_addons-0.11.2-cp37-cp37m-linux_x86_64.whl
tensorflow-cloud==0.1.10
tensorflow-datasets==3.0.0
tensorflow-estimator==2.3.0
tensorflow-gcs-config @ file:///tmp/tensorflow_gcs_config/tensorflow_gcs_config-2.1.7-py3-none-any.whl
tensorflow-hub==0.10.0
tensorflow-metadata==0.25.0
tensorflow-probability==0.11.1
tensorflow-text==2.3.0
tensorpack==0.10.1

$ pip freeze | grep jax
jax==0.2.7
jaxlib==0.1.57

$ python -V
Python 3.7.6

For bugs: reproduction and error logs

# Steps to reproduce:
from trax import layers as tl
tl.WeightedCategoryCrossEntropy()
tl.WeightedCategoryAccuracy()

# Error logs:
AttributeError                            Traceback (most recent call last)
<ipython-input-68-2e7eea5d8c65> in <module>
      4 train_task = training.TrainTask(
      5     labeled_data=train_batches_stream,
----> 6     loss_layer=tl.WeightedCategoryCrossEntropy(),
      7     optimizer=trax.optimizers.Adam(0.01),
      8     n_steps_per_checkpoint=500,

AttributeError: module 'trax.layers' has no attribute 'WeightedCategoryCrossEntropy'

[thoo]

I think it is only available in the development version. For Colab, you can try:

!pip install -q -U git+https://github.com/google/trax.git
!pip install -q -U git+https://github.com/google/gin-config.git

This document is for the latest version by default. You can switch it to stable version.

[AlyHdr]

Thanks for your reply, yes now it's working but why in the documentation only pip install trax is introduced and those metrics are used in the tutorial itself. If you run the notebook provided in Colab, it will fail with the same error I have mentioned before unless one installs Trax as you have said.

[pdubz-sudo]

I think it is only available in the development version. For Colab, you can try:

!pip install -q -U git+https://github.com/google/trax.git
!pip install -q -U git+https://github.com/google/gin-config.git

This document is for the latest version by default. You can switch it to stable version.

What do you mean "switch it to stable version"?