jitting the step() method in scheduler

Author: prishajain1

src/maxdiffusion/schedulers/scheduling_flow_match_flax.py

@@ -15,6 +15,7 @@
 # DISCLAIMER: This is a JAX/Flax conversion of a PyTorch implementation.
 # The original PyTorch code was provided by the user.
 
+from functools import partial
 from typing import Optional, Tuple, Union
 
 import flax
@@ -243,6 +244,7 @@ def _find_timestep_id(self, state: FlowMatchSchedulerState, timestep: jnp.ndarra
       diffs = jnp.abs(state.timesteps[None, :] - timestep[:, None])
       return jnp.argmin(diffs, axis=1)
 
+  @partial(jax.jit, static_argnums=(0, 5, 6))
   def step(
       self,
       state: FlowMatchSchedulerState,