def parse_args(parser):
    """
    Parse commandline arguments.
    """
    build_model = parser.add_argument_group('model setup')    
    build_model.add_argument( '--num_tokens',      type=int,  default=256, help='')
    build_model.add_argument( '--dim',             type=int,  default=512, help='')
    build_model.add_argument( '--depth',           type=int,  default=  8, help='')
    build_model.add_argument( '--heads',           type=int,  default=  8, help='')       
    build_model.add_argument( '--seq_len',         type=int,  default=512, help='')   
    build_model.add_argument( '--mem_len',         type=int,  default=512, help='')   
    build_model.add_argument( '--cmem_len',        type=int,  default=128, help='')           
    build_model.add_argument( '--num_mem_layers',  type=int,  default=  3, help='')     

    training = parser.add_argument_group('training setup')    

    training.add_argument( '--validate_every',  type=int,  default=100, help='')
    training.add_argument( '--generate_every',  type=int,  default=200, help='')
    training.add_argument( '--prime_length',    type=int,  default= 512, help='')
    training.add_argument( '--generate_length', type=int,  default=1024, help='')    


    optimization = parser.add_argument_group('optimization setup')
    optimization.add_argument( '--optimizer',      type=str,  default='adam',   help='Optimization algorithm')
    optimization.add_argument( '--learning_rate',  type=int,  default=1e-4,     help='learning rate')
    optimization.add_argument( '--num_batches',    type=int,  default=100000,   help='max iteration')
    optimization.add_argument( '--batch_size',     type=int,  default=16,       help='batch size')    
    optimization.add_argument( '--max_batch_size', type=int,  default=4,        help='gradient accumulation')    


    dataset = parser.add_argument_group('dataset parameters')
    dataset.add_argument('--zip_filename',  type=str, default='/content/enwik8.gz',  help='Path to training filelist')
    dataset.add_argument('--num_segments',  type=int, default =   4,  help='num_segments')        
     
    return parser