Merge pull request #160 from google-research/rajat_dev

Fixes build and upload to pypi

.github/workflows/main.yml

@@ -1,4 +1,3 @@
-
 name: Poetry publish
 
 on:
@@ -24,8 +23,8 @@ jobs:
       # e.g. poetry version 0.1.${{ github.run_number }}
       - name: Set Version number
         run: |
-          poetry version 1.2.0
+          poetry version 1.2.1
       - name: Build and Publish to PyPI
         run: |
           poetry config pypi-token.pypi ${{ secrets.PYPI_API_TOKEN }}
-          poetry publish --build
+          poetry publish --build

pyproject.toml

@@ -37,26 +37,35 @@ wandb = ">=0.17.5"
 ipython = "^8.27.0"
 absl-py = ">=1.4.0"
 
-[tool.poetry.group.pax]
-optional = true
+[tool.poetry.extras]
+pax = ["paxml", "lingvo", "jax", "jaxlib"]
+torch = ["torch",  "jax", "jaxlib"]  # jax & jaxlib are already in pax
 
-[tool.poetry.group.pax.dependencies]
-paxml = { version = ">=1.4.0", python = ">=3.10,<3.11" }
-lingvo = { version = ">=0.12.7", python = ">=3.10,<3.11" }
-jax = { version = ">=0.4.26", extras = ["cuda12"], python = ">=3.10,<3.11" }
-jaxlib = { version = ">=0.4.26", python = ">=3.10,<3.11" }
+[tool.poetry.dependencies.paxml]
+version = ">=1.4.0"
+python = ">=3.10,<3.11"
 
-[tool.poetry.group.torch]
-optional = true
+[tool.poetry.dependencies.lingvo]
+version = ">=0.12.7"
+python = ">=3.10,<3.11"
 
-[tool.poetry.group.torch.dependencies]
-torch = { version = ">=2.0.0", extras = ["cuda"], python = ">=3.11,<3.12" }
-jax = { version = ">=0.4.26", extras = ["cuda12"], python = ">=3.11,<3.12" }
-jaxlib = { version = ">=0.4.26", python = ">=3.11,<3.12" }
+[tool.poetry.dependencies.jax]
+version = ">=0.4.26"
+extras = ["cuda12"]
+python = ">=3.10,<3.12"  # Support both python versions
+
+[tool.poetry.dependencies.jaxlib]
+version = ">=0.4.26"
+python = ">=3.10,<3.12"  # Support both python versions
+
+[tool.poetry.dependencies.torch]
+version = ">=2.0.0"
+extras = ["cuda"]
+python = ">=3.11,<3.12"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=8.3.2"
 
 [build-system]
 requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+build-backend = "poetry.core.masonry.api"