Update all sanning utility with code from seismo-performer repo

geophystech · Aug 31, 2021 · 6a9533a · 6a9533a
1 parent 4d56ced
commit 6a9533a
Show file tree

Hide file tree

Showing 26 changed files with 2,217 additions and 1,393 deletions.
diff --git a/archive_scan.py b/archive_scan.py
diff --git a/fast_attention/README.md b/fast_attention/README.md
@@ -0,0 +1,14 @@
+# Tensorflow Variant of FAVOR+.
+
+The main attention function in the Tensorflow variant is `favor_attention`.
+
+* In order to use softmax attention, set its argument `kernel_transformation=softmax_kernel_transformation`.
+* In order to use generalized ReLU attention, set its argument `kernel_transformation=relu_kernel_transformation`. 
+
+To use as a `tf.keras.layers.Layer` module, use the FAVOR `Attention` class instead (after setting the FAVOR+ configuration). This has a similar API to `tf.keras.layers.Attention()`.
+
+## Notes:
+
+* Just like the Jax case, the causal (unidirectional) variant uses custom gradients via `@tf.custom_gradient`, in order to provide significant memory reductions.
+
+