diff --git a/mkdocs.yml b/mkdocs.yml index ab62f6f9..35697146 100644 --- 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 index 7a060a98..4dc76d47 100644 --- 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" -] \ No newline at end of file + "seek_all_images", +]