Coverage for src/arraybridge/array_operations.py: 58%
247 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:15 +0000
1"""Typed array-operation leaves carried by ``MemoryType`` declarations."""
3from collections.abc import Callable, Sequence
4from dataclasses import dataclass
5from typing import Any
7import numpy as np
9ToNumpy = Callable[[Any, Any], Any]
10FromNumpy = Callable[[Any, Any, int], Any]
11StackArrays = Callable[[Sequence[Any], Any], Any]
12ScaleDtype = Callable[[Any, Any, Any], Any]
13DtypeName = Callable[[Any], str]
14CastArray = Callable[[Any, Any, Any], Any]
15LogicalAnd = Callable[[Any, Any, Any], Any]
17_SCALING_RANGES: dict[str, float | tuple[float, float]] = {
18 "uint8": 255.0,
19 "uint16": 65535.0,
20 "uint32": 4294967295.0,
21 "int16": (65535.0, 32768.0),
22 "int32": (4294967295.0, 2147483648.0),
23}
26def _dtype_name(dtype: Any) -> str:
27 declared_name = getattr(dtype, "name", None)
28 if declared_name is not None:
29 return str(declared_name)
30 return getattr(dtype, "__name__", str(dtype).rsplit(".", maxsplit=1)[-1])
33def _numpy_dtype_name(dtype: Any) -> str:
34 return str(np.dtype(dtype).name)
37def _torch_dtype_name(dtype: Any) -> str:
38 return str(dtype).rsplit(".", maxsplit=1)[-1]
41def _tensorflow_dtype_name(dtype: Any) -> str:
42 numpy_dtype = getattr(dtype, "as_numpy_dtype", dtype)
43 return str(np.dtype(numpy_dtype).name)
46def _scaled_values(result: Any, result_min: Any, result_max: Any, target_dtype: Any) -> Any:
47 normalized = (result - result_min) / (result_max - result_min)
48 range_info = _SCALING_RANGES.get(_dtype_name(target_dtype))
49 if range_info is None:
50 return normalized
51 if isinstance(range_info, tuple):
52 scale, offset = range_info
53 return normalized * scale - offset
54 return normalized * range_info
57def _clamp_bounds(target_dtype: Any) -> tuple[float, float] | None:
58 range_info = _SCALING_RANGES.get(_dtype_name(target_dtype))
59 if range_info is None:
60 return None
61 if isinstance(range_info, tuple):
62 scale, offset = range_info
63 return -offset, scale - offset - 128
64 return 0, range_info
67def _identity_to_numpy(data: Any, module: Any) -> Any:
68 del module
69 return data
72def _identity_from_numpy(data: Any, module: Any, device_id: int) -> Any:
73 del module, device_id
74 return data
77def _numpy_stack(values: Sequence[Any], module: Any) -> Any:
78 return module.stack(values, axis=0)
81def _numpy_cast(data: Any, dtype: Any, module: Any) -> Any:
82 del module
83 return data.astype(dtype, copy=False)
86def _module_logical_and(left: Any, right: Any, module: Any) -> Any:
87 return module.logical_and(left, right)
90def _numpy_scale(result: Any, target_dtype: Any, module: Any) -> Any:
91 if not hasattr(result, "dtype"):
92 return result
93 if not (
94 module.issubdtype(result.dtype, module.floating)
95 and module.issubdtype(target_dtype, module.integer)
96 ):
97 return result.astype(target_dtype)
98 result_min = result.min()
99 result_max = result.max()
100 if result_max <= result_min:
101 return result.astype(target_dtype)
102 scaled = _scaled_values(result, result_min, result_max, target_dtype)
103 bounds = _clamp_bounds(target_dtype)
104 if bounds is not None:
105 scaled = module.clip(scaled, *bounds)
106 return scaled.astype(target_dtype)
109def _cupy_to_numpy(data: Any, module: Any) -> Any:
110 del module
111 return data.get()
114def _cupy_from_numpy(data: Any, module: Any, device_id: int) -> Any:
115 del device_id
116 return module.array(data)
119def _cupy_stack(values: Sequence[Any], module: Any) -> Any:
120 return module.stack(values, axis=0)
123def _cupy_scale(result: Any, target_dtype: Any, module: Any) -> Any:
124 if not hasattr(result, "dtype"):
125 return result
126 if not (
127 module.issubdtype(result.dtype, module.floating)
128 and not module.issubdtype(target_dtype, module.floating)
129 ):
130 return result.astype(target_dtype)
131 result_min = module.min(result)
132 result_max = module.max(result)
133 if result_max <= result_min:
134 return result.astype(target_dtype)
135 scaled = _scaled_values(result, result_min, result_max, target_dtype)
136 bounds = _clamp_bounds(target_dtype)
137 if bounds is not None:
138 scaled = module.clip(scaled, *bounds)
139 return scaled.astype(target_dtype)
142def _torch_to_numpy(data: Any, module: Any) -> Any:
143 del module
144 return data.cpu().numpy()
147def _torch_from_numpy(data: Any, module: Any, device_id: int) -> Any:
148 host_data = (
149 np.ascontiguousarray(data)
150 if any(stride < 0 for stride in getattr(data, "strides", ()))
151 else data
152 )
153 return module.from_numpy(host_data).to(f"cuda:{device_id}")
156def _torch_stack(values: Sequence[Any], module: Any) -> Any:
157 return module.stack(tuple(values), dim=0)
160def _torch_cast(data: Any, dtype: Any, module: Any) -> Any:
161 return data.to(dtype=_mapped_dtype(dtype, module))
164def _mapped_dtype(target_dtype: Any, module: Any) -> Any:
165 try:
166 dtype_name = np.dtype(target_dtype).name
167 except TypeError as error:
168 raise TypeError(f"Unsupported target dtype {target_dtype!r}") from error
169 mapped = getattr(module, dtype_name, None)
170 if mapped is None:
171 module_name = getattr(module, "__name__", type(module).__name__)
172 raise TypeError(f"{module_name} does not expose dtype {dtype_name}")
173 return mapped
176def _torch_scale(result: Any, target_dtype: Any, module: Any) -> Any:
177 if not hasattr(result, "dtype"):
178 return result
179 mapped = _mapped_dtype(target_dtype, module)
180 floats = (module.float16, module.float32, module.float64)
181 if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
182 return result.to(mapped)
183 result_min = result.min()
184 result_max = result.max()
185 if result_max <= result_min:
186 return result.to(mapped)
187 scaled = _scaled_values(result, result_min, result_max, target_dtype)
188 bounds = _clamp_bounds(target_dtype)
189 if bounds is not None:
190 scaled = module.clamp(scaled, min=bounds[0], max=bounds[1])
191 return scaled.to(mapped)
194def _tensorflow_to_numpy(data: Any, module: Any) -> Any:
195 del module
196 return data.numpy()
199def _tensorflow_from_numpy(data: Any, module: Any, device_id: int) -> Any:
200 del device_id
201 return module.convert_to_tensor(data)
204def _tensorflow_stack(values: Sequence[Any], module: Any) -> Any:
205 return module.stack(tuple(values), axis=0)
208def _tensorflow_cast(data: Any, dtype: Any, module: Any) -> Any:
209 return module.cast(data, _mapped_dtype(dtype, module))
212def _tensorflow_scale(result: Any, target_dtype: Any, module: Any) -> Any:
213 if not hasattr(result, "dtype"):
214 return result
215 mapped = _mapped_dtype(target_dtype, module)
216 floats = (module.float16, module.float32, module.float64)
217 if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
218 return module.cast(result, mapped)
219 result_min = module.reduce_min(result)
220 result_max = module.reduce_max(result)
221 if result_max <= result_min:
222 return module.cast(result, mapped)
223 scaled = _scaled_values(result, result_min, result_max, target_dtype)
224 bounds = _clamp_bounds(target_dtype)
225 if bounds is not None:
226 scaled = module.clip_by_value(scaled, *bounds)
227 return module.cast(scaled, mapped)
230def _jax_to_numpy(data: Any, module: Any) -> Any:
231 del module
232 return np.asarray(data)
235def _jax_from_numpy(data: Any, module: Any, device_id: int) -> Any:
236 devices = tuple(device for device in module.devices() if device.platform == "gpu")
237 return module.device_put(data, devices[device_id])
240def _jax_stack(values: Sequence[Any], module: Any) -> Any:
241 return module.numpy.stack(tuple(values), axis=0)
244def _jax_cast(data: Any, dtype: Any, module: Any) -> Any:
245 return data.astype(_mapped_dtype(dtype, module.numpy))
248def _jax_logical_and(left: Any, right: Any, module: Any) -> Any:
249 return module.numpy.logical_and(left, right)
252def _jax_scale(result: Any, target_dtype: Any, module: Any) -> Any:
253 if not hasattr(result, "dtype"):
254 return result
255 if np.dtype(target_dtype) == np.dtype(np.float64):
256 x64_enabled = getattr(module.config, "x64_enabled", None)
257 if x64_enabled is None:
258 x64_enabled = module.config.read("jax_enable_x64")
259 if not x64_enabled:
260 raise ValueError(
261 "JAX float64 output requires x64 mode; set JAX_ENABLE_X64=true before import"
262 )
263 jnp = module.numpy
264 mapped = _mapped_dtype(target_dtype, jnp)
265 floats = (jnp.float16, jnp.float32, jnp.float64)
266 if not (result.dtype in floats and np.issubdtype(np.dtype(target_dtype), np.integer)):
267 return result.astype(mapped)
268 result_min = jnp.min(result)
269 result_max = jnp.max(result)
270 if result_max <= result_min:
271 return result.astype(mapped)
272 scaled = _scaled_values(result, result_min, result_max, target_dtype)
273 bounds = _clamp_bounds(target_dtype)
274 if bounds is not None:
275 scaled = jnp.clip(scaled, *bounds)
276 return scaled.astype(mapped)
279def _pyclesperanto_to_numpy(data: Any, module: Any) -> Any:
280 return module.pull(data)
283def _pyclesperanto_from_numpy(data: Any, module: Any, device_id: int) -> Any:
284 del device_id
285 return module.push(data)
288def _pyclesperanto_stack(values: Sequence[Any], module: Any) -> Any:
289 if not values:
290 raise ValueError("Cannot stack an empty pyclesperanto sequence")
291 if len(values) == 1:
292 source = values[0]
293 result = module.create((1, *source.shape), dtype=source.dtype)
294 return module.copy_slice(source, result, 0)
295 result = values[0]
296 for value in values[1:]:
297 result = module.concatenate_along_z(result, value)
298 return result
301def _pyclesperanto_cast(data: Any, dtype: Any, module: Any) -> Any:
302 return module.push(module.pull(data).astype(dtype, copy=False))
305def _pyclesperanto_logical_and(left: Any, right: Any, module: Any) -> Any:
306 return module.push(np.logical_and(module.pull(left), module.pull(right)))
309def _pyclesperanto_scale(result: Any, target_dtype: Any, module: Any) -> Any:
310 if not hasattr(result, "dtype"):
311 return result
312 target_is_int = np.issubdtype(np.dtype(target_dtype), np.integer)
313 if not (np.issubdtype(result.dtype, np.floating) and target_is_int):
314 return module.push(module.pull(result).astype(target_dtype))
315 result_min = float(module.minimum_of_all_pixels(result))
316 result_max = float(module.maximum_of_all_pixels(result))
317 if result_max <= result_min:
318 return module.push(module.pull(result).astype(target_dtype))
319 normalized = module.subtract_image_from_scalar(result, scalar=result_min)
320 normalized = module.multiply_image_and_scalar(
321 normalized,
322 scalar=1.0 / (result_max - result_min),
323 )
324 range_info = _SCALING_RANGES.get(_dtype_name(target_dtype))
325 if isinstance(range_info, tuple):
326 scale, offset = range_info
327 scaled = module.multiply_image_and_scalar(normalized, scalar=scale)
328 scaled = module.subtract_image_from_scalar(scaled, scalar=offset)
329 elif range_info is not None:
330 scaled = module.multiply_image_and_scalar(normalized, scalar=range_info)
331 else:
332 scaled = normalized
333 host_values = module.pull(scaled)
334 bounds = _clamp_bounds(target_dtype)
335 if bounds is not None:
336 host_values = np.clip(host_values, *bounds)
337 return module.push(host_values.astype(target_dtype))
340@dataclass(frozen=True, slots=True)
341class ArrayOperations:
342 """Framework-specific array leaves referenced by one declaration."""
344 to_numpy: ToNumpy
345 from_numpy: FromNumpy
346 stack: StackArrays
347 scale_dtype: ScaleDtype
348 dtype_name: DtypeName = _numpy_dtype_name
349 cast: CastArray = _numpy_cast
350 logical_and: LogicalAnd = _module_logical_and
353NUMPY_OPERATIONS = ArrayOperations(
354 to_numpy=_identity_to_numpy,
355 from_numpy=_identity_from_numpy,
356 stack=_numpy_stack,
357 scale_dtype=_numpy_scale,
358)
359CUPY_OPERATIONS = ArrayOperations(
360 to_numpy=_cupy_to_numpy,
361 from_numpy=_cupy_from_numpy,
362 stack=_cupy_stack,
363 scale_dtype=_cupy_scale,
364)
365TORCH_OPERATIONS = ArrayOperations(
366 to_numpy=_torch_to_numpy,
367 from_numpy=_torch_from_numpy,
368 stack=_torch_stack,
369 scale_dtype=_torch_scale,
370 dtype_name=_torch_dtype_name,
371 cast=_torch_cast,
372)
373TENSORFLOW_OPERATIONS = ArrayOperations(
374 to_numpy=_tensorflow_to_numpy,
375 from_numpy=_tensorflow_from_numpy,
376 stack=_tensorflow_stack,
377 scale_dtype=_tensorflow_scale,
378 dtype_name=_tensorflow_dtype_name,
379 cast=_tensorflow_cast,
380)
381JAX_OPERATIONS = ArrayOperations(
382 to_numpy=_jax_to_numpy,
383 from_numpy=_jax_from_numpy,
384 stack=_jax_stack,
385 scale_dtype=_jax_scale,
386 cast=_jax_cast,
387 logical_and=_jax_logical_and,
388)
389PYCLESPERANTO_OPERATIONS = ArrayOperations(
390 to_numpy=_pyclesperanto_to_numpy,
391 from_numpy=_pyclesperanto_from_numpy,
392 stack=_pyclesperanto_stack,
393 scale_dtype=_pyclesperanto_scale,
394 cast=_pyclesperanto_cast,
395 logical_and=_pyclesperanto_logical_and,
396)