added debug for values

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -296,6 +296,7 @@ def wrap_flash_attention(query, key, value):
     kv_padded_len = key.shape[2]
     print("KV padded len:", kv_padded_len)
     kv_indices = jax.lax.broadcasted_iota(jnp.int32, (kv_padded_len,), 0)
+    print("KV indices:", kv_indices)
     kv_segment_ids = (kv_indices < key_seq_len).astype(jnp.int32)
     print("KV segment ids:", kv_segment_ids)
     segment_ids = splash_attention_kernel.SegmentIds(q=q_segment_ids, kv=kv_segment_ids)