Fix prk_timesteps and plms_timesteps dtype

Author: skirsten

src/diffusers/schedulers/scheduling_pndm_flax.py

@@ -183,19 +183,19 @@ def set_timesteps(self, state: PNDMSchedulerState, num_inference_steps: int, sha
             # produce better results. When using PNDM with `self.config.skip_prk_steps` the implementation
             # is based on crowsonkb's PLMS sampler implementation: https://github.com/CompVis/latent-diffusion/pull/51
 
-            prk_timesteps = jnp.array([], dtype=self.config.dtype)
+            prk_timesteps = jnp.array([], dtype=jnp.int32)
             plms_timesteps = jnp.concatenate([_timesteps[:-1], _timesteps[-2:-1], _timesteps[-1:]])[::-1]
 
         else:
             prk_timesteps = jnp.array(_timesteps[-self.pndm_order :]).repeat(2) + jnp.tile(
-                jnp.array([0, self.config.num_train_timesteps // num_inference_steps // 2], dtype=self.config.dtype),
+                jnp.array([0, self.config.num_train_timesteps // num_inference_steps // 2], dtype=jnp.int32),
                 self.pndm_order,
             )
 
             prk_timesteps = (prk_timesteps[:-1].repeat(2)[1:-1])[::-1]
             plms_timesteps = _timesteps[:-3][::-1]
 
-        timesteps = jnp.concatenate([prk_timesteps, plms_timesteps]).astype(jnp.int32)
+        timesteps = jnp.concatenate([prk_timesteps, plms_timesteps])
 
         # initial running values