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

        Args:
            state (`PNDMSchedulerState`):
                the `FlaxPNDMScheduler` state data class instance.
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
            shape (`Tuple`):
                the shape of the samples to be generated.
        """
        offset = self.config.steps_offset

        step_ratio = self.config.num_train_timesteps // num_inference_steps
        # creates integer timesteps by multiplying by ratio
        # rounding to avoid issues when num_inference_step is power of 3
        _timesteps = (jnp.arange(0, num_inference_steps) * step_ratio).round() + offset

        state = state.replace(num_inference_steps=num_inference_steps, _timesteps=_timesteps)

        if self.config.skip_prk_steps:
            # for some models like stable diffusion the prk steps can/should be skipped to
            # 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
            state = state.replace(
                prk_timesteps=jnp.array([]),
                plms_timesteps=jnp.concatenate(
                    [state._timesteps[:-1], state._timesteps[-2:-1], state._timesteps[-1:]]
                )[::-1],
            )
        else:
            prk_timesteps = jnp.array(state._timesteps[-self.pndm_order :]).repeat(2) + jnp.tile(
                jnp.array([0, self.config.num_train_timesteps // num_inference_steps // 2]), self.pndm_order
            )

            state = state.replace(
                prk_timesteps=(prk_timesteps[:-1].repeat(2)[1:-1])[::-1],
                plms_timesteps=state._timesteps[:-3][::-1],
            )

        return state.replace(
            timesteps=jnp.concatenate([state.prk_timesteps, state.plms_timesteps]).astype(jnp.int32),
            counter=0,
            # Reserve space for the state variables
            cur_model_output=jnp.zeros(shape),
            cur_sample=jnp.zeros(shape),
            ets=jnp.zeros((4,) + shape),
        )

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

        Args:
            state (`PNDMSchedulerState`):
                the `FlaxPNDMScheduler` state data class instance.
            num_inference_steps (`int`):
                the number of diffusion steps used when generating samples with a pre-trained model.
            shape (`Tuple`):
                the shape of the samples to be generated.
        
promptist/diffusers/src/diffusers/schedulers/scheduling_pndm_flax.py