MCPcopy
hub / github.com/NVIDIA/warp / numpy

Method numpy

warp/_src/types.py:4129–4178  ·  view source on GitHub ↗

Convert the array to a :class:`numpy.ndarray` (aliasing memory through the array interface protocol) If the array is on the GPU, a synchronous device-to-host copy (on the CUDA default stream) will be automatically performed to ensure that any outstanding work is completed.

(self, *, _suppress_bfloat16_warning=False)

Source from the content-addressed store, hash-verified

4127 warp.copy(self, array(data=src, dtype=self.dtype, copy=False, device="cpu"))
4128
4129 def numpy(self, *, _suppress_bfloat16_warning=False):
4130 """Convert the array to a :class:`numpy.ndarray` (aliasing memory through the array interface protocol)
4131 If the array is on the GPU, a synchronous device-to-host copy (on the CUDA default stream) will be
4132 automatically performed to ensure that any outstanding work is completed.
4133 """
4134 # Determine if this array contains bfloat16 data (scalar or compound type)
4135 is_bf16 = self.dtype is bfloat16 or (
4136 isinstance(self.dtype, type)
4137 and issubclass(self.dtype, ctypes.Array)
4138 and self.dtype._wp_scalar_type_ is bfloat16
4139 )
4140
4141 if is_bf16 and not _suppress_bfloat16_warning:
4142 ml_bf16 = _get_ml_dtypes_bfloat16()
4143 if ml_bf16 is None:
4144 log_warning(
4145 "bfloat16 arrays are returned as np.uint16 (raw bit representation) "
4146 "because NumPy does not natively support bfloat16. "
4147 "Use wp.to_torch() or wp.to_jax() for frameworks that support bfloat16 natively, "
4148 "or install ml-dtypes for a NumPy bfloat16 dtype.",
4149 once=True,
4150 )
4151
4152 if self.ptr:
4153 # use the CUDA default stream for synchronous behaviour with other streams
4154 with warp.ScopedStream(self.device.null_stream):
4155 a = self.to("cpu", requires_grad=False)
4156 # convert through __array_interface__
4157 # Note: this handles arrays of structs using `descr`, so the result will be a structured NumPy array
4158 result = np.asarray(a)
4159 else:
4160 # return an empty numpy array with the correct dtype and shape
4161 if isinstance(self.dtype, warp._src.codegen.Struct):
4162 npdtype = self.dtype.numpy_dtype()
4163 npshape = self.shape
4164 elif issubclass(self.dtype, ctypes.Array):
4165 npdtype = warp_type_to_np_dtype[self.dtype._wp_scalar_type_]
4166 npshape = (*self.shape, *self.dtype._shape_)
4167 else:
4168 npdtype = warp_type_to_np_dtype[self.dtype]
4169 npshape = self.shape
4170 result = np.empty(npshape, dtype=npdtype)
4171
4172 # If ml_dtypes is available, view as bfloat16 instead of returning raw uint16
4173 if is_bf16:
4174 ml_bf16 = _get_ml_dtypes_bfloat16()
4175 if ml_bf16 is not None:
4176 result = result.view(ml_bf16)
4177
4178 return result
4179
4180 def cptr(self):
4181 """Return a ctypes cast of the array address.

Calls 5

toMethod · 0.95
log_warningFunction · 0.90
_get_ml_dtypes_bfloat16Function · 0.85
viewMethod · 0.80
numpy_dtypeMethod · 0.45