Onboard Flux.2-klein-4B: Implement latent noise generation, Qwen3 prompt embeddings, JAX projection layer, Euler scheduler, 4D RoPE, and verified block parity

Author: amepas

new-test.txt

@@ -446,7 +446,7 @@ def __call__(self, ids):
     pos = ids.astype(self.dtype)
     freqs_dtype = self.dtype
     for i in range(n_axes):
-      out = get_1d_rotary_pos_embed(self.axes_dim[i], pos[..., i], freqs_dtype=freqs_dtype)
+      out = get_1d_rotary_pos_embed(self.axes_dim[i], pos[..., i], theta=self.theta, freqs_dtype=freqs_dtype)
       out_freqs.append(out)
 
     out_freqs = jnp.concatenate(out_freqs, axis=1)