gufunc-numba-axis-examples.ipynb

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": []
    }
  ]
}