Unable to train a reformer model using tensorflow-numpy due to incompatable types

[JEF1056]

Description

Unable to train a reformer model using tensorflow-numpy due to incompatable types

Environment information

OS: Colaboratory

$ pip freeze | grep trax
trax==1.3.6 (same issue in 1.3.7)

$ pip freeze | grep tensor
mesh-tensorflow==0.1.18
tensorboard==2.4.0
tensorboard-plugin-wit==1.7.0
tensorboardcolab==0.0.22
tensorflow==2.4.0
tensorflow-addons==0.8.3
tensorflow-datasets==4.0.1
tensorflow-estimator==2.4.0
tensorflow-gcs-config==2.4.0
tensorflow-hub==0.10.0
tensorflow-metadata==0.26.0
tensorflow-privacy==0.2.2
tensorflow-probability==0.11.0
tensorflow-text==2.4.2

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

$ python -V
Python 3.6.9

For bugs: reproduction and error logs

# Steps to reproduce:
1. Open https://colab.research.google.com/drive/1S8KjoG7vSeGA8yj1dYI3An9SpzcEShxl?usp=sharing (the reformer example but modified for tensorflow-numpy)
2. Run all the blocks down to the 9th block, where the first step is run
3. Observe the error

# Error logs:
WARNING:tensorflow:The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int64
WARNING:tensorflow:The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int64
WARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int64
WARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int64
WARNING:tensorflow:The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int64
WARNING:tensorflow:The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int64
WARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int64
WARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int64
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py in _apply_op_helper(op_type_name, name, **keywords)
    521                 as_ref=input_arg.is_ref,
--> 522                 preferred_dtype=default_dtype)
    523         except TypeError as err:

48 frames
ValueError: Tensor conversion requested dtype float32 for Tensor with dtype float64: <tf.Tensor 'truediv_3:0' shape=(256,) dtype=float64>

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py in _apply_op_helper(op_type_name, name, **keywords)
    556                 "%s type %s of argument '%s'." %
    557                 (prefix, dtypes.as_dtype(attrs[input_arg.type_attr]).name,
--> 558                  inferred_from[input_arg.type_attr]))
    559 
    560         types = [values.dtype]

TypeError: Input 'y' of 'Mul' Op has type float64 that does not match type float32 of argument 'x'.