| # zhangqiaorjc@google.com | |
| # `jax.distributed.initialize` is available in jax-0.2.25. | |
| # $ pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html # Note: wheels only available on linux. | |
| # Run this script on 2 GPU nodes, assuming 10.128.0.6 is the master node | |
| # python nvidia_gpu_pjit.py --server_addr="10.128.0.6:1456" --num_hosts=2 --host_idx=0 | |
| # python nvidia_gpu_pjit.py --server_addr="10.128.0.6:1456" --num_hosts=2 --host_idx=1 | |
| from absl import app | |
| from absl import flags |
| import torch | |
| import torch.utils.dlpack | |
| import jax | |
| import jax.dlpack | |
| # A generic mechanism for turning a JAX function into a PyTorch function. | |
| def j2t(x_jax): | |
| x_torch = torch.utils.dlpack.from_dlpack(jax.dlpack.to_dlpack(x_jax)) | |
| return x_torch |
| #! /usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| # vim:fenc=utf-8 | |
| # | |
| # Copyright © 2019 Pi-Yueh Chuang <pychuang@gwu.edu> | |
| # | |
| # Distributed under terms of the MIT license. | |
| """An example of using tfp.optimizer.lbfgs_minimize to optimize a TensorFlow model. |
This is a step-by-step tutorial for setting up your own conda feedstock from scratch. I started with the instructions here:
As a novice, I found that the above tutorial lacked sufficient detail. So, I decided to document the process that I used for my own future reference.
| cimport numpy as np | |
| import numpy as np | |
| import cython | |
| from cython.parallel import prange | |
| @cython.boundscheck(False) | |
| @cython.wraparound(False) | |
| @cython.cdivision(True) | |
| def p_arange(int n): | |
| ''' |
This is a short post that explains how to write a high-performance matrix multiplication program on modern processors. In this tutorial I will use a single core of the Skylake-client CPU with AVX2, but the principles in this post also apply to other processors with different instruction sets (such as AVX512).
Matrix multiplication is a mathematical operation that defines the product of
| import numpy as np | |
| from sklearn.utils import shuffle | |
| # Data comes from y = f(x) = [2, 3].x + [5, 7] | |
| X0 = np.random.randn(100, 2) - 1 | |
| X1 = np.random.randn(100, 2) + 1 | |
| X = np.vstack([X0, X1]) | |
| t = np.vstack([np.zeros([100, 1]), np.ones([100, 1])]) |
| import tensorflow as tf | |
| from tensorflow.python.framework import function | |
| @function.Defun() | |
| def my_op_grad(op, grad): ### instead of my_op_grad(x) | |
| return tf.sigmoid(op) | |
| @function.Defun(grad_func=my_op_grad) | |
| def my_op(a): | |
| return tf.identity(a) |
| """Convert jupyter notebook to sphinx gallery notebook styled examples. | |
| Usage: python ipynb_to_gallery.py <notebook.ipynb> | |
| Dependencies: | |
| pypandoc: install using `pip install pypandoc` | |
| """ | |
| import pypandoc as pdoc | |
| import json |
