Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/_memory/_device_memory_resource.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ class DeviceMemoryResource(_MemPool):
def __cinit__(self, *args, **kwargs) -> None:
...

def __init__(self, device_id: Device | int, options: DeviceMemoryResourceOptions | dict[str, object] | None=None) -> None:
def __init__(self, device_id: Device | int, options: DeviceMemoryResourceOptions | None=None) -> None:
...

def __reduce__(self) -> tuple[object, ...]:
Expand Down
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/_memory/_device_memory_resource.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ cdef class DeviceMemoryResource(_MemPool):
def __init__(
self,
device_id: Device | int,
options: DeviceMemoryResourceOptions | dict[str, object] | None = None
options: DeviceMemoryResourceOptions | None = None
) -> None:
_DMR_init(self, device_id, options)

Expand Down
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/_memory/_managed_memory_resource.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ class ManagedMemoryResource(_MemPool):
memory pools.
"""

def __init__(self, options: ManagedMemoryResourceOptions | dict[str, object] | None=None) -> None:
def __init__(self, options: ManagedMemoryResourceOptions | None=None) -> None:
...

def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> ManagedBuffer:
Expand Down
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/_memory/_managed_memory_resource.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ cdef class ManagedMemoryResource(_MemPool):
memory pools.
"""

def __init__(self, options: ManagedMemoryResourceOptions | dict[str, object] | None = None) -> None:
def __init__(self, options: ManagedMemoryResourceOptions | None = None) -> None:
_MMR_init(self, options)

def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> ManagedBuffer:
Expand Down
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ class PinnedMemoryResource(_MemPool):
See :class:`DeviceMemoryResource` for more details on IPC usage patterns.
"""

def __init__(self, options: PinnedMemoryResourceOptions | dict[str, object] | None=None) -> None:
def __init__(self, options: PinnedMemoryResourceOptions | None=None) -> None:
...

def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer:
Expand Down
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ cdef class PinnedMemoryResource(_MemPool):
See :class:`DeviceMemoryResource` for more details on IPC usage patterns.
"""

def __init__(self, options: PinnedMemoryResourceOptions | dict[str, object] | None = None) -> None:
def __init__(self, options: PinnedMemoryResourceOptions | None = None) -> None:
_PMR_init(self, options)

def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> Buffer:
Expand Down
12 changes: 6 additions & 6 deletions cuda_core/tests/test_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -1981,11 +1981,11 @@ def test_mempool_attributes_repr(memory_resource_factory):
device.set_current()

if MR is DeviceMemoryResource:
mr = MR(device, options={"max_size": 2048})
mr = MR(device, options=DeviceMemoryResourceOptions(max_size=2048))
elif MR is PinnedMemoryResource:
mr = MR(options={"max_size": 2048})
mr = MR(options=PinnedMemoryResourceOptions(max_size=2048))
elif MR is ManagedMemoryResource:
mr = create_managed_memory_resource_or_skip(options={})
mr = create_managed_memory_resource_or_skip(options=ManagedMemoryResourceOptions())

buffer1 = mr.allocate(64, stream=device.default_stream)
buffer2 = mr.allocate(64, stream=device.default_stream)
Expand Down Expand Up @@ -2018,11 +2018,11 @@ def test_mempool_attributes_ownership(memory_resource_factory):
device.set_current()

if MR is DeviceMemoryResource:
mr = MR(device, {"max_size": POOL_SIZE})
mr = MR(device, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
elif MR is PinnedMemoryResource:
mr = MR({"max_size": POOL_SIZE})
mr = MR(PinnedMemoryResourceOptions(max_size=POOL_SIZE))
elif MR is ManagedMemoryResource:
mr = create_managed_memory_resource_or_skip({})
mr = create_managed_memory_resource_or_skip(ManagedMemoryResourceOptions())

attributes = mr.attributes
mr.close()
Expand Down
Loading