[DOCS] CLEANUP]

kyegomez · Feb 22, 2024 · c351720 · c351720
1 parent dbd6a82
commit c351720
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 104 deletions.
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -182,116 +182,124 @@ nav:
  - MixtureOfAutoregressiveAttention: "zeta/nn/attention/mixture_of_attention_ar.md"
  - SparseAttention: "zeta/nn/attention/sparse_attn.md"
  - zeta.tokenizers:
- - MultiModalTokenizer: "zeta/tokenizers/multi_modal_tokenizer.md"
- - LanguageTokenizerGPTX: "zeta/tokenizers/language_tokenizer.md"
- - SentencePieceTokenizer: "zeta/tokenizers/sentencepiece.md"
- - TokenMonster: "zeta/tokenizers/token_monster.md"
+ - Language:
+ - LanguageTokenizerGPTX: "zeta/tokenizers/language_tokenizer.md"
+ - SentencePieceTokenizer: "zeta/tokenizers/sentencepiece.md"
+ - TokenMonster: "zeta/tokenizers/token_monster.md"
+ - MultiModal:
+ - MultiModalTokenizer: "zeta/tokenizers/multi_modal_tokenizer.md"
+
  - zeta.utils:
- - cast_tuple: "zeta/utils/cast_tuple.md"
- - group_by_key_prefix: "zeta/utils/group_by_key_prefix.md"
- - eval_decorator: "zeta/utils/eval_decorator.md"
- - print_cuda_memory_usage: "zeta/utils/print_cuda_memory_usage.md"
- - once: "zeta/utils/once.md"
- - default: "zeta/utils/default.md"
- - gumbel_noise: "zeta/utils/gumbel_noise.md"
- - pad_at_dim: "zeta/utils/pad_at_dim.md"
- - init_zero_: "zeta/utils/init_zero_.md"
- - top_p: "zeta/utils/top_p.md"
- - cast_if_src_dtype: "zeta/utils/cast_if_src_dtype.md"
- - disable_warnings_and_logs: "zeta/utils/disable_warnings_and_logs.md"
- - save_load_wrapper: "zeta/utils/save_load_wrapper.md"
- - get_sinusoid_encoding_table: "zeta/utils/get_sinusoid_encoding_table.md"
- - main: "zeta/utils/main.md"
- - string_begins_with: "zeta/utils/string_begins_with.md"
- - gif_to_tensor: "zeta/utils/gif_to_tensor.md"
- - l2norm: "zeta/utils/l2norm.md"
- - save_load: "zeta/utils/save_load.md"
- - log: "zeta/utils/log.md"
- - module_device: "zeta/utils/module_device.md"
- - print_num_params: "zeta/utils/print_num_params.md"
- - top_a: "zeta/utils/top_a.md"
- - interpolate_pos_encoding_2d: "zeta/utils/interpolate_pos_encoding_2d.md"
- - exists: "zeta/utils/exists.md"
- - cosine_beta_schedule: "zeta/utils/cosine_beta_schedule.md"
- - track_cuda_memory: "zeta/utils/track_cuda_memory.md"
- - maybe: "zeta/utils/maybe.md"
- - save_memory_snapshot: "zeta/utils/save_memory_snapshot.md"
- - top_k: "zeta/utils/top_k.md"
- - print_main: "zeta/utils/print_main.md"
- - pick_and_pop: "zeta/utils/pick_and_pop.md"
- - track_cuda_memory_usage: "zeta/utils/track_cuda_memory_usage.md"
- - group_dict_by_key: "zeta/utils/group_dict_by_key.md"
- - video_tensor_to_gift: "zeta/utils/video_tensor_to_gift.md"
+ - Misc:
+ - cast_tuple: "zeta/utils/cast_tuple.md"
+ - group_by_key_prefix: "zeta/utils/group_by_key_prefix.md"
+ - eval_decorator: "zeta/utils/eval_decorator.md"
+ - print_cuda_memory_usage: "zeta/utils/print_cuda_memory_usage.md"
+ - once: "zeta/utils/once.md"
+ - default: "zeta/utils/default.md"
+ - gumbel_noise: "zeta/utils/gumbel_noise.md"
+ - pad_at_dim: "zeta/utils/pad_at_dim.md"
+ - init_zero_: "zeta/utils/init_zero_.md"
+ - top_p: "zeta/utils/top_p.md"
+ - cast_if_src_dtype: "zeta/utils/cast_if_src_dtype.md"
+ - disable_warnings_and_logs: "zeta/utils/disable_warnings_and_logs.md"
+ - save_load_wrapper: "zeta/utils/save_load_wrapper.md"
+ - get_sinusoid_encoding_table: "zeta/utils/get_sinusoid_encoding_table.md"
+ - main: "zeta/utils/main.md"
+ - string_begins_with: "zeta/utils/string_begins_with.md"
+ - gif_to_tensor: "zeta/utils/gif_to_tensor.md"
+ - l2norm: "zeta/utils/l2norm.md"
+ - save_load: "zeta/utils/save_load.md"
+ - log: "zeta/utils/log.md"
+ - module_device: "zeta/utils/module_device.md"
+ - print_num_params: "zeta/utils/print_num_params.md"
+ - top_a: "zeta/utils/top_a.md"
+ - interpolate_pos_encoding_2d: "zeta/utils/interpolate_pos_encoding_2d.md"
+ - exists: "zeta/utils/exists.md"
+ - cosine_beta_schedule: "zeta/utils/cosine_beta_schedule.md"
+ - track_cuda_memory: "zeta/utils/track_cuda_memory.md"
+ - maybe: "zeta/utils/maybe.md"
+ - save_memory_snapshot: "zeta/utils/save_memory_snapshot.md"
+ - top_k: "zeta/utils/top_k.md"
+ - print_main: "zeta/utils/print_main.md"
+ - pick_and_pop: "zeta/utils/pick_and_pop.md"
+ - track_cuda_memory_usage: "zeta/utils/track_cuda_memory_usage.md"
+ - group_dict_by_key: "zeta/utils/group_dict_by_key.md"
+ - video_tensor_to_gift: "zeta/utils/video_tensor_to_gift.md"
  - zeta.ops:
- - img_compose_decompose: "zeta/ops/img_compose_decompose.md"
- - img_transpose_2daxis: "zeta/ops/img_transpose_2daxis.md"
- - img_transpose: "zeta/ops/img_transpose.md"
- - img_order_of_axes: "zeta/ops/img_order_of_axes.md"
- - mos: "zeta/ops/mos.md"
- - merge_small_dims: "zeta/ops/merge_small_dims.md"
- - multi_dim_cat: "zeta/ops/multi_dim_cat.md"
- - img_compose_bw: "zeta/ops/img_compose_bw.md"
- - squeeze_2d_new: "zeta/ops/squeeze_2d_new.md"
- - temp_softmax: "zeta/ops/temp_softmax.md"
- - gumbelmax: "zeta/ops/gumbelmax.md"
- - _matrix_inverse_root_newton: "zeta/ops/_matrix_inverse_root_newton.md"
- - compute_matrix_root_inverse_residuals: "zeta/ops/compute_matrix_root_inverse_residuals.md"
- - matrix_root_diagonal: "zeta/ops/matrix_root_diagonal.md"
- - sparse_softmax: "zeta/ops/sparse_softmax.md"
- - reshape_audio_to_text: "zeta/ops/reshape_audio_to_text.md"
- - local_softmax: "zeta/ops/local_softmax.md"
- - softmaxes: "zeta/ops/softmaxes.md"
- - _matrix_root_eigen: "zeta/ops/_matrix_root_eigen.md"
- - main: "zeta/ops/main.md"
- - norm_exp_softmax: "zeta/ops/norm_exp_softmax.md"
- - multi_dim_split: "zeta/ops/multi_dim_split.md"
- - img_width_to_height: "zeta/ops/img_width_to_height.md"
- - fast_softmax: "zeta/ops/fast_softmax.md"
- - standard_softmax: "zeta/ops/standard_softmax.md"
- - unitwise_norm: "zeta/ops/unitwise_norm.md"
- - reshape_video_to_text: "zeta/ops/reshape_video_to_text.md"
- - img_decompose: "zeta/ops/img_decompose.md"
- - unsqueeze_2d_new: "zeta/ops/unsqueeze_2d_new.md"
- - reshape_img_to_text: "zeta/ops/reshape_img_to_text.md"
- - channel_shuffle_new: "zeta/ops/channel_shuffle_new.md"
- - matrix_inverse_root: "zeta/ops/matrix_inverse_root.md"
- - sparsemax: "zeta/ops/sparsemax.md"
- - gram_matrix_new: "zeta/ops/gram_matrix_new.md"
- - logit_scaled_softmax: "zeta/ops/logit_scaled_softmax.md"
- - selu_softmax: "zeta/ops/selu_softmax.md"
- - reshape_text_to_img: "zeta/ops/reshape_text_to_img.md"
+ - Misc:
+ - img_compose_decompose: "zeta/ops/img_compose_decompose.md"
+ - img_transpose_2daxis: "zeta/ops/img_transpose_2daxis.md"
+ - img_transpose: "zeta/ops/img_transpose.md"
+ - img_order_of_axes: "zeta/ops/img_order_of_axes.md"
+ - mos: "zeta/ops/mos.md"
+ - merge_small_dims: "zeta/ops/merge_small_dims.md"
+ - multi_dim_cat: "zeta/ops/multi_dim_cat.md"
+ - img_compose_bw: "zeta/ops/img_compose_bw.md"
+ - squeeze_2d_new: "zeta/ops/squeeze_2d_new.md"
+ - temp_softmax: "zeta/ops/temp_softmax.md"
+ - gumbelmax: "zeta/ops/gumbelmax.md"
+ - _matrix_inverse_root_newton: "zeta/ops/_matrix_inverse_root_newton.md"
+ - compute_matrix_root_inverse_residuals: "zeta/ops/compute_matrix_root_inverse_residuals.md"
+ - matrix_root_diagonal: "zeta/ops/matrix_root_diagonal.md"
+ - sparse_softmax: "zeta/ops/sparse_softmax.md"
+ - reshape_audio_to_text: "zeta/ops/reshape_audio_to_text.md"
+ - local_softmax: "zeta/ops/local_softmax.md"
+ - softmaxes: "zeta/ops/softmaxes.md"
+ - _matrix_root_eigen: "zeta/ops/_matrix_root_eigen.md"
+ - main: "zeta/ops/main.md"
+ - norm_exp_softmax: "zeta/ops/norm_exp_softmax.md"
+ - multi_dim_split: "zeta/ops/multi_dim_split.md"
+ - img_width_to_height: "zeta/ops/img_width_to_height.md"
+ - fast_softmax: "zeta/ops/fast_softmax.md"
+ - standard_softmax: "zeta/ops/standard_softmax.md"
+ - unitwise_norm: "zeta/ops/unitwise_norm.md"
+ - reshape_video_to_text: "zeta/ops/reshape_video_to_text.md"
+ - img_decompose: "zeta/ops/img_decompose.md"
+ - unsqueeze_2d_new: "zeta/ops/unsqueeze_2d_new.md"
+ - reshape_img_to_text: "zeta/ops/reshape_img_to_text.md"
+ - channel_shuffle_new: "zeta/ops/channel_shuffle_new.md"
+ - matrix_inverse_root: "zeta/ops/matrix_inverse_root.md"
+ - sparsemax: "zeta/ops/sparsemax.md"
+ - gram_matrix_new: "zeta/ops/gram_matrix_new.md"
+ - logit_scaled_softmax: "zeta/ops/logit_scaled_softmax.md"
+ - selu_softmax: "zeta/ops/selu_softmax.md"
+ - reshape_text_to_img: "zeta/ops/reshape_text_to_img.md"
  - zeta.optim:
- - StableAdamWUnfused: "zeta/optims/adamw.md"
- - GradientAscent: "zeta/optims/ga.md"
- - DecoupledLionW: "zeta/training/optimizers/decoupled_lion.md"
- - SophiaG: "zeta/training/optimizers/sophia.md"
+ - Optimizers:
+ - StableAdamWUnfused: "zeta/optims/adamw.md"
+ - GradientAscent: "zeta/optims/ga.md"
+ - DecoupledLionW: "zeta/training/optimizers/decoupled_lion.md"
+ - SophiaG: "zeta/training/optimizers/sophia.md"
  - zeta.training:
  - fsdp: "zeta/training/fsdp.md"
  - ParallelWrapper: "zeta/training/parallel_wrapper.md"
  - train: "zeta/training/train.md"
  - zeta.models:
- - vit: "zeta/models/vit.md"
- - gpt4multimodal: "zeta/models/gpt4multimodal.md"
- - maxvit: "zeta/models/maxvit.md"
- - llama2: "zeta/models/llama2.md"
- - gpt4: "zeta/models/gpt4.md"
- - andromeda: "zeta/models/andromeda.md"
- - basemodel: "zeta/models/basemodel.md"
- - palme: "zeta/models/palme.md"
- - megavit: "zeta/models/megavit.md"
- - navit: "zeta/models/navit.md"
+ - Language and MultiModal:
+ - vit: "zeta/models/vit.md"
+ - gpt4multimodal: "zeta/models/gpt4multimodal.md"
+ - maxvit: "zeta/models/maxvit.md"
+ - llama2: "zeta/models/llama2.md"
+ - gpt4: "zeta/models/gpt4.md"
+ - andromeda: "zeta/models/andromeda.md"
+ - basemodel: "zeta/models/basemodel.md"
+ - palme: "zeta/models/palme.md"
+ - megavit: "zeta/models/megavit.md"
+ - navit: "zeta/models/navit.md"
  - zeta.structs:
- - Decoder: "zeta/nn/architecture/decoder.md"
- - Transformer: "zeta/nn/architecture/transformer.md"
- - TransformerBlock: "zeta/nn/architecture/transformerblock.md"
- - paralleltransformerblock: "paralleltransformerblock.md"
- - hierarchicalblock: "hierarchicalblock.md"
- - vitransformerwrapper: "vitransformerwrapper.md"
- - localtransformer: "localtransformer.md"
- - autoregressivewrapper: "autoregressivewrapper.md"
- - simpletransformer: "simpletransformer.md"
- - encoder: "encoder.md"
- - encoderdecoder: "encoderdecoder.md"
+ - Structures:
+ - Decoder: "zeta/nn/architecture/decoder.md"
+ - Transformer: "zeta/nn/architecture/transformer.md"
+ - TransformerBlock: "zeta/nn/architecture/transformerblock.md"
+ - paralleltransformerblock: "paralleltransformerblock.md"
+ - hierarchicalblock: "hierarchicalblock.md"
+ - vitransformerwrapper: "vitransformerwrapper.md"
+ - localtransformer: "localtransformer.md"
+ - autoregressivewrapper: "autoregressivewrapper.md"
+ - simpletransformer: "simpletransformer.md"
+ - encoder: "encoder.md"
+ - encoderdecoder: "encoderdecoder.md"
  - zeta.quant:
  - QUIK: "zeta/quant/quik.md"
  - BitLinear: "zeta/quant/bitlinear.md"

diff --git a/zeta/utils/__init__.py b/zeta/utils/__init__.py
@@ -1,4 +1,3 @@
-
 from zeta.utils.cuda_memory_wrapper import track_cuda_memory_usage
 
 from zeta.utils.benchmark import (
@@ -90,5 +89,5 @@
  "append_nvcc_threads",
  "check_cuda",
  "VerboseExecution",
- "seek_all_images"
-]
+ "seek_all_images",
+]