Reformer: Image Generation: module 'trax' has no attribute 'fastmath'

[asj100]

I've tried to run the Reformer: Image Generation colab a couple of times but get an error message on the last cell: module 'trax' has no attribute 'fastmath'. It may be something I'm doing or not but I can't get it to run.

The entire error message is: AttributeError Traceback (most recent call last) in () 1 flat_prompt = [] ----> 2 for i, img in enumerate(DATA[:trax.fastmath.device_count()]): 3 img = img.reshape((-1, 64, 3))[:32, :, :] 4 flat_prompt.append(img.reshape((-1,))) 5 prompt = np.stack(flat_prompt, 0)

AttributeError: module 'trax' has no attribute 'fastmath' ...

Environment information

OS: <your answer here>

$ pip freeze | grep trax
# your output here

$ pip freeze | grep tensor
# your output here

$ pip freeze | grep jax
# your output here

$ python -V
# your output here

For bugs: reproduction and error logs

# Steps to reproduce:
...

# Error logs:
...

[demdecuong]

I replace that line with for i, img in enumerate(DATA[:trax.math.device_count()]):

[asj100]

Thanks - works perfectly. Is anyone going to change the original code?