Updating Docs

erfanzar · Aug 12, 2024 · f02053e · f02053e
1 parent a4e6b8a
commit f02053e
Show file tree

Hide file tree

Showing 14 changed files with 47 additions and 39 deletions.
diff --git a/docs/api_docs/APIs.rst b/docs/api_docs/APIs.rst
@@ -7,7 +7,7 @@ FJFormer APIs 🔮
    Bit Quantization
    Checkpoint
    Core
-   Custom Array
+   Dtypes
    Functions
    Lora
    Monitor

diff --git a/docs/api_docs/Bit Quantization.rst b/docs/api_docs/Bit Quantization.rst
@@ -4,11 +4,11 @@ Bit Quantization
 .. toctree::
    :maxdepth: 2
 
-   bit_quantization.q_dot_general
-   bit_quantization.config
-   bit_quantization.no_numerics
-   bit_quantization.q_flax
-   bit_quantization.calibration
-   bit_quantization.stochastic_rounding
    bit_quantization.numerics
+   bit_quantization.calibration
+   bit_quantization.q_flax
+   bit_quantization.q_dot_general
    bit_quantization.int_numerics
+   bit_quantization.stochastic_rounding
+   bit_quantization.no_numerics
+   bit_quantization.config
diff --git a/docs/api_docs/Core.rst b/docs/api_docs/Core.rst
@@ -5,7 +5,7 @@ Core
    :maxdepth: 2
 
    core.symbols
-   core.errors
    core.types
    core.utilities
    core.implicit_array
+   core.errors
diff --git a/docs/api_docs/Custom Array.rst b/docs/api_docs/Custom Array.rst
diff --git a/docs/api_docs/Dtypes.rst b/docs/api_docs/Dtypes.rst
@@ -0,0 +1,9 @@
+Dtypes
+======
+
+.. toctree::
+   :maxdepth: 2
+
+   dtypes.arraynf4
+   dtypes.array8lt
+   dtypes.array8bit
diff --git a/docs/api_docs/Functions.rst b/docs/api_docs/Functions.rst
@@ -4,5 +4,5 @@ Functions
 .. toctree::
    :maxdepth: 2
 
-   functions._functions
    functions.loss_functions
+   functions._functions
diff --git a/docs/api_docs/Lora.rst b/docs/api_docs/Lora.rst
@@ -4,5 +4,5 @@ Lora
 .. toctree::
    :maxdepth: 2
 
-   lora.lora_core
    lora.rapture
+   lora.lora_core
diff --git a/docs/api_docs/Optimizers.rst b/docs/api_docs/Optimizers.rst
@@ -4,8 +4,8 @@ Optimizers
 .. toctree::
    :maxdepth: 2
 
+   optimizers.adafactor
    optimizers.adamw
+   optimizers.rmsprop
    optimizers.lion
    optimizers.optimizer_utils
-   optimizers.rmsprop
-   optimizers.adafactor
diff --git a/docs/api_docs/Pallas Operations.rst b/docs/api_docs/Pallas Operations.rst
@@ -4,15 +4,15 @@ Pallas Operations
 .. toctree::
    :maxdepth: 2
 
-   pallas_operations.efficient_attention.efficient_attention
-   pallas_operations.pallas_attention.attention
-   pallas_operations.gpu.flash_attention.mha
-   pallas_operations.gpu.layer_norm.layer_norm
-   pallas_operations.gpu.softmax.softmax
-   pallas_operations.gpu.rms_norm.rms_norm
+   pallas_operations.tpu.ring_attention.ring_attention
    pallas_operations.tpu.flash_attention.flash_attention
-   pallas_operations.tpu.paged_attention.paged_attention
+   pallas_operations.tpu.splash_attention.splash_attention_kernel
    pallas_operations.tpu.splash_attention.splash_attention_mask_info
    pallas_operations.tpu.splash_attention.splash_attention_mask
-   pallas_operations.tpu.splash_attention.splash_attention_kernel
-   pallas_operations.tpu.ring_attention.ring_attention
+   pallas_operations.tpu.paged_attention.paged_attention
+   pallas_operations.pallas_attention.attention
+   pallas_operations.gpu.softmax.softmax
+   pallas_operations.gpu.layer_norm.layer_norm
+   pallas_operations.gpu.rms_norm.rms_norm
+   pallas_operations.gpu.flash_attention.mha
+   pallas_operations.efficient_attention.efficient_attention
diff --git a/docs/api_docs/Sharding.rst b/docs/api_docs/Sharding.rst
@@ -4,5 +4,5 @@ Sharding
 .. toctree::
    :maxdepth: 2
 
-   sharding.sharding
    sharding.t5x_partitioning
+   sharding.sharding
diff --git a/docs/api_docs/custom_array.array4bit.rst b/docs/api_docs/custom_array.array4bit.rst
diff --git a/docs/api_docs/custom_array.array8bit.rst → docs/api_docs/dtypes.array8bit.rst b/docs/api_docs/custom_array.array8bit.rst → docs/api_docs/dtypes.array8bit.rst
@@ -1,7 +1,7 @@
-Array8Lt package
+array8bit package
 =================
 
-.. automodule:: fjformer.custom_array.Array8Lt
+.. automodule:: fjformer.dtypes.array8bit
     :members:
     :undoc-members:
     :show-inheritance:
diff --git a/docs/api_docs/dtypes.array8lt.rst b/docs/api_docs/dtypes.array8lt.rst
@@ -0,0 +1,7 @@
+array8lt package
+================
+
+.. automodule:: fjformer.dtypes.array8lt
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/api_docs/dtypes.arraynf4.rst b/docs/api_docs/dtypes.arraynf4.rst
@@ -0,0 +1,7 @@
+arraynf4 package
+================
+
+.. automodule:: fjformer.dtypes.arraynf4
+    :members:
+    :undoc-members:
+    :show-inheritance: