test publish

Author: homerjed

.github/workflows/publish.yml

@@ -0,0 +1,42 @@
+name: Publish 
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - '0.0.*'
+    paths:
+      - transformer_flows/**
+      - pyproject.toml
+
+jobs:
+  pypi:
+    name: Publish to PyPI
+    runs-on: ubuntu-latest
+    # Environment and permissions trusted publishing.
+    environment:
+      # Create this environment in the GitHub repository under Settings -> Environments
+      name: release
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v3
+      - run: uv build
+      # Check that basic features work and we didn't miss to include crucial files
+      - name: Smoke test (wheel)
+        run: uv run --isolated --no-project -p 3.13 --with dist/*.whl tests/smoke_test.py
+      - name: Smoke test (source distribution)
+        run: uv run --isolated --no-project -p 3.13 --with dist/*.tar.gz tests/smoke_test.py
+      - run: uv publish --trusted-publishing always
+
+  github-release:
+    name: GitHub Release
+    needs:
+    - pypi
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write  # IMPORTANT: mandatory for making GitHub Releases
+      id-token: write  # IMPORTANT: mandatory for sigstore

.gitignore

@@ -10,4 +10,5 @@ data/
 __pycache__/
 tests/
 grfs.py
-guidance.py
+guidance.py
+test.yml

transformer_flow.py

@@ -25,7 +25,7 @@
 from attention import MultiheadAttention, self_attention
 
 
-typecheck = lambda x: x #jaxtyped(typechecker=typechecker)
+typecheck = jaxtyped(typechecker=typechecker)
 
 
 MetricsDict = dict[
@@ -1652,7 +1652,7 @@ def train(
     # Sharding: data and model
     sharding: Optional[NamedSharding] = None,
     replicated_sharding: Optional[NamedSharding] = None,
-    save_fn: Callable[[Optional[str], TransformerFlow], None]
+    save_fn: Callable[[Optional[str], TransformerFlow], None] = None
 ) -> TransformerFlow:
 
     print("n_params={:.3E}".format(count_parameters(model)))