bartz 0.2.0__tar.gz → 0.2.1__tar.gz

{bartz-0.2.0 → bartz-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bartz
-Version: 0.2.0
+Version: 0.2.1
 Summary: A JAX implementation of BART
 Home-page: https://github.com/Gattocrucco/bartz
 License: MIT

{bartz-0.2.0 → bartz-0.2.1}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "bartz"
-version = "0.2.0"
+version = "0.2.1"
 description = "A JAX implementation of BART"
 authors = ["Giacomo Petrillo <info@giacomopetrillo.com>"]
 license = "MIT"
@@ -61,7 +61,7 @@ pytest = "^8.1.1"
 
 [tool.poetry.group.docs.dependencies]
 Sphinx = "^7.2.6"
-numpydoc = "^1.6.0,<1.7.0" # 1.7.0 breaks linkcode, it seems
+numpydoc = "^1.6.0"
 myst-parser = "^2.0.0"
 
 [tool.pytest.ini_options]

bartz-0.2.1/src/bartz/_version.py

@@ -0,0 +1 @@
+__version__ = '0.2.1'

{bartz-0.2.0 → bartz-0.2.1}/src/bartz/mcmcstep.py

@@ -84,7 +84,8 @@ def init(*,
         The minimum number of data points in a leaf node. 0 if not specified.
     suffstat_batch_size : int, None, str, default 'auto'
         The batch size for computing sufficient statistics. `None` for no
-        batching. If 'auto', pick a value based on the device of `y`.
+        batching. If 'auto', pick a value based on the device of `y`, or the
+        default device.
 
     Returns
     -------
@@ -188,7 +189,12 @@ def init(*,
 
 def _choose_suffstat_batch_size(size, y):
     if size == 'auto':
-        platform = y.devices().pop().platform
+        try:
+            device = y.devices().pop()
+        except jax.errors.ConcretizationTypeError:
+            device = jax.devices()[0]
+        platform = device.platform
+
         if platform == 'cpu':
             return None
                 # maybe I should batch residuals (not counts) for numerical
@@ -198,8 +204,10 @@ def _choose_suffstat_batch_size(size, y):
                        # 512 is good on T4, and V100 at low n
         else:
             raise KeyError(f'Unknown platform: {platform}')
+    
     elif size is not None:
         return int(size)
+    
     return size
 
 def step(bart, key):

bartz-0.2.0/src/bartz/_version.py

@@ -1 +0,0 @@
-__version__ = '0.2.0'