A reduction to count the number of elements. This has an additional kwarg in coerce_np_ndarray, which determines whether to ensure that the resulting array is a numpy.ndarray, or whether we allow it to be other array types via `np.full_like`.
(x, coerce_np_ndarray: bool, **kwargs)
| 276 | |
| 277 | |
| 278 | def _numel(x, coerce_np_ndarray: bool, **kwargs): |
| 279 | """ |
| 280 | A reduction to count the number of elements. |
| 281 | |
| 282 | This has an additional kwarg in coerce_np_ndarray, which determines |
| 283 | whether to ensure that the resulting array is a numpy.ndarray, or whether |
| 284 | we allow it to be other array types via `np.full_like`. |
| 285 | """ |
| 286 | shape = x.shape |
| 287 | keepdims = kwargs.get("keepdims", False) |
| 288 | axis = kwargs.get("axis") |
| 289 | dtype = kwargs.get("dtype", np.float64) |
| 290 | |
| 291 | if axis is None: |
| 292 | prod = np.prod(shape, dtype=dtype) |
| 293 | if keepdims is False: |
| 294 | return prod |
| 295 | |
| 296 | if coerce_np_ndarray: |
| 297 | return np.full(shape=(1,) * len(shape), fill_value=prod, dtype=dtype) |
| 298 | else: |
| 299 | return np.full_like(x, prod, shape=(1,) * len(shape), dtype=dtype) |
| 300 | |
| 301 | if not isinstance(axis, (tuple, list)): |
| 302 | axis = [axis] |
| 303 | |
| 304 | prod = math.prod(shape[dim] for dim in axis) |
| 305 | if keepdims is True: |
| 306 | new_shape = tuple( |
| 307 | shape[dim] if dim not in axis else 1 for dim in range(len(shape)) |
| 308 | ) |
| 309 | else: |
| 310 | new_shape = tuple(shape[dim] for dim in range(len(shape)) if dim not in axis) |
| 311 | |
| 312 | if coerce_np_ndarray: |
| 313 | return np.broadcast_to(np.array(prod, dtype=dtype), new_shape) |
| 314 | else: |
| 315 | return np.full_like(x, prod, shape=new_shape, dtype=dtype) |
| 316 | |
| 317 | |
| 318 | @nannumel_lookup.register((object, np.ndarray)) |
no test coverage detected