Changed rho type

xlb/operator/macroscopic/first_moment.py

@@ -7,6 +7,7 @@
 from xlb.compute_backend import ComputeBackend
 from xlb.operator.operator import Operator
 
+
 class FirstMoment(Operator):
     """A class to compute the first moment (velocity) of distribution functions."""
 
@@ -22,7 +23,10 @@ def _construct_warp(self):
         _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
 
         @wp.func
-        def functional(f: _f_vec, rho: float):
+        def functional(
+            f: _f_vec,
+            rho: Any,
+        ):
             u = _u_vec()
             for l in range(self.velocity_set.q):
                 for d in range(self.velocity_set.d):
@@ -80,4 +84,4 @@ def warp_implementation(self, f, rho, u):
             inputs=[f, rho, u],
             dim=u.shape[1:],
         )
-        return u
+        return u

xlb/operator/macroscopic/macroscopic.py

@@ -9,6 +9,7 @@
 from xlb.operator.macroscopic.zero_moment import ZeroMoment
 from xlb.operator.macroscopic.first_moment import FirstMoment
 
+
 class Macroscopic(Operator):
     """A class to compute both zero and first moments of distribution functions (rho, u)."""
 
@@ -82,4 +83,4 @@ def warp_implementation(self, f, rho, u):
             inputs=[f, rho, u],
             dim=rho.shape[1:],
         )
-        return rho, u
+        return rho, u

xlb/operator/macroscopic/second_moment.py

@@ -123,12 +123,5 @@ def kernel2d(
     @Operator.register_backend(ComputeBackend.WARP)
     def warp_implementation(self, f, pi):
         # Launch the warp kernel
-        wp.launch(
-            self.warp_kernel,
-            inputs=[
-                f,
-                pi,
-            ],
-            dim=pi.shape[1:],
-        )
+        wp.launch(self.warp_kernel, inputs=[f, pi], dim=pi.shape[1:])
         return pi

xlb/operator/macroscopic/zero_moment.py

@@ -7,6 +7,7 @@
 from xlb.compute_backend import ComputeBackend
 from xlb.operator.operator import Operator
 
+
 class ZeroMoment(Operator):
     """A class to compute the zeroth moment (density) of distribution functions."""
 
@@ -61,9 +62,5 @@ def kernel2d(
 
     @Operator.register_backend(ComputeBackend.WARP)
     def warp_implementation(self, f, rho):
-        wp.launch(
-            self.warp_kernel,
-            inputs=[f, rho],
-            dim=rho.shape[1:],
-        )
-        return rho
+        wp.launch(self.warp_kernel, inputs=[f, rho], dim=rho.shape[1:])
+        return rho