feat: make cross_entropy_chunk_size tunable, with an auto memory budget (#2190) [3/3]

feat: make cross_entropy_chunk_size tunable, with an auto memory budget (#2190) [3/3] #8662