tmp loc

arguments.py

data_utils/configure_data.py

draw_diff.py

env/baai_setup_connection.py

env/setup_connection.py

generate_samples.py

generation/__init__.py

generation/local_sampling.py

generation/sampling.py

mpu/local_attention_function.py

mpu/sparse_transformer.py

pretrain_gpt2.py

random_display.py

scripts/cuda_2d_text2image.sh

scripts/ds_config_zero.json

scripts/pretrain_multiple_nodes.sh

scripts/pretrain_single_node.sh

scripts/testnan.sh

scripts/text2image.sh

test_sparse_attention.py