From 6cf990847cf3daa62af55fdd0e69575bd4a61da9 Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashishkmr472@gmail.com>
Date: Wed, 22 Jul 2026 23:32:34 -0400
Subject: [PATCH 1/5] fix(samplers): native step domain, exact ancestral DDPM,
 video shapes, DDIM eta, RK4, DPM reuse

The sampler step domain was hardcoded to [0, 1000] and rescaled to the
schedule's domain internally, so the default start_step=max_timesteps was
wrong for every continuous schedule: the inference pipeline collapsed to a
single no-op step at sigma_min and returned the initial noise. Steps are
now in the schedule's own native domain and the defaults are correct by
construction. Linear spacing on a Karras schedule is already the EDM rho
spacing, so the broken quadratic/karras/exponential spacing modes
(log(0), duplicate int16 steps) are gone.

Also: DDPMSampler is now the algebraically exact ancestral posterior
(handles batches and arbitrary strides; the posterior-table variant
crashed on batch > 1 and ignored the stride), DDIM eta > 0 uses the paper
formula with the missing direction-term reduction, RK4 no longer jits
over a python callable, MultiStepDPM resets its history per trajectory,
Heun falls back to euler when sigma reaches 0, and every sampler reshapes
rates against the actual sample rank so 5D video works.

All samplers are now verified end to end against an analytic gaussian
oracle denoiser; the corresponding strict xfails are flipped to passing
tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 flaxdiff/inference/pipeline.py                |  3 -
 flaxdiff/samplers/common.py                   | 97 ++++---------------
 flaxdiff/samplers/ddim.py                     | 20 ++--
 flaxdiff/samplers/ddpm.py                     | 39 ++++----
 flaxdiff/samplers/euler.py                    | 16 +--
 flaxdiff/samplers/heun_sampler.py             | 19 +++-
 flaxdiff/samplers/multistep_dpm.py            | 11 ++-
 flaxdiff/samplers/rk4_sampler.py              | 12 ++-
 flaxdiff/trainer/diffusion_trainer.py         |  1 -
 flaxdiff/trainer/general_diffusion_trainer.py |  1 -
 tests/test_samplers.py                        | 22 +----
 11 files changed, 89 insertions(+), 152 deletions(-)

diff --git a/flaxdiff/inference/pipeline.py b/flaxdiff/inference/pipeline.py
index 27c7e15..acd72aa 100644
--- a/flaxdiff/inference/pipeline.py
+++ b/flaxdiff/inference/pipeline.py
@@ -224,7 +224,6 @@ def generate_samples(
         diffusion_steps: int = 50,
         guidance_scale: float = 1.0,
         sampler_class=EulerAncestralSampler,
-        timestep_spacing: str = 'linear',
         seed: Optional[int] = None,
         start_step: Optional[int] = None,
         end_step: int = 0,
@@ -241,8 +240,6 @@ def generate_samples(
             guidance_scale=guidance_scale,
             sampler_class=sampler_class,
         )
-        if hasattr(sampler, 'timestep_spacing'):
-            sampler.timestep_spacing = timestep_spacing
         print(f"Generating samples: steps={diffusion_steps}, num_samples={num_samples}, guidance={guidance_scale}")
         
         if use_best_params:
diff --git a/flaxdiff/samplers/common.py b/flaxdiff/samplers/common.py
index 69d1559..ee03df9 100644
--- a/flaxdiff/samplers/common.py
+++ b/flaxdiff/samplers/common.py
@@ -7,7 +7,7 @@
 from typing import List, Tuple, Dict, Any, Optional
 
 from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
-from ..schedulers import NoiseScheduler
+from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple
 from ..utils import RandomMarkovState, MarkovState, clip_images
 from jax.experimental.shard_map import shard_map
 from jax.sharding import Mesh, PartitionSpec as P
@@ -25,7 +25,6 @@ def __init__(
         input_config: DiffusionInputConfig,
         guidance_scale: float = 0.0,
         autoencoder: AutoEncoder = None,
-        timestep_spacing: str = 'linear',
     ):
         """Initialize the diffusion sampler.
         
@@ -36,27 +35,16 @@ def __init__(
             model_output_transform: Transform for model predictions
             guidance_scale: Scale for classifier-free guidance (0.0 means disabled)
             autoencoder: Optional autoencoder for latent diffusion
-            timestep_spacing: Strategy for timestep spacing in sampling
-                             'linear' - Default equal spacing
-                             'quadratic' - Emphasizes early steps
-                             'karras' - Based on EDM paper, better with fewer steps
-                             'exponential' - Concentrates steps near the end
         """
         self.model = model
         self.noise_schedule = noise_schedule
         self.model_output_transform = model_output_transform
         self.guidance_scale = guidance_scale
         self.autoencoder = autoencoder
-        self.timestep_spacing = timestep_spacing
         self.input_config = input_config
         
         unconditionals = input_config.get_unconditionals()
         
-        # For Karras spacing if needed
-        if hasattr(noise_schedule, 'min_inv_rho') and hasattr(noise_schedule, 'max_inv_rho'):
-            self.min_inv_rho = noise_schedule.min_inv_rho
-            self.max_inv_rho = noise_schedule.max_inv_rho
-        
         if self.guidance_scale > 0:
             # Classifier free guidance
             print("Using classifier-free guidance")
@@ -65,7 +53,7 @@ def sample_model(params, x_t, t, *conditioning_inputs):
                 # Concatenate unconditional and conditional inputs
                 x_t_cat = jnp.concatenate([x_t] * 2, axis=0)
                 t_cat = jnp.concatenate([t] * 2, axis=0)
-                rates_cat = self.noise_schedule.get_rates(t_cat)
+                rates_cat = self.noise_schedule.get_rates(t_cat, get_coeff_shapes_tuple(x_t_cat))
                 c_in_cat = self.model_output_transform.get_input_scale(rates_cat)
                 
                 final_conditionals = []
@@ -92,7 +80,7 @@ def sample_model(params, x_t, t, *conditioning_inputs):
         else:
             # Unconditional sampling
             def sample_model(params, x_t, t, *conditioning_inputs):
-                rates = self.noise_schedule.get_rates(t)
+                rates = self.noise_schedule.get_rates(t, get_coeff_shapes_tuple(x_t))
                 c_in = self.model_output_transform.get_input_scale(rates)
                 model_output = self.model.apply(
                     params, 
@@ -136,7 +124,7 @@ def sample_step(
         Returns:
             Tuple of (new samples, updated state)
         """
-        step_ones = jnp.ones((len(current_samples),), dtype=jnp.int32)
+        step_ones = jnp.ones((len(current_samples),), dtype=jnp.float32)
         current_step = step_ones * current_step
         next_step = step_ones * next_step
         
@@ -175,74 +163,30 @@ def take_next_step(
         raise NotImplementedError("Subclasses must implement take_next_step method")
 
 
-    def scale_steps(self, steps):
-        """Scale timesteps to match the noise schedule's range."""
-        scale_factor = self.noise_schedule.max_timesteps / 1000
-        return steps * scale_factor
+    def get_steps(self, start_step, end_step, diffusion_steps):
+        """Get the descending sequence of timesteps for the diffusion process.
 
+        Timesteps are in the noise schedule's own domain: [0, max_timesteps],
+        so [0, 1] for the continuous schedules and [0, T] for the discrete ones.
+        Note that for a Karras schedule, linear spacing in t is already the
+        rho-spacing in sigma from the EDM paper.
 
-    def get_steps(self, start_step, end_step, diffusion_steps):
-        """Get the sequence of timesteps for the diffusion process.
-        
         Args:
-            start_step: Starting timestep (typically the max)
+            start_step: Starting timestep (typically max_timesteps)
             end_step: Ending timestep (typically 0)
             diffusion_steps: Number of steps to use
-            
+
         Returns:
             Array of timesteps for sampling
         """
         step_range = start_step - end_step
         if diffusion_steps is None or diffusion_steps == 0:
-            diffusion_steps = step_range
-        diffusion_steps = min(diffusion_steps, step_range)
-        
-        # Linear spacing (default)
-        if getattr(self, 'timestep_spacing', 'linear') == 'linear':
-            steps = jnp.linspace(
-                end_step, start_step,
-                diffusion_steps, dtype=jnp.int16
-            )[::-1]
-        
-        # Quadratic spacing (emphasizes early steps)
-        elif self.timestep_spacing == 'quadratic':
-            steps = jnp.linspace(0, 1, diffusion_steps) ** 2
-            steps = (start_step - end_step) * steps + end_step
-            steps = jnp.asarray(steps, dtype=jnp.int16)[::-1]
-            
-        # Karras spacing from the EDM paper - often gives better results with fewer steps
-        elif self.timestep_spacing == 'karras':
-            # Implementation based on the EDM paper's recommendations
-            sigma_min = end_step / start_step
-            sigma_max = 1.0
-            rho = 7.0  # Karras paper default, controls the distribution
-            
-            # Create log-spaced steps in sigma space
-            sigmas = jnp.exp(jnp.linspace(
-                jnp.log(sigma_max), jnp.log(sigma_min), diffusion_steps
-            ))
-            steps = jnp.clip(
-                (sigmas ** (1 / rho) - self.min_inv_rho) / 
-                (self.max_inv_rho - self.min_inv_rho), 
-                0, 1
-            ) * start_step
-            steps = jnp.asarray(steps, dtype=jnp.int16)
-            
-        # Exponential spacing (concentrates steps near the end)
-        elif self.timestep_spacing == 'exponential':
-            steps = jnp.linspace(0, 1, diffusion_steps)
-            steps = jnp.exp(steps * jnp.log((start_step + 1) / (end_step + 1))) * (end_step + 1) - 1
-            steps = jnp.clip(steps, end_step, start_step)
-            steps = jnp.asarray(steps, dtype=jnp.int16)[::-1]
-        
-        # Fallback to linear spacing
-        else:
-            steps = jnp.linspace(
-                end_step, start_step,
-                diffusion_steps, dtype=jnp.int16
-            )[::-1]
-            
-        return steps
+            diffusion_steps = int(step_range)
+        if self.noise_schedule.max_timesteps > 1:
+            # Discrete schedules cannot take more steps than they have entries
+            diffusion_steps = min(diffusion_steps, int(step_range))
+
+        return jnp.linspace(start_step, end_step, diffusion_steps, dtype=jnp.float32)
 
 
     def generate_samples(
@@ -374,8 +318,8 @@ def sample_step(sample_model_fn, state: RandomMarkovState, samples, current_step
             steps = self.get_steps(start_step, end_step, diffusion_steps)
 
         for i in tqdm.tqdm(range(0, len(steps))):
-            current_step = self.scale_steps(steps[i])
-            next_step = self.scale_steps(steps[i+1] if i+1 < len(steps) else 0)
+            current_step = steps[i]
+            next_step = steps[i+1] if i+1 < len(steps) else end_step
             
             if i != len(steps) - 1:
                 samples, rngstate = sample_step(
@@ -397,7 +341,6 @@ def _get_noise_parameters(self, resolution, start_step):
         Returns:
             Tuple of (variance, image_size, image_channels)
         """
-        start_step = self.scale_steps(start_step)
         alpha_n, sigma_n = self.noise_schedule.get_rates(start_step)
         variance = jnp.sqrt(alpha_n ** 2 + sigma_n ** 2)
         
diff --git a/flaxdiff/samplers/ddim.py b/flaxdiff/samplers/ddim.py
index 4921952..6bcc235 100644
--- a/flaxdiff/samplers/ddim.py
+++ b/flaxdiff/samplers/ddim.py
@@ -32,18 +32,18 @@ def take_next_step(
         
         # Extract random noise if needed for stochastic sampling
         if self.eta > 0:
-            # For DDIM, we need to compute the variance coefficient
-            # This is based on the original DDIM paper's formula
-            # When eta=0, it's deterministic DDIM, when eta=1.0 it approaches DDPM
-            sigma_tilde = self.eta * sigma_next * (1 - alpha_t**2 / alpha_next**2).sqrt() / (1 - alpha_t**2).sqrt()
+            # DDIM paper eq. 16: eta=0 is deterministic DDIM, eta=1.0 approaches DDPM.
+            # The direction term must shrink to keep the marginal variance right.
+            sigma_tilde = self.eta * (sigma_next / sigma_t) * jnp.sqrt(
+                jnp.maximum(1 - alpha_t**2 / alpha_next**2, 0.0))
             state, noise_key = state.get_random_key()
             noise = jax.random.normal(noise_key, current_samples.shape)
-            # Add the stochastic component
-            stochastic_term = sigma_tilde * noise
+            direction_coeff = jnp.sqrt(jnp.maximum(sigma_next**2 - sigma_tilde**2, 0.0))
+            new_samples = (alpha_next * reconstructed_samples
+                           + direction_coeff * pred_noise
+                           + sigma_tilde * noise)
         else:
-            stochastic_term = 0
-            
-        # Direct DDIM update formula
-        new_samples = alpha_next * reconstructed_samples + sigma_next * pred_noise + stochastic_term
+            # Direct DDIM update formula
+            new_samples = alpha_next * reconstructed_samples + sigma_next * pred_noise
         
         return new_samples, state
diff --git a/flaxdiff/samplers/ddpm.py b/flaxdiff/samplers/ddpm.py
index 7c07d46..1f5da65 100644
--- a/flaxdiff/samplers/ddpm.py
+++ b/flaxdiff/samplers/ddpm.py
@@ -1,37 +1,34 @@
 import jax
 import jax.numpy as jnp
 from .common import DiffusionSampler
+from ..schedulers import get_coeff_shapes_tuple
 from ..utils import MarkovState, RandomMarkovState
+
 class DDPMSampler(DiffusionSampler):
-    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
-                 pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
-        mean = self.noise_schedule.get_posterior_mean(reconstructed_samples, current_samples, current_step)
-        variance = self.noise_schedule.get_posterior_variance(steps=current_step)
-        
-        state, rng = state.get_random_key()
-        # Now sample from the posterior
-        noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32)
+    """Exact ancestral sampler for the reverse diffusion SDE.
 
-        return mean + noise * variance, state
-    
-    def generate_images(self, num_images=16, diffusion_steps=1000, start_step: int = None, *args, **kwargs):
-        return super().generate_images(num_images=num_images, diffusion_steps=diffusion_steps, start_step=start_step, *args, **kwargs)
-    
-class SimpleDDPMSampler(DiffusionSampler):
+    This is the simplified (but algebraically exact) form of the DDPM
+    posterior, phrased purely in terms of signal/noise rates so it works for
+    any schedule and any step stride, not just t -> t-1.
+    """
     def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
         state, rng = state.get_random_key()
         noise = jax.random.normal(rng, reconstructed_samples.shape, dtype=jnp.float32)
 
-        # Compute noise rates and signal rates only once
-        current_signal_rate, current_noise_rate = self.noise_schedule.get_rates(current_step)
-        next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step)
-        
+        shape = get_coeff_shapes_tuple(current_samples)
+        current_signal_rate, current_noise_rate = self.noise_schedule.get_rates(current_step, shape)
+        next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step, shape)
+
         pred_noise_coeff = ((next_noise_rate ** 2) * current_signal_rate) / (current_noise_rate * next_signal_rate)
-        
+
         noise_ratio_squared = (next_noise_rate ** 2) / (current_noise_rate ** 2)
         signal_ratio_squared = (current_signal_rate ** 2) / (next_signal_rate ** 2)
         gamma = jnp.sqrt(noise_ratio_squared * (1 - signal_ratio_squared))
-        
+
         next_samples = next_signal_rate * reconstructed_samples + pred_noise_coeff * pred_noise + noise * gamma
-        return next_samples, state
\ No newline at end of file
+        return next_samples, state
+
+# The two used to be separate implementations; the posterior-table variant
+# crashed for batched steps and could not handle strided sampling
+SimpleDDPMSampler = DDPMSampler
diff --git a/flaxdiff/samplers/euler.py b/flaxdiff/samplers/euler.py
index 62438b7..d7011f9 100644
--- a/flaxdiff/samplers/euler.py
+++ b/flaxdiff/samplers/euler.py
@@ -1,14 +1,16 @@
 import jax
 import jax.numpy as jnp
 from .common import DiffusionSampler
+from ..schedulers import get_coeff_shapes_tuple
 from ..utils import RandomMarkovState
 
 class EulerSampler(DiffusionSampler):
     # Basically a DDIM Sampler but parameterized as an ODE
     def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
-        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
-        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+        shape = get_coeff_shapes_tuple(current_samples)
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape)
 
         dt = next_sigma - current_sigma
         
@@ -23,8 +25,9 @@ class SimplifiedEulerSampler(DiffusionSampler):
     """
     def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
-        _, current_sigma = self.noise_schedule.get_rates(current_step)
-        _, next_sigma = self.noise_schedule.get_rates(next_step)
+        shape = get_coeff_shapes_tuple(current_samples)
+        _, current_sigma = self.noise_schedule.get_rates(current_step, shape)
+        _, next_sigma = self.noise_schedule.get_rates(next_step, shape)
 
         dt = next_sigma - current_sigma
         
@@ -38,8 +41,9 @@ class EulerAncestralSampler(DiffusionSampler):
     """
     def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
-        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
-        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+        shape = get_coeff_shapes_tuple(current_samples)
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape)
 
         sigma_up = (next_sigma**2 * (current_sigma**2 - next_sigma**2) / current_sigma**2) ** 0.5
         sigma_down = (next_sigma**2 - sigma_up**2) ** 0.5
diff --git a/flaxdiff/samplers/heun_sampler.py b/flaxdiff/samplers/heun_sampler.py
index d553614..4baa1c3 100644
--- a/flaxdiff/samplers/heun_sampler.py
+++ b/flaxdiff/samplers/heun_sampler.py
@@ -1,14 +1,16 @@
 import jax
 import jax.numpy as jnp
 from .common import DiffusionSampler
+from ..schedulers import get_coeff_shapes_tuple
 from ..utils import RandomMarkovState
 
 class HeunSampler(DiffusionSampler):
     def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
         # Get the noise and signal rates for the current and next steps
-        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
-        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+        shape = get_coeff_shapes_tuple(current_samples)
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape)
 
         dt = next_sigma - current_sigma
         x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / dt
@@ -20,8 +22,15 @@ def take_next_step(self, current_samples, reconstructed_samples, model_condition
         estimated_x_0, _, _ = sample_model_fn(next_samples_0, next_step, *model_conditioning_inputs)
         
         # Estimate the refined derivative using the midpoint (Heun's method)
-        dx_1 = (next_samples_0 - x_0_coeff * estimated_x_0) / next_sigma
-        # Compute the final next samples by averaging the initial and refined derivatives
-        final_next_samples = current_samples + 0.5 * (dx_0 + dx_1) * dt
+        safe_next_sigma = jnp.where(next_sigma > 0, next_sigma, 1.0)
+        dx_1 = (next_samples_0 - x_0_coeff * estimated_x_0) / safe_next_sigma
+        # Compute the final next samples by averaging the initial and refined derivatives.
+        # When sigma reaches 0 there is no derivative there - fall back to the euler step
+        # (Karras et al. 2022, Algorithm 2)
+        final_next_samples = jnp.where(
+            next_sigma > 0,
+            current_samples + 0.5 * (dx_0 + dx_1) * dt,
+            next_samples_0,
+        )
         
         return final_next_samples, state
diff --git a/flaxdiff/samplers/multistep_dpm.py b/flaxdiff/samplers/multistep_dpm.py
index b3e7fe9..cb184e8 100644
--- a/flaxdiff/samplers/multistep_dpm.py
+++ b/flaxdiff/samplers/multistep_dpm.py
@@ -1,6 +1,7 @@
 import jax
 import jax.numpy as jnp
 from .common import DiffusionSampler
+from ..schedulers import get_coeff_shapes_tuple
 from ..utils import RandomMarkovState
 
 class MultiStepDPM(DiffusionSampler):
@@ -8,11 +9,17 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.history = []
 
+    def generate_samples(self, *args, **kwargs):
+        # The multistep history is only valid within a single trajectory
+        self.history = []
+        return super().generate_samples(*args, **kwargs)
+
     def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, sample_model_fn, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
         # Get the noise and signal rates for the current and next steps
-        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
-        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
+        shape = get_coeff_shapes_tuple(current_samples)
+        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step, shape)
+        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step, shape)
 
         dt = next_sigma - current_sigma
 
diff --git a/flaxdiff/samplers/rk4_sampler.py b/flaxdiff/samplers/rk4_sampler.py
index a4cdbe8..1763b88 100644
--- a/flaxdiff/samplers/rk4_sampler.py
+++ b/flaxdiff/samplers/rk4_sampler.py
@@ -2,13 +2,14 @@
 import jax.numpy as jnp
 from .common import DiffusionSampler
 from ..utils import RandomMarkovState, MarkovState
-from ..schedulers import GeneralizedNoiseScheduler
+from ..schedulers import GeneralizedNoiseScheduler, get_coeff_shapes_tuple
 
 class RK4Sampler(DiffusionSampler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         assert issubclass(type(self.noise_schedule), GeneralizedNoiseScheduler), "Noise schedule must be a GeneralizedNoiseScheduler"
-        @jax.jit
+        # Not jitted here: sample_model_fn is a python callable, the model call
+        # inside it is already jitted
         def get_derivative(sample_model_fn, x_t, sigma, state:RandomMarkovState, model_conditioning_inputs) -> tuple[jnp.ndarray, RandomMarkovState]:
             t = self.noise_schedule.get_timesteps(sigma)
             x_0, eps, _ = sample_model_fn(x_t, t, *model_conditioning_inputs)
@@ -17,11 +18,12 @@ def get_derivative(sample_model_fn, x_t, sigma, state:RandomMarkovState, model_c
         self.get_derivative = get_derivative
 
     def sample_step(self, sample_model_fn, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
-        step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.int32)
+        step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.float32)
         current_step = step_ones * current_step
         next_step = step_ones * next_step
-        _, current_sigma = self.noise_schedule.get_rates(current_step)
-        _, next_sigma = self.noise_schedule.get_rates(next_step)
+        shape = get_coeff_shapes_tuple(current_samples)
+        _, current_sigma = self.noise_schedule.get_rates(current_step, shape)
+        _, next_sigma = self.noise_schedule.get_rates(next_step, shape)
 
         dt = next_sigma - current_sigma
 
diff --git a/flaxdiff/trainer/diffusion_trainer.py b/flaxdiff/trainer/diffusion_trainer.py
index 2004da6..0b78783 100644
--- a/flaxdiff/trainer/diffusion_trainer.py
+++ b/flaxdiff/trainer/diffusion_trainer.py
@@ -300,7 +300,6 @@ def generate_samples(
                 resolution=image_size,
                 num_samples=len(labels_seq),
                 diffusion_steps=diffusion_steps,
-                start_step=1000,
                 end_step=0,
                 priors=None,
                 model_conditioning_inputs=(labels_seq,),
diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py
index 08c6e64..dab5544 100644
--- a/flaxdiff/trainer/general_diffusion_trainer.py
+++ b/flaxdiff/trainer/general_diffusion_trainer.py
@@ -392,7 +392,6 @@ def generate_samples(
                 num_samples=batch_size,
                 sequence_length=sequence_length,  # Will be None for images
                 diffusion_steps=diffusion_steps,
-                start_step=1000,
                 end_step=0,
                 priors=None,
                 model_conditioning_inputs=tuple(model_conditioning_inputs),
diff --git a/tests/test_samplers.py b/tests/test_samplers.py
index 0e9eac0..9156c91 100644
--- a/tests/test_samplers.py
+++ b/tests/test_samplers.py
@@ -105,7 +105,6 @@ def generate(model, sampler, **kwargs):
         num_samples=256,
         resolution=8,
         diffusion_steps=100,
-        start_step=1000,
         rngstate=RandomMarkovState(jax.random.PRNGKey(2)),
     )
     defaults.update(kwargs)
@@ -126,20 +125,14 @@ def test_karras_sampler_converges(sampler_class):
     assert_gaussian_stats(samples)
 
 
-@pytest.mark.xfail(strict=True, reason="bug: default start_step uses max_timesteps but the step domain is hardcoded to [0, 1000]")
 def test_karras_sampler_default_start_step():
-    """With no explicit start_step the sampler must still denoise from sigma_max.
-    Today max_timesteps=1 collapses the schedule to a single no-op step and
-    the 'samples' are just the initial noise."""
+    """With no explicit start_step the sampler must denoise from sigma_max."""
     model, sampler = make_karras_sampler(EulerSampler)
     samples = generate(model, sampler, start_step=None)
     assert_gaussian_stats(samples)
 
 
-@pytest.mark.xfail(strict=True, reason="bug: sample_model reshapes rates with the default (-1,1,1,1), breaking 5D video shapes for every sampler")
 def test_ddim_video_samples():
-    """DDIM's own take_next_step is 5D-safe, but the shared sample_model wrapper
-    still scales the input with 4D-reshaped rates, so video is broken everywhere."""
     model, sampler = make_karras_sampler(DDIMSampler)
     params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,)))
     samples = sampler.generate_samples(
@@ -148,14 +141,12 @@ def test_ddim_video_samples():
         resolution=8,
         sequence_length=3,
         diffusion_steps=50,
-        start_step=1000,
         rngstate=RandomMarkovState(jax.random.PRNGKey(2)),
     )
     assert samples.shape == (2, 3, 8, 8, 3)
     assert jnp.all(jnp.isfinite(samples))
 
 
-@pytest.mark.xfail(strict=True, reason="bug: samplers reshape rates with the default (-1,1,1,1), breaking 5D video shapes")
 def test_euler_video_samples():
     model, sampler = make_karras_sampler(EulerSampler)
     params = model.init(jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,)))
@@ -165,34 +156,29 @@ def test_euler_video_samples():
         resolution=8,
         sequence_length=3,
         diffusion_steps=50,
-        start_step=1000,
         rngstate=RandomMarkovState(jax.random.PRNGKey(2)),
     )
     assert samples.shape == (2, 3, 8, 8, 3)
 
 
-@pytest.mark.xfail(strict=True, reason="bug: DDPMSampler casts the step array through int(), crashing for batch > 1")
 def test_ddpm_sampler_converges():
     model, sampler = make_vp_sampler(DDPMSampler)
     samples = generate(model, sampler, diffusion_steps=1000)
     assert_gaussian_stats(samples)
 
 
-@pytest.mark.xfail(strict=True, reason="bug: DDIM eta > 0 calls .sqrt() on a jax array and uses the wrong variance split")
 def test_ddim_eta_converges():
     model, sampler = make_vp_sampler(DDIMSampler, eta=0.5)
     samples = generate(model, sampler)
     assert_gaussian_stats(samples)
 
 
-@pytest.mark.xfail(strict=True, reason="bug: RK4Sampler jits over its python-callable argument")
 def test_rk4_sampler_runs():
     model, sampler = make_karras_sampler(RK4Sampler)
     samples = generate(model, sampler, diffusion_steps=20)
     assert jnp.all(jnp.isfinite(samples))
 
 
-@pytest.mark.xfail(strict=True, reason="bug: MultiStepDPM keeps stale derivative history across generate calls")
 def test_multistep_dpm_reentrant():
     model, sampler = make_karras_sampler(MultiStepDPM)
     first = generate(model, sampler)
@@ -200,9 +186,3 @@ def test_multistep_dpm_reentrant():
     assert jnp.allclose(first, second, atol=1e-5)
 
 
-@pytest.mark.parametrize("spacing", ["quadratic", "karras", "exponential"])
-@pytest.mark.xfail(strict=True, reason="bug: non-linear timestep spacings produce log(0) or duplicate int16 steps (dt=0)")
-def test_timestep_spacing_produces_finite_samples(spacing):
-    model, sampler = make_karras_sampler(EulerSampler, timestep_spacing=spacing)
-    samples = generate(model, sampler, diffusion_steps=50)
-    assert_gaussian_stats(samples, tol=0.1)

From 30863e07c09ad8656482bdc4bbbf8af8f5be636f Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashishkmr472@gmail.com>
Date: Wed, 22 Jul 2026 23:37:03 -0400
Subject: [PATCH 2/5] fix: honor --noise_schedule, shared train/infer presets,
 UViT hilbert, EDM weights, stable hashing

- training.py silently ignored --noise_schedule and always trained EDM
  while inference honored the flag, so a cosine-flagged run trained one
  convention and sampled another. Both sides now build from a single
  get_diffusion_preset, with a test pinning their agreement.
- UViT applied the hilbert permutation twice on the way in but inverted
  it once on the way out, training on a spatial scramble (SimpleDiT was
  correct; UDiT recomputed what patchify already returns).
- The EDM loss weight's epsilon guard halved the weight at sigma_min
  ((sigma*sd)^2 == 1e-6 exactly); rewritten as 1/sd^2 + 1/sigma^2 which
  needs no guard.
- FourierEmbedding frequencies came from jax's PRNG, which changed
  defaults in 0.5.0 and silently altered every model's time conditioning;
  now fixed numpy values, identical across jax versions.
- arguments_hash used python hash(), randomized per process, so identical
  runs never mapped to the same experiment directory; now sha256.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 flaxdiff/inference/utils.py     | 21 +++------------
 flaxdiff/models/common.py       |  6 ++++-
 flaxdiff/models/simple_vit.py   | 13 +++-------
 flaxdiff/predictors/__init__.py | 45 +++++++++++++++++++++++++++++++--
 flaxdiff/schedulers/karras.py   |  7 ++---
 tests/test_config.py            | 23 +++++++++++++++++
 tests/test_models.py            | 28 +++++++++++---------
 tests/test_schedulers.py        |  1 -
 training.py                     | 19 ++++++--------
 9 files changed, 106 insertions(+), 57 deletions(-)

diff --git a/flaxdiff/inference/utils.py b/flaxdiff/inference/utils.py
index cdf598d..f25b3c4 100644
--- a/flaxdiff/inference/utils.py
+++ b/flaxdiff/inference/utils.py
@@ -1,14 +1,7 @@
 import jax
 import jax.numpy as jnp
 import json
-from flaxdiff.schedulers import (
-    CosineNoiseScheduler,
-    KarrasVENoiseScheduler,
-)
-from flaxdiff.predictors import (
-    VPredictionTransform,
-    KarrasPredictionTransform,
-)
+from flaxdiff.predictors import get_diffusion_preset
 from flaxdiff.models.common import kernel_init
 from flaxdiff.models.simple_unet import Unet
 from flaxdiff.models.simple_vit import UViT
@@ -236,17 +229,9 @@ def map_nested_config(config):
     if 'diffusers' in architecture:
         model = BCHWModelWrapper(model)
     
-    # Create noise scheduler based on configuration
+    # Same preset as training, so the sampling convention always matches
     noise_schedule_type = conf.get('noise_schedule', conf.get('arguments', {}).get('noise_schedule', 'edm'))
-    if noise_schedule_type in ['edm', 'karras']:
-        # For both EDM and karras, we use the karras scheduler for inference
-        noise_schedule = KarrasVENoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
-        prediction_transform = KarrasPredictionTransform(sigma_data=noise_schedule.sigma_data)
-    elif noise_schedule_type == 'cosine':
-        noise_schedule = CosineNoiseScheduler(1000, beta_end=1)
-        prediction_transform = VPredictionTransform()
-    else:
-        raise ValueError(f"Unknown noise schedule: {noise_schedule_type}")
+    _, noise_schedule, prediction_transform = get_diffusion_preset(noise_schedule_type)
     
     # Prepare return dictionary with all components
     result = {
diff --git a/flaxdiff/models/common.py b/flaxdiff/models/common.py
index 2b5b84f..50ea1dd 100644
--- a/flaxdiff/models/common.py
+++ b/flaxdiff/models/common.py
@@ -1,5 +1,6 @@
 import jax.numpy as jnp
 import jax
+import numpy as np
 from flax import linen as nn
 from typing import Optional, Any, Callable, Sequence, Union
 from flax.typing import Dtype, PrecisionLike
@@ -99,7 +100,10 @@ class FourierEmbedding(nn.Module):
     scale:int = 16
 
     def setup(self):
-        self.freqs = jax.random.normal(jax.random.PRNGKey(42), (self.features // 2, ), dtype=jnp.float32) * self.scale
+        # Fixed frequencies via numpy so they are identical across jax versions
+        # (jax 0.5.0 changed the default PRNG and silently altered these)
+        freqs = np.random.RandomState(42).normal(size=(self.features // 2,))
+        self.freqs = jnp.asarray(freqs, dtype=jnp.float32) * self.scale
 
     def __call__(self, x):
         x = jax.lax.convert_element_type(x, jnp.float32)
diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py
index 92f130c..4df02d1 100644
--- a/flaxdiff/models/simple_vit.py
+++ b/flaxdiff/models/simple_vit.py
@@ -182,13 +182,10 @@ def __call__(self, x, temb, textcontext=None):
 
         hilbert_inv_idx = None
         if self.use_hilbert:
-            patches_raw, hilbert_inv_idx_calc = hilbert_patchify(
-                x, self.patch_size)
+            # hilbert_patchify already returns the patches in hilbert order
+            # along with the inverse permutation for the output path
+            patches_raw, hilbert_inv_idx = hilbert_patchify(x, self.patch_size)
             x_patches = self.hilbert_proj(patches_raw)
-            idx = hilbert_indices(H_P, W_P)
-            hilbert_inv_idx = inverse_permutation(
-                idx, total_size=num_patches)
-            x_patches = x_patches[:, idx, :]
         else:
             x_patches = self.patch_embed(x)
 
@@ -396,10 +393,8 @@ def __call__(self, x, temb, textcontext=None):
 
         hilbert_inv_idx = None
         if self.use_hilbert:
-            patches_raw, _ = hilbert_patchify(x, self.patch_size)
+            patches_raw, hilbert_inv_idx = hilbert_patchify(x, self.patch_size)
             x_seq = self.hilbert_proj(patches_raw)
-            idx = hilbert_indices(H_P, W_P)
-            hilbert_inv_idx = inverse_permutation(idx, total_size=num_patches)
         else:
             x_seq = self.patch_embed(x)
 
diff --git a/flaxdiff/predictors/__init__.py b/flaxdiff/predictors/__init__.py
index 3aeb5d2..68390f2 100644
--- a/flaxdiff/predictors/__init__.py
+++ b/flaxdiff/predictors/__init__.py
@@ -1,6 +1,13 @@
 from typing import Union
 import jax.numpy as jnp
-from ..schedulers import NoiseScheduler, GeneralizedNoiseScheduler, get_coeff_shapes_tuple
+from ..schedulers import (
+    NoiseScheduler,
+    GeneralizedNoiseScheduler,
+    get_coeff_shapes_tuple,
+    CosineNoiseScheduler,
+    KarrasVENoiseScheduler,
+    EDMNoiseScheduler,
+)
 
 ############################################################################################################
 # Prediction Transforms
@@ -93,4 +100,38 @@ def pred_transform(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray], eps
     def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8) -> jnp.ndarray:
         _, sigma = rates
         c_in = 1 / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon)
-        return c_in
\ No newline at end of file
+        return c_in
+
+############################################################################################################
+# Noise schedule presets
+############################################################################################################
+
+def get_diffusion_preset(
+    name: str,
+    timesteps: int = 1000,
+    sigma_min: float = 0.002,
+    sigma_max: float = 80.0,
+    rho: float = 7.0,
+    sigma_data: float = 0.5,
+) -> tuple[NoiseScheduler, NoiseScheduler, DiffusionPredictionTransform]:
+    """Named (train schedule, sampling schedule, prediction transform) presets.
+
+    The single source of truth for which schedule pairs with which
+    parameterization. Both training and inference build from here, so a model
+    is always sampled with the same convention it was trained with.
+    """
+    if name == 'edm':
+        train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
+        sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
+        transform = KarrasPredictionTransform(sigma_data=sigma_data)
+    elif name == 'karras':
+        train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
+        sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
+        transform = KarrasPredictionTransform(sigma_data=sigma_data)
+    elif name == 'cosine':
+        train = CosineNoiseScheduler(timesteps, beta_end=1)
+        sample = CosineNoiseScheduler(timesteps, beta_end=1)
+        transform = VPredictionTransform()
+    else:
+        raise ValueError(f"Unknown noise schedule preset: {name}")
+    return train, sample, transform
diff --git a/flaxdiff/schedulers/karras.py b/flaxdiff/schedulers/karras.py
index b651c75..b9e09d5 100644
--- a/flaxdiff/schedulers/karras.py
+++ b/flaxdiff/schedulers/karras.py
@@ -19,9 +19,10 @@ def get_sigmas(self, steps) -> jnp.ndarray:
         
     def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
         sigma = self.get_sigmas(steps)
-        # Add epsilon for numerical stability
-        epsilon = 1e-6
-        weights = ((sigma ** 2 + self.sigma_data ** 2) / ((sigma * self.sigma_data) ** 2 + epsilon))
+        # EDM lambda(sigma) = (sigma^2 + sd^2) / (sigma * sd)^2, written in a
+        # form that needs no epsilon guard (the old guard halved the weight at
+        # sigma_min where (sigma * sd)^2 == 1e-6)
+        weights = 1 / self.sigma_data ** 2 + 1 / sigma ** 2
         return weights.reshape(shape)
     
     def transform_inputs(self, x, steps, num_discrete_chunks=1000) -> tuple[jnp.ndarray, jnp.ndarray]:
diff --git a/tests/test_config.py b/tests/test_config.py
index b7ef699..7d8b831 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -75,3 +75,26 @@ def test_parse_config_preserves_dotted_values():
     model = result["model"]
     # a dropped key falls back to the class default (swish) instead of erroring
     assert model.activation is jax.nn.mish, "activation was silently dropped"
+
+
+def test_training_and_inference_share_schedule_presets():
+    """--noise_schedule must mean the same thing at train and inference time.
+    Both sides now build from get_diffusion_preset."""
+    from flaxdiff.predictors import get_diffusion_preset
+
+    train, sample, transform = get_diffusion_preset("edm")
+    assert type(train).__name__ == "EDMNoiseScheduler"
+    assert type(sample).__name__ == "KarrasVENoiseScheduler"
+    assert type(transform).__name__ == "KarrasPredictionTransform"
+
+    train, sample, transform = get_diffusion_preset("cosine")
+    assert type(train).__name__ == "CosineNoiseScheduler"
+    assert type(sample).__name__ == "CosineNoiseScheduler"
+    assert type(transform).__name__ == "VPredictionTransform"
+
+    # parse_config must agree with the preset for every name
+    for name in ("edm", "karras", "cosine"):
+        _, sample, transform = get_diffusion_preset(name)
+        result = parse_config(make_config(arguments_overrides={"noise_schedule": name}))
+        assert type(result["noise_schedule"]) is type(sample)
+        assert type(result["prediction_transform"]) is type(transform)
diff --git a/tests/test_models.py b/tests/test_models.py
index c86f173..ca5cb13 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -101,17 +101,21 @@ def test_hybrid_dit_forward(rng, kwargs):
     assert out.shape == x.shape
 
 
-@pytest.mark.xfail(strict=True, reason="bug: UViT applies the hilbert permutation twice")
-def test_uvit_hilbert_matches_raster_information(rng):
-    """A zero-layer sanity check: with the permutation applied and inverted once,
-    patch content must land back at its own spatial position. UViT permutes twice
-    on the way in but inverts once on the way out, scrambling the output."""
-    from flaxdiff.models.hilbert import hilbert_patchify, hilbert_unpatchify, hilbert_indices, inverse_permutation
-
-    x = jax.random.normal(rng, (1, 16, 16, 3))
+def test_uvit_hilbert_forward(rng):
+    """UViT used to apply the hilbert permutation twice on the way in but
+    invert it only once on the way out, scrambling every output spatially."""
+    model = UViT(patch_size=4, emb_features=64, num_layers=4, num_heads=2, use_hilbert=True)
+    x, temb, textcontext = small_inputs(rng)
+    out = run_forward(model, rng, x, temb, textcontext)
+    assert out.shape == x.shape
+
+
+def test_hilbert_patchify_roundtrip(rng):
+    """patchify returns patches in hilbert order plus the inverse permutation;
+    unpatchify with that permutation must be an exact identity."""
+    from flaxdiff.models.hilbert import hilbert_patchify, hilbert_unpatchify
+
+    x = jax.random.normal(rng, (2, 16, 16, 3))
     patches, inv_idx = hilbert_patchify(x, 4)
-    # UViT's forward applies idx again on the already-permuted patches
-    idx = hilbert_indices(4, 4)
-    double_permuted = patches[:, idx, :]
-    rec = hilbert_unpatchify(double_permuted, inv_idx, 4, 16, 16, 3)
+    rec = hilbert_unpatchify(patches, inv_idx, 4, 16, 16, 3)
     assert jnp.allclose(rec, x)
diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py
index b54cc54..44e05fc 100644
--- a/tests/test_schedulers.py
+++ b/tests/test_schedulers.py
@@ -66,7 +66,6 @@ def test_karras_weights_match_edm_lambda():
     assert jnp.allclose(got, expected, rtol=1e-3)
 
 
-@pytest.mark.xfail(strict=True, reason="bug: epsilon guard distorts the weight near sigma_min")
 def test_karras_weights_at_sigma_min():
     schedule = KarrasVENoiseScheduler(1, sigma_min=0.002, sigma_max=80, rho=7, sigma_data=0.5)
     sigma = jnp.array([0.002])
diff --git a/training.py b/training.py
index 3b6606c..1000826 100644
--- a/training.py
+++ b/training.py
@@ -15,8 +15,7 @@
     import jax.experimental.pallas.ops.tpu.flash_attention
 except (ImportError, ModuleNotFoundError):
     pass
-from flaxdiff.predictors import VPredictionTransform, EpsilonPredictionTransform, DiffusionPredictionTransform, DirectPredictionTransform, KarrasPredictionTransform
-from flaxdiff.schedulers import CosineNoiseScheduler, NoiseScheduler, GeneralizedNoiseScheduler, KarrasVENoiseScheduler, EDMNoiseScheduler
+from flaxdiff.predictors import get_diffusion_preset
 from diffusers import FlaxUNet2DConditionModel
 from flaxdiff.samplers.euler import EulerAncestralSampler
 import struct as st
@@ -30,6 +29,7 @@
 from datetime import datetime
 
 import json
+import hashlib
 # For CLIP
 import argparse
 from dataclasses import dataclass
@@ -490,7 +490,8 @@ def main(args):
         model_config['emb_features'] = 768
         
     sorted_args_json = json.dumps(vars(args), sort_keys=True)
-    arguments_hash = hash(sorted_args_json)
+    # hash() is randomized per process; identical args must map to the same experiment
+    arguments_hash = hashlib.sha256(sorted_args_json.encode()).hexdigest()[:16]
     
     text_encoder = defaultTextEncodeModel()
     
@@ -543,10 +544,7 @@ def main(args):
     
     batches = batches if args.steps_per_epoch is None else args.steps_per_epoch
 
-    cosine_schedule = CosineNoiseScheduler(1000, beta_end=1)
-    karas_ve_schedule = KarrasVENoiseScheduler(
-        1, sigma_max=80, rho=7, sigma_data=0.5)
-    edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
+    train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset(args.noise_schedule)
     
     if args.experiment_name is not None:
         experiment_name = args.experiment_name
@@ -618,11 +616,10 @@ def main(args):
     trainer = GeneralDiffusionTrainer(
         model, optimizer=solver,
         input_config=input_config,
-        noise_schedule=edm_schedule,
+        noise_schedule=train_schedule,
         rngs=jax.random.PRNGKey(4),
         name=experiment_name,
-        model_output_transform=KarrasPredictionTransform(
-        sigma_data=edm_schedule.sigma_data),
+        model_output_transform=prediction_transform,
         load_from_checkpoint=args.load_from_checkpoint,
         wandb_config=wandb_config,
         distributed_training=args.distributed_training,
@@ -656,7 +653,7 @@ def get_val_dataset():
         training_steps_per_epoch=batches,
         epochs=CONFIG['epochs'], 
         sampler_class=EulerAncestralSampler, 
-        sampling_noise_schedule=karas_ve_schedule,
+        sampling_noise_schedule=sampling_schedule,
         val_steps_per_epoch=args.val_steps_per_epoch,
     )
     

From c46aa265a88515cce5556abbf502c75640a89e6d Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashishkmr472@gmail.com>
Date: Wed, 22 Jul 2026 23:40:43 -0400
Subject: [PATCH 3/5] fix(trainer): optional wandb, multi-host pbar, honest
 divergence handling, full state restore, checkpoint deletion

- self.wandb was only assigned when wandb_config was set on process 0, so
  training without wandb (or on any non-zero process) crashed on the first
  logging tick; it is now always initialized and the getattr patches at the
  call sites are gone.
- pbar was unbound on non-zero processes (UnboundLocalError at step 0 of
  any multi-host run since d50040d).
- Checkpoint restore rebuilt the TrainState via create(), silently zeroing
  the adam moments and resetting the optax step counter (restarting lr
  warmup) on every resume. Restore now goes through orbax with the freshly
  initialized states as a typed template, so optimizer state, step counter
  and EMA params all round-trip. The unused param_transforms hook is gone.
- The NaN 'recovery' block reset params and optimizer to a possibly
  epochs-old best_state mid-epoch and patched the loss to 1.0 to keep the
  average looking sane. A diverged run now fails loudly after 5
  consecutive non-finite steps.
- save() deleted the checkpoint it just wrote even when the registry push
  was skipped; deletion now only happens after a successful push.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 flaxdiff/trainer/diffusion_trainer.py         |  29 +---
 flaxdiff/trainer/general_diffusion_trainer.py |  14 +-
 flaxdiff/trainer/simple_trainer.py            | 124 +++++++-----------
 tests/test_trainer.py                         |  11 +-
 4 files changed, 69 insertions(+), 109 deletions(-)

diff --git a/flaxdiff/trainer/diffusion_trainer.py b/flaxdiff/trainer/diffusion_trainer.py
index 0b78783..99440dc 100644
--- a/flaxdiff/trainer/diffusion_trainer.py
+++ b/flaxdiff/trainer/diffusion_trainer.py
@@ -93,42 +93,25 @@ def generate_states(
         self,
         optimizer: optax.GradientTransformation,
         rngs: jax.random.PRNGKey,
-        existing_state: dict = None,
-        existing_best_state: dict = None,
         model: nn.Module = None,
-        param_transforms: Callable = None,
         use_dynamic_scale: bool = False
     ) -> Tuple[TrainState, TrainState]:
         print("Generating states for DiffusionTrainer")
         rngs, subkey = jax.random.split(rngs)
 
-        if existing_state == None:
-            input_vars = self.get_input_ones()
-            params = model.init(subkey, **input_vars)
-            new_state = {"params": params, "ema_params": params}
-        else:
-            new_state = existing_state
-
-        if param_transforms is not None:
-            params = param_transforms(params)
+        input_vars = self.get_input_ones()
+        params = model.init(subkey, **input_vars)
 
         state = TrainState.create(
             apply_fn=model.apply,
-            params=new_state['params'],
-            ema_params=new_state['ema_params'],
+            params=params,
+            ema_params=params,
             tx=optimizer,
             rngs=rngs,
             metrics=Metrics.empty(),
             dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None
         )
-            
-        if existing_best_state is not None:
-            best_state = state.replace(
-                params=existing_best_state['params'], ema_params=existing_best_state['ema_params'])
-        else:
-            best_state = state
-
-        return state, best_state
+        return state, state
 
     def _define_train_step(self, batch_size):
         noise_schedule: NoiseScheduler = self.noise_schedule
@@ -332,7 +315,7 @@ def validation_loop(
             )
             
             # Put each sample on wandb
-            if getattr(self, 'wandb', None) is not None and self.wandb:
+            if self.wandb is not None and self.wandb:
                 import numpy as np
                 from wandb import Image as wandbImage
                 wandb_images = []
diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py
index dab5544..2539530 100644
--- a/flaxdiff/trainer/general_diffusion_trainer.py
+++ b/flaxdiff/trainer/general_diffusion_trainer.py
@@ -466,7 +466,7 @@ def validation_loop(
                 if i == 0:
                     print(f"Evaluation started for process index {process_index}")
                     # Log samples to wandb
-                    if getattr(self, 'wandb', None) is not None and self.wandb:
+                    if self.wandb is not None and self.wandb:
                         import numpy as np
                         
                         # Process samples differently based on dimensionality
@@ -488,7 +488,7 @@ def validation_loop(
                         prev = self.best_val_metrics[final_key]
                         self.best_val_metrics[final_key] = max(prev, value) if higher_is_better else min(prev, value)
                 # Log the best validation metrics
-                if getattr(self, 'wandb', None) is not None and self.wandb:
+                if self.wandb is not None and self.wandb:
                     # Log the metrics
                     for key, value in metrics.items():
                         if isinstance(value, jnp.ndarray):
@@ -711,12 +711,12 @@ def save(self, epoch=0, step=0, state=None, rngstate=None):
                         aliases.append("best")
                     self.push_to_registry(aliases=aliases)
                     print("Model pushed to registry successfully with aliases:", aliases)
+                    # Only delete after a successful registry push - the local
+                    # checkpoint is the only copy otherwise
+                    shutil.rmtree(checkpoint, ignore_errors=True)
+                    print(f"Checkpoint deleted at {checkpoint}")
                 else:
-                    print("Current run is not one of the best runs. Not saving model.")
-                
-                # Only delete after successful registry push
-                shutil.rmtree(checkpoint, ignore_errors=True)
-                print(f"Checkpoint deleted at {checkpoint}")
+                    print("Current run is not one of the best runs. Keeping local checkpoint.")
             except Exception as e:
                 print(f"Error during registry operations: {e}")
                 print(f"Checkpoint preserved at {checkpoint}")
diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py
index b31a33c..1d23acd 100644
--- a/flaxdiff/trainer/simple_trainer.py
+++ b/flaxdiff/trainer/simple_trainer.py
@@ -161,7 +161,6 @@ def __init__(self,
                  name: str = "Simple",
                  load_from_checkpoint: str = None,
                  loss_fn=optax.l2_loss,
-                 param_transforms: Callable = None,
                  wandb_config: Dict[str, Any] = None,
                  distributed_training: bool = None,
                  checkpoint_base_path: str = "./checkpoints",
@@ -186,6 +185,7 @@ def __init__(self,
         
         load_directly_from_dir = False
         
+        self.wandb = None
         if wandb_config is not None and jax.process_index() == 0:
             import wandb
             run = wandb.init(resume='allow', **wandb_config)
@@ -234,27 +234,12 @@ def __init__(self,
         self.checkpointer = orbax.checkpoint.CheckpointManager(
             self.checkpoint_path(), async_checkpointer, options)
 
-        if load_from_checkpoint is not None:
-            latest_step, old_state, old_best_state, rngstate = self.load(load_from_checkpoint, checkpoint_step, load_directly_from_dir)
-        else:
-            latest_step, old_state, old_best_state, rngstate = 0, None, None, None
-
-        self.latest_step = latest_step
-        
-        if train_start_step_override is not None:
-            self.latest_step = train_start_step_override
-            print(f"Overriding start step to {self.latest_step}")
-        
-        if rngstate:
-            self.rngstate = RandomMarkovState(**rngstate)
-        else:
-            self.rngstate = RandomMarkovState(rngs)
-            
+        self.rngstate = RandomMarkovState(rngs)
         self.rngstate, subkey = self.rngstate.get_random_key()
 
         if train_state == None:
             state, best_state = self.generate_states(
-                optimizer, subkey, old_state, old_best_state, model, param_transforms, use_dynamic_scale
+                optimizer, subkey, model, use_dynamic_scale
             )
             self.init_state(state, best_state)
         else:
@@ -262,6 +247,14 @@ def __init__(self,
             self.best_state = train_state
             self.best_loss = 1e9
 
+        self.latest_step = 0
+        if load_from_checkpoint is not None:
+            self.latest_step = self.load(load_from_checkpoint, checkpoint_step, load_directly_from_dir)
+
+        if train_start_step_override is not None:
+            self.latest_step = train_start_step_override
+            print(f"Overriding start step to {self.latest_step}")
+
     def get_input_ones(self):
         return {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()}
 
@@ -269,20 +262,14 @@ def generate_states(
         self,
         optimizer: optax.GradientTransformation,
         rngs: jax.random.PRNGKey,
-        existing_state: dict = None,
-        existing_best_state: dict = None,
         model: nn.Module = None,
-        param_transforms: Callable = None,
         use_dynamic_scale: bool = False
     ) -> Tuple[SimpleTrainState, SimpleTrainState]:
         print("Generating states for SimpleTrainer")
         rngs, subkey = jax.random.split(rngs)
 
-        if existing_state == None:
-            input_vars = self.get_input_ones()
-            params = model.init(subkey, **input_vars)
-        else:
-            params = existing_state['params']
+        input_vars = self.get_input_ones()
+        params = model.init(subkey, **input_vars)
 
         state = SimpleTrainState.create(
             apply_fn=model.apply,
@@ -291,13 +278,7 @@ def generate_states(
             metrics=Metrics.empty(),
             dynamic_scale = dynamic_scale_lib.DynamicScale() if use_dynamic_scale else None
         )
-        if existing_best_state is not None:
-            best_state = state.replace(
-                params=existing_best_state['params'])
-        else:
-            best_state = state
-            
-        return state, best_state
+        return state, state
 
     def init_state(
         self,
@@ -353,18 +334,29 @@ def load(self, checkpoint_path, checkpoint_step=None, load_directly_from_dir=Fal
             checkpoint_path if checkpoint_path else self.checkpoint_path(),
             f"{step}")
         self.loaded_checkpoint_path = loaded_checkpoint_path
-        ckpt = checkpointer.restore(step) if not load_directly_from_dir else checkpointer.restore(checkpoint_path)
+
+        # Restore against the freshly-initialized states as a template so orbax
+        # rebuilds the exact pytree types - optimizer state and step included.
+        # Restoring untyped used to silently discard opt_state and reset the
+        # step counter (and with it the lr schedule) on every resume.
+        template = {
+            'rngs': self.get_rngstate(),
+            'state': self.get_state(),
+            'best_state': self.get_best_state(),
+            'best_loss': np.array(self.best_loss),
+            'epoch': 0,
+        }
+        ckpt = checkpointer.restore(step, items=template) if not load_directly_from_dir else checkpointer.restore(checkpoint_path, items=template)
         
-        state = ckpt['state']
-        best_state = ckpt['best_state']
-        rngstate = ckpt['rngs']
-        # Convert the state to a TrainState
-        self.best_loss = ckpt['best_loss']
+        self.state = ckpt['state']
+        self.best_state = ckpt['best_state']
+        self.rngstate = ckpt['rngs']
+        self.best_loss = float(ckpt['best_loss'])
         if self.best_loss == 0:
             # It cant be zero as that must have been some problem
             self.best_loss = 1e9
-        print(f"Loaded model from checkpoint at step {step}", ckpt['best_loss'])
-        return step, state, best_state, rngstate
+        print(f"Loaded model from checkpoint at step {step}", self.best_loss)
+        return step
 
     def save(self, epoch=0, step=0, state=None, rngstate=None):
         print(f"Saving model at epoch {epoch} step {step}")
@@ -516,11 +508,14 @@ def train_loop(
             global_device_indexes = 0
             
         epoch_loss = 0
+        bad_loss_steps = 0
         current_epoch = current_step // train_steps_per_epoch
         last_save_time = time.time()
         
         if process_index == 0:
             pbar = tqdm.tqdm(total=train_steps_per_epoch, desc=f'\t\tEpoch {current_epoch}', ncols=100, unit='step')
+        else:
+            pbar = None
             
         for i in range(train_steps_per_epoch):
             batch = next(train_ds)
@@ -539,41 +534,18 @@ def train_loop(
                 # loss = jax.experimental.multihost_utils.process_allgather(loss)
                 loss = jnp.mean(loss) # Just to make sure its a scaler value
                     
-            if loss <= 1e-8 or jnp.isnan(loss) or jnp.isinf(loss):
-                # If the loss is too low or NaN/Inf, log the issue and attempt recovery
-                print(colored(f"Abnormal loss at step {current_step}: {loss}", 'red'))
-                
-                # Check model parameters for NaN/Inf values
-                params = train_state.params
-                has_nan_or_inf = False
-                
-                if isinstance(params, dict):
-                    for key, value in params.items():
-                        if isinstance(value, jnp.ndarray):
-                            if jnp.isnan(value).any() or jnp.isinf(value).any():
-                                print(colored(f"NaN/inf values found in params[{key}] at step {current_step}", 'red'))
-                                has_nan_or_inf = True
-                                break
-                    
-                    if not has_nan_or_inf:
-                        print(colored(f"Model parameters seem valid despite abnormal loss", 'yellow'))
-                
-                # Try to recover - clear JAX caches and collect garbage
-                gc.collect()
-                if hasattr(jax, "clear_caches"):
-                    jax.clear_caches()
-                
-                # If we have a best state and the loss is truly invalid, consider restoring
-                if (loss <= 1e-8 or jnp.isnan(loss) or jnp.isinf(loss)) and self.best_state is not None:
-                    print(colored(f"Attempting recovery by resetting model to last best state", 'yellow'))
-                    train_state = self.best_state
-                    loss = self.best_loss
-                else:
-                    # If we can't recover, skip this step but continue training
-                    print(colored(f"Unable to recover - continuing with current state", 'yellow'))
-                    if loss <= 1e-8:
-                        loss = 1.0  # Set to a reasonable default to continue training
-                            
+            if not jnp.isfinite(loss):
+                # No silent recovery: a diverged run must fail loudly, not be
+                # papered over with a stale best_state and a cosmetic loss value
+                print(colored(f"Non-finite loss at step {current_step}: {loss}", 'red'))
+                bad_loss_steps += 1
+                if bad_loss_steps >= 5:
+                    raise RuntimeError(
+                        f"Loss has been non-finite for {bad_loss_steps} consecutive steps, stopping"
+                    )
+            else:
+                bad_loss_steps = 0
+
             epoch_loss += loss
             current_step += 1
             if i % 100 == 0:
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index f8d86f6..a33c15f 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -39,11 +39,17 @@ def batch_iterator():
         yield {"image": jnp.ones((8, 4)), "label": jnp.zeros((8, 2))}
 
 
-@pytest.mark.xfail(strict=True, reason="bug: self.wandb is only assigned when wandb_config is set, training without wandb crashes")
 def test_fit_without_wandb(tmp_path):
     trainer = make_trainer(tmp_path)
     data = {"train": batch_iterator, "train_len": 32}
-    trainer.fit(data, train_steps_per_epoch=4, epochs=1, val_steps_per_epoch=0)
+    state = trainer.fit(data, train_steps_per_epoch=4, epochs=2, val_steps_per_epoch=0)
+    # the tiny regression must actually learn something
+    final_loss = float(jnp.mean(optax.l2_loss(
+        state.apply_fn(state.params, jnp.ones((8, 4))), jnp.zeros((8, 2)))))
+    initial = make_trainer(tmp_path, name="fresh")
+    initial_loss = float(jnp.mean(optax.l2_loss(
+        initial.state.apply_fn(initial.state.params, jnp.ones((8, 4))), jnp.zeros((8, 2)))))
+    assert final_loss < initial_loss
 
 
 def test_save_writes_checkpoint(tmp_path):
@@ -53,7 +59,6 @@ def test_save_writes_checkpoint(tmp_path):
     assert trainer.checkpointer.latest_step() == 1
 
 
-@pytest.mark.xfail(strict=True, reason="bug: restore rebuilds the TrainState with tx.init, discarding opt_state and the step counter")
 def test_restore_preserves_optimizer_state(tmp_path):
     trainer = make_trainer(tmp_path)
 

From 31576f7672028d94f72211435693f06342ea9e6f Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashishkmr472@gmail.com>
Date: Wed, 22 Jul 2026 23:54:29 -0400
Subject: [PATCH 4/5] fix(models): implement dropout for real, delete the
 broken flash-attention path, dead flags, bf16 output heads

--dropout_rate was threaded into every DiT-family block and applied
nowhere (any sweep over it measured noise). nn.Dropout now sits on both
residual branches of DiTBlock/SSMDiTBlock/MMDiTBlock with a uniform
train flag across every model, and the trainers thread a dedicated
dropout rng. UViT's variant was never even forwarded, so that field is
gone.

--flash_attention is deleted along with EfficientAttention: the pallas
call passed no sm_scale (inflating logits 8x at head_dim 64 vs the
normal path), crashed on GPU/CPU where the import guard swallowed the
dependency but the call remained, and used a different param tree so
checkpoints weren't interchangeable. On every DiT-family model the flag
was accepted and silently ignored anyway. A proper fused-attention path
(cudnn/splash with identical params) comes with the parallelism work.

Also: dead norm_groups/activation fields removed from the DiT family
(training.py now scopes those args to the unet/uvit which read them),
S5Layer validates its features field, parse_config filters config keys
against the model's actual fields (old runs logged since-removed flags)
and resolves dotted function paths instead of silently dropping them,
the activation string is no longer mutated into a live function object
inside the logged config, DiT/MMDiT/hybrid output heads are fp32 (the
loss is fp32; bf16 heads quantized early-training residuals), the
attention-block FFN no longer silently runs fp32 inside bf16 models,
and jax.tree_map (removed in jax 0.6) is jax.tree.map.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 flaxdiff/inference/utils.py                   |  21 +++-
 flaxdiff/models/attention.py                  | 112 +-----------------
 flaxdiff/models/common.py                     |   1 -
 flaxdiff/models/general.py                    |   2 +-
 flaxdiff/models/simple_dit.py                 |  17 ++-
 flaxdiff/models/simple_mmdit.py               |  34 +++---
 flaxdiff/models/simple_unet.py                |   5 +-
 flaxdiff/models/simple_vit.py                 |  22 ++--
 flaxdiff/models/ssm_dit.py                    |  20 ++--
 flaxdiff/samplers/common.py                   |   2 -
 flaxdiff/trainer/diffusion_trainer.py         |   6 +-
 flaxdiff/trainer/general_diffusion_trainer.py |  11 +-
 tests/test_config.py                          |  32 +++--
 tests/test_models.py                          |  22 +++-
 training.py                                   |  35 +++---
 15 files changed, 130 insertions(+), 212 deletions(-)

diff --git a/flaxdiff/inference/utils.py b/flaxdiff/inference/utils.py
index f25b3c4..e45036c 100644
--- a/flaxdiff/inference/utils.py
+++ b/flaxdiff/inference/utils.py
@@ -140,9 +140,15 @@ def map_nested_config(config):
                     new_config[key] = ACTIVATION_MAP[value]
                 elif value == 'None':
                     new_config[key] = None
-                elif '.'in value:
-                    # Ignore any other string that contains a dot
-                    print(f"Ignoring key {key} with value {value} as it contains a dot.")
+                elif value.startswith('jax.') or value.startswith('jax._src.'):
+                    # Resolve function paths like 'jax.nn.mish' by attribute walk;
+                    # jax._src paths leak from old configs that stored the live
+                    # function object instead of its name
+                    attr_path = value.replace('jax._src.nn.functions', 'jax.nn')
+                    obj = jax
+                    for part in attr_path.split('.')[1:]:
+                        obj = getattr(obj, part)
+                    new_config[key] = obj
                 else:
                     new_config[key] = value
             else:
@@ -222,6 +228,15 @@ def map_nested_config(config):
     if not model_class:
         raise ValueError(f"Unknown architecture: {architecture}. Supported architectures: {', '.join(MODEL_CLASSES.keys())}")
     
+    # Drop config keys the model no longer has fields for (older runs logged
+    # since-removed flags like use_flash_attention)
+    import dataclasses
+    valid_fields = {f.name for f in dataclasses.fields(model_class)}
+    dropped = sorted(set(model_kwargs) - valid_fields)
+    if dropped:
+        print(f"Dropping config keys not accepted by {model_class.__name__}: {dropped}")
+    model_kwargs = {k: v for k, v in model_kwargs.items() if k in valid_fields}
+    
     # Instantiate the model
     model = model_class(**model_kwargs)
     
diff --git a/flaxdiff/models/attention.py b/flaxdiff/models/attention.py
index a7f761e..8106bf5 100644
--- a/flaxdiff/models/attention.py
+++ b/flaxdiff/models/attention.py
@@ -11,108 +11,6 @@
 import functools
 import math
 from .common import kernel_init
-try:
-    import jax.experimental.pallas.ops.tpu.flash_attention
-except (ImportError, ModuleNotFoundError):
-    pass
-
-class EfficientAttention(nn.Module):
-    """
-    Based on the pallas attention implementation.
-    """
-    query_dim: int
-    heads: int = 4
-    dim_head: int = 64
-    dtype: Optional[Dtype] = None
-    precision: PrecisionLike = None
-    use_bias: bool = True
-    # kernel_init: Callable = kernel_init(1.0)
-    force_fp32_for_softmax: bool = True
-
-    def setup(self):
-        inner_dim = self.dim_head * self.heads
-        # Weights were exported with old names {to_q, to_k, to_v, to_out}
-        dense = functools.partial(
-            nn.Dense,
-            self.heads * self.dim_head,
-            precision=self.precision, 
-            use_bias=self.use_bias, 
-            # kernel_init=self.kernel_init, 
-            dtype=self.dtype
-        )
-        self.query = dense(name="to_q")
-        self.key = dense(name="to_k")
-        self.value = dense(name="to_v")
-        
-        self.proj_attn = nn.DenseGeneral(
-            self.query_dim, 
-            use_bias=False, 
-            precision=self.precision, 
-            # kernel_init=self.kernel_init,
-            dtype=self.dtype,
-            name="to_out_0"
-        )
-        # self.attnfn = make_fast_generalized_attention(qkv_dim=inner_dim, lax_scan_unroll=16)
-    
-    def _reshape_tensor_to_head_dim(self, tensor):
-        batch_size, _, seq_len, dim = tensor.shape
-        head_size = self.heads
-        tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size)
-        tensor = jnp.transpose(tensor, (0, 2, 1, 3))
-        return tensor
-    
-    def _reshape_tensor_from_head_dim(self, tensor):
-        batch_size, _, seq_len, dim = tensor.shape
-        head_size = self.heads
-        tensor = jnp.transpose(tensor, (0, 2, 1, 3))
-        tensor = tensor.reshape(batch_size, 1, seq_len, dim * head_size)
-        return tensor
-
-    @nn.compact
-    def __call__(self, x:jax.Array, context=None):
-        # print(x.shape)
-        # x has shape [B, H * W, C]
-        context = x if context is None else context
-        
-        orig_x_shape = x.shape
-        if len(x.shape) == 4:
-            B, H, W, C = x.shape
-            x = x.reshape((B, 1, H * W, C))
-        else:
-            B, SEQ, C = x.shape
-            x = x.reshape((B, 1, SEQ, C))
-        
-        if len(context.shape) == 4:
-            B, _H, _W, _C = context.shape
-            context = context.reshape((B, 1, _H * _W, _C))
-        else:
-            B, SEQ, _C = context.shape
-            context = context.reshape((B, 1, SEQ, _C))
-        
-        query = self.query(x)
-        key = self.key(context)
-        value = self.value(context)
-        
-        query = self._reshape_tensor_to_head_dim(query)
-        key = self._reshape_tensor_to_head_dim(key)
-        value = self._reshape_tensor_to_head_dim(value)
-        
-        hidden_states = jax.experimental.pallas.ops.tpu.flash_attention.flash_attention(
-            query, key, value, None
-        )
-        
-        hidden_states = self._reshape_tensor_from_head_dim(hidden_states)
-        
-        
-        # hidden_states = nn.dot_product_attention(
-        #     query, key, value, dtype=self.dtype, broadcast_dropout=False, dropout_rng=None, precision=self.precision
-        # )
-        
-        proj = self.proj_attn(hidden_states)
-        
-        proj = proj.reshape(orig_x_shape)
-        
-        return proj
 
 class NormalAttention(nn.Module):
     """
@@ -246,17 +144,13 @@ class BasicTransformerBlock(nn.Module):
     precision: PrecisionLike = None
     use_bias: bool = True
     # kernel_init: Callable = kernel_init(1.0)
-    use_flash_attention:bool = False
     use_cross_only:bool = False
     only_pure_attention:bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-4
     
     def setup(self):
-        if self.use_flash_attention:
-            attenBlock = EfficientAttention
-        else:
-            attenBlock = NormalAttention
+        attenBlock = NormalAttention
             
         self.attention1 = attenBlock(
          query_dim=self.query_dim,
@@ -281,7 +175,7 @@ def setup(self):
             force_fp32_for_softmax=self.force_fp32_for_softmax
         )
         
-        self.ff = FlaxFeedForward(dim=self.query_dim)
+        self.ff = FlaxFeedForward(dim=self.query_dim, dtype=self.dtype, precision=self.precision)
         self.norm1 = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)
         self.norm2 = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)
         self.norm3 = nn.RMSNorm(epsilon=self.norm_epsilon, dtype=self.dtype)
@@ -309,7 +203,6 @@ class TransformerBlock(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     use_projection: bool = False
-    use_flash_attention:bool = False
     use_self_and_cross:bool = True
     only_pure_attention:bool = False
     force_fp32_for_softmax: bool = True
@@ -351,7 +244,6 @@ def __call__(self, x, context=None):
             precision=self.precision,
             use_bias=False,
             dtype=self.dtype,
-            use_flash_attention=self.use_flash_attention,
             use_cross_only=(not self.use_self_and_cross),
             only_pure_attention=self.only_pure_attention,
             force_fp32_for_softmax=self.force_fp32_for_softmax,
diff --git a/flaxdiff/models/common.py b/flaxdiff/models/common.py
index 50ea1dd..b732eae 100644
--- a/flaxdiff/models/common.py
+++ b/flaxdiff/models/common.py
@@ -180,7 +180,6 @@ def setup(self):
                 kernel_size=self.kernel_size,
                 strides=self.strides,
                 padding="SAME",
-                param_dtype=self.dtype,
                 dtype=self.dtype,
                 precision=self.precision
             )
diff --git a/flaxdiff/models/general.py b/flaxdiff/models/general.py
index 19d2f71..cf7b09c 100644
--- a/flaxdiff/models/general.py
+++ b/flaxdiff/models/general.py
@@ -6,7 +6,7 @@ class BCHWModelWrapper(nn.Module):
     model: nn.Module
 
     @nn.compact
-    def __call__(self, x, temb, textcontext):
+    def __call__(self, x, temb, textcontext, train: bool = False):
         # Reshape the input to BCHW format from BHWC
         x = jnp.transpose(x, (0, 3, 1, 2))
         # Pass the input through the UNet model
diff --git a/flaxdiff/models/simple_dit.py b/flaxdiff/models/simple_dit.py
index d7624f0..70fa752 100644
--- a/flaxdiff/models/simple_dit.py
+++ b/flaxdiff/models/simple_dit.py
@@ -28,7 +28,6 @@ class DiTBlock(nn.Module):
     dropout_rate: float = 0.0
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    use_flash_attention: bool = False # Keep placeholder
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-5
     use_gating: bool = True # Add flag to easily disable gating
@@ -60,8 +59,10 @@ def setup(self):
             nn.Dense(features=self.features, dtype=self.dtype, precision=self.precision)
         ])
 
+        self.dropout = nn.Dropout(rate=self.dropout_rate)
+
     @nn.compact
-    def __call__(self, x, conditioning, freqs_cis):
+    def __call__(self, x, conditioning, freqs_cis, train: bool = False):
         # Get scale/shift/gate parameters
         # Shape: [B, 1, 6*F] -> split into 6 of [B, 1, F]
         scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split(
@@ -74,6 +75,7 @@ def __call__(self, x, conditioning, freqs_cis):
         # Modulate after norm
         x_attn_modulated = norm_x_attn * (1 + scale_attn) + shift_attn
         attn_output = self.attention(x_attn_modulated, context=None, freqs_cis=freqs_cis)
+        attn_output = self.dropout(attn_output, deterministic=not train)
 
         if self.use_gating:
              x = residual + gate_attn * attn_output
@@ -86,6 +88,7 @@ def __call__(self, x, conditioning, freqs_cis):
         # Modulate after norm
         x_mlp_modulated = norm_x_mlp * (1 + scale_mlp) + shift_mlp
         mlp_output = self.mlp(x_mlp_modulated)
+        mlp_output = self.dropout(mlp_output, deterministic=not train)
 
         if self.use_gating:
              x = residual + gate_mlp * mlp_output
@@ -111,14 +114,11 @@ class SimpleDiT(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     # Passed down, but RoPEAttention uses NormalAttention
-    use_flash_attention: bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-5
     learn_sigma: bool = False  # Option to predict sigma like in DiT paper
     use_hilbert: bool = False  # Toggle Hilbert patch reorder
     use_zigzag: bool = False  # Toggle zigzag (serpentine) patch reorder
-    norm_groups: int = 0
-    activation: Callable = jax.nn.swish
 
     def setup(self):
         self.patch_embed = PatchEmbedding(
@@ -178,7 +178,6 @@ def setup(self):
                 dropout_rate=self.dropout_rate,
                 dtype=self.dtype,
                 precision=self.precision,
-                use_flash_attention=self.use_flash_attention,
                 force_fp32_for_softmax=self.force_fp32_for_softmax,
                 norm_epsilon=self.norm_epsilon,
                 rope_emb=self.rope,  # Pass RoPE instance
@@ -198,14 +197,14 @@ def setup(self):
 
         self.final_proj = nn.Dense(
             features=output_dim,
-            dtype=self.dtype,
+            dtype=jnp.float32,  # fp32 output head - the loss is computed in fp32
             precision=self.precision,
             kernel_init=nn.initializers.zeros,  # Initialize final layer to zero
             name="final_proj"
         )
 
     @nn.compact
-    def __call__(self, x, temb, textcontext=None):
+    def __call__(self, x, temb, textcontext=None, train: bool = False):
         B, H, W, C = x.shape
         assert H % self.patch_size == 0 and W % self.patch_size == 0, "Image dimensions must be divisible by patch size"
 
@@ -272,7 +271,7 @@ def __call__(self, x, temb, textcontext=None):
 
         # 4. Apply Transformer Blocks with adaLN-Zero conditioning
         for block in self.blocks:
-            x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin))
+            x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin), train=train)
 
         # 5. Final Layer
         x_out = self.final_norm(x_seq)
diff --git a/flaxdiff/models/simple_mmdit.py b/flaxdiff/models/simple_mmdit.py
index 5f27bf8..904701e 100644
--- a/flaxdiff/models/simple_mmdit.py
+++ b/flaxdiff/models/simple_mmdit.py
@@ -103,7 +103,6 @@ class MMDiTBlock(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     # Keep option, though RoPEAttention doesn't use it
-    use_flash_attention: bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-5
 
@@ -134,8 +133,10 @@ def setup(self):
                      precision=self.precision)
         ])
 
+        self.dropout = nn.Dropout(rate=self.dropout_rate)
+
     @nn.compact
-    def __call__(self, x, t_emb, text_emb, freqs_cis):
+    def __call__(self, x, t_emb, text_emb, freqs_cis, train: bool = False):
         # x shape: [B, S, F]
         # t_emb shape: [B, D_t] or [B, 1, D_t]
         # text_emb shape: [B, D_text] or [B, 1, D_text]
@@ -149,10 +150,12 @@ def __call__(self, x, t_emb, text_emb, freqs_cis):
         # Attention block (remains the same)
         attn_output = self.attention(
             x_attn, context=None, freqs_cis=freqs_cis)  # Self-attention only
+        attn_output = self.dropout(attn_output, deterministic=not train)
         x = residual + gate_attn * attn_output
 
         # MLP block (remains the same)
         mlp_output = self.mlp(x_mlp)
+        mlp_output = self.dropout(mlp_output, deterministic=not train)
         x = x + gate_mlp * mlp_output
 
         return x
@@ -176,13 +179,10 @@ class SimpleMMDiT(nn.Module):
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
     # Passed down, but RoPEAttention uses NormalAttention
-    use_flash_attention: bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-5
     learn_sigma: bool = False  # Option to predict sigma like in DiT paper
     use_hilbert: bool = False  # Toggle Hilbert patch reorder
-    norm_groups: int = 0
-    activation: Callable = jax.nn.swish
 
     def setup(self):
         self.patch_embed = PatchEmbedding(
@@ -228,7 +228,6 @@ def setup(self):
                 dropout_rate=self.dropout_rate,
                 dtype=self.dtype,
                 precision=self.precision,
-                use_flash_attention=self.use_flash_attention,
                 force_fp32_for_softmax=self.force_fp32_for_softmax,
                 norm_epsilon=self.norm_epsilon,
                 rope_emb=self.rope,  # Pass RoPE instance
@@ -248,14 +247,14 @@ def setup(self):
 
         self.final_proj = nn.Dense(
             features=output_dim,
-            dtype=self.dtype,
+            dtype=jnp.float32,  # fp32 output head - the loss is computed in fp32
             precision=self.precision,
             kernel_init=nn.initializers.zeros,  # Initialize final layer to zero
             name="final_proj"
         )
 
     @nn.compact
-    def __call__(self, x, temb, textcontext):  # textcontext is required
+    def __call__(self, x, temb, textcontext, train: bool = False):  # textcontext is required
         B, H, W, C = x.shape
         assert H % self.patch_size == 0 and W % self.patch_size == 0, "Image dimensions must be divisible by patch size"
         assert textcontext is not None, "textcontext must be provided for SimpleMMDiT"
@@ -291,7 +290,7 @@ def __call__(self, x, temb, textcontext):  # textcontext is required
         for block in self.blocks:
             # Pass t_emb and text_emb separately to the block
             x_seq = block(x_seq, t_emb, text_emb,
-                          freqs_cis=(freqs_cos, freqs_sin))
+                          freqs_cis=(freqs_cos, freqs_sin), train=train)
 
         # 5. Final Layer
         x_seq = self.final_norm(x_seq)
@@ -446,13 +445,10 @@ class HierarchicalMMDiT(nn.Module):
     dropout_rate: float = 0.0
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    use_flash_attention: bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-5
     learn_sigma: bool = False
     use_hilbert: bool = False
-    norm_groups: int = 0 # Not used in this structure, maybe remove later
-    activation: Callable = jax.nn.swish # Not used directly here, used in MLP inside MMDiTBlock
 
     def setup(self):
         assert len(self.emb_features) == len(self.num_layers) == len(self.num_heads), \
@@ -525,8 +521,7 @@ def setup(self):
                     dropout_rate=self.dropout_rate,
                     dtype=self.dtype,
                     precision=self.precision,
-                    use_flash_attention=self.use_flash_attention,
-                    force_fp32_for_softmax=self.force_fp32_for_softmax,
+                        force_fp32_for_softmax=self.force_fp32_for_softmax,
                     norm_epsilon=self.norm_epsilon,
                     rope_emb=self.ropes[stage],
                     name=f"encoder_block_stage{stage}_{i}"
@@ -589,8 +584,7 @@ def setup(self):
                     dropout_rate=self.dropout_rate,
                     dtype=self.dtype,
                     precision=self.precision,
-                    use_flash_attention=self.use_flash_attention,
-                    force_fp32_for_softmax=self.force_fp32_for_softmax,
+                        force_fp32_for_softmax=self.force_fp32_for_softmax,
                     norm_epsilon=self.norm_epsilon,
                     rope_emb=self.ropes[stage],
                     name=f"decoder_block_stage{stage}_{i}"
@@ -618,13 +612,13 @@ def setup(self):
 
         self.final_proj = nn.Dense(
             features=output_dim,
-            dtype=self.dtype,
+            dtype=jnp.float32,  # fp32 output head - the loss is computed in fp32
             precision=self.precision,
             kernel_init=nn.initializers.zeros,  # Zero init
             name="final_proj"
         )
 
-    def __call__(self, x, temb, textcontext):
+    def __call__(self, x, temb, textcontext, train: bool = False):
         B, H, W, C = x.shape
         num_stages = len(self.emb_features)
         finest_patch_size = self.base_patch_size
@@ -668,7 +662,7 @@ def __call__(self, x, temb, textcontext):
 
             # Apply blocks for this stage
             for block in self.encoder_blocks[stage]:
-                x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin))
+                x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin), train=train)
 
             # Store skip features (before merging)
             skip_features[stage] = x_seq
@@ -701,7 +695,7 @@ def __call__(self, x, temb, textcontext):
 
             # Apply blocks for this stage
             for block in self.decoder_blocks[i]: # Use index i for the decoder block list
-                x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin))
+                x_seq = block(x_seq, t_embs[stage], text_embs[stage], freqs_cis=(freqs_cos, freqs_sin), train=train)
 
         # --- Final Layer ---
         # x_seq should now be at the finest resolution (stage 0 features)
diff --git a/flaxdiff/models/simple_unet.py b/flaxdiff/models/simple_unet.py
index ae9e3c4..90cc982 100644
--- a/flaxdiff/models/simple_unet.py
+++ b/flaxdiff/models/simple_unet.py
@@ -30,7 +30,7 @@ def setup(self):
             self.conv_out_norm = norm()
         
     @nn.compact
-    def __call__(self, x, temb, textcontext):
+    def __call__(self, x, temb, textcontext, train: bool = False):
         # print("embedding features", self.emb_features)
         temb = FourierEmbedding(features=self.emb_features)(temb)
         temb = TimeProjection(features=self.emb_features)(temb)
@@ -74,7 +74,6 @@ def __call__(self, x, temb, textcontext):
                 if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
                     x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32),
                                         dim_head=dim_in // attention_config['heads'],
-                                        use_flash_attention=attention_config.get("flash_attention", False),
                                         use_projection=attention_config.get("use_projection", False),
                                         use_self_and_cross=attention_config.get("use_self_and_cross", True),
                                         precision=attention_config.get("precision", self.precision),
@@ -115,7 +114,6 @@ def __call__(self, x, temb, textcontext):
             if middle_attention is not None and j == self.num_middle_res_blocks - 1:   # Apply attention only on the last block
                 x = TransformerBlock(heads=middle_attention['heads'], dtype=middle_attention.get('dtype', jnp.float32), 
                                     dim_head=middle_dim_out // middle_attention['heads'],
-                                    use_flash_attention=middle_attention.get("flash_attention", False),
                                     use_linear_attention=False,
                                     use_projection=middle_attention.get("use_projection", False),
                                     use_self_and_cross=False,
@@ -161,7 +159,6 @@ def __call__(self, x, temb, textcontext):
                 if attention_config is not None and j == self.num_res_blocks - 1:   # Apply attention only on the last block
                     x = TransformerBlock(heads=attention_config['heads'], dtype=attention_config.get('dtype', jnp.float32), 
                                         dim_head=dim_out // attention_config['heads'],
-                                        use_flash_attention=attention_config.get("flash_attention", False),
                                         use_projection=attention_config.get("use_projection", False),
                                         use_self_and_cross=attention_config.get("use_self_and_cross", True),
                                         precision=attention_config.get("precision", self.precision),
diff --git a/flaxdiff/models/simple_vit.py b/flaxdiff/models/simple_vit.py
index 4df02d1..0aca9fe 100644
--- a/flaxdiff/models/simple_vit.py
+++ b/flaxdiff/models/simple_vit.py
@@ -21,9 +21,7 @@ class UViT(nn.Module):
     emb_features: int = 768
     num_layers: int = 12  # Should be even for U-Net structure
     num_heads: int = 12
-    dropout_rate: float = 0.1  # Dropout is often 0 in diffusion models
     use_projection: bool = False  # In TransformerBlock MLP
-    use_flash_attention: bool = False  # Passed to TransformerBlock
     # Passed to TransformerBlock (likely False for UViT)
     use_self_and_cross: bool = False
     force_fp32_for_softmax: bool = True  # Passed to TransformerBlock
@@ -94,7 +92,7 @@ def setup(self):
                 heads=self.num_heads,
                 dim_head=self.emb_features // self.num_heads,
                 dtype=self.dtype, precision=self.precision, use_projection=self.use_projection,
-                use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross,
+                use_self_and_cross=self.use_self_and_cross,
                 force_fp32_for_softmax=self.force_fp32_for_softmax,
                 only_pure_attention=False, norm_inputs=self.norm_inputs,
                 explicitly_add_residual=self.explicitly_add_residual,
@@ -107,7 +105,7 @@ def setup(self):
             heads=self.num_heads,
             dim_head=self.emb_features // self.num_heads,
             dtype=self.dtype, precision=self.precision, use_projection=self.use_projection,
-            use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross,
+            use_self_and_cross=self.use_self_and_cross,
             force_fp32_for_softmax=self.force_fp32_for_softmax,
             only_pure_attention=False, norm_inputs=self.norm_inputs,
             explicitly_add_residual=self.explicitly_add_residual,
@@ -128,7 +126,7 @@ def setup(self):
                 heads=self.num_heads,
                 dim_head=self.emb_features // self.num_heads,
                 dtype=self.dtype, precision=self.precision, use_projection=self.use_projection,
-                use_flash_attention=self.use_flash_attention, use_self_and_cross=self.use_self_and_cross,
+                use_self_and_cross=self.use_self_and_cross,
                 force_fp32_for_softmax=self.force_fp32_for_softmax,
                 only_pure_attention=False, norm_inputs=self.norm_inputs,
                 explicitly_add_residual=self.explicitly_add_residual,
@@ -172,7 +170,7 @@ def setup(self):
             )
 
     @nn.compact
-    def __call__(self, x, temb, textcontext=None):
+    def __call__(self, x, temb, textcontext=None, train: bool = False):
         original_img = x
         B, H, W, C = original_img.shape
         H_P = H // self.patch_size
@@ -264,7 +262,6 @@ class SimpleUDiT(nn.Module):
     dropout_rate: float = 0.0  # Typically 0 for diffusion
     dtype: Optional[Dtype] = None # e.g., jnp.float32 or jnp.bfloat16
     precision: PrecisionLike = None
-    use_flash_attention: bool = False # Passed to DiTBlock -> RoPEAttention
     force_fp32_for_softmax: bool = True # Passed to DiTBlock -> RoPEAttention
     norm_epsilon: float = 1e-5
     learn_sigma: bool = False
@@ -320,7 +317,6 @@ def setup(self):
                 dropout_rate=self.dropout_rate,
                 dtype=self.dtype,
                 precision=self.precision,
-                use_flash_attention=self.use_flash_attention,
                 force_fp32_for_softmax=self.force_fp32_for_softmax,
                 norm_epsilon=self.norm_epsilon,
                 rope_emb=self.rope,
@@ -335,7 +331,6 @@ def setup(self):
             dropout_rate=self.dropout_rate,
             dtype=self.dtype,
             precision=self.precision,
-            use_flash_attention=self.use_flash_attention,
             force_fp32_for_softmax=self.force_fp32_for_softmax,
             norm_epsilon=self.norm_epsilon,
             rope_emb=self.rope,
@@ -358,7 +353,6 @@ def setup(self):
                 dropout_rate=self.dropout_rate,
                 dtype=self.dtype,
                 precision=self.precision,
-                use_flash_attention=self.use_flash_attention,
                 force_fp32_for_softmax=self.force_fp32_for_softmax,
                 norm_epsilon=self.norm_epsilon,
                 rope_emb=self.rope,
@@ -382,7 +376,7 @@ def setup(self):
         )
 
     @nn.compact
-    def __call__(self, x, temb, textcontext=None):
+    def __call__(self, x, temb, textcontext=None, train: bool = False):
         B, H, W, C = x.shape
         H_P = H // self.patch_size
         W_P = W // self.patch_size
@@ -410,16 +404,16 @@ def __call__(self, x, temb, textcontext=None):
 
         skips = []
         for i in range(self.num_layers // 2):
-            x_seq = self.down_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None)
+            x_seq = self.down_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None, train=train)
             skips.append(x_seq)
 
-        x_seq = self.mid_block(x_seq, conditioning=cond_emb, freqs_cis=None)
+        x_seq = self.mid_block(x_seq, conditioning=cond_emb, freqs_cis=None, train=train)
 
         for i in range(self.num_layers // 2):
             skip_conn = skips.pop()
             x_seq = jnp.concatenate([x_seq, skip_conn], axis=-1)
             x_seq = self.up_dense[i](x_seq)
-            x_seq = self.up_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None)
+            x_seq = self.up_blocks[i](x_seq, conditioning=cond_emb, freqs_cis=None, train=train)
 
         x_out = self.final_norm(x_seq)
         x_out = self.final_proj(x_out)
diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py
index 3d2c2c8..94930d7 100644
--- a/flaxdiff/models/ssm_dit.py
+++ b/flaxdiff/models/ssm_dit.py
@@ -55,6 +55,7 @@ class S5Layer(nn.Module):
     def __call__(self, u):
         # u: [B, S, F]
         B, S, F = u.shape
+        assert F == self.features, f"S5Layer built for {self.features} features, got {F}"
 
         # A: diagonal complex state matrix, HiPPO init, parameterized as
         # log of the negative real part for stability
@@ -257,7 +258,6 @@ class SSMDiTBlock(nn.Module):
     dropout_rate: float = 0.0
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    use_flash_attention: bool = False  # Ignored, interface compat
     force_fp32_for_softmax: bool = True  # Ignored, interface compat
     norm_epsilon: float = 1e-5
     use_gating: bool = True
@@ -308,6 +308,8 @@ def setup(self):
             nn.Dense(features=self.features, dtype=self.dtype, precision=self.precision)
         ])
 
+        self.dropout = nn.Dropout(rate=self.dropout_rate)
+
     def _apply_2d_fusion(self, ssm_output):
         """Un-permute scan-ordered SSM output to a 2D grid, fuse, re-permute back."""
         B, S, F = ssm_output.shape
@@ -351,7 +353,7 @@ def _apply_2d_fusion(self, ssm_output):
         return y_fused
 
     @nn.compact
-    def __call__(self, x, conditioning, freqs_cis):
+    def __call__(self, x, conditioning, freqs_cis, train: bool = False):
         # Get scale/shift/gate parameters
         scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split(
             self.ada_params_module(conditioning), 6, axis=-1
@@ -365,6 +367,7 @@ def __call__(self, x, conditioning, freqs_cis):
 
         if self.use_2d_fusion:
             ssm_output = self._apply_2d_fusion(ssm_output)
+        ssm_output = self.dropout(ssm_output, deterministic=not train)
 
         if self.use_gating:
             x = residual + gate_attn * ssm_output
@@ -376,6 +379,7 @@ def __call__(self, x, conditioning, freqs_cis):
         norm_x_mlp = self.norm2(x)
         x_mlp_modulated = norm_x_mlp * (1 + scale_mlp) + shift_mlp
         mlp_output = self.mlp(x_mlp_modulated)
+        mlp_output = self.dropout(mlp_output, deterministic=not train)
 
         if self.use_gating:
             x = residual + gate_mlp * mlp_output
@@ -401,14 +405,11 @@ class HybridSSMAttentionDiT(nn.Module):
     dropout_rate: float = 0.0
     dtype: Optional[Dtype] = None
     precision: PrecisionLike = None
-    use_flash_attention: bool = False
     force_fp32_for_softmax: bool = True
     norm_epsilon: float = 1e-5
     learn_sigma: bool = False
     use_hilbert: bool = False
     use_zigzag: bool = False  # ZigMa-style serpentine scan
-    norm_groups: int = 0
-    activation: Callable = jax.nn.swish
     block_pattern: Optional[Sequence[str]] = None  # e.g., ['ssm','ssm','ssm','attn']
     ssm_attention_ratio: str = "3:1"  # e.g., "3:1", "1:1", "all-ssm", "all-attn"
     bidirectional_ssm: bool = True
@@ -500,8 +501,7 @@ def setup(self):
                     dropout_rate=self.dropout_rate,
                     dtype=self.dtype,
                     precision=self.precision,
-                    use_flash_attention=self.use_flash_attention,
-                    force_fp32_for_softmax=self.force_fp32_for_softmax,
+                        force_fp32_for_softmax=self.force_fp32_for_softmax,
                     norm_epsilon=self.norm_epsilon,
                     name=f"dit_block_{i}"
                 ))
@@ -517,14 +517,14 @@ def setup(self):
 
         self.final_proj = nn.Dense(
             features=output_dim,
-            dtype=self.dtype,
+            dtype=jnp.float32,  # fp32 output head - the loss is computed in fp32
             precision=self.precision,
             kernel_init=nn.initializers.zeros,
             name="final_proj"
         )
 
     @nn.compact
-    def __call__(self, x, temb, textcontext=None):
+    def __call__(self, x, temb, textcontext=None, train: bool = False):
         B, H, W, C = x.shape
         assert H % self.patch_size == 0 and W % self.patch_size == 0
 
@@ -579,7 +579,7 @@ def __call__(self, x, temb, textcontext=None):
 
         # 4. Hybrid blocks (SSM and attention interleaved)
         for block in self.blocks:
-            x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin))
+            x_seq = block(x_seq, conditioning=cond_emb, freqs_cis=(freqs_cos, freqs_sin), train=train)
 
         # 5. Final output
         x_out = self.final_norm(x_seq)
diff --git a/flaxdiff/samplers/common.py b/flaxdiff/samplers/common.py
index ee03df9..910202c 100644
--- a/flaxdiff/samplers/common.py
+++ b/flaxdiff/samplers/common.py
@@ -9,8 +9,6 @@
 from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
 from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple
 from ..utils import RandomMarkovState, MarkovState, clip_images
-from jax.experimental.shard_map import shard_map
-from jax.sharding import Mesh, PartitionSpec as P
 from flaxdiff.models.autoencoder import AutoEncoder
 from flaxdiff.inputs import DiffusionInputConfig
 
diff --git a/flaxdiff/trainer/diffusion_trainer.py b/flaxdiff/trainer/diffusion_trainer.py
index 99440dc..adbb834 100644
--- a/flaxdiff/trainer/diffusion_trainer.py
+++ b/flaxdiff/trainer/diffusion_trainer.py
@@ -175,6 +175,7 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc
             
             local_rng_state, rngs = local_rng_state.get_random_key()
             noise: jax.Array = jax.random.normal(rngs, shape=images.shape, dtype=jnp.float32)
+            local_rng_state, dropout_key = local_rng_state.get_random_key()
             
             # Make sure image is also float32
             images = images.astype(jnp.float32)
@@ -183,7 +184,10 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc
             noisy_images, c_in, expected_output = model_output_transform.forward_diffusion(images, noise, rates)
 
             def model_loss(params):
-                preds = model.apply(params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq)
+                preds = model.apply(
+                    params, *noise_schedule.transform_inputs(noisy_images*c_in, noise_level), label_seq,
+                    train=True, rngs={'dropout': dropout_key},
+                )
                 preds = model_output_transform.pred_transform(noisy_images, preds, rates)
                 nloss = loss_fn(preds, expected_output)
                 # Ignore the loss contribution of images with zero standard deviation
diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py
index 2539530..147bc22 100644
--- a/flaxdiff/trainer/general_diffusion_trainer.py
+++ b/flaxdiff/trainer/general_diffusion_trainer.py
@@ -279,6 +279,8 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc
             # Generate noise
             local_rng_state, noise_key = local_rng_state.get_random_key()
             noise = jax.random.normal(noise_key, shape=data.shape, dtype=jnp.float32)
+
+            local_rng_state, dropout_key = local_rng_state.get_random_key()
             
             # Forward diffusion process
             rates = noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(data))
@@ -288,7 +290,10 @@ def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch, loc
             def model_loss(params):
                 # Apply model
                 inputs = noise_schedule.transform_inputs(noisy_data * c_in, noise_level)
-                preds = model.apply(params, *inputs, *all_conditional_inputs)
+                preds = model.apply(
+                    params, *inputs, *all_conditional_inputs,
+                    train=True, rngs={'dropout': dropout_key},
+                )
                 
                 # Transform predictions and calculate loss
                 preds = model_output_transform.pred_transform(noisy_data, preds, rates)
@@ -312,8 +317,8 @@ def model_loss(params):
                 # Handle NaN/Inf gradients
                 select_fn = functools.partial(jnp.where, is_finite)
                 new_state = new_state.replace(
-                    opt_state=jax.tree_map(select_fn, new_state.opt_state, train_state.opt_state),
-                    params=jax.tree_map(select_fn, new_state.params, train_state.params)
+                    opt_state=jax.tree.map(select_fn, new_state.opt_state, train_state.opt_state),
+                    params=jax.tree.map(select_fn, new_state.params, train_state.params)
                 )
             else:
                 # Standard gradient computation
diff --git a/tests/test_config.py b/tests/test_config.py
index 7d8b831..dea975d 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -18,9 +18,7 @@ def make_config(model_overrides=None, arguments_overrides=None):
         "emb_features": 64,
         "dtype": "bfloat16",
         "precision": "high",
-        "activation": "swish",
         "output_channels": 3,
-        "norm_groups": 4,
         "patch_size": 4,
         "num_layers": 2,
         "num_heads": 2,
@@ -56,7 +54,14 @@ def test_parse_config_rebuilds_model():
     assert model.num_layers == 2
     assert model.dtype == jnp.bfloat16
     assert model.precision == jax.lax.Precision.HIGH
-    assert model.activation is jax.nn.swish
+
+
+def test_parse_config_drops_removed_fields():
+    """Configs from older runs carry since-removed flags; reconstruction must
+    drop them instead of crashing."""
+    result = parse_config(make_config(model_overrides={
+        "use_flash_attention": False, "norm_groups": 8, "activation": "swish"}))
+    assert type(result["model"]).__name__ == "SimpleDiT"
 
 
 def test_parse_config_noise_schedule_selection():
@@ -67,14 +72,19 @@ def test_parse_config_noise_schedule_selection():
     assert type(cosine["noise_schedule"]).__name__ == "CosineNoiseScheduler"
 
 
-@pytest.mark.xfail(strict=True, reason="bug: parse_config silently drops any string value containing a dot")
-def test_parse_config_preserves_dotted_values():
-    """String config values with a dot in them (module paths, filenames, version
-    strings) must survive the round trip, not vanish into class defaults."""
-    result = parse_config(make_config(model_overrides={"activation": "jax.nn.mish"}))
-    model = result["model"]
-    # a dropped key falls back to the class default (swish) instead of erroring
-    assert model.activation is jax.nn.mish, "activation was silently dropped"
+def test_parse_config_resolves_dotted_values():
+    """Function paths like 'jax.nn.mish' (and the 'jax._src.nn.functions.silu'
+    that old configs contain from the aliasing bug) must resolve to the actual
+    function, not silently vanish into class defaults."""
+    config = make_config(arguments_overrides={"architecture": "uvit"})
+    config["model"]["activation"] = "jax.nn.mish"
+    result = parse_config(config)
+    assert result["model"].activation is jax.nn.mish
+
+    config = make_config(arguments_overrides={"architecture": "uvit"})
+    config["model"]["activation"] = "jax._src.nn.functions.silu"
+    result = parse_config(config)
+    assert result["model"].activation is jax.nn.silu
 
 
 def test_training_and_inference_share_schedule_presets():
diff --git a/tests/test_models.py b/tests/test_models.py
index ca5cb13..7747b49 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -34,7 +34,7 @@ def test_unet_forward(rng):
     model = Unet(
         emb_features=64,
         feature_depths=[16, 32],
-        attention_configs=[None, {"heads": 2, "dtype": jnp.float32, "flash_attention": False,
+        attention_configs=[None, {"heads": 2, "dtype": jnp.float32,
                                   "use_projection": False, "use_self_and_cross": False}],
         num_res_blocks=1,
         num_middle_res_blocks=1,
@@ -53,7 +53,6 @@ def test_simple_dit_forward(rng):
     assert out.dtype == jnp.float32
 
 
-@pytest.mark.xfail(strict=True, reason="bug: DiT final_proj keeps compute dtype, bf16 output vs fp32 loss")
 def test_simple_dit_bf16_outputs_fp32(rng):
     model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2,
                       mlp_ratio=2, dtype=jnp.bfloat16)
@@ -119,3 +118,22 @@ def test_hilbert_patchify_roundtrip(rng):
     patches, inv_idx = hilbert_patchify(x, 4)
     rec = hilbert_unpatchify(patches, inv_idx, 4, 16, 16, 3)
     assert jnp.allclose(rec, x)
+
+
+def test_dropout_is_active_in_train_mode(rng):
+    """--dropout_rate used to be threaded into every block and applied nowhere."""
+    model = SimpleDiT(patch_size=4, emb_features=64, num_layers=2, num_heads=2,
+                      mlp_ratio=2, dropout_rate=0.5)
+    x, temb, textcontext = small_inputs(rng)
+    params = model.init(rng, x, temb, textcontext)
+    # nudge every param off init: the zero-initialized adaLN gates and final
+    # projection make a fresh DiT output exactly zero, hiding dropout
+    params = jax.tree.map(lambda p: p + 0.02, params)
+
+    d0 = model.apply(params, x, temb, textcontext, train=True, rngs={"dropout": jax.random.PRNGKey(1)})
+    d1 = model.apply(params, x, temb, textcontext, train=True, rngs={"dropout": jax.random.PRNGKey(2)})
+    assert not jnp.allclose(d0, d1), "different dropout rngs must give different outputs"
+
+    e0 = model.apply(params, x, temb, textcontext)
+    e1 = model.apply(params, x, temb, textcontext)
+    assert jnp.allclose(e0, e1), "eval mode must be deterministic"
diff --git a/training.py b/training.py
index 1000826..834380a 100644
--- a/training.py
+++ b/training.py
@@ -11,10 +11,6 @@
 from flaxdiff.models.simple_dit import SimpleDiT
 from flaxdiff.models.simple_mmdit import SimpleMMDiT, HierarchicalMMDiT
 from flaxdiff.models.ssm_dit import HybridSSMAttentionDiT
-try:
-    import jax.experimental.pallas.ops.tpu.flash_attention
-except (ImportError, ModuleNotFoundError):
-    pass
 from flaxdiff.predictors import get_diffusion_preset
 from diffusers import FlaxUNet2DConditionModel
 from flaxdiff.samplers.euler import EulerAncestralSampler
@@ -139,7 +135,6 @@ def boolean_string(s):
 parser.add_argument('--emb_features', type=int, default=256, help='Embedding features')
 parser.add_argument('--feature_depths', type=int, nargs='+', default=[64, 128, 256, 512], help='Feature depths')
 parser.add_argument('--attention_heads', type=int, default=8, help='Number of attention heads')
-parser.add_argument('--flash_attention', type=boolean_string, default=False, help='Use Flash Attention')
 parser.add_argument('--use_projection', type=boolean_string, default=False, help='Use projection')
 parser.add_argument('--use_self_and_cross', type=boolean_string, default=True, help='Use self and cross attention')
 parser.add_argument('--only_pure_attention', type=boolean_string, default=True, help='Use only pure attention or proper transformer in the attention blocks') 
@@ -314,14 +309,14 @@ def main(args):
     if args.attention_heads > 0:
         attention_configs += [
             {
-                "heads": args.attention_heads, "dtype": DTYPE, "flash_attention": args.flash_attention,
+                "heads": args.attention_heads, "dtype": DTYPE,
                 "use_projection": args.use_projection, "use_self_and_cross": args.use_self_and_cross,
                 "only_pure_attention": args.only_pure_attention,    
             },
         ] * (len(args.feature_depths) - 2)
         attention_configs += [
             {
-                "heads": args.attention_heads, "dtype": DTYPE, "flash_attention": False,
+                "heads": args.attention_heads, "dtype": DTYPE,
                 "use_projection": False, "use_self_and_cross": args.use_self_and_cross,
                 "only_pure_attention": args.only_pure_attention
             },
@@ -370,9 +365,7 @@ def main(args):
             "emb_features": args.emb_features,
             "dtype": DTYPE,
             "precision": PRECISION,
-            "activation": args.activation,
             "output_channels": INPUT_CHANNELS,
-            "norm_groups": args.norm_groups,
         }
     
     
@@ -385,6 +378,8 @@ def main(args):
                 "num_res_blocks": args.num_res_blocks,
                 "num_middle_res_blocks": args.num_middle_res_blocks,
                 "named_norms": args.named_norms,
+                "activation": args.activation,
+                "norm_groups": args.norm_groups,
             },
         },
         "uvit": {
@@ -393,9 +388,9 @@ def main(args):
                 "patch_size":  args.patch_size,
                 "num_layers":  args.num_layers,
                 "num_heads":  args.num_heads,
-                "dropout_rate": args.dropout_rate,
                 "add_residualblock_output": args.add_residualblock_output,
-                "use_flash_attention": args.flash_attention,
+                "activation": args.activation,
+                "norm_groups": args.norm_groups,
                 "use_self_and_cross": args.use_self_and_cross,
                 "use_hilbert": use_hilbert,
             },
@@ -407,7 +402,6 @@ def main(args):
                 "num_layers":  args.num_layers,
                 "num_heads":  args.num_heads,
                 "dropout_rate": args.dropout_rate,
-                "use_flash_attention": args.flash_attention,
                 "mlp_ratio": args.mlp_ratio,
                 "use_hilbert": use_hilbert,
             },
@@ -419,7 +413,6 @@ def main(args):
                 "num_layers":  args.num_layers,
                 "num_heads":  args.num_heads,
                 "dropout_rate": args.dropout_rate,
-                "use_flash_attention": args.flash_attention,
                 "mlp_ratio": args.mlp_ratio,
                 "use_hilbert": use_hilbert,
                 "use_zigzag": use_zigzag,
@@ -432,7 +425,6 @@ def main(args):
                 "num_layers":  args.num_layers,
                 "num_heads":  args.num_heads,
                 "dropout_rate": args.dropout_rate,
-                "use_flash_attention": args.flash_attention,
                 "mlp_ratio": args.mlp_ratio,
                 "use_hilbert": use_hilbert,
             },
@@ -445,7 +437,6 @@ def main(args):
                 "num_layers": (args.num_layers // 3, args.num_layers // 2, args.num_layers),  # Default layers per stage
                 "num_heads": (args.num_heads - 2, args.num_heads, args.num_heads + 2),  # Default heads per stage
                 "dropout_rate": args.dropout_rate,
-                "use_flash_attention": args.flash_attention,
                 "mlp_ratio": args.mlp_ratio,
                 "use_hilbert": use_hilbert,
             },
@@ -457,7 +448,6 @@ def main(args):
                 "num_layers": args.num_layers,
                 "num_heads": args.num_heads,
                 "dropout_rate": args.dropout_rate,
-                "use_flash_attention": args.flash_attention,
                 "mlp_ratio": args.mlp_ratio,
                 "use_hilbert": use_hilbert,
                 "use_zigzag": use_zigzag,
@@ -476,7 +466,6 @@ def main(args):
                 "block_out_channels":args.feature_depths,
                 "cross_attention_dim":args.emb_features,
                 "dtype": DTYPE,
-                "use_memory_efficient_attention": args.flash_attention,
                 "attention_head_dim": args.attention_heads,
                 "only_cross_attention": not args.use_self_and_cross,
             },
@@ -575,10 +564,14 @@ def main(args):
         
     print("Experiment_Name:", experiment_name)
 
-    if 'activation' in model_config:
-        model_config['activation'] = ACTIVATION_MAP[model_config['activation']]
-        
-    model = model_architecture(**model_config)
+    # Map the activation string to the actual function only for construction.
+    # model_config is aliased into CONFIG (already logged to wandb), so mutating
+    # it here used to leak a live function object into the stored config.
+    construction_config = dict(model_config)
+    if 'activation' in construction_config:
+        construction_config['activation'] = ACTIVATION_MAP[construction_config['activation']]
+
+    model = model_architecture(**construction_config)
     
     # If using the Diffusers UNet, we need to wrap it 
     if 'diffusers' in architecture_name:

From 897c6013ac9fdbe9f4fc0db755b2dd6b5fa00714 Mon Sep 17 00:00:00 2001
From: Ashish Kumar Singh <ashishkmr472@gmail.com>
Date: Wed, 22 Jul 2026 23:55:24 -0400
Subject: [PATCH 5/5] docs: fix the inference example to match the sampler API

The old example called generate_images(num_images=...) which never
matched the actual signature, and showed a custom sampler with a stale
take_next_step contract.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 README.md | 32 +++++++++++++++-----------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 3eee41c..f3d6beb 100644
--- a/README.md
+++ b/README.md
@@ -210,23 +210,21 @@ final_state = trainer.fit(data, batches, epochs=2000, sampler_class=EulerAncestr
 Here is a simplified example for generating images using a trained model:
 
 ```python
-from flaxdiff.samplers import DiffusionSampler
-
-class EulerSampler(DiffusionSampler):
-    def take_next_step(self, current_samples, reconstructed_samples, pred_noise, current_step, state, next_step=None):
-        current_alpha, current_sigma = self.noise_schedule.get_rates(current_step)
-        next_alpha, next_sigma = self.noise_schedule.get_rates(next_step)
-        dt = next_sigma - current_sigma
-        x_0_coeff = (current_alpha * next_sigma - next_alpha * current_sigma) / dt
-        dx = (current_samples - x_0_coeff * reconstructed_samples) / current_sigma
-        next_samples = current_samples + dx * dt
-        return next_samples, state
-
-# Create sampler
-sampler = EulerSampler(trainer.model, trainer.state.ema_params, edm_schedule, model_output_transform=trainer.model_output_transform)
-
-# Generate images
-samples = sampler.generate_images(num_images=64, diffusion_steps=100, start_step=1000, end_step=0)
+from flaxdiff.samplers.euler import EulerSampler
+
+sampler = EulerSampler(
+    trainer.model,
+    noise_schedule=karas_ve_schedule,
+    model_output_transform=trainer.model_output_transform,
+    input_config=input_config,
+)
+
+samples = sampler.generate_samples(
+    params=trainer.state.ema_params,
+    num_samples=64,
+    resolution=128,
+    diffusion_steps=100,
+)
 plotImages(samples, dpi=300)
 ```