Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
gufunc-numba-axis-examples.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "gufunc-numba-axis-examples.ipynb",
"version": "0.3.2",
"provenance": [],
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"[View in Colaboratory](https://colab.research.google.com/gist/shoyer/cce8e042213313685000374200ae498e/gufunc-numba-axis-examples.ipynb)"
]
},
{
"metadata": {
"id": "5BFPnuQkJjZd",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 289
},
"outputId": "3e52be55-4a10-4629-b171-b2f00812498f"
},
"cell_type": "code",
"source": [
"! pip install -U numba numpy"
],
"execution_count": 1,
"outputs": [
{
"output_type": "stream",
"text": [
"Collecting numba\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/24/89/6f1755892d60ddd528090dc313349e7cc491170d6737f6b3a7a5b317ef81/numba-0.39.0-cp36-cp36m-manylinux1_x86_64.whl (1.9MB)\n",
"\u001b[K 100% |████████████████████████████████| 1.9MB 9.3MB/s \n",
"\u001b[?25hCollecting numpy\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/fe/94/7049fed8373c52839c8cde619acaf2c9b83082b935e5aa8c0fa27a4a8bcc/numpy-1.15.1-cp36-cp36m-manylinux1_x86_64.whl (13.9MB)\n",
"\u001b[K 100% |████████████████████████████████| 13.9MB 2.2MB/s \n",
"\u001b[?25hCollecting llvmlite>=0.24.0dev0 (from numba)\n",
"\u001b[?25l Downloading https://files.pythonhosted.org/packages/be/05/c1b933d6b3dd6a234b681605e4154c44d21ee22cbef315c4eb9d64e6ab6a/llvmlite-0.24.0-cp36-cp36m-manylinux1_x86_64.whl (15.8MB)\n",
"\u001b[K 100% |████████████████████████████████| 15.9MB 2.8MB/s \n",
"\u001b[?25hInstalling collected packages: numpy, llvmlite, numba\n",
" Found existing installation: numpy 1.14.5\n",
" Uninstalling numpy-1.14.5:\n",
" Successfully uninstalled numpy-1.14.5\n",
"Successfully installed llvmlite-0.24.0 numba-0.39.0 numpy-1.15.1\n"
],
"name": "stdout"
}
]
},
{
"metadata": {
"id": "GwfpM9WvKHc-",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 35
},
"outputId": "d86a20e8-6233-4c84-c4c2-a81a5a318fee"
},
"cell_type": "code",
"source": [
"float64[:]"
],
"execution_count": 3,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"array(float64, 1d, A)"
]
},
"metadata": {
"tags": []
},
"execution_count": 3
}
]
},
{
"metadata": {
"id": "NTsIntQFJlWK",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"from numba import guvectorize, float64\n",
"\n",
"@guvectorize([(float64[:], float64[:])], '(n)->(n)')\n",
"def f(x, res):\n",
" for i in range(x.shape[0]):\n",
" res[i] = x[i]\n",
"\n",
"@guvectorize([(float64[:], float64[:], float64[:])], '(n),(n)->(n)')\n",
"def g(x, y, res):\n",
" for i in range(x.shape[0]):\n",
" res[i] = x[i] + y[i]\n"
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "UUP00n_GKa9W",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"import numpy as np"
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "DYivaC45KALG",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 71
},
"outputId": "7135c06a-c74e-4d15-cfd0-14dc0fbacdfa"
},
"cell_type": "code",
"source": [
"f(np.zeros((3, 4)), axes=[0, 0])"
],
"execution_count": 10,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"array([[0., 0., 0., 0.],\n",
" [0., 0., 0., 0.],\n",
" [0., 0., 0., 0.]])"
]
},
"metadata": {
"tags": []
},
"execution_count": 10
}
]
},
{
"metadata": {
"id": "vJE9blh9LLbO",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
"a = np.random.randn(5, 1, 6)\n",
"b = np.random.randn(3, 5)\n"
],
"execution_count": 0,
"outputs": []
},
{
"metadata": {
"id": "KqVJy96QKaTu",
"colab_type": "code",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 192
},
"outputId": "f1ad5380-f0f7-4cbc-ee99-6e19cc482ece"
},
"cell_type": "code",
"source": [
"g(a, b, axes=[0, 2])"
],
"execution_count": 14,
"outputs": [
{
"output_type": "error",
"ename": "ValueError",
"evalue": "ignored",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mValueError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m<ipython-input-14-186f936f07ba>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;31mValueError\u001b[0m: axes should be a list with an entry for all 3 inputs and outputs; entries for outputs can only be omitted if none of them has core axes."
]
}
]
},
{
"metadata": {
"id": "pcOIVlBiLNWN",
"colab_type": "code",
"colab": {}
},
"cell_type": "code",
"source": [
""
],
"execution_count": 0,
"outputs": []
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.