Update cola version to prepare for jax>0.5.0

Author: theo-brown

pyproject.toml

@@ -28,12 +28,15 @@ dependencies = [
   "tqdm>4.66.2",
   "tensorflow-probability>=0.24.0",
   "beartype>0.16.1",
-  "cola-ml==0.0.5",
+  "cola-ml>=0.0.7",
   "jaxopt==0.8.2",
   "flax<0.10.0",
   "numpy<2.0.0",
 ]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.hatch.envs.default]
 installer = "uv"
 python = "3.10"