fix test_serialization.py

Signed-off-by: Hexu Zhao <hexuz@nvidia.com>

Author: Hexu Zhao

tests/unit/test_serialization.py

@@ -8,7 +8,7 @@
 
 from fvdb import GridBatch, JaggedTensor
 from fvdb.utils.tests import dtype_to_atol
-from fvdb.utils.tests.grid_utils import create_test_grid_batch
+from fvdb.utils.tests.grid_utils import make_grid_batch_and_jagged_point_data
 
 all_device_dtype_combos = [
     ["cuda", torch.float16],
@@ -28,15 +28,11 @@ def setUp(self):
     def test_morton_permutation(self, device, dtype):
         """Test Morton order permutations (xyz and zyx variants)."""
         # Create a test grid batch with known active voxels
-        batch_size = 2
-        grid_size = 8
-
         # Create a grid batch with some active voxels
-        grid_batch = create_test_grid_batch(
-            batch_size=batch_size,
-            grid_size=grid_size,
+        grid_batch, _, _ = make_grid_batch_and_jagged_point_data(
             device=device,
-            density=0.3  # 30% of voxels will be active
+            dtype=dtype,
+            include_boundary_points=True
         )
 
         # Get permutation indices for both Morton orderings
@@ -77,15 +73,11 @@ def test_morton_permutation(self, device, dtype):
     def test_hilbert_permutation(self, device, dtype):
         """Test Hilbert curve permutations (xyz and zyx variants)."""
         # Create a test grid batch with known active voxels
-        batch_size = 2
-        grid_size = 8
-
         # Create a grid batch with some active voxels
-        grid_batch = create_test_grid_batch(
-            batch_size=batch_size,
-            grid_size=grid_size,
+        grid_batch, _, _ = make_grid_batch_and_jagged_point_data(
             device=device,
-            density=0.3  # 30% of voxels will be active
+            dtype=dtype,
+            include_boundary_points=True
         )
 
         # Get permutation indices for both Hilbert orderings
@@ -125,17 +117,12 @@ def test_hilbert_permutation(self, device, dtype):
     @parameterized.expand(all_device_dtype_combos)
     def test_permutation_validity(self, device, dtype):
         """Test that permutation indices are valid (complete and unique)."""
-        batch_size = 3
-        grid_size = 16
-
-        # Create grid batch with varying densities
-        densities = [0.1, 0.5, 0.8]  # Test different sparsity levels
-        for density in densities:
-            grid_batch = create_test_grid_batch(
-                batch_size=batch_size,
-                grid_size=grid_size,
+        # Create multiple grid batches to test different configurations
+        for _ in range(3):
+            grid_batch, _, _ = make_grid_batch_and_jagged_point_data(
                 device=device,
-                density=density
+                dtype=dtype,
+                include_boundary_points=True
             )
 
             # Test all permutation types