Minor corrections

Author: dirmeier

.github/workflows/ci.yaml

@@ -8,8 +8,6 @@ on:
 
 jobs:
   precommit:
-    needs:
-      - format
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3

ebm/__init__.py

@@ -1,7 +1,6 @@
-"""ebm: ebm."""
-
-__version__ = "0.0.1"
+"""ebm: Implementation of an EBM trained using NCE."""
 
+__version__ = "0.1.0"
 
 from ebm.ebm import ebm_training_fns_and_sample_fn, EBMConfig
 

pyproject.toml

@@ -4,11 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name =  "ebm"
-description = "ebm"
+description = "Implementation of an EBM trained using NCE."
 authors = [{name = "Simon Dirmeier", email = "simd23@pm.me"}]
 readme = "README.md"
 license = {file = "LICENSE"}
-keywords = []
 classifiers = [
     "Development Status :: 1 - Planning",
     "Intended Audience :: Science/Research",
@@ -20,19 +19,16 @@ classifiers = [
 requires-python = ">=3.11"
 dependencies = [
     "jax-ai-stack",
-    "tensorflow-probability>=0.25.0",
+    "tensorflow-probability[jax]>=0.25.0",
     "tensorflow>=2.20.0",
     "tensorflow-datasets>=4.9.9",
     "ml-collections>=1.1.0",
-    "tf-keras>=2.20.1",
     "flax>=0.11.1",
 ]
 dynamic = ["version"]
 
 [dependency-groups]
 dev = [
-    "gitlint",
-    "jupyter",
     "pre-commit",
     "ruff",
 ]