Tutorial on how to extract reformer layers as tf.layers

[gaceladri]

hello, do you have any tutorial on how to extract the reversible self attention layer as a tf.layer? Is it possible? Could it be possible to just take the self attention layer and integrate it to Bert? It would be amazing! Or any tutorial on bow to integrate jax with tf also will be amazing. Thanks!

[lukaszkaiser]

We don't have a tutorial yet and it only works with reference_code=True for LSH attention, but you can already get a Keras layer out of it using this file: https://github.com/google/trax/blob/master/trax/trax2keras.py#L60

More support and tests coming!