def __getattribute__(self, name): """ Decorate callable functions of RandomState such that they return chunks of the total `N` random numbers generated """ attr = RandomState.__getattribute__(self, name) if callable(attr) and not getattr(attr, 'mpi_enabled', False): attr = _mpi_enabled_rng(attr) return attr