diff --git a/xlb/operator/macroscopic/first_moment.py b/xlb/operator/macroscopic/first_moment.py
index 218458e..561fe7a 100644
--- a/xlb/operator/macroscopic/first_moment.py
+++ b/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
\ No newline at end of file
+        return u
diff --git a/xlb/operator/macroscopic/macroscopic.py b/xlb/operator/macroscopic/macroscopic.py
index b585b5e..495a6a0 100644
--- a/xlb/operator/macroscopic/macroscopic.py
+++ b/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
\ No newline at end of file
+        return rho, u
diff --git a/xlb/operator/macroscopic/second_moment.py b/xlb/operator/macroscopic/second_moment.py
index bda2a25..917c86a 100644
--- a/xlb/operator/macroscopic/second_moment.py
+++ b/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
diff --git a/xlb/operator/macroscopic/zero_moment.py b/xlb/operator/macroscopic/zero_moment.py
index a37ede7..d0fbf51 100644
--- a/xlb/operator/macroscopic/zero_moment.py
+++ b/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
\ No newline at end of file
+        wp.launch(self.warp_kernel, inputs=[f, rho], dim=rho.shape[1:])
+        return rho