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: $ 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?