Skip to content

Commit

Permalink
DOC: Add guvectorize scalar return example
Browse files Browse the repository at this point in the history
  • Loading branch information
Matt711 committed Nov 28, 2022
1 parent 7282635 commit b236c93
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 2 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
repos:
- repo: https://gitlab.com/pycqa/flake8
- repo: https://github.com/PyCQA/flake8
rev: 3.7.8
hooks:
- id: flake8
29 changes: 28 additions & 1 deletion docs/source/user/vectorize.rst
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,33 @@ complicated inputs, depending on their shapes::
Use it to ensure the generated code does not fallback to
:term:`object mode`.


.. _scalar-return-values:

Scalar return values
------------------------

Now suppose we want to return a scalar value from :func:`~numba.guvectorize`::

@guvectorize([(int64[:], int64, int64[:])], '(n),()->()')
def g(x, y, res):
acc=0
for i in range(x.shape[0]):
acc += x[i] + y
res[0]=acc

The modified function takes the sum of the 1-dimensional array (``x``) plus the scalar (``y``).
There are a few important changes that allow the wrapped function to return a scalar:

* in the signatures, declare the scalar return with ``[:]`` like a 1-dimensional array (eg. ``int64[:]``)
* in the layout, declare it as ``()``
* in the implementation, write the return value to output. (e.g. ``res[0]=acc``)

>>> a = np.arange(5)
>>> g(a,2)
20


.. _overwriting-input-values:

Overwriting input values
Expand Down Expand Up @@ -388,7 +415,7 @@ compilation works for a :class:`~numba.GUFunc`::
>>> res = np.zeros_like(x)
>>> g(x, y, res)
>>> res
array([5, 6, 7, 8, 9])
array([10, 11, 12, 13, 14])
>>> g.types
['ll->l']

Expand Down

0 comments on commit b236c93

Please sign in to comment.