def set_timesteps(
        self, state: DDPMSchedulerState, num_inference_steps: int, shape: Tuple = ()
    ) -> DDPMSchedulerState:
        """
        Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.

        Args:
            state (`DDIMSchedulerState`):
                the `FlaxDDPMScheduler` state data class instance.
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
        """
        num_inference_steps = min(self.config.num_train_timesteps, num_inference_steps)
        timesteps = jnp.arange(
            0, self.config.num_train_timesteps, self.config.num_train_timesteps // num_inference_steps
        )[::-1]
        return state.replace(num_inference_steps=num_inference_steps, timesteps=timesteps)

        Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.

        Args:
            state (`DDIMSchedulerState`):
                the `FlaxDDPMScheduler` state data class instance.
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
        
promptist/diffusers/src/diffusers/schedulers/scheduling_ddpm_flax.py