• Kashif Rasul's avatar
    Karras VE, DDIM and DDPM flax schedulers (#508) · b34be039
    Kashif Rasul authored
    * beta never changes removed from state
    
    * fix typos in docs
    
    * removed unused var
    
    * initial ddim flax scheduler
    
    * import
    
    * added dummy objects
    
    * fix style
    
    * fix typo
    
    * docs
    
    * fix typo in comment
    
    * set return type
    
    * added flax ddom
    
    * fix style
    
    * remake
    
    * pass PRNG key as argument and split before use
    
    * fix doc string
    
    * use config
    
    * added flax Karras VE scheduler
    
    * make style
    
    * fix dummy
    
    * fix ndarray type annotation
    
    * replace returns a new state
    
    * added lms_discrete scheduler
    
    * use self.config
    
    * add_noise needs state
    
    * use config
    
    * use config
    
    * docstring
    
    * added flax score sde ve
    
    * fix imports
    
    * fix typos
    b34be039
__init__.py 2.3 KB