added the import at the top

Author: cheachu

keras_hub/src/models/deepseek_v31/deepseek_v31_attention.py

@@ -1,5 +1,7 @@
 """DeepSeek V31 Multi-head Latent Attention layer."""
 
+import math
+
 import keras
 from keras import ops
 
@@ -232,7 +234,6 @@ def _yarn_inv_freq(self, dtype):
 
         # Wavelength = 2π / freq. High-freq → small wavelength, low-freq →
         # large wavelength. YaRN applies more scaling to low-freq dimensions.
-        import math
 
         wavelengths = 2.0 * math.pi / freqs
         old_ctx = float(self.yarn_original_max_position_embeddings)