Update cos_dataset.py

Added make_rng("uniform", "randn")