fix tests

GridTools · Dec 19, 2024 · 6709b04 · 6709b04
1 parent 7817ec3
commit 6709b04
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 15 deletions.
diff --git a/src/gt4py/next/_allocators.py b/src/gt4py/next/_allocators.py
@@ -520,7 +520,7 @@ def empty(
                 device: Optional[core_defs.Device] = None,
                 aligned_index: Optional[Sequence[common.NamedIndex]] = None,
             ) -> core_defs.NDArrayObject:
-                assert is_field_allocator(allocator)
+                assert is_field_allocation_tool(allocator)
                 return allocate(
                     domain=common.domain(domain),
                     dtype=core_defs.dtype(dtype),
@@ -537,7 +537,7 @@ def zeros(
                 device: Optional[core_defs.Device] = None,
                 aligned_index: Optional[Sequence[common.NamedIndex]] = None,
             ) -> core_defs.NDArrayObject:
-                assert is_field_allocator(allocator)
+                assert is_field_allocation_tool(allocator)
                 buffer = allocate(
                     domain=common.domain(domain),
                     dtype=core_defs.dtype(dtype),
@@ -556,7 +556,7 @@ def ones(
                 device: Optional[core_defs.Device] = None,
                 aligned_index: Optional[Sequence[common.NamedIndex]] = None,
             ) -> core_defs.NDArrayObject:
-                assert is_field_allocator(allocator)
+                assert is_field_allocation_tool(allocator)
                 buffer = allocate(
                     domain=common.domain(domain),
                     dtype=core_defs.dtype(dtype),
@@ -576,7 +576,7 @@ def full(
                 device: Optional[core_defs.Device] = None,
                 aligned_index: Optional[Sequence[common.NamedIndex]] = None,
             ) -> core_defs.NDArrayObject:
-                assert is_field_allocator(allocator)
+                assert is_field_allocation_tool(allocator)
                 buffer = allocate(
                     domain=common.domain(domain),
                     dtype=core_defs.dtype(dtype),  # TODO check all dtypes
@@ -597,7 +597,7 @@ def asarray(
                 copy: Optional[bool] = None,
                 aligned_index: Optional[Sequence[common.NamedIndex]] = None,
             ) -> core_defs.NDArrayObject:
-                assert is_field_allocator(allocator)
+                assert is_field_allocation_tool(allocator)
                 if not copy:
                     raise NotImplementedError("Zero-copy construction is not yet supported.")
                 dtype = core_defs.dtype(data.dtype) if dtype is None else core_defs.dtype(dtype)

diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace.py
@@ -173,13 +173,13 @@ def verify_testee(offset_provider):
         )
         mock_fast_call.assert_called_once()
 
-    if gtx.allocators.is_field_allocator_for(
+    if gtx._allocators.is_field_allocator_for(
         unstructured_case.backend.allocator, core_defs.DeviceType.CPU
     ):
         offset_provider = unstructured_case.offset_provider
     else:
-        assert gtx.allocators.is_field_allocator_for(
-            unstructured_case.backend.allocator, gtx.allocators.CUPY_DEVICE
+        assert gtx._allocators.is_field_allocator_for(
+            unstructured_case.backend.allocator, gtx._allocators.CUPY_DEVICE
         )
 
         import cupy as cp

diff --git a/tests/next_tests/unit_tests/test_allocators.py b/tests/next_tests/unit_tests/test_allocators.py
@@ -152,7 +152,7 @@ def test_allocate(self):
 
 
 def test_allocate():
-    from gt4py.next._allocators import StandardCPUFieldBufferAllocator, make_concrete_allocator
+    from gt4py.next._allocators import StandardCPUFieldBufferAllocator, allocate
 
     I = common.Dimension("I")
     J = common.Dimension("J")
@@ -161,25 +161,25 @@ def test_allocate():
 
     # Test with a explicit field allocator
     allocator = StandardCPUFieldBufferAllocator()
-    tensor_buffer = make_concrete_allocator(domain, dtype, allocator=allocator)()
+    tensor_buffer = allocate(domain=domain, dtype=dtype, allocator=allocator)
     assert tensor_buffer.shape == domain.shape
     assert tensor_buffer.dtype == dtype
 
     # Test with a device
     device = core_defs.Device(core_defs.DeviceType.CPU, 0)
-    tensor_buffer = make_concrete_allocator(domain, dtype, device=device)()
+    tensor_buffer = allocate(domain=domain, dtype=dtype, device=device)
     assert tensor_buffer.shape == domain.shape
     assert tensor_buffer.dtype == dtype
 
     # Test with both allocator and device
     with pytest.raises(ValueError, match="are incompatible"):
-        make_concrete_allocator(
-            domain,
-            dtype,
+        allocate(
+            domain=domain,
+            dtype=dtype,
             allocator=allocator,
             device=core_defs.Device(core_defs.DeviceType.CUDA, 0),
         )
 
     # Test with no device or allocator
     with pytest.raises(ValueError, match="No 'device' or 'allocator' specified"):
-        make_concrete_allocator(domain, dtype)
+        allocate(domain=domain, dtype=dtype)