Remove jaxtyping import.

astro-informatics · May 10, 2024 · da4b31f · da4b31f
1 parent 321f345
commit da4b31f
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/harmonic/flows.py b/harmonic/flows.py
@@ -2,7 +2,6 @@
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array
 import tensorflow_probability as tfp
 import distrax
 
@@ -211,7 +210,7 @@ class RQSpline(nn.Module):
     num_bins: int
     spline_range: Sequence[float] = (-10.0, 10.0)
     multimodal_base: bool = False
-    base_centers: Sequence[Array] = None
+    base_centers: Sequence[jnp.ndarray]] = None
 
     def setup(self):
         conditioner = []