Fix numerical instability that prevents AF3 from running without a GPU

Discovered by @wuooo339 in #273 (comment)

PiperOrigin-RevId: 948202523
Change-Id: I7566974d2dbd04abdabfcaa471a01b8ca3e12f33

Author: Augustin-Zidek

src/alphafold3/model/network/diffusion_head.py

@@ -341,7 +341,10 @@ def apply_denoising_step(carry, noise_level):
     gamma = config.gamma_0 * (noise_level > config.gamma_min)
     t_hat = noise_level_prev * (1 + gamma)
 
-    noise_scale = config.noise_scale * jnp.sqrt(t_hat**2 - noise_level_prev**2)
+    noise_scale = config.noise_scale * jnp.sqrt(
+        # Don't take sqrt of tiny negative number (happens when running on CPU).
+        jnp.maximum(t_hat**2 - noise_level_prev**2, 0.0)
+    )
     noise = noise_scale * jax.random.normal(key_noise, positions.shape)
     positions_noisy = positions + noise