Failed to sample from the Reformer model about trax

[2877992943]

got the same issue here

[nitinnairk]

Same issue here too!

Environment information

$ pip freeze | grep tensor
mesh-tensorflow==0.1.9
tensor2tensor==1.15.4
tensorboard==2.0.2
tensorflow-datasets==2.0.0
tensorflow-estimator==2.0.1
tensorflow-gan==2.0.0
tensorflow-hub==0.7.0
tensorflow-metadata==0.21.1
tensorflow-probability==0.7.0

$ pip freeze | grep jax
jax==0.1.59
jaxlib==0.1.39

$ python -V
Python 3.7.4

Changes

From

# Prepare a jitted copy of the model.
jit_model_infer = trax.layers.base._accelerate(
    model_infer._forward_internal, trax.math.device_count())
# Set up the initial state for sampling.
infer_state = model_infer.new_weights_and_state(
    trax.supervised.trainer_lib.ShapeDtype((1,1), dtype=np.int32))[1]
infer_state = trainer._for_n_devices(infer_state)

To

# Prepare a jitted copy of the model.
jit_model_infer = trax.layers.base._accelerate(
    model_infer.pure_fn, trax.math.device_count())
# Set up the initial state for sampling.
infer_state = model_infer.new_weights_and_state(
    trax.supervised.trainer_lib.ShapeDtype((1,1), dtype=np.int32))[1]
infer_state = trainer._for_n_devices(infer_state)

[nitinnairk]

Changing the below line in the Sample function fixed the issue

From

rngs = trax.math.random.split(trax.math.random.get_prng(0), trax.math.device_count())

To

rngs = trax.math.random.get_prng(0)

[lukaszkaiser]

Thanks for letting us know! Should be corrected at head now.