Removing Documents

erfanzar · Nov 27, 2023 · 857c393 · 857c393
1 parent 1a5cdc9
commit 857c393
Show file tree

Hide file tree

Showing 47 changed files with 577 additions and 40,242 deletions.
diff --git a/documents/fjformer.html b/documents/fjformer.html
diff --git a/documents/fjformer/attention.html b/documents/fjformer/attention.html
diff --git a/documents/fjformer/attention/blockwise_attention.html b/documents/fjformer/attention/blockwise_attention.html
diff --git a/documents/fjformer/attention/efficient_attention.html b/documents/fjformer/attention/efficient_attention.html
diff --git a/documents/fjformer/attention/flash_attention.html b/documents/fjformer/attention/flash_attention.html
diff --git a/documents/fjformer/attention/flash_attention_0.html b/documents/fjformer/attention/flash_attention_0.html
diff --git a/documents/fjformer/attention/flash_attention_gpu.html b/documents/fjformer/attention/flash_attention_gpu.html
diff --git a/documents/fjformer/attention/flash_attention_tpu.html b/documents/fjformer/attention/flash_attention_tpu.html
diff --git a/documents/fjformer/attention/jax_flash_attn_tpu.html b/documents/fjformer/attention/jax_flash_attn_tpu.html
diff --git a/documents/fjformer/bits.html b/documents/fjformer/bits.html
diff --git a/documents/fjformer/bits/bits.html b/documents/fjformer/bits/bits.html
diff --git a/documents/fjformer/bits/calibration.html b/documents/fjformer/bits/calibration.html
diff --git a/documents/fjformer/bits/config.html b/documents/fjformer/bits/config.html
diff --git a/documents/fjformer/bits/int_numerics.html b/documents/fjformer/bits/int_numerics.html
diff --git a/documents/fjformer/bits/q_dot_general.html b/documents/fjformer/bits/q_dot_general.html
diff --git a/documents/fjformer/bits/q_flax.html b/documents/fjformer/bits/q_flax.html
diff --git a/documents/fjformer/bits/stochastic_rounding.html b/documents/fjformer/bits/stochastic_rounding.html
diff --git a/documents/fjformer/datasets.html b/documents/fjformer/datasets.html
diff --git a/documents/fjformer/datasets/datasets.html b/documents/fjformer/datasets/datasets.html
diff --git a/documents/fjformer/func.html b/documents/fjformer/func.html
diff --git a/documents/fjformer/func/_func.html b/documents/fjformer/func/_func.html
diff --git a/documents/fjformer/func/loss_func.html b/documents/fjformer/func/loss_func.html
diff --git a/documents/fjformer/gpu_pallas.html b/documents/fjformer/gpu_pallas.html
diff --git a/documents/fjformer/gpu_pallas/attention.html b/documents/fjformer/gpu_pallas/attention.html
diff --git a/documents/fjformer/gpu_pallas/layer_norm.html b/documents/fjformer/gpu_pallas/layer_norm.html
diff --git a/documents/fjformer/gpu_pallas/rms_norm.html b/documents/fjformer/gpu_pallas/rms_norm.html
diff --git a/documents/fjformer/gpu_pallas/softmax.html b/documents/fjformer/gpu_pallas/softmax.html
diff --git a/documents/fjformer/load.html b/documents/fjformer/load.html
diff --git a/documents/fjformer/load/_load.html b/documents/fjformer/load/_load.html
diff --git a/documents/fjformer/load/streamer.html b/documents/fjformer/load/streamer.html
diff --git a/documents/fjformer/monitor.html b/documents/fjformer/monitor.html
diff --git a/documents/fjformer/monitor/tracker.html b/documents/fjformer/monitor/tracker.html
diff --git a/documents/fjformer/optimizers.html b/documents/fjformer/optimizers.html
diff --git a/documents/fjformer/optimizers/adafactor.html b/documents/fjformer/optimizers/adafactor.html
diff --git a/documents/fjformer/optimizers/adamw.html b/documents/fjformer/optimizers/adamw.html
diff --git a/documents/fjformer/optimizers/lion.html b/documents/fjformer/optimizers/lion.html
diff --git a/documents/fjformer/optimizers/optimizer_utils.html b/documents/fjformer/optimizers/optimizer_utils.html
diff --git a/documents/fjformer/partition_utils.html b/documents/fjformer/partition_utils.html
diff --git a/documents/fjformer/partition_utils/mesh_utils.html b/documents/fjformer/partition_utils/mesh_utils.html
diff --git a/documents/fjformer/partition_utils/t5x_partitioning.html b/documents/fjformer/partition_utils/t5x_partitioning.html
diff --git a/documents/fjformer/utils.html b/documents/fjformer/utils.html
diff --git a/documents/index.html b/documents/index.html
diff --git a/documents/search.js b/documents/search.js
diff --git a/fjformer/__init__.py b/fjformer/__init__.py
@@ -151,4 +151,4 @@
     count_num_params
 )
 
-__version__ = '0.0.14'
+__version__ = '0.0.15'
diff --git a/fjformer/attention/__init__.py b/fjformer/attention/__init__.py
@@ -16,10 +16,14 @@
  `blockwise_attn` from `flash_attention` module: These imports bring in various attention mechanisms such as
  ring attention, flash attention for GPU and TPU, and blockwise feed-forward network and attention mechanisms.
 
+4. `tpu_flash_attention` `gpu_flash_attention` are imported from jax_flash_attn_tpu/gpu
+
 """
 
 from .efficient_attention import efficient_attention
 from .flash_attention_0 import dot_product_attention_multihead, dot_product_attention_multiquery, \
     dot_product_attention_queries_per_head
 from .flash_attention import ring_attention, ring_attention_standard, ring_flash_attention_gpu, \
     ring_flash_attention_tpu, blockwise_ffn, blockwise_attn
+from .jax_flash_attn_tpu import flash_attention as tpu_flash_attention
+from .jax_flash_attn_gpu import mha as gpu_flash_attention