precommit

Author: cetagostini

python/nutpie/compile_pymc.py

@@ -576,7 +576,7 @@ def expand(_x, **shared):
         dims=dims,
         coords=coords,
         raw_logp_fn=orig_logp_fn,
-        force_single_core=(gradient_backend == "mlx")
+        force_single_core=(gradient_backend == "mlx"),
     )
 
 

python/nutpie/compiled_pyfunc.py

@@ -43,15 +43,16 @@ def with_data(self, **updates):
                 raise ValueError(f"Unknown data variable: {name}")
 
         updated = self._shared_data.copy()
-        
+
         # Convert to MLX arrays if using MLX backend (indicated by force_single_core)
         if self._force_single_core:
             import mlx.core as mx
+
             for name, value in updates.items():
                 updated[name] = mx.array(value)
         else:
             updated.update(**updates)
-        
+
         return dataclasses.replace(self, _shared_data=updated)
 
     def with_transform_adapt(self, **kwargs):