Created
October 7, 2019 02:49
-
-
Save fehiepsi/7f70265780aaabfa32c5819b479399fd to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
{ | |
"cells": [ | |
{ | |
"cell_type": "code", | |
"execution_count": 1, | |
"metadata": {}, | |
"outputs": [], | |
"source": [ | |
"from jax import jit" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 2, | |
"metadata": {}, | |
"outputs": [], | |
"source": [ | |
"@jit\n", | |
"def f(x, *args, **kwargs):\n", | |
" return x + args[0] + (list(kwargs.values())[0] if kwargs else 0)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 3, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"CPU times: user 136 ms, sys: 163 ms, total: 300 ms\n", | |
"Wall time: 456 ms\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/plain": [ | |
"DeviceArray(10, dtype=int32)" | |
] | |
}, | |
"execution_count": 3, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"%%time\n", | |
"f(1, 2, a=7)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 4, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"CPU times: user 1.59 ms, sys: 125 µs, total: 1.72 ms\n", | |
"Wall time: 887 µs\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/plain": [ | |
"DeviceArray(11, dtype=int32)" | |
] | |
}, | |
"execution_count": 4, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"%%time\n", | |
"f(1, 2, a=8)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 5, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"CPU times: user 8.01 ms, sys: 9.38 ms, total: 17.4 ms\n", | |
"Wall time: 65.3 ms\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/plain": [ | |
"DeviceArray(4, dtype=int32)" | |
] | |
}, | |
"execution_count": 5, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"%%time\n", | |
"f(1, 3)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 6, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"CPU times: user 1.27 ms, sys: 44 µs, total: 1.31 ms\n", | |
"Wall time: 824 µs\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/plain": [ | |
"DeviceArray(5, dtype=int32)" | |
] | |
}, | |
"execution_count": 6, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"%%time\n", | |
"f(1, 4)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 7, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"CPU times: user 10.6 ms, sys: 0 ns, total: 10.6 ms\n", | |
"Wall time: 9.59 ms\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/plain": [ | |
"DeviceArray(13, dtype=int32)" | |
] | |
}, | |
"execution_count": 7, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"%%time\n", | |
"f(1, 2, b=10)" | |
] | |
}, | |
{ | |
"cell_type": "code", | |
"execution_count": 8, | |
"metadata": {}, | |
"outputs": [ | |
{ | |
"name": "stdout", | |
"output_type": "stream", | |
"text": [ | |
"CPU times: user 1.09 ms, sys: 0 ns, total: 1.09 ms\n", | |
"Wall time: 561 µs\n" | |
] | |
}, | |
{ | |
"data": { | |
"text/plain": [ | |
"DeviceArray(103, dtype=int32)" | |
] | |
}, | |
"execution_count": 8, | |
"metadata": {}, | |
"output_type": "execute_result" | |
} | |
], | |
"source": [ | |
"%%time\n", | |
"f(1, 2, b=100)" | |
] | |
} | |
], | |
"metadata": { | |
"kernelspec": { | |
"display_name": "Python 3", | |
"language": "python", | |
"name": "python3" | |
}, | |
"language_info": { | |
"codemirror_mode": { | |
"name": "ipython", | |
"version": 3 | |
}, | |
"file_extension": ".py", | |
"mimetype": "text/x-python", | |
"name": "python", | |
"nbconvert_exporter": "python", | |
"pygments_lexer": "ipython3", | |
"version": "3.6.9" | |
} | |
}, | |
"nbformat": 4, | |
"nbformat_minor": 4 | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment