fix static linkage and make THD statically linked

Previously, these Functions passed themselves, i.e. the saved_for from
ConvForward would be ConvForward.

.gitignore

@@ -5,6 +5,7 @@ torch.egg-info/
 torch/version.py
 torch/csrc/generic/TensorMethods.cpp
 torch/lib/*.so*
+torch/lib/*.a*
 torch/lib/*.dylib*
 torch/lib/*.h
 torch/lib/build
@@ -19,8 +20,10 @@ torch/csrc/nn/THCUNN.cpp
 torch/csrc/nn/THNN_generic.cwrap
 torch/csrc/nn/THNN_generic.cpp
 torch/csrc/nn/THNN_generic.h
+torch/csrc/generated
 docs/src/**/*
 test/data/legacy_modules.t7
+test/data/gpu_tensors.pt
 test/htmlcov
 test/.coverage
 */*.pyc
@@ -31,3 +34,18 @@ test/.coverage
 */*.so*
 */**/*.so*
 */**/*.dylib*
+test/data/legacy_serialized.pt
+test/data/linear.pt
+
+# IPython notebook checkpoints
+.ipynb_checkpoints
+
+# Editor temporaries
+*.swn
+*.swo
+*.swp
+*~
+
+# OSX dir files
+.DS_Store
+

.travis.yml

@@ -1,7 +1,8 @@
 # https://travis-ci.org/pytorch/pytorch
 language: python
+dist: trusty
 python:
-    - 2.7.8
+    - 2.7.9
     - 2.7
     - 3.5
     - 3.6
@@ -18,7 +19,8 @@ install:
     - export CC="ccache gcc-4.8"
     - export CXX="ccache g++-4.8"
     - ccache --show-stats
-    - travis_retry pip install -r requirements.txt
+    - travis_retry pip install --upgrade pip setuptools wheel
+    - travis_retry pip install -r requirements.txt --only-binary=scipy
     - python setup.py install
 
 script:
@@ -43,5 +45,5 @@ matrix:
         env: LINT_CHECK
         python: "2.7"
         addons: true
-        install: pip install pep8
-        script: pep8
+        install: pip install flake8
+        script: flake8

CONTRIBUTING.md

@@ -0,0 +1,185 @@
+## Contributing to PyTorch
+
+If you are interested in contributing to PyTorch, your contributions will fall
+into two categories:
+1. You want to propose a new Feature and implement it
+    - post about your intended feature, and we shall discuss the design and
+    implementation. Once we agree that the plan looks good, go ahead and implement it.
+2. You want to implement a feature or bug-fix for an outstanding issue
+    - Look at the outstanding issues here: https://github.com/pytorch/pytorch/issues
+    - Especially look at the Low Priority and Medium Priority issues
+    - Pick an issue and comment on the task that you want to work on this feature
+    - If you need more context on a particular issue, please ask and we shall provide.
+
+Once you finish implementing a feature or bugfix, please send a Pull Request to
+https://github.com/pytorch/pytorch
+
+If you are not familiar with creating a Pull Request, here are some guides:
+- http://stackoverflow.com/questions/14680711/how-to-do-a-github-pull-request
+- https://help.github.com/articles/creating-a-pull-request/
+
+
+## Developing locally with PyTorch
+
+To locally develop with PyTorch, here are some tips:
+
+1. Uninstall all existing pytorch installs
+```
+conda uninstall pytorch
+pip uninstall torch
+pip uninstall torch # run this command twice
+```
+
+2. Locally clone a copy of PyTorch from source:
+
+```
+git clone https://github.com/pytorch/pytorch
+cd pytorch
+```
+
+3. Install PyTorch in `build develop` mode:
+
+A full set of instructions on installing PyTorch from Source are here:
+https://github.com/pytorch/pytorch#from-source
+
+The change you have to make is to replace
+
+```
+python setup.py install
+```
+
+with
+
+```
+python setup.py build develop
+```
+
+This is especially useful if you are only changing Python files.
+
+This mode will symlink the python files from the current local source tree into the
+python install.
+
+Hence, if you modify a python file, you do not need to reinstall pytorch again and again.
+
+For example:
+- Install local pytorch in `build develop` mode
+- modify your python file `torch/__init__.py` (for example)
+- test functionality
+- modify your python file `torch/__init__.py`
+- test functionality
+- modify your python file `torch/__init__.py`
+- test functionality
+
+You do not need to repeatedly install after modifying python files.
+
+
+## Writing documentation
+
+PyTorch uses [Google style](http://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)
+for formatting docstrings. Length of line inside docstrings block must be limited to 80 characters to
+fit into Jupyter documentation popups.
+
+
+## Managing multiple build trees
+
+One downside to using `python setup.py develop` is that your development
+version of pytorch will be installed globally on your account (e.g., if
+you run `import torch` anywhere else, the development version will be
+used.
+
+If you want to manage multiple builds of PyTorch, you can make use of
+[conda environments](https://conda.io/docs/using/envs.html) to maintain
+separate Python package environments, each of which can be tied to a
+specific build of PyTorch.  To set one up:
+
+```
+conda create -n pytorch-myfeature
+source activate pytorch-myfeature
+# if you run python now, torch will NOT be installed
+python setup.py build develop
+```
+
+## C++ Development tips
+
+If you are working on the C++ code, there are a few important things that you
+will want to keep in mind:
+
+1. How to rebuild only the code you are working on, and
+2. How to make rebuilds in the absence of changes go faster.
+
+### Build only what you need.
+
+`python setup.py build` will build everything, but since our build system is
+not very optimized for incremental rebuilds, this will actually be very slow.
+Far better is to only request rebuilds of the parts of the project you are
+working on:
+
+- Working on `torch/csrc`?  Run `python setup.py develop` to rebuild
+  (NB: no `build` here!)
+
+- Working on `torch/lib/TH`, did not make any cmake changes, and just want to
+  see if it compiles?  Run `(cd torch/lib/build/TH && make install -j$(getconf _NPROCESSORS_ONLN))`.  This
+  applies for any other subdirectory of `torch/lib`.  **Warning: Changes you
+  make here will not be visible from Python.**  See below.
+
+- Working on `torch/lib` and want to run your changes / rerun cmake?  Run
+  `python setup.py build_deps`.  Note that this will rerun cmake for
+  every subdirectory in TH; if you are only working on one project,
+  consider editing `torch/lib/build_all.sh` and commenting out the
+  `build` lines of libraries you are not working on.
+
+On the initial build, you can also speed things up with the environment
+variables `DEBUG` and `NO_CUDA`.
+
+- `DEBUG=1` will enable debug builds (-g -O0)
+- `NO_CUDA=1` will disable compiling CUDA (in case you are developing on something not CUDA related), to save compile time.
+
+For example:
+```
+NO_CUDA=1 DEBUG=1 python setup.py build develop
+```
+
+Make sure you continue to pass these flags on subsequent builds.
+
+### Make no-op build fast.
+
+Python `setuptools` is pretty dumb, and always rebuilds every C file in a
+project. Using ccache in a situation like this is a real time-saver. However, by
+default, ccache does not properly support CUDA stuff, so here are the
+instructions for installing a custom `ccache` fork that has CUDA support:
+
+```
+# install and export ccache
+if ! ls ~/ccache/bin/ccache
+then
+    sudo apt-get update
+    sudo apt-get install -y automake autoconf
+    sudo apt-get install -y asciidoc
+    mkdir -p ~/ccache
+    pushd /tmp
+    rm -rf ccache
+    git clone https://github.com/colesbury/ccache -b ccbin
+    pushd ccache
+    ./autogen.sh
+    ./configure
+    make install prefix=~/ccache
+    popd
+    popd
+
+    mkdir -p ~/ccache/lib
+    mkdir -p ~/ccache/cuda
+    ln -s ~/ccache/bin/ccache ~/ccache/lib/cc
+    ln -s ~/ccache/bin/ccache ~/ccache/lib/c++
+    ln -s ~/ccache/bin/ccache ~/ccache/lib/gcc
+    ln -s ~/ccache/bin/ccache ~/ccache/lib/g++
+    ln -s ~/ccache/bin/ccache ~/ccache/cuda/nvcc
+
+    ~/ccache/bin/ccache -M 25Gi
+fi
+
+export PATH=~/ccache/lib:$PATH
+export CUDA_NVCC_EXECUTABLE=~/ccache/cuda/nvcc
+```
+
+
+Hope this helps, and thanks for considering to contribute.

Dockerfile

@@ -1,10 +1,13 @@
-FROM nvidia/cuda:8.0-cudnn5-devel-ubuntu14.04 
+FROM nvidia/cuda:8.0-cudnn6-devel-ubuntu16.04 
+
+RUN echo "deb http://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64 /" > /etc/apt/sources.list.d/nvidia-ml.list
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
          build-essential \
          cmake \
          git \
          curl \
+         vim \
          ca-certificates \
          libjpeg-dev \
          libpng-dev &&\
@@ -15,7 +18,7 @@ RUN curl -o ~/miniconda.sh -O  https://repo.continuum.io/miniconda/Miniconda3-4.
      ~/miniconda.sh -b -p /opt/conda && \     
      rm ~/miniconda.sh && \
      /opt/conda/bin/conda install conda-build && \
-     /opt/conda/bin/conda create -y --name pytorch-py35 python=3.5.2 numpy scipy ipython mkl&& \
+     /opt/conda/bin/conda create -y --name pytorch-py35 python=3.5.2 numpy pyyaml scipy ipython mkl&& \
      /opt/conda/bin/conda clean -ya 
 ENV PATH /opt/conda/envs/pytorch-py35/bin:$PATH
 RUN conda install --name pytorch-py35 -c soumith magma-cuda80
@@ -23,11 +26,11 @@ RUN conda install --name pytorch-py35 -c soumith magma-cuda80
 WORKDIR /opt/pytorch
 COPY . .
 
-RUN cat requirements.txt | xargs -n1 pip install --no-cache-dir && \
-    TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1+PTX" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
-    CMAKE_LIBRARY_PATH=/opt/conda/envs/pytorch-py35/lib \
-    CMAKE_INCLUDE_PATH=/opt/conda/envs/pytorch-py35/include \
+RUN TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1+PTX" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
+    CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
     pip install -v .
 
+RUN git clone https://github.com/pytorch/vision.git && cd vision && pip install -v .
+
 WORKDIR /workspace
 RUN chmod -R a+w /workspace

README.md

@@ -2,54 +2,72 @@
 
 --------------------------------------------------------------------------------
 
-PyTorch is a python package that provides two high-level features:
-- Tensor computation (like numpy) with strong GPU acceleration
-- Deep Neural Networks built on a tape-based autograd system
+PyTorch is a Python package that provides two high-level features:
+- Tensor computation (like NumPy) with strong GPU acceleration
+- Deep neural networks built on a tape-based autograd system
 
-You can reuse your favorite python packages such as numpy, scipy and Cython to extend PyTorch when needed.
+You can reuse your favorite Python packages such as NumPy, SciPy and Cython to extend PyTorch when needed.
 
-We are in an early-release Beta. Expect some adventures and rough edges.
+We are in an early-release beta. Expect some adventures and rough edges.
 
-- [More About PyTorch](#more-about-pytorch)
+- [More about PyTorch](#more-about-pytorch)
 - [Installation](#installation)
   - [Binaries](#binaries)
-  - [From source](#from-source)
-  - [Docker image](#docker-image)
+  - [From Source](#from-source)
+  - [Docker Image](#docker-image)
 - [Getting Started](#getting-started)
 - [Communication](#communication)
 - [Releases and Contributing](#releases-and-contributing)
 - [The Team](#the-team)
 
-| System | Python | Status |
+| System | 2.7 | 3.5 |
 | --- | --- | --- |
-| Linux CPU | 2.7.8, 2.7, 3.5, nightly | [![Build Status](https://travis-ci.org/pytorch/pytorch.svg?branch=master)](https://travis-ci.org/pytorch/pytorch) |
-| Linux GPU | 2.7 | [![Build Status](http://build.pytorch.org:8080/buildStatus/icon?job=pytorch-master-py2)](https://build.pytorch.org/job/pytorch-master-py2) |
-| Linux GPU | 3.5 | [![Build Status](http://build.pytorch.org:8080/buildStatus/icon?job=pytorch-master-py3)](https://build.pytorch.org/job/pytorch-master-py3) |
+| Linux CPU | [![Build Status](https://travis-ci.org/pytorch/pytorch.svg?branch=master)](https://travis-ci.org/pytorch/pytorch) | [![Build Status](https://travis-ci.org/pytorch/pytorch.svg?branch=master)](https://travis-ci.org/pytorch/pytorch) |
+| Linux GPU | [![Build Status](http://build.pytorch.org:8080/buildStatus/icon?job=pytorch-master-py2-linux)](https://build.pytorch.org/job/pytorch-master-py2-linux) | [![Build Status](http://build.pytorch.org:8080/buildStatus/icon?job=pytorch-master-py3-linux)](https://build.pytorch.org/job/pytorch-master-py3-linux) |
+| macOS CPU | [![Build Status](http://build.pytorch.org:8080/buildStatus/icon?job=pytorch-master-py2-osx-cpu)](https://build.pytorch.org/job/pytorch-master-py2-osx-cpu) | [![Build Status](http://build.pytorch.org:8080/buildStatus/icon?job=pytorch-master-py3-osx-cpu)](https://build.pytorch.org/job/pytorch-master-py3-osx-cpu) |
+
 
 ## More about PyTorch
 
 At a granular level, PyTorch is a library that consists of the following components:
 
-| \_                       | \_ |
-| ------------------------ | --- |
-| torch                    | a Tensor library like NumPy, with strong GPU support |
-| torch.autograd           | a tape based automatic differentiation library that supports all differentiable Tensor operations in torch |
-| torch.nn                 | a neural networks library deeply integrated with autograd designed for maximum flexibility |
-| torch.optim              | an optimization package to be used with torch.nn with standard optimization methods such as SGD, RMSProp, LBFGS, Adam etc. |
-| torch.multiprocessing    | python multiprocessing, but with magical memory sharing of torch Tensors across processes. Useful for data loading and hogwild training. |
-| torch.utils              | DataLoader, Trainer and other utility functions for convenience |
-| torch.legacy(.nn/.optim) | legacy code that has been ported over from torch for backward compatibility reasons |
+<table>
+<tr>
+    <td><b> torch </b></td>
+    <td> a Tensor library like NumPy, with strong GPU support </td>
+</tr>
+<tr>
+    <td><b> torch.autograd </b></td>
+    <td> a tape-based automatic differentiation library that supports all differentiable Tensor operations in torch </td>
+</tr>
+<tr>
+    <td><b> torch.nn </b></td>
+    <td> a neural networks library deeply integrated with autograd designed for maximum flexibility </td>
+</tr>
+<tr>
+    <td><b> torch.multiprocessing  </b></td>
+    <td> Python multiprocessing, but with magical memory sharing of torch Tensors across processes. Useful for data loading and Hogwild training. </td>
+</tr>
+<tr>
+    <td><b> torch.utils </b></td>
+    <td> DataLoader, Trainer and other utility functions for convenience </td>
+</tr>
+<tr>
+    <td><b> torch.legacy(.nn/.optim) </b></td>
+    <td> legacy code that has been ported over from torch for backward compatibility reasons </td>
+</tr>
+</table>
 
 Usually one uses PyTorch either as:
 
-- A replacement for numpy to use the power of GPUs.
+- a replacement for NumPy to use the power of GPUs.
 - a deep learning research platform that provides maximum flexibility and speed
 
 Elaborating further:
 
-### A GPU-ready Tensor library
+### A GPU-Ready Tensor Library
 
-If you use numpy, then you have used Tensors (a.k.a ndarray).
+If you use NumPy, then you have used Tensors (a.k.a ndarray).
 
 <p align=center><img width="30%" src="docs/source/_static/img/tensor_illustration.png" /></p>
 
@@ -60,15 +78,15 @@ We provide a wide variety of tensor routines to accelerate and fit your scientif
 such as slicing, indexing, math operations, linear algebra, reductions.
 And they are fast!
 
-### Dynamic Neural Networks: Tape based Autograd
+### Dynamic Neural Networks: Tape-Based Autograd
 
 PyTorch has a unique way of building neural networks: using and replaying a tape recorder.
 
-Most frameworks such as `TensorFlow`, `Theano`, `Caffe` and `CNTK` have a static view of the world.
+Most frameworks such as TensorFlow, Theano, Caffe and CNTK have a static view of the world.
 One has to build a neural network, and reuse the same structure again and again.
 Changing the way the network behaves means that one has to start from scratch.
 
-With PyTorch, we use a technique called Reverse-mode auto-differentiation, which allows you to
+With PyTorch, we use a technique called reverse-mode auto-differentiation, which allows you to
 change the way your network behaves arbitrarily with zero lag or overhead. Our inspiration comes
 from several research papers on this topic, as well as current and past work such as
 [autograd](https://github.com/twitter/torch-autograd),
@@ -80,69 +98,68 @@ You get the best of speed and flexibility for your crazy research.
 
 <p align=center><img width="80%" src="docs/source/_static/img/dynamic_graph.gif" /></p>
 
-### Python first
+### Python First
 
-PyTorch is not a Python binding into a monolothic C++ framework.
+PyTorch is not a Python binding into a monolithic C++ framework.
 It is built to be deeply integrated into Python.
-You can use it naturally like you would use numpy / scipy / scikit-learn etc.
+You can use it naturally like you would use NumPy / SciPy / scikit-learn etc.
 You can write your new neural network layers in Python itself, using your favorite libraries
 and use packages such as Cython and Numba.
 Our goal is to not reinvent the wheel where appropriate.
 
-### Imperative experiences
+### Imperative Experiences
 
 PyTorch is designed to be intuitive, linear in thought and easy to use.
 When you execute a line of code, it gets executed. There isn't an asynchronous view of the world.
-When you drop into a debugger, or receive error messages and stack traces, understanding them is straight-forward.
-The stack-trace points to exactly where your code was defined.
+When you drop into a debugger, or receive error messages and stack traces, understanding them is straightforward.
+The stack trace points to exactly where your code was defined.
 We hope you never spend hours debugging your code because of bad stack traces or asynchronous and opaque execution engines.
 
 ### Fast and Lean
 
-PyTorch has minimal framework overhead. We integrate acceleration libraries 
-such as Intel MKL and NVIDIA (CuDNN, NCCL) to maximize speed. 
-At the core, its CPU and GPU Tensor and Neural Network backends 
+PyTorch has minimal framework overhead. We integrate acceleration libraries
+such as Intel MKL and NVIDIA (cuDNN, NCCL) to maximize speed.
+At the core, its CPU and GPU Tensor and neural network backends
 (TH, THC, THNN, THCUNN) are written as independent libraries with a C99 API.  
 They are mature and have been tested for years.
 
-Hence, PyTorch is quite fast -- whether you run small or large neural networks.
+Hence, PyTorch is quite fast – whether you run small or large neural networks.
 
 The memory usage in PyTorch is extremely efficient compared to Torch or some of the alternatives.
 We've written custom memory allocators for the GPU to make sure that
 your deep learning models are maximally memory efficient.
 This enables you to train bigger deep learning models than before.
 
-### Extensions without pain
+### Extensions without Pain
 
-Writing new neural network modules, or interfacing with PyTorch's Tensor API was designed to be straight-forward
+Writing new neural network modules, or interfacing with PyTorch's Tensor API was designed to be straightforward
 and with minimal abstractions.
 
 You can write new neural network layers in Python using the torch API
-[or your favorite numpy based libraries such as SciPy](https://github.com/pytorch/tutorials/blob/master/Creating%20extensions%20using%20numpy%20and%20scipy.ipynb).
+[or your favorite NumPy-based libraries such as SciPy](http://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html).
 
 If you want to write your layers in C/C++, we provide an extension API based on
-[cffi](http://cffi.readthedocs.io/en/latest/) that is efficient and with minimal boilerplate.  
-There is no wrapper code that needs to be written. [You can see an example here](https://github.com/pytorch/extension-ffi).
+[cffi](http://cffi.readthedocs.io/en/latest/) that is efficient and with minimal boilerplate.
+There is no wrapper code that needs to be written. You can see [a tutorial here](http://pytorch.org/tutorials/advanced/c_extension.html) and [an example here](https://github.com/pytorch/extension-ffi).
 
 
 ## Installation
 
 ### Binaries
-- Anaconda
-```bash
-conda install pytorch torchvision -c soumith
-```
+Commands to install from binaries via Conda or pip wheels are on our website:
 
-### From source
+[http://pytorch.org](http://pytorch.org)
+
+### From Source
 
 If you are installing from source, we highly recommend installing an [Anaconda](https://www.continuum.io/downloads) environment.
 You will get a high-quality BLAS library (MKL) and you get a controlled compiler version regardless of your Linux distro.
 
-Once you have [anaconda](https://www.continuum.io/downloads) installed, here are the instructions.
+Once you have [Anaconda](https://www.continuum.io/downloads) installed, here are the instructions.
 
 If you want to compile with CUDA support, install
 - [NVIDIA CUDA](https://developer.nvidia.com/cuda-downloads) 7.5 or above
-- [NVIDIA CuDNN](https://developer.nvidia.com/cudnn) v5.x
+- [NVIDIA cuDNN](https://developer.nvidia.com/cudnn) v5.x or above
 
 If you want to disable CUDA support, export environment variable `NO_CUDA=1`.
 
@@ -150,64 +167,70 @@ If you want to disable CUDA support, export environment variable `NO_CUDA=1`.
 
 On Linux
 ```bash
-export CMAKE_PREFIX_PATH=[anaconda root directory]
+export CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" # [anaconda root directory]
 
 # Install basic dependencies
-conda install numpy mkl setuptools cmake gcc cffi
+conda install numpy pyyaml mkl setuptools cmake gcc cffi
 
 # Add LAPACK support for the GPU
-conda install -c soumith magma-cuda75 # or magma-cuda80 if CUDA 8.0
+conda install -c soumith magma-cuda80 # or magma-cuda75 if CUDA 7.5
 ```
 
 On OSX
 ```bash
 export CMAKE_PREFIX_PATH=[anaconda root directory]
-conda install numpy setuptools cmake cffi
+conda install numpy pyyaml setuptools cmake cffi
 ```
 
 #### Install PyTorch
+On Linux
 ```bash
-export MACOSX_DEPLOYMENT_TARGET=10.9 # if OSX
-pip install -r requirements.txt
 python setup.py install
 ```
 
+On OSX
+```bash
+MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ python setup.py install
+```
+
 ### Docker image
 
-Dockerfiles are supplied to build images with cuda support and cudnn v5 and cudnn v6 RC. Build them as usual
+Dockerfile is supplied to build images with cuda support and cudnn v6. Build as usual
 ```
-docker build . -t pytorch-cudnnv5 
+docker build -t pytorch .
 ```
-or 
+Alternatively, if you want a runtime image, build with
+
 ```
-docker build . -t pytorch-cudnnv6 -f tools/docker/Dockerfile-v6
+docker build -t pytorch . -f tools/docker/Dockerfile_runtime
+
 ```
-and run them with nvidia-docker:
+and run with nvidia-docker:
 ```
-nvidia-docker run --rm -ti --ipc=host pytorch-cudnnv5
+nvidia-docker run --rm -ti --ipc=host pytorch
 ```
-Please note that pytorch uses shared memory to share data between processes, so if torch multiprocessing is used (e.g.
+Please note that PyTorch uses shared memory to share data between processes, so if torch multiprocessing is used (e.g.
 for multithreaded data loaders) the default shared memory segment size that container runs with is not enough, and you
-should increase shared memory size either with --ipc=host or --shm-size command line options to nvidia-docker run. 
+should increase shared memory size either with `--ipc=host` or `--shm-size` command line options to `nvidia-docker run`.
 
 
 ## Getting Started
 
 Three pointers to get you started:
-- [Tutorials: notebooks to get you started with understanding and using PyTorch](https://github.com/pytorch/tutorials)
+- [Tutorials: get you started with understanding and using PyTorch](http://pytorch.org/tutorials/)
 - [Examples: easy to understand pytorch code across all domains](https://github.com/pytorch/examples)
 - The API Reference: [http://pytorch.org/docs/](http://pytorch.org/docs/)
 
 ## Communication
 * forums: discuss implementations, research, etc. http://discuss.pytorch.org
-* github issues: bug reports, feature requests, install issues, RFCs, thoughts, etc.
-* slack: general chat, online discussions, collaboration etc. https://pytorch.slack.com/ . If you need a slack invite, ping us at soumith@pytorch.org
+* GitHub issues: bug reports, feature requests, install issues, RFCs, thoughts, etc.
+* Slack: general chat, online discussions, collaboration etc. https://pytorch.slack.com/ . If you need a slack invite, ping us at soumith@pytorch.org
 * newsletter: no-noise, one-way email newsletter with important announcements about pytorch. You can sign-up here: http://eepurl.com/cbG0rv
 
 ## Releases and Contributing
 
-PyTorch has a 90 day release cycle (major releases). 
-It's current state is Beta (v0.1.6), we expect no obvious bugs. Please let us know if you encounter a bug by [filing an issue](https://github.com/pytorch/pytorch/issues).
+PyTorch has a 90 day release cycle (major releases).
+It's current state is Beta, we expect no obvious bugs. Please let us know if you encounter a bug by [filing an issue](https://github.com/pytorch/pytorch/issues).
 
 We appreciate all contributions. If you are planning to contribute back bug-fixes, please do so without any further discussion.
 

cmake/FindCUDA/FindCUDA/select_compute_arch.cmake

@@ -63,11 +63,16 @@ function(CUDA_DETECT_INSTALLED_GPUS OUT_VARIABLE)
       "}\n")
 
     execute_process(COMMAND "${CUDA_NVCC_EXECUTABLE}" "--run" "${cufile}"
+                    "-ccbin" ${CMAKE_CXX_COMPILER}
                     WORKING_DIRECTORY "${PROJECT_BINARY_DIR}/CMakeFiles/"
                     RESULT_VARIABLE nvcc_res OUTPUT_VARIABLE nvcc_out
                     ERROR_QUIET OUTPUT_STRIP_TRAILING_WHITESPACE)
 
     if(nvcc_res EQUAL 0)
+      # only keep the last line of nvcc_out
+      STRING(REGEX REPLACE ";" "\\\\;" nvcc_out "${nvcc_out}")
+      STRING(REGEX REPLACE "\n" ";" nvcc_out "${nvcc_out}")
+      list(GET nvcc_out -1 nvcc_out)
       string(REPLACE "2.1" "2.1(2.0)" nvcc_out "${nvcc_out}")
       set(CUDA_GPU_DETECT_OUTPUT ${nvcc_out} CACHE INTERNAL "Returned GPU architetures from detect_gpus tool" FORCE)
     endif()
@@ -116,13 +121,13 @@ function(CUDA_SELECT_NVCC_ARCH_FLAGS out_variable)
       set(add_ptx TRUE)
       set(arch_name ${CMAKE_MATCH_1})
     endif()
-    if(arch_name MATCHES "([0-9]\\.[0-9])$")
+    if(arch_name MATCHES "(^[0-9]\\.[0-9](\\([0-9]\\.[0-9]\\))?)$")
       set(arch_bin ${CMAKE_MATCH_1})
       set(arch_ptx ${arch_bin})
     else()
       # Look for it in our list of known architectures
       if(${arch_name} STREQUAL "Fermi")
-        set(arch_bin 2.0 "2.1(2.0)")
+        set(arch_bin "2.0 2.1(2.0)")
       elseif(${arch_name} STREQUAL "Kepler+Tegra")
         set(arch_bin 3.2)
       elseif(${arch_name} STREQUAL "Kepler+Tesla")
@@ -173,11 +178,11 @@ function(CUDA_SELECT_NVCC_ARCH_FLAGS out_variable)
   # Tell NVCC to add binaries for the specified GPUs
   foreach(arch ${cuda_arch_bin})
     if(arch MATCHES "([0-9]+)\\(([0-9]+)\\)")
-      # User explicitly specified PTX for the concrete BIN
+      # User explicitly specified ARCH for the concrete CODE
       list(APPEND nvcc_flags -gencode arch=compute_${CMAKE_MATCH_2},code=sm_${CMAKE_MATCH_1})
       list(APPEND nvcc_archs_readable sm_${CMAKE_MATCH_1})
     else()
-      # User didn't explicitly specify PTX for the concrete BIN, we assume PTX=BIN
+      # User didn't explicitly specify ARCH for the concrete CODE, we assume ARCH=CODE
       list(APPEND nvcc_flags -gencode arch=compute_${arch},code=sm_${arch})
       list(APPEND nvcc_archs_readable sm_${arch})
     endif()

docs/Makefile

@@ -12,7 +12,14 @@ BUILDDIR      = build
 help:
 	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
-.PHONY: help Makefile
+docset: html
+	doc2dash --name $(SPHINXPROJ) --icon $(SOURCEDIR)/_static/img/pytorch-logo-flame.png --enable-js --online-redirect-url http://pytorch.org/docs/ --force $(BUILDDIR)/html/
+
+	# Manually fix because Zeal doesn't deal well with `icon.png`-only at 2x resolution.
+	cp $(SPHINXPROJ).docset/icon.png $(SPHINXPROJ).docset/icon@2x.png
+	convert $(SPHINXPROJ).docset/icon@2x.png -resize 16x16 $(SPHINXPROJ).docset/icon.png
+
+.PHONY: help Makefile docset
 
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).

docs/source/_static/css/pytorch_theme.css

@@ -112,3 +112,7 @@ footer p {
 nav .hidden-section {
     display: inherit;
 }
+
+.wy-side-nav-search>div.version {
+    color: #000;
+}

docs/source/_static/img/pytorch-logo-flame.svg

@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<svg
+   xmlns:dc="http://purl.org/dc/elements/1.1/"
+   xmlns:cc="http://creativecommons.org/ns#"
+   xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
+   xmlns:svg="http://www.w3.org/2000/svg"
+   xmlns="http://www.w3.org/2000/svg"
+   height="40.200001"
+   width="40.200001"
+   xml:space="preserve"
+   viewBox="0 0 40.200002 40.2"
+   y="0px"
+   x="0px"
+   id="Layer_1"
+   version="1.1"><metadata
+     id="metadata4717"><rdf:RDF><cc:Work
+         rdf:about=""><dc:format>image/svg+xml</dc:format><dc:type
+           rdf:resource="http://purl.org/dc/dcmitype/StillImage" /><dc:title></dc:title></cc:Work></rdf:RDF></metadata><defs
+     id="defs4715" /><style
+     id="style4694"
+     type="text/css">
+	.st0{fill:#F05732;}
+	.st1{fill:#9E529F;}
+	.st2{fill:#333333;}
+</style><path
+     style="fill:#f05732"
+     id="path4696"
+     d="m 26.975479,12.199999 c -1.3,-1 -1.8,3.9 -4.4,3.9 -3,0 -4,-12.9999998 -6.3,-12.9999998 -0.7,0 -0.8,-0.4 -7.9000003,21.2999998 -2.9000001,9 4.4000003,15.8 11.8000003,15.8 4.6,0 12.3,-3 12.3,-12.6 0,-7.1 -3.5,-13.9 -5.5,-15.4 z m -6.9,23.1 c -3.7,0 -6.7,-3.1 -6.7,-7 0,-3.9 3,-7 6.7,-7 3.7,0 6.7,3.1 6.7,7 0,3.8 -3,7 -6.7,7 z"
+     class="st0" /><path
+     style="fill:#9e529f"
+     id="path4698"
+     d="m 24.075479,-7.6293945e-7 c -0.5,0 -1.8,2.49999996293945 -1.8,3.59999996293945 0,1.5 1,2 1.8,2 0.8,0 1.8,-0.5 1.8,-2 -0.1,-1.1 -1.4,-3.59999996293945 -1.8,-3.59999996293945 z"
+     class="st1" /></svg>

docs/source/autograd.rst

@@ -9,6 +9,8 @@ Automatic differentiation package - torch.autograd
 
 .. autofunction:: backward
 
+.. autofunction:: grad
+
 Variable
 --------
 
@@ -20,7 +22,7 @@ of a couple in-place methods, that would overwrite inputs required for
 gradient computation). In most cases Tensors can be safely replaced with
 Variables and the code will remain to work just fine. Because of this,
 we're not documenting all the operations on variables, and you should
-refere to :class:`torch.Tensor` docs for this purpose.
+refer to :class:`torch.Tensor` docs for this purpose.
 
 In-place operations on Variables
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -38,8 +40,8 @@ All :class:`Variable` s keep track of in-place operations applied to them, and
 if the implementation detects that a variable was saved for backward in one of
 the functions, but it was modified in-place afterwards, an error will be raised
 once backward pass is started. This ensures that if you're using in-place
-functions and not seing any errors, you can be sure that the computed gradients
-are correct.
+functions and not seeing any errors, you can be sure that the computed
+gradients are correct.
 
 
 .. autoclass:: Variable

docs/source/conf.py

@@ -74,9 +74,11 @@ author = 'Torch Contributors'
 # built documents.
 #
 # The short X.Y version.
-version = '0.1.6'
+# TODO: change to [:2] at v1.0
+version = 'master (' + torch.__version__ + ' )'
 # The full version, including alpha/beta/rc tags.
-release = '0.1.6'
+# TODO: verify this works as expected
+release = 'master'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
@@ -111,7 +113,7 @@ html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
 #
 html_theme_options = {
     'collapse_navigation': False,
-    'display_version': False,
+    'display_version': True,
     'logo_only': True,
 }
 
@@ -202,7 +204,10 @@ from sphinx.util.docfields import TypedField
 from sphinx import addnodes
 
 
-def patched_make_field(self, types, domain, items):
+def patched_make_field(self, types, domain, items, **kw):
+    # `kw` catches `env=None` needed for newer sphinx while maintaining
+    #  backwards compatibility when passed along further down!
+
     # type: (List, unicode, Tuple) -> nodes.field
     def handle_item(fieldarg, content):
         par = nodes.paragraph()
@@ -222,7 +227,7 @@ def patched_make_field(self, types, domain, items):
                 typename = typename.replace('float', 'python:float')
                 typename = typename.replace('type', 'python:type')
                 par.extend(self.make_xrefs(self.typerolename, domain, typename,
-                                           addnodes.literal_emphasis))
+                                           addnodes.literal_emphasis, **kw))
             else:
                 par += fieldtype
             par += nodes.Text(')')

docs/source/cuda.rst

@@ -25,3 +25,10 @@ Streams and events
 
 .. autoclass:: Event
    :members:
+
+NVIDIA Tools Extension (NVTX)
+-----------------------------
+
+.. autofunction:: torch.cuda.nvtx.mark
+.. autofunction:: torch.cuda.nvtx.range_push
+.. autofunction:: torch.cuda.nvtx.range_pop

docs/source/data.rst

@@ -5,3 +5,9 @@ torch.utils.data
 .. autoclass:: Dataset
 .. autoclass:: TensorDataset
 .. autoclass:: DataLoader
+.. autoclass:: torch.utils.data.sampler.Sampler
+.. autoclass:: torch.utils.data.sampler.SequentialSampler
+.. autoclass:: torch.utils.data.sampler.RandomSampler
+.. autoclass:: torch.utils.data.sampler.SubsetRandomSampler
+.. autoclass:: torch.utils.data.sampler.WeightedRandomSampler
+.. autoclass:: torch.utils.data.distributed.DistributedSampler

docs/source/distributed.rst

@@ -0,0 +1,165 @@
+.. role:: hidden
+    :class: hidden-section
+
+Distributed communication package - torch.distributed
+=====================================================
+
+.. automodule:: torch.distributed
+.. currentmodule:: torch.distributed
+
+Currently torch.distributed supports three backends, each with
+different capabilities. The table below shows which functions are available
+for use with CPU / CUDA tensors.
+MPI supports cuda only iff the implementation used to build PyTorch supports it.
+
++------------+-----------+-----------+-----------+
+| Backend    | ``tcp``   | ``gloo``  | ``mpi``   |
++------------+-----+-----+-----+-----+-----+-----+
+| Device     | CPU | GPU | CPU | GPU | CPU | GPU |
++============+=====+=====+=====+=====+=====+=====+
+| send       | ✓   | ✘   | ✘   | ✘   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| recv       | ✓   | ✘   | ✘   | ✘   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| broadcast  | ✓   | ✘   | ✓   | ✓   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| all_reduce | ✓   | ✘   | ✓   | ✓   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| reduce     | ✓   | ✘   | ✘   | ✘   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| all_gather | ✓   | ✘   | ✘   | ✘   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| gather     | ✓   | ✘   | ✘   | ✘   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| scatter    | ✓   | ✘   | ✘   | ✘   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+| barrier    | ✓   | ✘   | ✓   | ✓   | ✓   | ?   |
++------------+-----+-----+-----+-----+-----+-----+
+
+Initialization
+--------------
+
+The package needs to be initialized using the :func:`torch.distributed.init_process_group`
+function before calling any other methods.
+
+.. autofunction:: init_process_group
+
+.. autofunction:: get_rank
+
+.. autofunction:: get_world_size
+
+--------------------------------------------------------------------------------
+
+Currently three initialization methods are supported:
+
+TCP initialization
+^^^^^^^^^^^^^^^^^^
+
+Initialization will utilize a network address reachable from all processes.
+If the address belongs to one of the machines, initialization requires that all processes
+have manually specified ranks.
+
+Alternatively, the address has to be a valid IP multicast address, in which case,
+ranks can be assigned automatically. Multicast initialization also supports
+a ``group_name`` argument, which allows you to use the same address for multiple jobs,
+as long as they use different group names.
+
+::
+
+    import torch.distributed as dist
+
+    # Use address of one of the machines
+    dist.init_process_group(init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4)
+
+    # or a multicast address - rank will be assigned automatically if unspecified
+    dist.init_process_group(init_method='tcp://[ff15:1e18:5d4c:4cf0:d02d:b659:53ba:b0a7]:23456',
+                            world_size=4)
+
+Shared file-system initialization
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Another initialization method makes use of a file system shared and visible from
+all machines in a group. The URL should start with ``file://`` and contain a path
+to a non-existent file (in an existing directory) on a shared file system.
+This initialization method also supports a ``group_name`` argument, which allows you to
+use the same shared file path for multiple jobs, as long as they use different
+group names.
+
+.. warning::
+    This method assumes that the file system supports locking using ``fcntl`` - most
+    local systems and NFS support it.
+
+::
+
+    import torch.distributed as dist
+
+    # Rank will be assigned automatically if unspecified
+    dist.init_process_group(init_method='file:///mnt/nfs/sharedfile', world_size=4,
+                            group_name=args.group)
+
+Environment variable initialization
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This method will read the configuration from environment variables, allowing
+one to fully customize how the information is obtained. The variables to be set
+are:
+
+* ``MASTER_PORT`` - required; has to be a free port on machine with rank 0
+* ``MASTER_ADDR`` - required (except for rank 0); address of rank 0 node
+* ``WORLD_SIZE`` - required; can be set either here, or in a call to init function
+* ``RANK`` - required; can be set either here, or in a call to init function
+
+The machine with rank 0 will be used to set up all connections.
+
+This is the default method, meaning that ``init_method`` does not have to be specified (or
+can be ``env://``).
+
+Groups
+------
+
+By default collectives operate on the default group (also called the world) and
+require all processes to enter the distributed function call. However, some workloads can benefit
+from more fine-grained communication. This is where distributed groups come
+into play. :func:`~torch.distributed.new_group` function can be
+used to create new groups, with arbitrary subsets of all processes. It returns
+an opaque group handle that can be given as a ``group`` argument to all collectives
+(collectives are distributed functions to exchange information in certain well-known programming patterns).
+
+.. autofunction:: new_group
+
+Point-to-point communication
+----------------------------
+
+.. autofunction:: send
+
+.. autofunction:: recv
+
+:func:`~torch.distributed.isend` and :func:`~torch.distributed.irecv`
+return distributed request objects when used. In general, the type of this object is unspecified
+as they should never be created manually, but they are guaranteed to support two methods:
+
+* ``is_completed()`` - returns True if the operation has finished
+* ``wait()`` - will block the process until the operation is finished.
+  ``is_completed()`` is guaranteed to return True once it returns.
+
+.. autofunction:: isend
+
+.. autofunction:: irecv
+
+Collective functions
+--------------------
+
+.. autofunction:: broadcast
+
+.. autofunction:: all_reduce
+
+.. autofunction:: reduce
+
+.. autofunction:: all_gather
+
+.. autofunction:: gather
+
+.. autofunction:: scatter
+
+.. autofunction:: barrier
+

docs/source/index.rst

@@ -24,11 +24,13 @@ PyTorch is an optimized tensor library for deep learning using GPUs and CPUs.
 
    torch
    tensors
+   sparse
    storage
    nn
    optim
    torch.autograd <autograd>
    torch.multiprocessing <multiprocessing>
+   torch.distributed <distributed>
    torch.legacy <legacy>
    cuda
    ffi

docs/source/multiprocessing.rst

@@ -83,6 +83,6 @@ the current process group, and will keep track of all shared memory allocations.
 Once all processes connected to it exit, it will wait a moment to ensure there
 will be no new connections, and will iterate over all shared memory files
 allocated by the group. If it finds that any of them still exist, they will be
-deallocated. We've tested this method and it prooved to be robust to various
+deallocated. We've tested this method and it proved to be robust to various
 failures. Still, if your system has high enough limits, and ``file_descriptor``
 is a supported strategy, we do not recommend switching to this one.

docs/source/nn.rst

@@ -22,6 +22,24 @@ Containers
 .. autoclass:: Module
     :members:
 
+:hidden:`Sequential`
+~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: Sequential
+    :members:
+
+:hidden:`ModuleList`
+~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ModuleList
+    :members:
+
+:hidden:`ParameterList`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ParameterList
+    :members:
+
 Convolution Layers
 ----------------------------------
 
@@ -132,6 +150,65 @@ Pooling Layers
 .. autoclass:: LPPool2d
     :members:
 
+:hidden:`AdaptiveMaxPool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AdaptiveMaxPool1d
+    :members:
+
+:hidden:`AdaptiveMaxPool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AdaptiveMaxPool2d
+    :members:
+
+:hidden:`AdaptiveAvgPool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AdaptiveAvgPool1d
+    :members:
+
+:hidden:`AdaptiveAvgPool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AdaptiveAvgPool2d
+    :members:
+
+
+Padding Layers
+--------------
+
+:hidden:`ReflectionPad2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ReflectionPad2d
+    :members:
+
+:hidden:`ReplicationPad2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ReplicationPad2d
+    :members:
+
+:hidden:`ReplicationPad3d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ReplicationPad3d
+    :members:
+
+:hidden:`ZeroPad2d`
+~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ZeroPad2d
+    :members:
+
+:hidden:`ConstantPad2d`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ConstantPad2d
+    :members:
+
+
 Non-linear Activations
 ----------------------------------
 
@@ -153,6 +230,12 @@ Non-linear Activations
 .. autoclass:: ELU
     :members:
 
+:hidden:`SELU`
+~~~~~~~~~~~~~~
+
+.. autoclass:: SELU
+    :members:
+
 :hidden:`PReLU`
 ~~~~~~~~~~~~~~~
 
@@ -259,6 +342,23 @@ Normalization layers
 .. autoclass:: BatchNorm3d
     :members:
 
+:hidden:`InstanceNorm1d`
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: InstanceNorm1d
+    :members:
+
+:hidden:`InstanceNorm2d`
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: InstanceNorm2d
+    :members:
+
+:hidden:`InstanceNorm3d`
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: InstanceNorm3d
+    :members:
 
 Recurrent layers
 ----------------------------------
@@ -330,6 +430,12 @@ Dropout layers
 .. autoclass:: Dropout3d
     :members:
 
+:hidden:`AlphaDropout`
+~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: AlphaDropout
+    :members:
+
 
 Sparse layers
 ----------------------------------
@@ -340,6 +446,27 @@ Sparse layers
 .. autoclass:: Embedding
     :members:
 
+:hidden:`EmbeddingBag`
+~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: EmbeddingBag
+    :members:
+
+Distance functions
+----------------------------------
+
+:hidden:`CosineSimilarity`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: CosineSimilarity
+    :members:
+
+:hidden:`PairwiseDistance`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: PairwiseDistance
+    :members:
+
 
 Loss functions
 ----------------------------------
@@ -368,6 +495,12 @@ Loss functions
 .. autoclass:: NLLLoss
     :members:
 
+:hidden:`PoissonNLLLoss`
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: PoissonNLLLoss
+    :members:
+
 :hidden:`NLLLoss2d`
 ~~~~~~~~~~~~~~~~~~~
 
@@ -386,6 +519,12 @@ Loss functions
 .. autoclass:: BCELoss
     :members:
 
+:hidden:`BCEWithLogitsLoss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: BCEWithLogitsLoss
+    :members:
+
 :hidden:`MarginRankingLoss`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -434,6 +573,12 @@ Loss functions
 .. autoclass:: MultiMarginLoss
     :members:
 
+:hidden:`TripletMarginLoss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: TripletMarginLoss
+    :members:
+
 
 Vision layers
 ----------------
@@ -444,28 +589,80 @@ Vision layers
 .. autoclass:: PixelShuffle
     :members:
 
+:hidden:`Upsample`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: Upsample
+    :members:
+
 :hidden:`UpsamplingNearest2d`
-~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: UpsamplingNearest2d
     :members:
 
 :hidden:`UpsamplingBilinear2d`
-~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: UpsamplingBilinear2d
     :members:
 
 
-Multi-GPU layers
-----------------
+DataParallel layers (multi-GPU, distributed)
+--------------------------------------------
 
 :hidden:`DataParallel`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: DataParallel
     :members:
 
+:hidden:`DistributedDataParallel`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: torch.nn.parallel.DataParallel
+    :members:
+
+
+Utilities
+---------
+
+:hidden:`clip_grad_norm`
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: torch.nn.utils.clip_grad_norm
+
+:hidden:`weight_norm`
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: torch.nn.utils.weight_norm
+
+:hidden:`remove_weight_norm`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: torch.nn.utils.remove_weight_norm
+
+
+.. currentmodule:: torch.nn.utils.rnn
+
+:hidden:`PackedSequence`
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: torch.nn.utils.rnn.PackedSequence
+
+
+:hidden:`pack_padded_sequence`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: torch.nn.utils.rnn.pack_padded_sequence
+
+
+:hidden:`pad_packed_sequence`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: torch.nn.utils.rnn.pad_packed_sequence
+
+
 torch.nn.functional
 ===================
 
@@ -557,6 +754,27 @@ Pooling functions
 
 .. autofunction:: lp_pool2d
 
+:hidden:`adaptive_max_pool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: adaptive_max_pool1d
+
+:hidden:`adaptive_max_pool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: adaptive_max_pool2d
+
+:hidden:`adaptive_avg_pool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: adaptive_avg_pool1d
+
+:hidden:`adaptive_avg_pool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: adaptive_avg_pool2d
+
+
 Non-linear activation functions
 -------------------------------
 
@@ -586,6 +804,11 @@ Non-linear activation functions
 
 .. autofunction:: elu
 
+:hidden:`selu`
+~~~~~~~~~~~~~~
+
+.. autofunction:: selu
+
 :hidden:`leaky_relu`
 ~~~~~~~~~~~~~~~~~~~~
 
@@ -664,6 +887,11 @@ Normalization functions
 
 .. autofunction:: batch_norm
 
+:hidden:`normalize`
+~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: normalize
+
 Linear functions
 ----------------
 
@@ -680,35 +908,123 @@ Dropout functions
 
 .. autofunction:: dropout
 
-Loss functions
---------------
-
-:hidden:`nll_loss`
-~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: nll_loss
-
-
-:hidden:`kl_div`
-~~~~~~~~~~~~~~~~
-
-.. autofunction:: kl_div
-
-:hidden:`cross_entropy`
+:hidden:`alpha_dropout`
 ~~~~~~~~~~~~~~~~~~~~~~~
 
-.. autofunction:: cross_entropy
+.. autofunction:: alpha_dropout
+
+:hidden:`dropout2d`
+~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: dropout2d
+
+:hidden:`dropout3d`
+~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: dropout3d
+
+Distance functions
+----------------------------------
+
+:hidden:`pairwise_distance`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: pairwise_distance
+
+:hidden:`cosine_similarity`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: cosine_similarity
+
+
+Loss functions
+--------------
 
 :hidden:`binary_cross_entropy`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: binary_cross_entropy
 
+:hidden:`poisson_nll_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: poisson_nll_loss
+
+:hidden:`cosine_embedding_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: cosine_embedding_loss
+
+:hidden:`cross_entropy`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: cross_entropy
+
+:hidden:`hinge_embedding_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: hinge_embedding_loss
+
+:hidden:`kl_div`
+~~~~~~~~~~~~~~~~
+
+.. autofunction:: kl_div
+
+:hidden:`l1_loss`
+~~~~~~~~~~~~~~~~~
+
+.. autofunction:: l1_loss
+
+:hidden:`mse_loss`
+~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: mse_loss
+
+:hidden:`margin_ranking_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: margin_ranking_loss
+
+:hidden:`multilabel_margin_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: multilabel_margin_loss
+
+:hidden:`multilabel_soft_margin_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: multilabel_soft_margin_loss
+
+:hidden:`multi_margin_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: multi_margin_loss
+
+:hidden:`nll_loss`
+~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: nll_loss
+
+:hidden:`binary_cross_entropy_with_logits`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: binary_cross_entropy_with_logits
+
 :hidden:`smooth_l1_loss`
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: smooth_l1_loss
 
+:hidden:`soft_margin_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: soft_margin_loss
+
+:hidden:`triplet_margin_loss`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: triplet_margin_loss
+
 Vision functions
 ----------------
 
@@ -716,3 +1032,51 @@ Vision functions
 ~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autofunction:: pixel_shuffle
+
+:hidden:`pad`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: pad
+
+:hidden:`upsample`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: upsample
+
+:hidden:`upsample_nearest`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: upsample_nearest
+
+:hidden:`upsample_bilinear`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: upsample_bilinear
+
+:hidden:`grid_sample`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: grid_sample
+
+:hidden:`affine_grid`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: affine_grid
+
+
+torch.nn.init
+=============
+
+.. currentmodule:: torch.nn.init
+.. autofunction:: calculate_gain
+.. autofunction:: uniform
+.. autofunction:: normal
+.. autofunction:: constant
+.. autofunction:: eye
+.. autofunction:: dirac
+.. autofunction:: xavier_uniform
+.. autofunction:: xavier_normal
+.. autofunction:: kaiming_uniform
+.. autofunction:: kaiming_normal
+.. autofunction:: orthogonal
+.. autofunction:: sparse

docs/source/notes/autograd.rst

@@ -67,18 +67,18 @@ model. ``volatile`` also determines that ``requires_grad is False``.
 
 Volatile differs from :ref:`excluding-requires_grad` in how the flag propagates.
 If there's even a single volatile input to an operation, its output is also
-going to be volatile. Volatility spreads accross the graph much easier than
+going to be volatile. Volatility spreads across the graph much easier than
 non-requiring gradient - you only need a **single** volatile leaf to have a
 volatile output, while you need **all** leaves to not require gradient to
-have an output the doesn't require gradient. Using volatile flag you don't
+have an output that doesn't require gradient. Using volatile flag you don't
 need to change any settings of your model parameters to use it for
 inference. It's enough to create a volatile input, and this will ensure that
 no intermediate states are saved.
 
 .. code::
 
-    >>> regular_input = Variable(torch.randn(5, 5))
-    >>> volatile_input = Variable(torch.randn(5, 5), volatile=True)
+    >>> regular_input = Variable(torch.randn(1, 3, 227, 227))
+    >>> volatile_input = Variable(torch.randn(1, 3, 227, 227), volatile=True)
     >>> model = torchvision.models.resnet18(pretrained=True)
     >>> model(regular_input).requires_grad
     True
@@ -86,21 +86,28 @@ no intermediate states are saved.
     False
     >>> model(volatile_input).volatile
     True
-    >>> model(volatile_input).creator is None
+    >>> model(volatile_input).grad_fn is None
     True
 
 How autograd encodes the history
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Each Variable has a ``.creator`` attribute, that points to the function, of
-which it is an output. This is an entry point to a directed acyclic graph (DAG)
-consisting of :class:`Function` objects as nodes, and references between them
-being the edges. Every time an operation is performed, a new :class:`Function`
-representing it is instantiated, its :meth:`~torch.autograd.Function.forward`
-method is called, and its output :class:`Variable` s creators are set to it.
-Then, by following the path from any :class:`Variable` to the leaves, it is
-possible to reconstruct the sequence of operations that has created the data,
-and automatically compute the gradients.
+Autograd is reverse automatic differentiation system.  Conceptually,
+autograd records a graph recording all of the operations that created
+the data as you execute operations, giving you a directed acyclic graph
+whose leaves are the input variables and roots are the output variables.
+By tracing this graph from roots to leaves, you can automatically
+compute the gradients using the chain rule.
+
+Internally, autograd represents this graph as a graph of
+:class:`Function` objects (really expressions), which can be
+:meth:`~torch.autograd.Function.apply` ed to compute the result of
+evaluating the graph.  When computing the forwards pass, autograd
+simultaneously performs the requested computations and builds up a graph
+representing the function that computes the gradient (the ``.grad_fn``
+attribute of each :class:`Variable` is an entry point into this graph).
+When the forwards pass is completed, we evaluate this graph in the
+backwards pass to compute the gradients.
 
 An important thing to note is that the graph is recreated from scratch at every
 iteration, and this is exactly what allows for using arbitrary Python control

docs/source/notes/broadcasting.rst

@@ -0,0 +1,113 @@
+.. _broadcasting-semantics:
+
+Broadcasting semantics
+======================
+
+Many PyTorch operations support :any:`NumPy Broadcasting Semantics <numpy.doc.broadcasting>`.
+
+In short, if a PyTorch operation supports broadcast, then its Tensor arguments can be
+automatically expanded to be of equal sizes (without making copies of the data).
+
+General semantics
+-----------------
+Two tensors are "broadcastable" if the following rules hold:
+
+- Each tensor has at least one dimension.
+- When iterating over the dimension sizes, starting at the trailing dimension,
+  the dimension sizes must either be equal, one of them is 1, or one of them
+  does not exist.
+
+For Example::
+
+    >>> x=torch.FloatTensor(5,7,3)
+    >>> y=torch.FloatTensor(5,7,3)
+    # same shapes are always broadcastable (i.e. the above rules always hold)
+
+    >>> x=torch.FloatTensor()
+    >>> y=torch.FloatTensor(2,2)
+    # x and y are not broadcastable, because x does not have at least 1 dimension
+
+    # can line up trailing dimensions
+    >>> x=torch.FloatTensor(5,3,4,1)
+    >>> y=torch.FloatTensor(  3,1,1)
+    # x and y are broadcastable.
+    # 1st trailing dimension: both have size 1
+    # 2nd trailing dimension: y has size 1
+    # 3rd trailing dimension: x size == y size
+    # 4th trailing dimension: y dimension doesn't exist
+
+    # but:
+    >>> x=torch.FloatTensor(5,2,4,1)
+    >>> y=torch.FloatTensor(  3,1,1)
+    # x and y are not broadcastable, because in the 3rd trailing dimension 2 != 3
+
+If two tensors :attr:`x`, :attr:`y` are "broadcastable", the resulting tensor size
+is calculated as follows:
+
+- If the number of dimensions of :attr:`x` and :attr:`y` are not equal, prepend 1
+  to the dimensions of the tensor with fewer dimensions to make them equal length.
+- Then, for each dimension size, the resulting dimension size is the max of the sizes of
+  :attr:`x` and :attr:`y` along that dimension.
+
+For Example::
+
+    # can line up trailing dimensions to make reading easier
+    >>> x=torch.FloatTensor(5,1,4,1)
+    >>> y=torch.FloatTensor(  3,1,1)
+    >>> (x+y).size()
+    torch.Size([5, 3, 4, 1])
+
+    # but not necessary:
+    >>> x=torch.FloatTensor(1)
+    >>> y=torch.FloatTensor(3,1,7)
+    >>> (x+y).size()
+    torch.Size([3, 1, 7])
+
+    >>> x=torch.FloatTensor(5,2,4,1)
+    >>> y=torch.FloatTensor(3,1,1)
+    >>> (x+y).size()
+    RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1
+
+In-place semantics
+------------------
+One complication is that in-place operations do not allow the in-place tensor to change shape
+as a result of the broadcast.
+
+For Example::
+
+    >>> x=torch.FloatTensor(5,3,4,1)
+    >>> y=torch.FloatTensor(3,1,1)
+    >>> (x.add_(y)).size()
+    torch.Size([5, 3, 4, 1])
+
+    # but:
+    >>> x=torch.FloatTensor(1,3,1)
+    >>> y=torch.FloatTensor(3,1,7)
+    >>> (x.add_(y)).size()
+    RuntimeError: The expanded size of the tensor (1) must match the existing size (7) at non-singleton dimension 2.
+
+Backwards compatibility
+-----------------------
+Prior versions of PyTorch allowed certain pointwise functions to execute on tensors with different shapes,
+as long as the number of elements in each tensor was equal.  The pointwise operation would then be carried
+out by viewing each tensor as 1-dimensional.  PyTorch now supports broadcasting and the "1-dimensional"
+pointwise behavior is considered deprecated and will generate a Python warning in cases where tensors are
+not broadcastable, but have the same number of elements.
+
+Note that the introduction of broadcasting can cause backwards incompatible changes in the case where
+two tensors do not have the same shape, but are broadcastable and have the same number of elements.
+For Example::
+
+    >>> torch.add(torch.ones(4,1), torch.randn(4))
+
+would previously produce a Tensor with size: torch.Size([4,1]), but now produces a Tensor with size: torch.Size([4,4]).
+In order to help identify cases in your code where backwards incompatibilities introduced by broadcasting may exist,
+you may set `torch.utils.backcompat.broadcast_warning.enabled` to `True`, which will generate a python warning
+in such cases.
+
+For Example::
+
+    >>> torch.utils.backcompat.broadcast_warning.enabled=True
+    >>> torch.add(torch.ones(4,1), torch.ones(4))
+    __main__:1: UserWarning: self and other do not have the same shape, but are broadcastable, and have the same number of elements.
+    Changing behavior in a backwards incompatible manner to broadcasting rather than viewing as 1-dimensional.

docs/source/notes/cuda.rst

@@ -1,3 +1,5 @@
+.. _cuda-semantics:
+
 CUDA semantics
 ==============
 
@@ -10,7 +12,7 @@ of your selected device, and the results will be always placed in on the same
 device as the tensor.
 
 Cross-GPU operations are not allowed by default, with the only exception of
-:meth:`~torch.Tensor.copy_`. Unless you enable peer-to-peer memory accesses
+:meth:`~torch.Tensor.copy_`. Unless you enable peer-to-peer memory accesses,
 any attempts to launch ops on tensors spread across different devices will
 raise an error.
 
@@ -60,4 +62,22 @@ Just pass an additional ``async=True`` argument to a :meth:`~torch.Tensor.cuda`
 call. This can be used to overlap data transfers with computation.
 
 You can make the :class:`~torch.utils.data.DataLoader` return batches placed in
-pinned memory by passing ``pinned=True`` to its constructor.
+pinned memory by passing ``pin_memory=True`` to its constructor.
+
+.. _cuda-nn-dataparallel-instead:
+
+Use nn.DataParallel instead of multiprocessing
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Most use cases involving batched input and multiple GPUs should default to using
+:class:`~torch.nn.DataParallel` to utilize more than one GPU. Even with the GIL,
+a single python process can saturate multiple GPUs.
+
+As of version 0.1.9, large numbers of GPUs (8+) might not be fully utilized.
+However, this is a known issue that is under active development. As always,
+test your use case.
+
+There are significant caveats to using CUDA models with
+:mod:`~torch.multiprocessing`; unless care is taken to meet the data handling
+requirements exactly, it is likely that your program will have incorrect or
+undefined behavior.

docs/source/notes/extending.rst

@@ -13,31 +13,28 @@ Extending :mod:`torch.autograd`
 Adding operations to :mod:`~torch.autograd` requires implementing a new
 :class:`Function` subclass for each operation. Recall that :class:`Function` s
 are what :mod:`~torch.autograd` uses to compute the results and gradients, and
-encode the operation history. Every new function requires you to implement 3
+encode the operation history. Every new function requires you to implement 2
 methods:
 
-- ``__init__`` (*optional*) - if your operation is parametrized by/uses
-  objects different than :class:`Variable` s, you should pass them as arguments
-  to ``__init__``. For example, ``AddConstant`` function takes a scalar to add,
-  while ``Transpose`` requires specifying which two dimensions to swap. If your
-  function doesn't require any additional parameters, you can skip it.
 - :meth:`~Function.forward` - the code that performs the operation. It can take
-  as many arguments as you want, with some of them being
-  optional, if you specify the default values. Keep in mind that only
-  :class:`Variable` s will be passed in here. You can return either a single
-  :class:`Variable` output, or a :class:`tuple` of :class:`Variable` s if there
-  are multiple. Also, please refer to the docs of :class:`Function` to find
-  descriptions of useful methods that can be called only from
-  :meth:`~Function.forward`.
+  as many arguments as you want, with some of them being optional, if you
+  specify the default values. All kinds of Python objects are accepted here.
+  :class:`Variable` arguments will be converted to :class:`Tensor` s before the
+  call, and their use will be registered in the graph. Note that this logic won't
+  traverse lists/dicts/any other data structures and will only consider Variables
+  that are direct arguments to the call. You can return either a single
+  :class:`Tensor` output, or a :class:`tuple` of :class:`Tensor` s if there are
+  multiple outputs. Also, please refer to the docs of :class:`Function` to find
+  descriptions of useful methods that can be called only from :meth:`~Function.forward`.
 - :meth:`~Function.backward` - gradient formula. It will be given
-  as many arguments as there were outputs, with each of them representing
-  gradient w.r.t. that output. It should return as many :class:`Tensor` s as
-  there were inputs, with each of them containing the gradient w.r.t.
-  corresponding input. If you inputs didn't require gradient (see
-  :attr:`~Variable.needs_input_grad`), or it was non-differentiable, you
-  can return :class:`None`. Also, if you have optional arguments to
-  :meth:`~Variable.forward` you can return more gradients than there were
-  inputs, as long as they're all :any:`python:None`.
+  as many :class:`Variable` arguments as there were outputs, with each of them
+  representing gradient w.r.t. that output. It should return as many
+  :class:`Variable` s as there were inputs, with each of them containing the
+  gradient w.r.t. its corresponding input. If your inputs didn't require
+  gradient (see :attr:`~Variable.needs_input_grad`), or were non-:class:`Variable`
+  objects, you can return :class:`python:None`. Also, if you have optional
+  arguments to :meth:`~Variable.forward` you can return more gradients than there
+  were inputs, as long as they're all :any:`python:None`.
 
 Below you can find code for a ``Linear`` function from :mod:`torch.nn`, with
 additional comments::
@@ -45,22 +42,25 @@ additional comments::
     # Inherit from Function
     class Linear(Function):
 
+        # Note that both forward and backward are @staticmethods
+        @staticmethod
         # bias is an optional argument
-        def forward(self, input, weight, bias=None):
-            self.save_for_backward(input, weight, bias)
+        def forward(ctx, input, weight, bias=None):
+            ctx.save_for_backward(input, weight, bias)
             output = input.mm(weight.t())
             if bias is not None:
                 output += bias.unsqueeze(0).expand_as(output)
             return output
 
         # This function has only a single output, so it gets only one gradient
-        def backward(self, grad_output):
+        @staticmethod
+        def backward(ctx, grad_output):
             # This is a pattern that is very convenient - at the top of backward
             # unpack saved_tensors and initialize all gradients w.r.t. inputs to
             # None. Thanks to the fact that additional trailing Nones are
             # ignored, the return statement is simple even when the function has
             # optional inputs.
-            input, weight, bias = self.saved_tensors
+            input, weight, bias = ctx.saved_variables
             grad_input = grad_weight = grad_bias = None
 
             # These needs_input_grad checks are optional and there only to
@@ -76,15 +76,40 @@ additional comments::
 
             return grad_input, grad_weight, grad_bias
 
-Now, to make it easier to use these custom ops, we recommend wrapping them in
-small helper functions::
+Now, to make it easier to use these custom ops, we recommend aliasing their
+``apply`` method::
 
-    def linear(input, weight, bias=None):
-        # First braces create a Function object. Any arguments given here
-        # will be passed to __init__. Second braces will invoke the __call__
-        # operator, that will then use forward() to compute the result and
-        # return it.
-        return Linear()(input, weight, bias)
+    linear = Linear.aply
+
+Here, we give an additional example of a function that is parametrized by
+non-Variable arguments::
+
+    class MulConstant(Function):
+        @staticmethod
+        def forward(ctx, tensor, constant):
+            # ctx is a context object that can be used to stash information
+            for backward computation
+            ctx.constant = constant
+            return tensor * constant
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            # We return as many input gradients as there were arguments.
+            # Gradients of non-Tensor arguments to forward must be None.
+            return grad_output * ctx.constant, None
+
+You probably want to check if the backward method you implemented actually
+computes the derivatives of your function. It is possible by comparing with
+numerical approximations using small finite differences::
+
+    from torch.autograd import gradcheck
+
+    # gradchek takes a tuple of tensor as input, check if your gradient
+    # evaluated with these tensors are close enough to numerical
+    # approximations and returns True if they all verify this condition.
+    input = (Variable(torch.randn(20,20).double(), requires_grad=True), Variable(torch.randn(30,20).double(), requires_grad=True),)
+    test = gradcheck(Linear.apply, input, eps=1e-6, atol=1e-4)
+    print(test)
 
 Extending :mod:`torch.nn`
 -------------------------
@@ -132,7 +157,7 @@ This is how a ``Linear`` module can be implemented::
             # nn.Parameters can never be volatile and, different than Variables,
             # they require gradients by default.
             self.weight = nn.Parameter(torch.Tensor(input_features, output_features))
-            if bias is not None:
+            if bias:
                 self.bias = nn.Parameter(torch.Tensor(output_features))
             else:
                 # You should always register all possible parameters, but the

docs/source/notes/multiprocessing.rst

@@ -33,6 +33,8 @@ by the CUDA runtime.
     kinds of data should be done with care. Note that this restriction doesn't
     apply to shared CPU memory.
 
+See also: :ref:`cuda-nn-dataparallel-instead`
+
 
 Best practices and tips
 -----------------------
@@ -100,11 +102,6 @@ example below as well::
     from model import MyModel
 
     def train(model):
-        # This for loop will break sharing of gradient buffers. It's not
-        # necessary but it reduces the contention, and has a small memory cost
-        # (equal to the total size of parameters).
-        for param in model.parameters():
-            param.grad.data = param.grad.data.clone()
         # Construct data_loader, optimizer, etc.
         for data, labels in data_loader:
             optimizer.zero_grad()

docs/source/notes/serialization.rst

@@ -0,0 +1,34 @@
+
+Serialization semantics
+=======================
+
+Best practices
+--------------
+
+.. _recommend-saving-models:
+
+Recommended approach for saving a model
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+There are two main approaches for serializing and restoring a model.
+
+The first (recommended) saves and loads only the model parameters::
+
+    torch.save(the_model.state_dict(), PATH)
+
+Then later::
+
+    the_model = TheModelClass(*args, **kwargs)
+    the_model.load_state_dict(torch.load(PATH))
+
+The second saves and loads the entire model::
+
+    torch.save(the_model, PATH)
+
+Then later::
+
+    the_model = torch.load(PATH)
+
+However in this case, the serialized data is bound to the specific classes
+and the exact directory structure used, so it can break in various ways when
+used in other projects, or after some serious refactors.

docs/source/optim.rst

@@ -114,3 +114,21 @@ Algorithms
     :members:
 .. autoclass:: SGD
     :members:
+
+How to adjust Learning Rate
+---------------------------
+
+:mod:`torch.optim.lr_scheduler` provides several methods to adjust the learning
+rate based on the number of epoches. :class:`torch.optim.lr_scheduler.ReduceLROnPlateau`
+allows dynamic learning rate reducing based on some validation measurements.
+
+.. autoclass:: torch.optim.lr_scheduler.LambdaLR
+    :members:
+.. autoclass:: torch.optim.lr_scheduler.StepLR
+    :members:
+.. autoclass:: torch.optim.lr_scheduler.MultiStepLR
+    :members:
+.. autoclass:: torch.optim.lr_scheduler.ExponentialLR
+    :members:
+.. autoclass:: torch.optim.lr_scheduler.ReduceLROnPlateau
+    :members:

docs/source/sparse.rst

@@ -0,0 +1,114 @@
+.. currentmodule:: torch.sparse
+
+torch.sparse
+============
+
+.. warning::
+
+    This API is currently experimental and may change in the near future.
+
+Torch supports sparse tensors in COO(rdinate) format, which can
+efficiently store and process tensors for which the majority of elements
+are zeros.
+
+A sparse tensor is represented as a pair of dense tensors: a tensor
+of values and a tensor of indices.  A sparse tensor can be constructed
+by providing these two tensors, as well as the size of the sparse tensor
+(which cannot be inferred from these tensors!)
+
+    >>> i = torch.LongTensor([[0, 1], [2, 0]])
+    >>> v = torch.FloatTensor([3, 4])
+    >>> torch.sparse.FloatTensor(i, v, torch.Size([2,3])).to_dense()
+     0  0  3
+     4  0  0
+    [torch.FloatTensor of size 2x2]
+
+You can also construct hybrid sparse tensors, where only the first n
+dimensions are sparse, and the rest of the dimensions are dense.
+
+    >>> i = torch.LongTensor([[2, 4]])
+    >>> v = torch.FloatTensor([[1, 3], [5, 7]])
+    >>> torch.sparse.FloatTensor(i, v).to_dense()
+     0  0
+     0  0
+     1  3
+     0  0
+     5  7
+    [torch.FloatTensor of size 5x2]
+
+An empty sparse tensor can be constructed by specifying its size:
+
+    >>> torch.sparse.FloatTensor(2, 3)
+    SparseFloatTensor of size 2x3 with indices:
+    [torch.LongTensor with no dimension]
+    and values:
+    [torch.FloatTensor with no dimension]
+
+.. note::
+
+    Our sparse tensor format permits *uncoalesced* sparse tensors, where
+    there may be duplicate coordinates in the indices; in this case,
+    the interpretation is that the value at that index is the sum of all
+    duplicate value entries. Uncoalesced tensors permit us to implement
+    certain operators more efficiently.
+
+    For the most part, you shouldn't have to care whether or not a
+    sparse tensor is coalesced or not, as most operations will work
+    identically given a coalesced or uncoalesced sparse tensor.
+    However, there are two cases in which you may need to care.
+
+    First, if you repeatedly perform an operation that can produce
+    duplicate entries (e.g., :func:`torch.sparse.FloatTensor.add`), you
+    should occasionally coalesce your sparse tensors to prevent
+    them from growing too large.
+
+    Second, some operators will produce different values depending on
+    whether or not they are coalesced or not (e.g.,
+    :func:`torch.sparse.FloatTensor._values` and
+    :func:`torch.sparse.FloatTensor._indices`, as well as
+    :func:`torch.Tensor._sparse_mask`).  These operators are
+    prefixed by an underscore to indicate that they reveal internal
+    implementation details and should be used with care, since code
+    that works with coalesced sparse tensors may not work with
+    uncoalesced sparse tensors; generally speaking, it is safest
+    to explicitly coalesce before working with these operators.
+
+    For example, suppose that we wanted to implement an operator
+    by operating directly on :func:`torch.sparse.FloatTensor._values`.
+    Multiplication by a scalar can be implemented in the obvious way,
+    as multiplication distributes over addition; however, square root
+    cannot be implemented directly, since ``sqrt(a + b) != sqrt(a) +
+    sqrt(b)`` (which is what would be computed if you were given an
+    uncoalesced tensor.)
+
+.. class:: FloatTensor()
+
+    .. method:: add
+    .. method:: add_
+    .. method:: clone
+    .. method:: dim
+    .. method:: div
+    .. method:: div_
+    .. method:: get_device
+    .. method:: hspmm
+    .. method:: mm
+    .. method:: mul
+    .. method:: mul_
+    .. method:: resizeAs_
+    .. method:: size
+    .. method:: spadd
+    .. method:: spmm
+    .. method:: sspaddmm
+    .. method:: sspmm
+    .. method:: sub
+    .. method:: sub_
+    .. method:: t_
+    .. method:: toDense
+    .. method:: transpose
+    .. method:: transpose_
+    .. method:: zero_
+    .. method:: coalesce
+    .. method:: is_coalesced
+    .. method:: _indices
+    .. method:: _values
+    .. method:: _nnz

docs/source/tensors.rst

@@ -13,7 +13,7 @@ Data type                CPU tensor                    GPU tensor
 ======================== ===========================   ================================
 32-bit floating point    :class:`torch.FloatTensor`    :class:`torch.cuda.FloatTensor`
 64-bit floating point    :class:`torch.DoubleTensor`   :class:`torch.cuda.DoubleTensor`
-16-bit floating point    N/A                           :class:`torch.cuda.HalfTensor`
+16-bit floating point    :class:`torch.HalfTensor`     :class:`torch.cuda.HalfTensor`
 8-bit integer (unsigned) :class:`torch.ByteTensor`     :class:`torch.cuda.ByteTensor`
 8-bit integer (signed)   :class:`torch.CharTensor`     :class:`torch.cuda.CharTensor`
 16-bit integer (signed)  :class:`torch.ShortTensor`    :class:`torch.cuda.ShortTensor`
@@ -196,9 +196,10 @@ view of a storage and defines numeric operations on it.
    .. automethod:: lt
    .. automethod:: lt_
    .. automethod:: map_
-   .. automethod:: masked_copy_
+   .. automethod:: masked_scatter_
    .. automethod:: masked_fill_
    .. automethod:: masked_select
+   .. automethod:: matmul
    .. automethod:: max
    .. automethod:: mean
    .. automethod:: median

docs/source/torch.rst

@@ -8,6 +8,7 @@ Tensors
 .. autofunction:: is_storage
 .. autofunction:: set_default_tensor_type
 .. autofunction:: numel
+.. autofunction:: set_printoptions
 
 
 Creation Ops
@@ -20,6 +21,7 @@ Creation Ops
 .. autofunction:: rand
 .. autofunction:: randn
 .. autofunction:: randperm
+.. autofunction:: arange
 .. autofunction:: range
 .. autofunction:: zeros
 
@@ -38,6 +40,7 @@ Indexing, Slicing, Joining, Mutating Ops
 .. autofunction:: t
 .. autofunction:: transpose
 .. autofunction:: unbind
+.. autofunction:: unsqueeze
 
 
 Random sampling
@@ -158,6 +161,8 @@ BLAS and LAPACK Operations
 .. autofunction:: addr
 .. autofunction:: baddbmm
 .. autofunction:: bmm
+.. autofunction:: btrifact
+.. autofunction:: btrisolve
 .. autofunction:: dot
 .. autofunction:: eig
 .. autofunction:: gels
@@ -165,6 +170,7 @@ BLAS and LAPACK Operations
 .. autofunction:: ger
 .. autofunction:: gesv
 .. autofunction:: inverse
+.. autofunction:: matmul
 .. autofunction:: mm
 .. autofunction:: mv
 .. autofunction:: orgqr

docs/source/torchvision/datasets.rst

@@ -1,109 +1,112 @@
 torchvision.datasets
 ====================
 
-The following dataset loaders are available:
+All datasets are subclasses of :class:`torch.utils.data.Dataset`
+i.e, they have ``__getitem__`` and ``__len__`` methods implemented.
+Hence, they can all be passed to a :class:`torch.utils.data.DataLoader`
+which can load multiple samples parallelly using ``torch.multiprocessing`` workers. 
+For example: ::
+    
+    imagenet_data = torchvision.datasets.ImageFolder('path/to/imagenet_root/')
+    data_loader = torch.utils.data.DataLoader(imagenet_data, 
+                                              batch_size=4,
+                                              shuffle=True,
+                                              num_workers=args.nThreads)
 
--  `COCO (Captioning and Detection)`_
--  `LSUN Classification`_
--  `ImageFolder`_
--  `Imagenet-12`_
--  `CIFAR10 and CIFAR100`_
+The following datasets are available:
 
-Datasets have the API:
+.. contents:: Datasets
+    :local:
 
--  ``__getitem__``
--  ``__len__``
-   They all subclass from ``torch.utils.data.Dataset``
-   Hence, they can all be multi-threaded (python multiprocessing) using
-   standard torch.utils.data.DataLoader.
+All the datasets have almost similar API. They all have two common arguments:
+``transform`` and  ``target_transform`` to transform the input and target respectively.
 
-For example:
 
-``torch.utils.data.DataLoader(coco_cap, batch_size=args.batchSize, shuffle=True, num_workers=args.nThreads)``
+.. currentmodule:: torchvision.datasets 
 
-In the constructor, each dataset has a slightly different API as needed,
-but they all take the keyword args:
 
--  ``transform`` - a function that takes in an image and returns a
-   transformed version
--  common stuff like ``ToTensor``, ``RandomCrop``, etc. These can be
-   composed together with ``transforms.Compose`` (see transforms section
-   below)
--  ``target_transform`` - a function that takes in the target and
-   transforms it. For example, take in the caption string and return a
-   tensor of word indices.
+MNIST
+~~~~~
+
+.. autoclass:: MNIST
 
 COCO
 ~~~~
 
-This requires the `COCO API to be installed`_
+.. note ::
+    These require the `COCO API to be installed`_
 
-Captions:
+.. _COCO API to be installed: https://github.com/pdollar/coco/tree/master/PythonAPI
+
+
+Captions
+^^^^^^^^
+
+.. autoclass:: CocoCaptions
+  :members: __getitem__
+  :special-members:
+
+
+Detection
 ^^^^^^^^^
 
-``dset.CocoCaptions(root="dir where images are", annFile="json annotation file", [transform, target_transform])``
-
-Example:
-
-.. code:: python
-
-    import torchvision.datasets as dset
-    import torchvision.transforms as transforms
-    cap = dset.CocoCaptions(root = 'dir where images are',
-                            annFile = 'json annotation file',
-                            transform=transforms.ToTensor())
-
-    print('Number of samples: ', len(cap))
-    img, target = cap[3] # load 4th sample
-
-    print("Image Size: ", img.size())
-    print(target)
-
-Output:
-
-::
-
-    Number of samples: 82783
-    Image Size: (3L, 427L, 640L)
-    [u'A plane emitting smoke stream flying over a mountain.',
-    u'A plane darts across a bright blue sky behind a mountain covered in snow',
-    u'A plane leaves a contrail above the snowy mountain top.',
-    u'A mountain that has a plane flying overheard in the distance.',
-    u'A mountain view with a plume of smoke in the background']
-
-Detection:
-^^^^^^^^^^
-
-``dset.CocoDetection(root="dir where images are", annFile="json annotation file", [transform, target_transform])``
+.. autoclass:: CocoDetection
+  :members: __getitem__
+  :special-members:
 
 LSUN
 ~~~~
 
-``dset.LSUN(db_path, classes='train', [transform, target_transform])``
+.. autoclass:: LSUN
+  :members: __getitem__
+  :special-members:
 
--  db\_path = root directory for the database files
--  classes =
--  ‘train’ - all categories, training set
--  ‘val’ - all categories, validation set
--  ‘test’ - all categories, test set
--  [‘bedroom\_train’, ‘church\_train’, …] : a list of categories to load
+ImageFolder
+~~~~~~~~~~~
+
+.. autoclass:: ImageFolder
+  :members: __getitem__
+  :special-members:
+
+
+Imagenet-12
+~~~~~~~~~~~
+
+This should simply be implemented with an ``ImageFolder`` dataset.
+The data is preprocessed `as described
+here <https://github.com/facebook/fb.resnet.torch/blob/master/INSTALL.md#download-the-imagenet-dataset>`__
+
+`Here is an
+example <https://github.com/pytorch/examples/blob/27e2a46c1d1505324032b1d94fc6ce24d5b67e97/imagenet/main.py#L48-L62>`__.
 
 CIFAR
 ~~~~~
 
-``dset.CIFAR10(root, train=True, transform=None, target_transform=None, download=False)``
+.. autoclass:: CIFAR10
+  :members: __getitem__
+  :special-members:
 
-``dset.CIFAR100(root, train=True, transform=None, target_transform=None, download=False)``
+STL10
+~~~~~
 
--  ``root`` : root directory of dataset where there is folder
-   ``cifar-10-batches-py``
--  ``train`` : ``True`` = Training set, ``False`` = Test set
--  ``download`` : ``True`` = downloads the dataset from the internet and
-   puts it in root directory. If dataset already downloaded, do
 
-.. _COCO (Captioning and Detection): #coco
-.. _LSUN Classification: #lsun
-.. _ImageFolder: #imagefolder
-.. _Imagenet-12: #imagenet-12
-.. _CIFAR10 and CIFAR100: #cifar
-.. _COCO API to be installed: https://github.com/pdollar/coco/tree/master/PythonAPI
+.. autoclass:: STL10
+  :members: __getitem__
+  :special-members:
+
+SVHN
+~~~~~
+
+
+.. autoclass:: SVHN
+  :members: __getitem__
+  :special-members:
+
+PhotoTour
+~~~~~~~~~
+
+
+.. autoclass:: PhotoTour
+  :members: __getitem__
+  :special-members:
+

docs/source/torchvision/models.rst

@@ -1,11 +1,12 @@
 torchvision.models
 ===================
 
-.. currentmodule:: torchvision.models
 
+.. currentmodule:: torchvision.models
 
 .. automodule:: torchvision.models
    :members: alexnet, resnet18, resnet34, resnet50, resnet101, resnet152,
              vgg11, vgg11_bn, vgg13, vgg13_bn, vgg16, vgg16_bn, vgg19,
-             vgg19_bn
+             vgg19_bn, inception_v3, squeezenet1_0, squeezenet1_1, densenet121, 
+             densenet169, densenet201, densenet161
    :undoc-members:

docs/source/torchvision/torchvision.rst

@@ -3,3 +3,6 @@ torchvision
 
 The :mod:`torchvision` package consists of popular datasets, model
 architectures, and common image transformations for computer vision.
+
+.. automodule:: torchvision
+   :members:

docs/source/torchvision/transforms.rst

@@ -3,6 +3,8 @@ torchvision.transforms
 
 .. currentmodule:: torchvision.transforms
 
+Transforms are common image transforms. They can be chained together using :class:`Compose`
+
 .. autoclass:: Compose
 
 Transforms on PIL.Image
@@ -24,16 +26,22 @@ Transforms on torch.\*Tensor
 ----------------------------
 
 .. autoclass:: Normalize
+	:members: __call__
+	:special-members:
 
 
 Conversion Transforms
 ---------------------
 
 .. autoclass:: ToTensor
+	:members: __call__
+	:special-members:
 
 .. autoclass:: ToPILImage
+	:members: __call__
+	:special-members:
 
-Generic Transofrms
+Generic Transforms
 ------------------
 
 .. autoclass:: Lambda

setup.py

@@ -15,9 +15,23 @@ import os
 from tools.setup_helpers.env import check_env_flag
 from tools.setup_helpers.cuda import WITH_CUDA, CUDA_HOME
 from tools.setup_helpers.cudnn import WITH_CUDNN, CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR
+from tools.setup_helpers.split_types import split_types
 DEBUG = check_env_flag('DEBUG')
-WITH_DISTRIBUTED = check_env_flag('WITH_DISTRIBUTED')
+WITH_DISTRIBUTED = not check_env_flag('NO_DISTRIBUTED')
 WITH_DISTRIBUTED_MW = WITH_DISTRIBUTED and check_env_flag('WITH_DISTRIBUTED_MW')
+WITH_NCCL = WITH_CUDA and platform.system() != 'Darwin'
+SYSTEM_NCCL = False
+
+
+################################################################################
+# Workaround setuptools -Wstrict-prototypes warnings
+# I lifted this code from https://stackoverflow.com/a/29634231/23845
+################################################################################
+import distutils.sysconfig
+cfg_vars = distutils.sysconfig.get_config_vars()
+for key, value in cfg_vars.items():
+    if type(value) == str:
+            cfg_vars[key] = value.replace("-Wstrict-prototypes", "")
 
 ################################################################################
 # Monkey-patch setuptools to compile in parallel
@@ -75,6 +89,8 @@ class build_deps(Command):
         build_all_cmd = ['bash', 'torch/lib/build_all.sh']
         if WITH_CUDA:
             build_all_cmd += ['--with-cuda']
+        if WITH_NCCL and not SYSTEM_NCCL:
+            build_all_cmd += ['--with-nccl']
         if WITH_DISTRIBUTED:
             build_all_cmd += ['--with-distributed']
         if subprocess.call(build_all_cmd) != 0:
@@ -134,6 +150,16 @@ class build_ext(setuptools.command.build_ext.build_ext):
             print('-- Detected CUDA at ' + CUDA_HOME)
         else:
             print('-- Not using CUDA')
+        if WITH_NCCL and SYSTEM_NCCL:
+            print('-- Using system provided NCCL library')
+        elif WITH_NCCL:
+            print('-- Building NCCL library')
+        else:
+            print('-- Not using NCCL')
+        if WITH_DISTRIBUTED:
+            print('-- Building with distributed package ')
+        else:
+            print('-- Building without distributed package')
 
         # cwrap depends on pyyaml, so we can't import it earlier
         from tools.cwrap import cwrap
@@ -144,10 +170,15 @@ class build_ext(setuptools.command.build_ext.build_ext):
         from tools.cwrap.plugins.KwargsPlugin import KwargsPlugin
         from tools.cwrap.plugins.NullableArguments import NullableArguments
         from tools.cwrap.plugins.CuDNNPlugin import CuDNNPlugin
+        from tools.cwrap.plugins.WrapDim import WrapDim
+        from tools.cwrap.plugins.AssertNDim import AssertNDim
+        from tools.cwrap.plugins.Broadcast import Broadcast
+        from tools.cwrap.plugins.ProcessorSpecificPlugin import ProcessorSpecificPlugin
         thp_plugin = THPPlugin()
         cwrap('torch/csrc/generic/TensorMethods.cwrap', plugins=[
-            BoolOption(), thp_plugin, AutoGPU(condition='IS_CUDA'),
-            ArgcountSortPlugin(), KwargsPlugin()
+            ProcessorSpecificPlugin(), BoolOption(), thp_plugin,
+            AutoGPU(condition='IS_CUDA'), ArgcountSortPlugin(), KwargsPlugin(),
+            AssertNDim(), WrapDim(), Broadcast()
         ])
         cwrap('torch/csrc/cudnn/cuDNN.cwrap', plugins=[
             CuDNNPlugin(), NullableArguments()
@@ -192,12 +223,12 @@ class clean(distutils.command.clean.clean):
 ################################################################################
 
 include_dirs = []
+library_dirs = []
 extra_link_args = []
-extra_compile_args = ['-std=c++11', '-Wno-write-strings']
-if os.getenv('PYTORCH_BINARY_BUILD') and platform.system() == 'Linux':
-    print('PYTORCH_BINARY_BUILD found. Static linking libstdc++ on Linux')
-    extra_compile_args += ['-static-libstdc++']
-    extra_link_args += ['-static-libstdc++']
+extra_compile_args = ['-std=c++11', '-Wno-write-strings',
+                      # Python 2.6 requires -fno-strict-aliasing, see
+                      # http://legacy.python.org/dev/peps/pep-3123/
+                      '-fno-strict-aliasing']
 
 cwd = os.path.dirname(os.path.abspath(__file__))
 lib_path = os.path.join(cwd, "torch", "lib")
@@ -210,9 +241,10 @@ include_dirs += [
     tmp_install_path + "/include/TH",
     tmp_install_path + "/include/THPP",
     tmp_install_path + "/include/THNN",
+    tmp_install_path + "/include/ATen",
 ]
 
-extra_link_args.append('-L' + lib_path)
+library_dirs.append(lib_path)
 
 # we specify exact lib names to avoid conflict with lua-torch installs
 TH_LIB = os.path.join(lib_path, 'libTH.so.1')
@@ -222,7 +254,11 @@ THCS_LIB = os.path.join(lib_path, 'libTHCS.so.1')
 THNN_LIB = os.path.join(lib_path, 'libTHNN.so.1')
 THCUNN_LIB = os.path.join(lib_path, 'libTHCUNN.so.1')
 THPP_LIB = os.path.join(lib_path, 'libTHPP.so.1')
-THD_LIB = os.path.join(lib_path, 'libTHD.so.1')
+ATEN_LIB = os.path.join(lib_path, 'libATen.so.1')
+GLOO_LIB = os.path.join(lib_path, 'libgloo.a')
+GLOO_CUDA_LIB = os.path.join(lib_path, 'libgloo_cuda.a')
+THD_LIB = os.path.join(lib_path, 'libTHD.a')
+NCCL_LIB = os.path.join(lib_path, 'libnccl.so.1')
 if platform.system() == 'Darwin':
     TH_LIB = os.path.join(lib_path, 'libTH.1.dylib')
     THS_LIB = os.path.join(lib_path, 'libTHS.1.dylib')
@@ -231,38 +267,50 @@ if platform.system() == 'Darwin':
     THNN_LIB = os.path.join(lib_path, 'libTHNN.1.dylib')
     THCUNN_LIB = os.path.join(lib_path, 'libTHCUNN.1.dylib')
     THPP_LIB = os.path.join(lib_path, 'libTHPP.1.dylib')
-    THD_LIB = os.path.join(lib_path, 'libTHD.1.dylib')
+    ATEN_LIB = os.path.join(lib_path, 'libATen.1.dylib')
+    NCCL_LIB = os.path.join(lib_path, 'libnccl.1.dylib')
+
+if WITH_NCCL and subprocess.call('ldconfig -p | grep libnccl >/dev/null', shell=True) == 0:
+    SYSTEM_NCCL = True
 
 main_compile_args = ['-D_THP_CORE']
 main_libraries = ['shm']
-main_link_args = [TH_LIB, THS_LIB, THPP_LIB, THNN_LIB]
+main_link_args = [TH_LIB, THS_LIB, THPP_LIB, THNN_LIB, ATEN_LIB]
 main_sources = [
     "torch/csrc/PtrWrapper.cpp",
     "torch/csrc/Module.cpp",
     "torch/csrc/Generator.cpp",
     "torch/csrc/Size.cpp",
     "torch/csrc/Exceptions.cpp",
-    "torch/csrc/Tensor.cpp",
     "torch/csrc/Storage.cpp",
     "torch/csrc/DynamicTypes.cpp",
     "torch/csrc/byte_order.cpp",
     "torch/csrc/utils.cpp",
+    "torch/csrc/expand_utils.cpp",
     "torch/csrc/utils/object_ptr.cpp",
+    "torch/csrc/utils/tuple_parser.cpp",
     "torch/csrc/allocators.cpp",
     "torch/csrc/serialization.cpp",
     "torch/csrc/autograd/init.cpp",
     "torch/csrc/autograd/engine.cpp",
     "torch/csrc/autograd/function.cpp",
     "torch/csrc/autograd/variable.cpp",
-    "torch/csrc/autograd/grad_buffer.cpp",
+    "torch/csrc/autograd/input_buffer.cpp",
     "torch/csrc/autograd/python_function.cpp",
     "torch/csrc/autograd/python_cpp_function.cpp",
     "torch/csrc/autograd/python_variable.cpp",
     "torch/csrc/autograd/python_engine.cpp",
+    "torch/csrc/autograd/python_hook.cpp",
     "torch/csrc/autograd/functions/batch_normalization.cpp",
+    "torch/csrc/autograd/functions/convolution.cpp",
+    "torch/csrc/autograd/functions/basic_ops.cpp",
+    "torch/csrc/autograd/functions/tensor.cpp",
+    "torch/csrc/autograd/functions/accumulate_grad.cpp",
+    "torch/csrc/autograd/functions/utils.cpp",
     "torch/csrc/autograd/functions/init.cpp",
     "torch/csrc/nn/THNN_generic.cpp",
 ]
+main_sources += split_types("torch/csrc/Tensor.cpp")
 
 try:
     import numpy as np
@@ -283,8 +331,11 @@ if WITH_DISTRIBUTED:
             "torch/csrc/distributed/Tensor.cpp",
             "torch/csrc/distributed/Storage.cpp",
         ]
+        extra_compile_args += ['-DWITH_DISTRIBUTED_MW']
     include_dirs += [tmp_install_path + "/include/THD"]
     main_link_args += [THD_LIB]
+    if platform.system() == 'Linux':
+        main_link_args += [GLOO_LIB]
 
 if WITH_CUDA:
     cuda_lib_dirs = ['lib64', 'lib']
@@ -295,30 +346,42 @@ if WITH_CUDA:
             break
     include_dirs.append(cuda_include_path)
     include_dirs.append(tmp_install_path + "/include/THCUNN")
-    extra_link_args.append('-L' + cuda_lib_path)
+    library_dirs.append(cuda_lib_path)
     extra_link_args.append('-Wl,-rpath,' + cuda_lib_path)
     extra_compile_args += ['-DWITH_CUDA']
     extra_compile_args += ['-DCUDA_LIB_PATH=' + cuda_lib_path]
-    main_libraries += ['cudart']
+    main_libraries += ['cudart', 'nvToolsExt']
     main_link_args += [THC_LIB, THCS_LIB, THCUNN_LIB]
+    if platform.system() == 'Linux':
+        main_link_args += [GLOO_CUDA_LIB]
     main_sources += [
         "torch/csrc/cuda/Module.cpp",
         "torch/csrc/cuda/Storage.cpp",
         "torch/csrc/cuda/Stream.cpp",
-        "torch/csrc/cuda/Tensor.cpp",
         "torch/csrc/cuda/AutoGPU.cpp",
         "torch/csrc/cuda/utils.cpp",
+        "torch/csrc/cuda/expand_utils.cpp",
         "torch/csrc/cuda/serialization.cpp",
     ]
+    main_sources += split_types("torch/csrc/cuda/Tensor.cpp")
+
+if WITH_NCCL:
+    if SYSTEM_NCCL:
+        main_libraries += ['nccl']
+    else:
+        main_link_args += [NCCL_LIB]
+    extra_compile_args += ['-DWITH_NCCL']
 
 if WITH_CUDNN:
     main_libraries += ['cudnn']
     include_dirs.append(CUDNN_INCLUDE_DIR)
-    extra_link_args.append('-L' + CUDNN_LIB_DIR)
+    library_dirs.append(CUDNN_LIB_DIR)
     main_sources += [
         "torch/csrc/cudnn/BatchNorm.cpp",
         "torch/csrc/cudnn/Conv.cpp",
         "torch/csrc/cudnn/cuDNN.cpp",
+        "torch/csrc/cudnn/GridSampler.cpp",
+        "torch/csrc/cudnn/AffineGridGenerator.cpp",
         "torch/csrc/cudnn/Types.cpp",
         "torch/csrc/cudnn/Handles.cpp",
     ]
@@ -328,6 +391,18 @@ if DEBUG:
     extra_compile_args += ['-O0', '-g']
     extra_link_args += ['-O0', '-g']
 
+if os.getenv('PYTORCH_BINARY_BUILD') and platform.system() == 'Linux':
+    print('PYTORCH_BINARY_BUILD found. Static linking libstdc++ on Linux')
+    # get path of libstdc++ and link manually.
+    # for reasons unknown, -static-libstdc++ doesn't fully link some symbols
+    CXXNAME = os.getenv('CXX', 'g++')
+    STDCPP_LIB = subprocess.check_output([CXXNAME, '-print-file-name=libstdc++.a'])
+    STDCPP_LIB = STDCPP_LIB[:-1]
+    if type(STDCPP_LIB) != str:  # python 3
+        STDCPP_LIB = STDCPP_LIB.decode(sys.stdout.encoding)
+    main_link_args += [STDCPP_LIB]
+    version_script = os.path.abspath("tools/pytorch.version")
+    extra_link_args += ['-Wl,--version-script=' + version_script]
 
 def make_relative_rpath(path):
     if platform.system() == 'Darwin':
@@ -340,7 +415,7 @@ def make_relative_rpath(path):
 ################################################################################
 
 extensions = []
-packages = find_packages(exclude=('tools.*',))
+packages = find_packages(exclude=('tools', 'tools.*',))
 
 C = Extension("torch._C",
               libraries=main_libraries,
@@ -348,6 +423,7 @@ C = Extension("torch._C",
               language='c++',
               extra_compile_args=main_compile_args + extra_compile_args,
               include_dirs=include_dirs,
+              library_dirs=library_dirs,
               extra_link_args=extra_link_args + main_link_args + [make_relative_rpath('lib')],
               )
 extensions.append(C)
@@ -386,7 +462,7 @@ if WITH_CUDA:
                        )
     extensions.append(THCUNN)
 
-version = '0.1.9'
+version = '0.2.0'
 if os.getenv('PYTORCH_BUILD_VERSION'):
     assert os.getenv('PYTORCH_BUILD_NUMBER') is not None
     version = os.getenv('PYTORCH_BUILD_VERSION') \
@@ -400,6 +476,7 @@ else:
 
 
 setup(name="torch", version=version,
+      description="Tensors and Dynamic neural networks in Python with strong GPU acceleration",
       ext_modules=extensions,
       cmdclass={
           'build': build,
@@ -418,5 +495,5 @@ setup(name="torch", version=version,
           'lib/*.h',
           'lib/include/TH/*.h', 'lib/include/TH/generic/*.h',
           'lib/include/THC/*.h', 'lib/include/THC/generic/*.h']},
-      install_requires=['pyyaml'],
+      install_requires=['pyyaml', 'numpy'],
       )

test/common.py

@@ -1,26 +1,42 @@
 import sys
+import os
 import argparse
 import unittest
+import warnings
 import contextlib
+from functools import wraps
 from itertools import product
 from copy import deepcopy
 
 import torch
 import torch.cuda
-from torch.autograd import Variable, Function
+from torch.autograd import Variable
 
 
 torch.set_default_tensor_type('torch.DoubleTensor')
 
+SEED = 0
+SEED_SET = 0
 
-def run_tests():
+
+def parse_set_seed_once():
+    global SEED
+    global SEED_SET
     parser = argparse.ArgumentParser(add_help=False)
     parser.add_argument('--seed', type=int, default=123)
     args, remaining = parser.parse_known_args()
-    torch.manual_seed(args.seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(args.seed)
+    if SEED_SET == 0:
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(args.seed)
+        SEED = args.seed
+        SEED_SET = 1
     remaining = [sys.argv[0]] + remaining
+    return remaining
+
+
+def run_tests():
+    remaining = parse_set_seed_once()
     unittest.main(argv=remaining)
 
 
@@ -30,6 +46,32 @@ try:
 except ImportError:
     TEST_NUMPY = False
 
+TEST_SCIPY = True
+try:
+    import scipy
+except ImportError:
+    TEST_SCIPY = False
+
+
+def skipIfNoLapack(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        try:
+            fn(*args, **kwargs)
+        except Exception as e:
+            if 'Lapack library not found' in e.args[0]:
+                raise unittest.SkipTest('Compiled without Lapack')
+            raise
+    return wrapper
+
+
+def suppress_warnings(fn):
+    def wrapper(*args, **kwargs):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            fn(*args, **kwargs)
+    return wrapper
+
 
 def get_cpu_type(t):
     assert t.__module__ == 'torch.cuda'
@@ -48,7 +90,7 @@ def to_gpu(obj, type_map={}):
     elif torch.is_storage(obj):
         return obj.new().resize_(obj.size()).copy_(obj)
     elif isinstance(obj, Variable):
-        assert obj.creator is None
+        assert obj.is_leaf
         t = type_map.get(type(obj.data), get_gpu_type(type(obj.data)))
         return Variable(obj.data.clone().type(t), requires_grad=obj.requires_grad)
     elif isinstance(obj, list):
@@ -89,23 +131,86 @@ def is_iterable(obj):
 class TestCase(unittest.TestCase):
     precision = 1e-5
 
+    def setUp(self):
+        torch.manual_seed(SEED)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(SEED)
+
+    def assertTensorsSlowEqual(self, x, y, prec=None, message=''):
+        max_err = 0
+        self.assertEqual(x.size(), y.size())
+        for index in iter_indices(x):
+            max_err = max(max_err, abs(x[index] - y[index]))
+        self.assertLessEqual(max_err, prec, message)
+
+    def safeCoalesce(self, t):
+        tc = t.coalesce()
+
+        value_map = {}
+        for idx, val in zip(t._indices().t(), t._values()):
+            idx_tup = tuple(idx)
+            if idx_tup in value_map:
+                value_map[idx_tup] += val
+            else:
+                value_map[idx_tup] = val.clone() if torch.is_tensor(val) else val
+
+        new_indices = sorted(list(value_map.keys()))
+        new_values = [value_map[idx] for idx in new_indices]
+        if t._values().ndimension() < 2:
+            new_values = t._values().new(new_values)
+        else:
+            new_values = torch.stack(new_values)
+
+        new_indices = t._indices().new(new_indices).t()
+        tg = t.new(new_indices, new_values, t.size())
+
+        self.assertEqual(tc._indices(), tg._indices())
+        self.assertEqual(tc._values(), tg._values())
+
+        return tg
+
+    def unwrapVariables(self, x, y):
+        if isinstance(x, Variable) and isinstance(y, Variable):
+            return x.data, y.data
+        elif isinstance(x, Variable) or isinstance(y, Variable):
+            raise AssertionError("cannot compare {} and {}".format(type(x), type(y)))
+        return x, y
+
     def assertEqual(self, x, y, prec=None, message=''):
         if prec is None:
             prec = self.precision
 
-        if isinstance(x, Variable) and isinstance(y, Variable):
-            x = x.data
-            y = y.data
+        x, y = self.unwrapVariables(x, y)
 
         if torch.is_tensor(x) and torch.is_tensor(y):
-            max_err = 0
-            super(TestCase, self).assertEqual(x.size(), y.size())
-            for index in iter_indices(x):
-                max_err = max(max_err, abs(x[index] - y[index]))
-            self.assertLessEqual(max_err, prec, message)
+            def assertTensorsEqual(a, b):
+                super(TestCase, self).assertEqual(a.size(), b.size())
+                if a.numel() > 0:
+                    b = b.type_as(a)
+                    b = b.cuda(device=a.get_device()) if a.is_cuda else b.cpu()
+                    # check that NaNs are in the same locations
+                    nan_mask = a != a
+                    self.assertTrue(torch.equal(nan_mask, b != b))
+                    diff = a - b
+                    diff[nan_mask] = 0
+                    if diff.is_signed():
+                        diff = diff.abs()
+                    max_err = diff.max()
+                    self.assertLessEqual(max_err, prec, message)
+            self.assertEqual(x.is_sparse, y.is_sparse, message)
+            if x.is_sparse:
+                x = self.safeCoalesce(x)
+                y = self.safeCoalesce(y)
+                assertTensorsEqual(x._indices(), y._indices())
+                assertTensorsEqual(x._values(), y._values())
+            else:
+                assertTensorsEqual(x, y)
         elif type(x) == str and type(y) == str:
             super(TestCase, self).assertEqual(x, y)
+        elif type(x) == set and type(y) == set:
+            super(TestCase, self).assertEqual(x, y)
         elif is_iterable(x) and is_iterable(y):
+            super(TestCase, self).assertEqual(len(x), len(y))
             for x_, y_ in zip(x, y):
                 self.assertEqual(x_, y_, prec, message)
         else:
@@ -120,17 +225,22 @@ class TestCase(unittest.TestCase):
         if prec is None:
             prec = self.precision
 
-        if isinstance(x, Variable) and isinstance(y, Variable):
-            x = x.data
-            y = y.data
+        x, y = self.unwrapVariables(x, y)
 
         if torch.is_tensor(x) and torch.is_tensor(y):
-            max_err = 0
             if x.size() != y.size():
                 super(TestCase, self).assertNotEqual(x.size(), y.size())
-            for index in iter_indices(x):
-                max_err = max(max_err, abs(x[index] - y[index]))
-            self.assertGreaterEqual(max_err, prec, message)
+            self.assertGreater(x.numel(), 0)
+            y = y.type_as(x)
+            y = y.cuda(device=x.get_device()) if x.is_cuda else y.cpu()
+            nan_mask = x != x
+            if torch.equal(nan_mask, y != y):
+                diff = x - y
+                if diff.is_signed():
+                    diff = diff.abs()
+                diff[nan_mask] = 0
+                max_err = diff.max()
+                self.assertGreaterEqual(max_err, prec, message)
         elif type(x) == str and type(y) == str:
             super(TestCase, self).assertNotEqual(x, y)
         elif is_iterable(x) and is_iterable(y):
@@ -149,66 +259,33 @@ class TestCase(unittest.TestCase):
                 return
         raise AssertionError("object not found in iterable")
 
+    if sys.version_info < (3, 2):
+        # assertRaisesRegexp renamed assertRaisesRegex in 3.2
+        assertRaisesRegex = unittest.TestCase.assertRaisesRegexp
 
-def make_jacobian(input, num_out):
-    if isinstance(input, Variable) and not input.requires_grad:
-        return None
-    if torch.is_tensor(input) or isinstance(input, Variable):
-        return torch.zeros(input.nelement(), num_out)
+
+def download_file(url, binary=True):
+    if sys.version_info < (3,):
+        from urlparse import urlsplit
+        import urllib2
+        request = urllib2
+        error = urllib2
     else:
-        return type(input)(filter(lambda x: x is not None,
-                                  (make_jacobian(elem, num_out) for elem in input)))
+        from urllib.parse import urlsplit
+        from urllib import request, error
 
+    filename = os.path.basename(urlsplit(url)[2])
+    data_dir = os.path.join(os.path.dirname(__file__), 'data')
+    path = os.path.join(data_dir, filename)
 
-def iter_tensors(x, only_requiring_grad=False):
-    if torch.is_tensor(x):
-        yield x
-    elif isinstance(x, Variable):
-        if x.requires_grad or not only_requiring_grad:
-            yield x.data
-    else:
-        for elem in x:
-            for result in iter_tensors(elem, only_requiring_grad):
-                yield result
-
-
-def contiguous(input):
-    if torch.is_tensor(input):
-        return input.contiguous()
-    elif isinstance(input, Variable):
-        return input.contiguous()
-    else:
-        return type(input)(contiguous(e) for e in input)
-
-
-def get_numerical_jacobian(fn, input, target):
-    perturbation = 1e-6
-    # To be able to use .view(-1) input must be contiguous
-    input = contiguous(input)
-    output_size = fn(input).numel()
-    jacobian = make_jacobian(target, output_size)
-
-    # It's much easier to iterate over flattened lists of tensors.
-    # These are reference to the same objects in jacobian, so any changes
-    # will be reflected in it as well.
-    x_tensors = [t for t in iter_tensors(target, True)]
-    j_tensors = [t for t in iter_tensors(jacobian)]
-
-    outa = torch.DoubleTensor(output_size)
-    outb = torch.DoubleTensor(output_size)
-
-    # TODO: compare structure
-    for x_tensor, d_tensor in zip(x_tensors, j_tensors):
-        flat_tensor = x_tensor.view(-1)
-        for i in range(flat_tensor.nelement()):
-            orig = flat_tensor[i]
-            flat_tensor[i] = orig - perturbation
-            outa.copy_(fn(input))
-            flat_tensor[i] = orig + perturbation
-            outb.copy_(fn(input))
-            flat_tensor[i] = orig
-
-            outb.add_(-1, outa).div_(2 * perturbation)
-            d_tensor[i] = outb
-
-    return jacobian
+    if os.path.exists(path):
+        return path
+    try:
+        data = request.urlopen(url, timeout=15).read()
+        with open(path, 'wb' if binary else 'w') as f:
+            f.write(data)
+        return path
+    except error.URLError:
+        msg = "could not download test file '{}'".format(url)
+        warnings.warn(msg, RuntimeWarning)
+        raise unittest.SkipTest(msg)

test/common_nn.py

@@ -7,8 +7,8 @@ from itertools import product
 import torch
 import torch.cuda
 from torch.autograd import Variable
-from common import TestCase, to_gpu, get_numerical_jacobian, iter_tensors, contiguous, \
-    freeze_rng_state
+from common import TestCase, to_gpu, freeze_rng_state
+from torch.autograd.gradcheck import get_numerical_jacobian, iter_tensors, contiguous
 import torch.backends.cudnn
 
 # tarfile module tries to obtain a file object name in python 3.3
@@ -53,29 +53,31 @@ module_tests = [
     dict(
         module_name='ReLU',
         input_size=(2, 3, 4, 5),
-        check_inplace=True
+        check_inplace=True,
     ),
     dict(
         module_name='ReLU6',
         input_size=(2, 3, 4, 5),
-        check_inplace=True
+        check_inplace=True,
     ),
     dict(
         module_name='RReLU',
         input_size=(1, 2, 2),
-        test_cuda=False
+        test_cuda=False,
+        check_gradgrad=False,
     ),
     dict(
         module_name='RReLU',
         constructor_args=(0.1, 0.9),
         input_size=(4, 4, 5),
         desc='with_up_down',
-        test_cuda=False
+        test_cuda=False,
+        check_gradgrad=False,
     ),
     dict(
         module_name='Hardtanh',
         input_size=(3, 2, 5),
-        reference_fn=lambda i, _: i.clamp(-1, 1)
+        reference_fn=lambda i, _: i.clamp(-1, 1),
     ),
     dict(
         module_name='Sigmoid',
@@ -88,35 +90,35 @@ module_tests = [
     dict(
         module_name='Softmax',
         input_size=(10, 20),
-        reference_fn=lambda i, _: torch.exp(i).div(torch.exp(i).sum(1).expand(10, 20))
+        reference_fn=lambda i, _: torch.exp(i).div(torch.exp(i).sum(1, True).expand(10, 20)),
     ),
     dict(
         module_name='Softmax2d',
         input_size=(1, 3, 10, 20),
-        reference_fn=lambda i, _: torch.exp(i).div(torch.exp(i).sum(1).expand_as(i))
+        reference_fn=lambda i, _: torch.exp(i).div(torch.exp(i).sum(1, False)),
     ),
     dict(
         module_name='LogSoftmax',
         input_size=(10, 20),
-        reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1).expand(10, 20)).log_()
+        reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1, True).expand(10, 20)).log_(),
     ),
     dict(
         module_name='LogSoftmax',
         input_size=(1, 3, 10, 20),
-        reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1).expand_as(i)).log_(),
-        desc='multiparam'
+        reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1, False)).log_(),
+        desc='multiparam',
     ),
     dict(
         module_name='ELU',
         constructor_args=(2.,),
         input_size=(3, 2, 5),
-        check_inplace=True
     ),
     # TODO: reference function
     dict(
         module_name='Hardshrink',
         constructor_args=(2.,),
-        input_size=(4, 3, 2, 4)
+        input_size=(4, 3, 2, 4),
+        check_gradgrad=False,
     ),
     dict(
         module_name='LeakyReLU',
@@ -133,53 +135,89 @@ module_tests = [
     dict(
         module_name='LogSigmoid',
         input_size=(2, 3, 4),
-        reference_fn=lambda i, _: i.sigmoid().log()
+        reference_fn=lambda i, _: i.sigmoid().log(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='Softplus',
         input_size=(10, 20),
-        reference_fn=lambda i, _: torch.log(1 + torch.exp(i))
+        reference_fn=lambda i, _: torch.log(1 + torch.exp(i)),
+        check_gradgrad=False,
     ),
     dict(
         module_name='Softplus',
         constructor_args=(2,),
         input_size=(10, 20),
         reference_fn=lambda i, _: 1. / 2. * torch.log(1 + torch.exp(2 * i)),
-        desc='beta'
+        desc='beta',
+        check_gradgrad=False,
     ),
     dict(
         module_name='Softshrink',
-        input_size=(3, 2, 5)
+        input_size=(3, 2, 5),
+        check_gradgrad=False,
     ),
     dict(
         module_name='Softshrink',
         constructor_args=(1,),
         input_size=(3, 2, 5),
-        desc='lambda'
+        desc='lambda',
+        check_gradgrad=False,
     ),
     dict(
         module_name='CrossMapLRN2d',
         constructor_args=(5, 5e-3, 1e-3, 2),
-        input_size=(2, 3, 6, 6)
+        input_size=(2, 3, 6, 6),
+        check_gradgrad=False,
     ),
     dict(
         module_name='PReLU',
-        input_size=(2, 3, 4, 5)
+        input_size=(2, 3, 4),
+        reference_fn=lambda i, p: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0],
+        desc='1d',
+    ),
+    dict(
+        module_name='PReLU',
+        constructor_args=(3,),
+        input_size=(2, 3, 4),
+        desc='1d_multiparam',
+        reference_fn=lambda i, p: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0],
+    ),
+    dict(
+        module_name='PReLU',
+        input_size=(2, 3, 4, 5),
+        desc='2d',
+        reference_fn=lambda i, p: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0],
     ),
     dict(
         module_name='PReLU',
         constructor_args=(3,),
         input_size=(2, 3, 4, 5),
-        desc='multiparam'
+        desc='2d_multiparam',
+        reference_fn=lambda i, p: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0],
+    ),
+    dict(
+        module_name='PReLU',
+        input_size=(2, 3, 4, 5, 6),
+        reference_fn=lambda i, p: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0],
+        desc='3d',
+    ),
+    dict(
+        module_name='PReLU',
+        constructor_args=(3,),
+        input_size=(2, 3, 4, 5, 6),
+        desc='3d_multiparam',
+        reference_fn=lambda i, p: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0],
     ),
     dict(
         module_name='Softsign',
         input_size=(3, 2, 5),
-        reference_fn=lambda i, _: i.div(1 + torch.abs(i))
+        reference_fn=lambda i, _: i.div(1 + torch.abs(i)),
     ),
     dict(
         module_name='Softmin',
-        input_size=(10, 20)
+        input_size=(10, 20),
+        check_gradgrad=False,
     ),
     dict(
         module_name='Tanhshrink',
@@ -187,19 +225,32 @@ module_tests = [
     ),
 ]
 
-
 criterion_tests = [
     dict(module_name='L1Loss',
          input_size=(2, 3, 4),
          target=torch.randn(2, 3, 4),
          reference_fn=lambda i, t, _: 1. / i.numel() *
-         sum((a - b).abs().sum() for a, b in zip(i, t))
+         sum((a - b).abs().sum() for a, b in zip(i, t)),
          ),
     dict(
         module_name='NLLLoss',
         input=torch.rand(15, 10).log(),
         target=torch.Tensor(15).uniform_().mul(10).floor().long(),
     ),
+    dict(
+        module_name='NLLLoss',
+        constructor_args=(None, False),
+        input=torch.rand(15, 10).log(),
+        target=torch.Tensor(15).uniform_().mul(10).floor().long(),
+        desc='no_size_average'
+    ),
+    dict(
+        module_name='NLLLoss',
+        constructor_args=(None, True, 2),
+        input=torch.rand(15, 10).log(),
+        target=torch.Tensor(15).uniform_().mul(10).floor().long(),
+        desc='ignore_index'
+    ),
     dict(
         module_name='NLLLoss',
         constructor_args=(torch.rand(10),),
@@ -207,113 +258,159 @@ criterion_tests = [
         target=torch.Tensor(15).uniform_().mul(10).floor().long(),
         desc='weights',
     ),
+    dict(
+        module_name='NLLLoss',
+        constructor_args=(torch.rand(10), True, 2),
+        input=torch.rand(15, 10).add(1e-2).log(),
+        target=torch.Tensor(15).uniform_().mul(10).floor().long(),
+        desc='weights_ignore_index'
+    ),
+    dict(
+        module_name='NLLLoss',
+        constructor_args=(torch.rand(10), True, -1),
+        input=torch.rand(15, 10).add(1e-2).log(),
+        target=torch.Tensor(15).uniform_().mul(10 + 1).floor().long() - 1,
+        desc='weights_ignore_index_neg'
+    ),
     dict(
         module_name='KLDivLoss',
         input=torch.rand(10, 10).log(),
-        target=torch.rand(10, 10)
+        target=torch.rand(10, 10),
+        check_gradgrad=False,
     ),
     dict(
         module_name='MSELoss',
         input=torch.randn(2, 3, 4, 5),
         target=torch.randn(2, 3, 4, 5),
-        reference_fn=lambda i, t, _: (i - t).abs().pow(2).sum() / i.numel()
+        reference_fn=lambda i, t, _: (i - t).abs().pow(2).sum() / i.numel(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='BCELoss',
         input=torch.rand(15, 10).clamp_(1e-2, 1 - 1e-2),
-        target=torch.randn(15, 10).gt(0).double()
+        target=torch.randn(15, 10).gt(0).double(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='BCELoss',
         constructor_args=(torch.rand(10),),
         input=torch.rand(15, 10).clamp_(1e-2, 1 - 1e-2),
         target=torch.randn(15, 10).gt(0).double(),
-        desc='weights'
+        desc='weights',
+        check_gradgrad=False,
     ),
     dict(
         module_name='CrossEntropyLoss',
         input=torch.randn(15, 10),
-        target=torch.Tensor(15).uniform_().mul(10).floor().long()
+        target=torch.Tensor(15).uniform_().mul(10).floor().long(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='CrossEntropyLoss',
         constructor_args=(torch.rand(10),),
         input=torch.randn(15, 10),
         target=torch.Tensor(15).uniform_().mul(10).floor().long(),
-        desc='weights'
+        desc='weights',
+        check_gradgrad=False,
     ),
     dict(
         module_name='NLLLoss2d',
         input_size=(2, 3, 5, 5),
-        target=torch.rand(2, 5, 5).mul(3).floor().long()
+        target=torch.rand(2, 5, 5).mul(3).floor().long(),
+    ),
+    dict(
+        module_name='NLLLoss2d',
+        constructor_args=(torch.rand(3),),
+        input_size=(2, 3, 5, 5),
+        target=torch.rand(2, 5, 5).mul(3).floor().long(),
+        desc='weights',
+    ),
+    dict(
+        module_name='NLLLoss2d',
+        constructor_args=(None, True, 3),
+        input_size=(2, 3, 5, 5),
+        target=torch.rand(2, 5, 5).mul(4).floor().long(),
+        desc='ignore_index',
     ),
     dict(
         module_name='HingeEmbeddingLoss',
         input=torch.rand(10),
-        target=torch.randn(10).gt(0).double().mul_(2).sub(1)
+        target=torch.randn(10).gt(0).double().mul_(2).sub(1),
+        check_gradgrad=False,
     ),
     dict(
         module_name='HingeEmbeddingLoss',
         constructor_args=(0.5,),
         input=torch.rand(10),
         target=torch.randn(10).gt(0).double().mul_(2).sub(1),
-        desc='margin'
+        desc='margin',
+        check_gradgrad=False,
     ),
     dict(
         module_name='MultiLabelMarginLoss',
         input_size=(5, 10),
-        target=torch.rand(5, 10).mul(10).floor().long()
+        target=torch.rand(5, 10).mul(10).floor().long(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='MultiLabelSoftMarginLoss',
         input_size=(5, 10),
-        target=torch.rand(5, 10).mul(2).floor()
+        target=torch.rand(5, 10).mul(2).floor(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='MultiLabelSoftMarginLoss',
         constructor_args=(torch.rand(10),),
         input_size=(5, 10),
         target=torch.rand(5, 10).mul(2).floor(),
-        desc='weights'
+        desc='weights',
+        check_gradgrad=False,
     ),
     dict(
         module_name='MultiMarginLoss',
         input_size=(5, 10),
-        target=torch.rand(5).mul(8).floor().long()
+        target=torch.rand(5).mul(8).floor().long(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='SmoothL1Loss',
         input_size=(5, 10),
-        target=torch.randn(5, 10)
+        target=torch.randn(5, 10),
+        check_gradgrad=False,
     ),
     dict(
         module_name='SoftMarginLoss',
         input_size=(5, 5),
-        target=torch.randn(5, 5).sign()
+        target=torch.randn(5, 5).sign(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='CosineEmbeddingLoss',
         input=(torch.rand(15, 10), torch.rand(15, 10)),
-        target=torch.randn(15).sign()
+        target=torch.randn(15).sign(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='CosineEmbeddingLoss',
         constructor_args=(0.7,),
         input=(torch.rand(15, 10), torch.rand(15, 10)),
         target=torch.randn(15).sign(),
-        desc='margin'
+        desc='margin',
+        check_gradgrad=False,
     ),
     dict(
         module_name='MarginRankingLoss',
         input=(torch.randn(50).mul(10), torch.randn(50).mul(10)),
-        target=torch.randn(50).sign()
+        target=torch.randn(50).sign(),
+        check_gradgrad=False,
     ),
     dict(
         module_name='MarginRankingLoss',
         constructor_args=(2,),
         input=(torch.randn(50).mul(10), torch.randn(50).mul(10)),
         target=torch.randn(50).sign(),
-        desc='margin'
+        desc='margin',
+        check_gradgrad=False,
     ),
 ]
 
@@ -330,16 +427,20 @@ class NNTestCase(TestCase):
 
     def _flatten_tensors(self, x):
         if torch.is_tensor(x):
-            return x.view(-1)
+            if x.is_sparse:
+                return x.to_dense().view(-1)
+            else:
+                return x.view(-1)
         elif isinstance(x, Variable):
-            return x.data.view(-1)
+            return self._flatten_tensors(x.data)
         else:
             return tuple(self._flatten_tensors(a) for a in x)
 
     def _zero_grad_input(self, input):
         if isinstance(input, Variable):
-            if input.requires_grad:
+            if input.requires_grad and input.grad is not None:
                 input.grad.data.zero_()
+                input.grad.detach_()
         elif torch.is_tensor(input):
             return
         else:
@@ -400,12 +501,11 @@ class NNTestCase(TestCase):
             return out
 
         res = tuple()
-        # TODO: enable non-contig tests
         input = contiguous(input)
         if jacobian_input:
-            res += get_numerical_jacobian(fw, input, input),
+            res += get_numerical_jacobian(fw, input, input, eps=1e-6),
         if jacobian_parameters:
-            res += torch.cat(list(get_numerical_jacobian(fw, input, p) for p in param), 0),
+            res += torch.cat(list(get_numerical_jacobian(fw, input, p, eps=1e-6) for p in param), 0),
         return res
 
     def check_jacobian(self, module, input, jacobian_input=True):
@@ -652,6 +752,7 @@ class CriterionTest(TestBase):
             test_case.assertEqual(out, expected_out)
 
         test_case.check_criterion_jacobian(module, input, self.target)
+        self._do_extra_tests(test_case, module, input, self.target)
 
     def test_cuda(self, test_case):
         if not TEST_CUDA or not self.should_test_cuda:
@@ -678,3 +779,6 @@ class CriterionTest(TestBase):
             test_case.assertEqual(cpu_gradInput, gpu_gradInput, 4e-4)
         except NotImplementedError:
             pass
+
+    def _do_extra_tests(self, test_case, module, input, target):
+        pass

test/run_test.sh

@@ -6,7 +6,7 @@ COVERAGE=0
 while [[ "$#" -gt 0 ]]; do
     case "$1" in
         -p|--python) PYCMD=$2; shift 2 ;;
-        -c|--coverage) COVERAGE=1; shift 2 ;;
+        -c|--coverage) COVERAGE=1; shift 1;;
         --) shift; break ;;
         *) echo "Invalid argument: $1!" ; exit 1 ;;
     esac
@@ -55,35 +55,39 @@ $PYCMD test_cuda.py $@
 echo "Running NCCL tests"
 $PYCMD test_nccl.py $@
 
-################################################################################
-if [[ "$TEST_DISTRIBUTED" -eq 1 ]]; then
-    distributed_set_up() {
-        export TEMP_DIR="$(mktemp -d)"
-        rm -rf "$TEMP_DIR/"*
-        mkdir "$TEMP_DIR/barrier"
-        mkdir "$TEMP_DIR/test_dir"
-    }
+distributed_set_up() {
+  export TEMP_DIR="$(mktemp -d)"
+  rm -rf "$TEMP_DIR/"*
+  mkdir "$TEMP_DIR/barrier"
+  mkdir "$TEMP_DIR/test_dir"
+}
 
-    distributed_tear_down() {
-        rm -rf "$TEMP_DIR"
-    }
+distributed_tear_down() {
+  rm -rf "$TEMP_DIR"
+}
 
-    trap distributed_tear_down EXIT SIGHUP SIGINT SIGTERM
+trap distributed_tear_down EXIT SIGHUP SIGINT SIGTERM
 
-    echo "Running distributed tests for the TCP backend"
-    distributed_set_up
-    BACKEND=tcp WORLD_SIZE=3 $PYCMD ./test_distributed.py
-    distributed_tear_down
+echo "Running distributed tests for the TCP backend"
+distributed_set_up
+BACKEND=tcp WORLD_SIZE=3 $PYCMD ./test_distributed.py
+distributed_tear_down
 
-    echo "Running distributed tests for the MPI backend"
-    distributed_set_up
-    BACKEND=mpi mpiexec -n 3 $PYCMD ./test_distributed.py
-    distributed_tear_down
+echo "Running distributed tests for the Gloo backend"
+distributed_set_up
+BACKEND=gloo WORLD_SIZE=3 $PYCMD ./test_distributed.py
+distributed_tear_down
+
+if [ -x "$(command -v mpiexec)" ]; then
+  echo "Running distributed tests for the MPI backend"
+  distributed_set_up
+  BACKEND=mpi mpiexec -n 3 $PYCMD ./test_distributed.py
+  distributed_tear_down
+else
+  echo "Skipping MPI backend tests (MPI not found)"
 fi
-################################################################################
 
-if [ "$1" == "coverage" ];
-then
+if [[ $COVERAGE -eq 1 ]]; then
     coverage combine
     coverage html
 fi

test/test_cuda.py

@@ -7,12 +7,14 @@ import torch
 import torch.cuda
 import torch.cuda.comm as comm
 
+from test_torch import TestTorch
 from common import TestCase, get_gpu_type, to_gpu, freeze_rng_state, run_tests
 
+HAS_CUDA = True
 if not torch.cuda.is_available():
     print('CUDA not available, skipping tests')
-    import sys
-    sys.exit()
+    TestCase = object  # noqa: F811
+    HAS_CUDA = False
 
 
 def is_floating(t):
@@ -59,6 +61,13 @@ def small_2d_scaled(t, scale=10):
     return make_tensor(t, S, S).mul(scale)
 
 
+def small_2d_oneish(t):
+    if is_floating(t):
+        return make_tensor(t, S, S).clamp(min=0.99, max=1.01)
+    else:
+        return t(S, S).fill_(1)
+
+
 def small_3d(t):
     return make_tensor(t, S, S, S)
 
@@ -85,23 +94,27 @@ def small_3d_positive(t):
 
 
 def small_3d_unique(t):
-    return t(S, S, S).copy_(torch.range(1, S * S * S))
+    return t(S, S, S).copy_(torch.arange(1, S * S * S + 1).view(S, S, S))
 
 
 def small_1d_lapack(t):
-    return t(1, 3).copy_(torch.range(1, 3).view(3))
+    return t(1, 3).copy_(torch.arange(1, 4).view(3))
 
 
 def small_2d_lapack(t):
-    return t(3, 3).copy_(torch.range(1, 9).view(3, 3))
+    return t(3, 3).copy_(torch.arange(1, 10).view(3, 3))
 
 
 def small_2d_lapack_skinny(t):
-    return t(3, 4).copy_(torch.range(1, 12).view(3, 4))
+    return t(3, 4).copy_(torch.arange(1, 13).view(3, 4))
 
 
 def small_2d_lapack_fat(t):
-    return t(4, 3).copy_(torch.range(1, 12).view(4, 3))
+    return t(4, 3).copy_(torch.arange(1, 13).view(4, 3))
+
+
+def large_2d_lapack(t):
+    return t(1000, 1000).normal_()
 
 
 def new_t(*sizes):
@@ -146,12 +159,15 @@ tests = [
     ('fmod', small_3d, lambda t: [small_3d_positive(t)], 'tensor'),
     ('chunk', medium_2d, lambda t: [4],),
     ('chunk', medium_2d, lambda t: [4, 1], 'dim'),
+    ('chunk', medium_2d, lambda t: [4, -2], 'neg_dim'),
     ('clamp', medium_2d_scaled, lambda t: [-1, 5],),
     ('clone', medium_2d, lambda t: [],),
     ('contiguous', medium_2d, lambda t: [],),
     ('cross', new_t(M, 3, M), lambda t: [new_t(M, 3, M)(t)],),
     ('cumprod', small_3d, lambda t: [1],),
+    ('cumprod', small_3d, lambda t: [-1], 'neg_dim'),
     ('cumsum', small_3d, lambda t: [1],),
+    ('cumsum', small_3d, lambda t: [-1], 'neg_dim'),
     ('dim', small_3d, lambda t: [],),
     ('dist', small_2d, lambda t: [small_2d(t)],),
     ('dist', small_2d, lambda t: [small_2d(t), 3], '3_norm'),
@@ -179,53 +195,75 @@ tests = [
     # TODO: positive case
     ('kthvalue', small_3d_unique, lambda t: [3],),
     ('kthvalue', small_3d_unique, lambda t: [3, 1], 'dim'),
+    ('kthvalue', small_3d_unique, lambda t: [3, -1], 'neg_dim'),
     ('lerp', small_3d, lambda t: [small_3d(t), 0.3],),
     ('max', small_3d_unique, lambda t: [],),
     ('max', small_3d_unique, lambda t: [1], 'dim'),
+    ('max', small_3d_unique, lambda t: [-1], 'neg_dim'),
     ('max', medium_2d, lambda t: [medium_2d(t)], 'elementwise'),
     ('min', small_3d_unique, lambda t: [],),
     ('min', small_3d_unique, lambda t: [1], 'dim'),
+    ('min', small_3d_unique, lambda t: [-1], 'neg_dim'),
     ('min', medium_2d, lambda t: [medium_2d(t)], 'elementwise'),
     ('mean', small_3d, lambda t: [],),
+    ('mean', small_3d, lambda t: [-1], 'neg_dim'),
     ('mean', small_3d, lambda t: [1], 'dim'),
     ('mode', small_3d, lambda t: [],),
     ('mode', small_3d, lambda t: [1], 'dim'),
+    ('mode', small_3d, lambda t: [-1], 'neg_dim'),
     ('remainder', small_3d, lambda t: [3], 'value'),
+    ('remainder', small_3d, lambda t: [-3], 'negative_value'),
     ('remainder', small_3d, lambda t: [small_3d_positive(t)], 'tensor'),
+    ('remainder', small_3d, lambda t: [0 - small_3d_positive(t)], 'negative_tensor'),
     ('std', small_3d, lambda t: [],),
     ('std', small_3d, lambda t: [1], 'dim'),
+    ('std', small_3d, lambda t: [-1], 'neg_dim'),
     ('var', small_3d, lambda t: [],),
     ('var', small_3d, lambda t: [1], 'dim'),
+    ('var', small_3d, lambda t: [-1], 'neg_dim'),
     ('ndimension', small_3d, lambda t: [],),
     ('nelement', small_3d, lambda t: [],),
     ('numel', small_3d, lambda t: [],),
     ('narrow', small_3d, lambda t: [1, 3, 2],),
+    ('narrow', small_3d, lambda t: [-1, 3, 2], 'neg_dim'),
     ('nonzero', small_3d, lambda t: [],),
     ('norm', small_3d, lambda t: [],),
     ('norm', small_3d, lambda t: [3], '3_norm'),
     ('norm', small_3d, lambda t: [3, 0], '3_norm_dim'),
+    ('norm', small_3d, lambda t: [3, -2], '3_norm_neg_dim'),
     ('ones', small_3d, lambda t: [1, 2, 3, 4, 5],),
     ('permute', new_t(1, 2, 3, 4), lambda t: [2, 1, 3, 0],),
-    ('prod', small_3d, lambda t: [],),
+    ('prod', small_2d_oneish, lambda t: [],),
     ('prod', small_3d, lambda t: [1], 'dim'),
+    ('prod', small_3d, lambda t: [-1], 'neg_dim'),
     ('sum', small_2d, lambda t: [],),
     ('sum', small_3d, lambda t: [1], 'dim'),
+    ('sum', small_3d, lambda t: [-1], 'neg_dim'),
     ('renorm', small_3d, lambda t: [2, 1, 1], '2_norm'),
+    ('renorm', small_3d, lambda t: [2, -1, 1], '2_norm_neg_dim'),
     ('renorm', small_3d, lambda t: [1.5, 1, 1], '1_5_norm'),
     ('repeat', small_2d, lambda t: [2, 2, 2],),
     ('size', new_t(1, 2, 3, 4), lambda t: [],),
+    ('size', new_t(1, 2, 3, 4), lambda t: [1], 'dim'),
+    ('size', new_t(1, 2, 3, 4), lambda t: [-2], 'neg_dim'),
     ('sort', small_3d_unique, lambda t: [],),
     ('sort', small_3d_unique, lambda t: [1], 'dim'),
+    ('sort', small_3d_unique, lambda t: [-1], 'neg_dim'),
     ('sort', small_3d_unique, lambda t: [1, True], 'dim_descending'),
+    ('sort', small_3d_unique, lambda t: [-1, True], 'neg_dim_descending'),
     ('split', small_3d, lambda t: [2],),
     ('split', small_3d, lambda t: [2, 1], 'dim'),
+    ('split', small_3d, lambda t: [2, -3], 'neg_dim'),
     ('squeeze', new_t(1, 2, 1, 4), lambda t: [],),
     ('squeeze', new_t(1, 2, 1, 4), lambda t: [2], 'dim'),
+    ('squeeze', new_t(1, 2, 1, 4), lambda t: [-2], 'neg_dim'),
     ('t', new_t(1, 2), lambda t: [],),
     ('transpose', new_t(1, 2, 3, 4), lambda t: [1, 2],),
+    ('transpose', new_t(1, 2, 3, 4), lambda t: [-1, -2], 'neg_dim'),
     ('to_list', small_3d, lambda t: [],),
-    ('topk', small_3d, lambda t: [2, 1, False, True], 'dim_sort'),
-    ('topk', small_3d, lambda t: [2, 1, True, True], 'dim_desc_sort'),
+    ('topk', small_3d_unique, lambda t: [2, 1, False, True], 'dim_sort'),
+    ('topk', small_3d_unique, lambda t: [2, -1, False, True], 'neg_dim_sort'),
+    ('topk', small_3d_unique, lambda t: [2, 1, True, True], 'dim_desc_sort'),
     ('trace', medium_2d, lambda t: [],),
     ('tril', medium_2d, lambda t: [],),
     ('tril', medium_2d, lambda t: [2], 'positive'),
@@ -234,6 +272,7 @@ tests = [
     ('triu', medium_2d, lambda t: [2], 'positive'),
     ('triu', medium_2d, lambda t: [-2], 'negative'),
     ('unsqueeze', new_t(2, 3, 4), lambda t: [2],),
+    ('unsqueeze', new_t(2, 3, 4), lambda t: [-2], 'neg_dim'),
     ('view', small_3d, lambda t: [100, 10],),
     ('view_as', small_3d, lambda t: [t(100, 10)],),
     ('zero', small_3d, lambda t: [],),
@@ -245,6 +284,8 @@ tests = [
     ('qr', small_2d_lapack, lambda t: [], 'square', float_types),
     ('qr', small_2d_lapack_skinny, lambda t: [], 'skinny', float_types),
     ('qr', small_2d_lapack_fat, lambda t: [], 'fat', float_types),
+    ('qr', large_2d_lapack, lambda t: [], 'big', float_types),
+    ('inverse', new_t(20, 20), lambda t: [], None, float_types),
 
 ]
 
@@ -259,6 +300,7 @@ custom_precision = {
     'baddbmm': 1e-4,
     'rsqrt': 1e-4,
     'cumprod': 1e-4,
+    'qr': 3e-4,
 }
 
 simple_pointwise = [
@@ -375,7 +417,7 @@ class TestCuda(TestCase):
             self.assertEqual(z.get_device(), 0)
             self.assertIs(z.cuda(0), z)
 
-    def test_serialization(self):
+    def test_serialization_array_with_storage(self):
         x = torch.randn(5, 5).cuda()
         y = torch.IntTensor(2, 5).fill_(0).cuda()
         q = [x, y, x, y.storage()]
@@ -427,6 +469,32 @@ class TestCuda(TestCase):
     def test_broadcast_gpu(self):
         self._test_broadcast(torch.randn(5, 5))
 
+    @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
+    def test_broadcast_coalesced(self):
+        numel = 5
+        num_bytes = numel * 8
+        tensors = [
+            torch.randn(numel).long().cuda(),
+            torch.randn(numel).cuda(),
+            torch.randn(numel).long().cuda(),
+            torch.randn(numel).long().cuda(),
+            torch.randn(numel * 2).int().cuda(),  # int is 2x shorter
+            torch.randn(numel).cuda(),
+        ]
+
+        b_tensors = [comm.broadcast(t, (0, 1)) for t in tensors]
+        for (_, bt), t in zip(b_tensors, tensors):
+            self.assertEqual(bt.get_device(), 1)
+            self.assertEqual(bt, t)
+            self.assertIsInstance(bt, type(t))
+
+        bc_tensors = comm.broadcast_coalesced(tensors, (0, 1), buffer_size=num_bytes * 5 // 2)
+        bc_tensors_t = list(zip(*bc_tensors))
+        self.assertEqual(b_tensors, bc_tensors_t)
+        for (_, bt), (_, bct) in zip(b_tensors, bc_tensors_t):
+            self.assertEqual(bt.get_device(), bct.get_device())
+            self.assertIsInstance(bct, type(bt))
+
     @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
     def test_reduce_add(self):
         x = torch.randn(5, 5)
@@ -437,6 +505,32 @@ class TestCuda(TestCase):
         self.assertEqual(result.get_device(), 0)
         self.assertEqual(result.cpu(), x + y)
 
+    @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
+    def test_reduce_add_coalesced(self):
+        numel = 5
+        num_bytes = numel * 8
+        tensors = [
+            torch.randn(numel).long().cuda(),
+            torch.randn(numel).cuda(),
+            torch.randn(numel).long().cuda(),
+            torch.randn(numel).long().cuda(),
+            torch.randn(numel * 2).int().cuda(),  # int is 2x shorter
+            torch.randn(numel).cuda(),
+        ]
+        dup_tensors = [tensors, list(map(lambda t: t.cuda(1), tensors))]
+
+        r_tensors = list(map(comm.reduce_add, zip(*dup_tensors)))
+        for r, t in zip(r_tensors, tensors):
+            self.assertEqual(r.get_device(), t.get_device())
+            self.assertEqual(r, t * 2)
+            self.assertIsInstance(r, type(t))
+
+        rc_tensors = comm.reduce_add_coalesced(dup_tensors, buffer_size=num_bytes * 5 // 2)
+        self.assertEqual(r_tensors, rc_tensors)
+        for r, rc in zip(r_tensors, rc_tensors):
+            self.assertEqual(rc.get_device(), r.get_device())
+            self.assertIsInstance(rc, type(r))
+
     def _test_scatter(self, input, chunk_sizes=None, dim=0):
         if torch.cuda.device_count() < 2:
             raise unittest.SkipTest("only one GPU detected")
@@ -458,6 +552,9 @@ class TestCuda(TestCase):
     def test_scatter_cpu_dim(self):
         self._test_scatter(torch.randn(4, 4), dim=1)
 
+    def test_scatter_cpu_neg_dim(self):
+        self._test_scatter(torch.randn(4, 4), dim=-2)
+
     def test_scatter_cpu_sizes(self):
         self._test_scatter(torch.randn(6, 4), chunk_sizes=(2, 4))
 
@@ -467,6 +564,9 @@ class TestCuda(TestCase):
     def test_scatter_gpu_dim(self):
         self._test_scatter(torch.randn(4, 4).cuda(), dim=1)
 
+    def test_scatter_gpu_neg_dim(self):
+        self._test_scatter(torch.randn(4, 4).cuda(), dim=-2)
+
     def test_scatter_gpu_sizes(self):
         self._test_scatter(torch.randn(6, 4).cuda(), chunk_sizes=(2, 4))
 
@@ -497,11 +597,22 @@ class TestCuda(TestCase):
 
     def test_from_sequence(self):
         seq = [list(range(i * 4, i * 4 + 4)) for i in range(5)]
-        reference = torch.range(0, 19).resize_(5, 4)
+        reference = torch.arange(0, 20).resize_(5, 4)
         for t in types:
             cuda_type = get_gpu_type(t)
             self.assertEqual(cuda_type(seq), reference)
 
+    def test_torch_manual_seed_seeds_cuda_devices(self):
+        with freeze_rng_state():
+            x = torch.zeros(4, 4).float().cuda()
+            torch.manual_seed(2)
+            self.assertEqual(torch.cuda.initial_seed(), 2)
+            x.uniform_()
+            torch.manual_seed(2)
+            y = x.clone().uniform_()
+            self.assertEqual(x, y)
+            self.assertEqual(torch.cuda.initial_seed(), 2)
+
     def test_manual_seed(self):
         with freeze_rng_state():
             x = torch.zeros(4, 4).float().cuda()
@@ -530,7 +641,7 @@ class TestCuda(TestCase):
         self.assertIs(type(x_copy), type(x))
         self.assertEqual(x_copy.get_device(), x.get_device())
 
-    def test_serialization_empty(self):
+    def test_serialization_array_with_empty(self):
         x = [torch.randn(4, 4).cuda(), torch.cuda.FloatTensor()]
         with tempfile.NamedTemporaryFile() as f:
             torch.save(x, f)
@@ -654,6 +765,38 @@ class TestCuda(TestCase):
         self.assertTrue(event.query())
         self.assertGreater(start_event.elapsed_time(event), 0)
 
+    def test_record_stream(self):
+        cycles_per_ms = get_cycles_per_ms()
+
+        t = torch.FloatTensor([1, 2, 3, 4]).pin_memory()
+        result = torch.cuda.FloatTensor(t.size())
+        stream = torch.cuda.Stream()
+        ptr = [None]
+
+        # Performs the CPU->GPU copy in a background stream
+        def perform_copy():
+            with torch.cuda.stream(stream):
+                tmp = t.cuda(async=True)
+                ptr[0] = tmp.data_ptr()
+            torch.cuda.current_stream().wait_stream(stream)
+            tmp.record_stream(torch.cuda.current_stream())
+            torch.cuda._sleep(int(50 * cycles_per_ms))  # delay the copy
+            result.copy_(tmp)
+
+        perform_copy()
+        with torch.cuda.stream(stream):
+            tmp2 = torch.cuda.FloatTensor(t.size())
+            tmp2.zero_()
+            self.assertNotEqual(tmp2.data_ptr(), ptr[0], 'allocation re-used to soon')
+
+        self.assertEqual(result.tolist(), [1, 2, 3, 4])
+
+        # Check that the block will be re-used after the main stream finishes
+        torch.cuda.current_stream().synchronize()
+        with torch.cuda.stream(stream):
+            tmp3 = torch.cuda.FloatTensor(t.size())
+            self.assertEqual(tmp3.data_ptr(), ptr[0], 'allocation not re-used')
+
     def test_caching_pinned_memory(self):
         cycles_per_ms = get_cycles_per_ms()
 
@@ -673,40 +816,121 @@ class TestCuda(TestCase):
         self.assertNotEqual(t.data_ptr(), ptr, 'allocation re-used too soon')
         self.assertEqual(list(gpu_tensor), [1])
 
+    @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
+    def test_caching_pinned_memory_multi_gpu(self):
+        # checks that the events preventing pinned memory from being re-used
+        # too early are recorded on the correct GPU
+        cycles_per_ms = get_cycles_per_ms()
 
-for decl in tests:
-    for t in types:
-        tensor = t()
-        gpu_tensor = get_gpu_type(t)()
-        if len(decl) == 3:
-            name, constr, arg_constr = decl
-            desc = ''
-        elif len(decl) == 4:
-            name, constr, arg_constr, desc = decl
-        elif len(decl) == 5:
-            name, constr, arg_constr, desc, type_subset = decl
-            if t not in type_subset:
-                continue
+        t = torch.FloatTensor([1]).pin_memory()
+        ptr = t.data_ptr()
+        gpu_tensor0 = torch.cuda.FloatTensor([0], device=0)
+        gpu_tensor1 = torch.cuda.FloatTensor([0], device=1)
 
-        precision = custom_precision.get(name, TestCuda.precision)
-        for inplace in (True, False):
-            if inplace:
-                name_inner = name + '_'
-            else:
-                name_inner = name
-            if not hasattr(tensor, name_inner):
-                continue
-            if not hasattr(gpu_tensor, name_inner):
-                print("Ignoring {}, because it's not implemented by torch.cuda.{}".format(
-                    name_inner, gpu_tensor.__class__.__name__))
-                continue
+        with torch.cuda.device(1):
+            torch.cuda._sleep(int(50 * cycles_per_ms))  # delay the copy
+            gpu_tensor1.copy_(t, async=True)
 
-            test_name = 'test_' + t.__name__ + '_' + name_inner
-            if desc:
-                test_name += '_' + desc
+        del t
+        t = torch.FloatTensor([2]).pin_memory()
+        self.assertNotEqual(t.data_ptr(), ptr, 'allocation re-used too soon')
+
+        with torch.cuda.device(0):
+            gpu_tensor0.copy_(t, async=True)
+
+        self.assertEqual(gpu_tensor1[0], 1)
+        self.assertEqual(gpu_tensor0[0], 2)
+
+    @staticmethod
+    def _select_broadcastable_dims(dims_full=None):
+        return TestTorch._select_broadcastable_dims(dims_full)
+
+    def test_broadcast(self):
+        TestTorch._test_broadcast(self, lambda t: t.cuda())
+
+    def test_broadcast_fallback(self):
+        TestTorch._test_broadcast_fallback(self, lambda t: t.cuda())
+
+    def test_broadcast_fused_matmul(self):
+        TestTorch._test_broadcast_fused_matmul(self, lambda t: t.cuda())
+
+    def test_broadcast_batched_matmul(self):
+        TestTorch._test_broadcast_batched_matmul(self, lambda t: t.cuda())
+
+    def test_advancedindex(self):
+        TestTorch._test_advancedindex(self, lambda t: t.cuda())
+
+    def test_advancedindex_big(self):
+        TestTorch._test_advancedindex_big(self, lambda t: t.cuda())
+
+    def test_btrifact(self):
+        TestTorch._test_btrifact(self, lambda t: t.cuda())
+
+    def test_btrisolve(self):
+        TestTorch._test_btrisolve(self, lambda t: t.cuda())
+
+    def test_tensor_gather(self):
+        TestTorch._test_gather(self, lambda t: t.cuda(), False)
+
+    def test_tensor_scatter(self):
+        TestTorch._test_scatter_base(self, lambda t: t.cuda(), 'scatter_', test_bounds=False)
+
+    def test_tensor_scatterAdd(self):
+        TestTorch._test_scatter_base(self, lambda t: t.cuda(), 'scatter_add_', test_bounds=False)
+
+    def test_tensor_scatterFill(self):
+        TestTorch._test_scatter_base(self, lambda t: t.cuda(), 'scatter_', True, test_bounds=False)
+
+    def test_arange(self):
+        for t in ['IntTensor', 'LongTensor', 'FloatTensor', 'DoubleTensor']:
+            a = torch.cuda.__dict__[t]()
+            torch.arange(0, 10, out=a)
+            b = torch.__dict__[t]()
+            torch.arange(0, 10, out=b)
+            self.assertEqual(a, b.cuda())
+
+    def test_nvtx(self):
+        # Just making sure we can see the symbols
+        torch.cuda.nvtx.range_push("foo")
+        torch.cuda.nvtx.mark("bar")
+        torch.cuda.nvtx.range_pop()
+
+
+if HAS_CUDA:
+    for decl in tests:
+        for t in types:
+            tensor = t()
+            gpu_tensor = get_gpu_type(t)()
+            if len(decl) == 3:
+                name, constr, arg_constr = decl
+                desc = ''
+            elif len(decl) == 4:
+                name, constr, arg_constr, desc = decl
+            elif len(decl) == 5:
+                name, constr, arg_constr, desc, type_subset = decl
+                if t not in type_subset:
+                    continue
+
+            precision = custom_precision.get(name, TestCuda.precision)
+            for inplace in (True, False):
+                if inplace:
+                    name_inner = name + '_'
+                else:
+                    name_inner = name
+                if not hasattr(tensor, name_inner):
+                    continue
+                if not hasattr(gpu_tensor, name_inner):
+                    print("Ignoring {}, because it's not implemented by torch.cuda.{}".format(
+                        name_inner, gpu_tensor.__class__.__name__))
+                    continue
+
+                test_name = 'test_' + t.__name__ + '_' + name_inner
+                if desc:
+                    test_name += '_' + desc
+
+                assert not hasattr(TestCuda, test_name), "Duplicated test name: " + test_name
+                setattr(TestCuda, test_name, compare_cpu_gpu(constr, arg_constr, name_inner, t, precision))
 
-            assert not hasattr(TestCuda, test_name), "Duplicated test name: " + test_name
-            setattr(TestCuda, test_name, compare_cpu_gpu(constr, arg_constr, name_inner, t, precision))
 
 if __name__ == '__main__':
     run_tests()

test/test_dataloader.py

@@ -3,8 +3,8 @@ import sys
 import torch
 import traceback
 import unittest
-from torch.utils.data import Dataset, TensorDataset, DataLoader
-from common import TestCase, run_tests
+from torch.utils.data import Dataset, TensorDataset, DataLoader, ConcatDataset
+from common import TestCase, run_tests, TEST_NUMPY
 from common_nn import TEST_CUDA
 
 
@@ -31,6 +31,38 @@ class TestTensorDataset(TestCase):
             self.assertEqual(l[i], source[i][1])
 
 
+class TestConcatDataset(TestCase):
+
+    def test_concat_two_singletons(self):
+        result = ConcatDataset([[0], [1]])
+        self.assertEqual(2, len(result))
+        self.assertEqual(0, result[0])
+        self.assertEqual(1, result[1])
+
+    def test_concat_two_non_singletons(self):
+        result = ConcatDataset([[0, 1, 2, 3, 4],
+                                [5, 6, 7, 8, 9]])
+        self.assertEqual(10, len(result))
+        self.assertEqual(0, result[0])
+        self.assertEqual(5, result[5])
+
+    def test_concat_two_non_singletons_with_empty(self):
+        # Adding an empty dataset somewhere is correctly handled
+        result = ConcatDataset([[0, 1, 2, 3, 4],
+                                [],
+                                [5, 6, 7, 8, 9]])
+        self.assertEqual(10, len(result))
+        self.assertEqual(0, result[0])
+        self.assertEqual(5, result[5])
+
+    def test_concat_raises_index_error(self):
+        result = ConcatDataset([[0, 1, 2, 3, 4],
+                                [5, 6, 7, 8, 9]])
+        with self.assertRaises(IndexError):
+            # this one goes to 11
+            result[11]
+
+
 class ErrorDataset(Dataset):
 
     def __init__(self, size):
@@ -77,7 +109,7 @@ class TestDataLoader(TestCase):
         errors = 0
         while True:
             try:
-                it.next()
+                next(it)
             except NotImplementedError:
                 errors += 1
             except StopIteration:
@@ -91,6 +123,14 @@ class TestDataLoader(TestCase):
     def test_sequential_batch(self):
         self._test_sequential(DataLoader(self.dataset, batch_size=2))
 
+    def test_growing_dataset(self):
+        dataset = [torch.ones(4) for _ in range(4)]
+        dataloader_seq = DataLoader(dataset, shuffle=False)
+        dataloader_shuffle = DataLoader(dataset, shuffle=True)
+        dataset.append(torch.ones(4))
+        self.assertEqual(len(dataloader_seq), 5)
+        self.assertEqual(len(dataloader_shuffle), 5)
+
     @unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
     def test_sequential_pin_memory(self):
         loader = DataLoader(self.dataset, batch_size=2, pin_memory=True)
@@ -116,6 +156,29 @@ class TestDataLoader(TestCase):
     def test_shuffle_batch_workers(self):
         self._test_shuffle(DataLoader(self.dataset, batch_size=2, shuffle=True, num_workers=4))
 
+    def _test_batch_sampler(self, **kwargs):
+        # [(0, 1), (2, 3, 4), (5, 6), (7, 8, 9), ...]
+        batches = []
+        for i in range(0, 100, 5):
+            batches.append(tuple(range(i, i + 2)))
+            batches.append(tuple(range(i + 2, i + 5)))
+
+        dl = DataLoader(self.dataset, batch_sampler=batches, **kwargs)
+        self.assertEqual(len(dl), 40)
+        for i, (input, _target) in enumerate(dl):
+            if i % 2 == 0:
+                offset = i * 5 // 2
+                self.assertEqual(len(input), 2)
+                self.assertEqual(input, self.data[offset:offset + 2])
+            else:
+                offset = i * 5 // 2
+                self.assertEqual(len(input), 3)
+                self.assertEqual(input, self.data[offset:offset + 3])
+
+    def test_batch_sampler(self):
+        self._test_batch_sampler()
+        self._test_batch_sampler(num_workers=4)
+
     @unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
     def test_shuffle_pin_memory(self):
         loader = DataLoader(self.dataset, batch_size=2, shuffle=True, num_workers=4, pin_memory=True)
@@ -123,6 +186,22 @@ class TestDataLoader(TestCase):
             self.assertTrue(input.is_pinned())
             self.assertTrue(target.is_pinned())
 
+    @unittest.skipIf(not TEST_NUMPY, "numpy unavailable")
+    def test_numpy(self):
+        import numpy as np
+
+        class TestDataset(torch.utils.data.Dataset):
+            def __getitem__(self, i):
+                return np.ones((2, 3, 4)) * i
+
+            def __len__(self):
+                return 1000
+
+        loader = DataLoader(TestDataset(), batch_size=12)
+        batch = next(iter(loader))
+        self.assertIsInstance(batch, torch.DoubleTensor)
+        self.assertEqual(batch.size(), torch.Size([12, 2, 3, 4]))
+
     def test_error(self):
         self._test_error(DataLoader(ErrorDataset(100), batch_size=2, shuffle=True))
 
@@ -157,6 +236,102 @@ class TestDataLoader(TestCase):
         check_len(DataLoader(self.dataset, batch_size=2), 50)
         check_len(DataLoader(self.dataset, batch_size=3), 34)
 
+    @unittest.skipIf(not TEST_NUMPY, "numpy unavailable")
+    def test_numpy_scalars(self):
+        import numpy as np
+
+        class ScalarDataset(torch.utils.data.Dataset):
+            def __init__(self, dtype):
+                self.dtype = dtype
+
+            def __getitem__(self, i):
+                return self.dtype()
+
+            def __len__(self):
+                return 4
+
+        dtypes = {
+            np.float64: torch.DoubleTensor,
+            np.float32: torch.FloatTensor,
+            np.float16: torch.HalfTensor,
+            np.int64: torch.LongTensor,
+            np.int32: torch.IntTensor,
+            np.int16: torch.ShortTensor,
+            np.int8: torch.CharTensor,
+            np.uint8: torch.ByteTensor,
+        }
+        for dt, tt in dtypes.items():
+            dset = ScalarDataset(dt)
+            loader = DataLoader(dset, batch_size=2)
+            batch = next(iter(loader))
+            self.assertIsInstance(batch, tt)
+
+
+class StringDataset(Dataset):
+    def __init__(self):
+        self.s = '12345'
+
+    def __len__(self):
+        return len(self.s)
+
+    def __getitem__(self, ndx):
+        return (self.s[ndx], ndx)
+
+
+class TestStringDataLoader(TestCase):
+    def setUp(self):
+        self.dataset = StringDataset()
+
+    @unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
+    def test_shuffle_pin_memory(self):
+        loader = DataLoader(self.dataset, batch_size=2, shuffle=True, num_workers=4, pin_memory=True)
+        for batch_ndx, (s, n) in enumerate(loader):
+            self.assertIsInstance(s[0], str)
+            self.assertTrue(n.is_pinned())
+
+
+class DictDataset(Dataset):
+    def __len__(self):
+        return 4
+
+    def __getitem__(self, ndx):
+        return {
+            'a_tensor': torch.Tensor(4, 2).fill_(ndx),
+            'another_dict': {
+                'a_number': ndx,
+            },
+        }
+
+
+class TestDictDataLoader(TestCase):
+    def setUp(self):
+        self.dataset = DictDataset()
+
+    def test_sequential_batch(self):
+        loader = DataLoader(self.dataset, batch_size=2, shuffle=False)
+        batch_size = loader.batch_size
+        for i, sample in enumerate(loader):
+            idx = i * batch_size
+            self.assertEqual(set(sample.keys()), {'a_tensor', 'another_dict'})
+            self.assertEqual(set(sample['another_dict'].keys()), {'a_number'})
+
+            t = sample['a_tensor']
+            self.assertEqual(t.size(), torch.Size([batch_size, 4, 2]))
+            self.assertTrue((t[0] == idx).all())
+            self.assertTrue((t[1] == idx + 1).all())
+
+            n = sample['another_dict']['a_number']
+            self.assertEqual(n.size(), torch.Size([batch_size]))
+            self.assertEqual(n[0], idx)
+            self.assertEqual(n[1], idx + 1)
+
+    @unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
+    def test_pin_memory(self):
+        loader = DataLoader(self.dataset, batch_size=2, pin_memory=True)
+        for batch_ndx, sample in enumerate(loader):
+            self.assertTrue(sample['a_tensor'].is_pinned())
+            self.assertTrue(sample['another_dict']['a_number'].is_pinned())
+
 
 if __name__ == '__main__':
     run_tests()

test/test_distributed.py

@@ -14,7 +14,12 @@ from common import TestCase
 BACKEND = os.environ['BACKEND']
 TEMP_DIR = os.environ['TEMP_DIR']
 MASTER_PORT = '29500'
-MASTER_ADDR = '127.0.0.1:' + MASTER_PORT
+MASTER_ADDR = '127.0.0.1'
+
+
+if not dist.is_available():
+    print('Distributed not available, skipping tests')
+    sys.exit(0)
 
 
 @contextmanager
@@ -64,7 +69,7 @@ class Barrier(object):
                         data = f.read()
                         if int(data) >= cls.barrier_id:
                             arrived += 1
-            if arrived == dist.get_num_processes():
+            if arrived == dist.get_world_size():
                 break
 
             if time.time() - start_time > timeout:
@@ -87,7 +92,7 @@ class _DistTestBase(object):
         return (group, group_id, rank)
 
     def _init_global_test(self):
-        group = [i for i in range(0, dist.get_num_processes())]
+        group = [i for i in range(0, dist.get_world_size())]
         group_id = dist.group.WORLD
         rank = dist.get_rank()
         return (group, group_id, rank)
@@ -96,7 +101,7 @@ class _DistTestBase(object):
     def test_get_rank(self):
         test_dir = os.path.join(TEMP_DIR, 'test_dir')
         pid = str(os.getpid())
-        num_processes = dist.get_num_processes()
+        num_processes = dist.get_world_size()
         with open(os.path.join(test_dir, pid), 'w') as f:
             f.write(str(dist.get_rank()))
 
@@ -117,15 +122,16 @@ class _DistTestBase(object):
         self._barrier()
 
     # SEND RECV
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support send/recv")
     def test_send_recv(self):
         rank = dist.get_rank()
         tensor = _build_tensor(rank + 1)
-        for dest in range(0, dist.get_num_processes()):
+        for dest in range(0, dist.get_world_size()):
             if dest == rank:
                 continue
             dist.send(tensor, dest)
 
-        for src in range(0, dist.get_num_processes()):
+        for src in range(0, dist.get_world_size()):
             if src == rank:
                 continue
             tensor = _build_tensor(src + 1, value=-1)
@@ -136,29 +142,32 @@ class _DistTestBase(object):
         self._barrier()
 
     # SEND RECV ANY SOURCE
+    @unittest.skipIf(BACKEND == 'gloo',
+                     "Gloo does not support send/recv from any source")
     def test_send_recv_any_source(self):
         rank = dist.get_rank()
         tensor = _build_tensor(10, rank)
-        for dest in range(0, dist.get_num_processes()):
+        for dest in range(0, dist.get_world_size()):
             if dest == rank:
                 continue
             dist.send(tensor, dest)
 
         recv_ranks = set()
-        for src in range(0, dist.get_num_processes()):
+        for src in range(0, dist.get_world_size()):
             if src == rank:
                 continue
             tensor = _build_tensor(10, value=-1)
             dist.recv(tensor)
             recv_ranks.add(tensor.resize_(1)[0])
 
-        self.assertEqual(len(recv_ranks), dist.get_num_processes() - 1)
+        self.assertEqual(len(recv_ranks), dist.get_world_size() - 1)
         self._barrier()
 
     # ISEND
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support isend")
     def test_isend(self):
         rank = dist.get_rank()
-        world_size = dist.get_num_processes()
+        world_size = dist.get_world_size()
 
         if rank == 0:
             requests = [
@@ -175,9 +184,10 @@ class _DistTestBase(object):
         self._barrier()
 
     # IRECV
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support irecv")
     def test_irecv(self):
         rank = dist.get_rank()
-        world_size = dist.get_num_processes()
+        world_size = dist.get_world_size()
 
         if rank == 0:
             expected_tensors = [_build_tensor(src, -1) for src in range(1, world_size)]
@@ -196,13 +206,17 @@ class _DistTestBase(object):
         self._barrier()
 
     # BROADCAST
-    def _test_broadcast_helper(self, group, group_id, rank):
+    def _test_broadcast_helper(self, group, group_id, rank, cuda=False):
         for src in group:
             expected_tensor = _build_tensor(src + 1)
+            if cuda:
+                expected_tensor = expected_tensor.cuda()
             if rank == src:
                 dist.broadcast(expected_tensor, src, group_id)
             else:
                 tensor = _build_tensor(src + 1, -1)
+                if cuda:
+                    tensor = tensor.cuda()
                 dist.broadcast(tensor, src, group_id)
                 self.assertEqual(tensor, expected_tensor)
 
@@ -212,6 +226,11 @@ class _DistTestBase(object):
         group, group_id, rank = self._init_global_test()
         self._test_broadcast_helper(group, group_id, rank)
 
+    @unittest.skipIf(BACKEND != 'gloo', "Only Gloo backend supports CUDA allReduce")
+    def test_broadcast_cuda(self):
+        group, group_id, rank = self._init_global_test()
+        self._test_broadcast_helper(group, group_id, rank, True)
+
     def test_broadcast_group(self):
         group, group_id, rank = self._init_group_test()
         self._test_broadcast_helper(group, group_id, rank)
@@ -229,12 +248,14 @@ class _DistTestBase(object):
 
         self._barrier()
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_sum(self):
         group, group_id, rank = self._init_global_test()
         self._test_reduce_helper(
             group, group_id, rank, dist.reduce_op.SUM, 2, 10, 2 + (10 * (len(group) - 1))
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_product(self):
         group, group_id, rank = self._init_global_test()
         self._test_reduce_helper(
@@ -242,24 +263,28 @@ class _DistTestBase(object):
             2, 10, reduce((lambda x, y: x * y), [10] * (len(group) - 1), 2)
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_min(self):
         group, group_id, rank = self._init_global_test()
         self._test_reduce_helper(
             group, group_id, rank, dist.reduce_op.MIN, 1010, 1, 1
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_max(self):
         group, group_id, rank = self._init_global_test()
         self._test_reduce_helper(
             group, group_id, rank, dist.reduce_op.MAX, -1, 10, 10
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_group_sum(self):
         group, group_id, rank = self._init_group_test()
         self._test_reduce_helper(
             group, group_id, rank, dist.reduce_op.SUM, 2, 10, 2 + (10 * (len(group) - 1))
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_group_product(self):
         group, group_id, rank = self._init_group_test()
         self._test_reduce_helper(
@@ -267,12 +292,14 @@ class _DistTestBase(object):
             2, 10, reduce((lambda x, y: x * y), [10] * (len(group) - 1), 2)
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_group_min(self):
         group, group_id, rank = self._init_group_test()
         self._test_reduce_helper(
             group, group_id, rank, dist.reduce_op.MIN, 1010, 1, 1
         )
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support reduce")
     def test_reduce_group_max(self):
         group, group_id, rank = self._init_group_test()
         self._test_reduce_helper(
@@ -280,14 +307,19 @@ class _DistTestBase(object):
         )
 
     # ALL REDUCE
-    def _test_all_reduce_helper(self, group, group_id, rank, op, master_value, worker_value, expected_value):
+    def _test_all_reduce_helper(self, group, group_id, rank, op, master_value,
+                                worker_value, expected_value, cuda=False):
         for src in group:
             if rank == src:
                 tensor = _build_tensor(src + 1).fill_(master_value)
+                if cuda:
+                    tensor = tensor.cuda()
                 dist.all_reduce(tensor, op, group_id)
                 self.assertEqual(tensor, _build_tensor(src + 1, expected_value))
             else:
                 tensor = _build_tensor(src + 1).fill_(worker_value)
+                if cuda:
+                    tensor = tensor.cuda()
                 dist.all_reduce(tensor, op, group_id)
                 self.assertEqual(tensor, _build_tensor(src + 1, expected_value))
 
@@ -299,6 +331,13 @@ class _DistTestBase(object):
             group, group_id, rank, dist.reduce_op.SUM, 2, 10, 2 + (10 * (len(group) - 1))
         )
 
+    @unittest.skipIf(BACKEND != 'gloo', "Only Gloo backend supports CUDA allReduce")
+    def test_all_reduce_sum_cuda(self):
+        group, group_id, rank = self._init_global_test()
+        self._test_all_reduce_helper(
+            group, group_id, rank, dist.reduce_op.SUM, 2, 10, 2 + (10 * (len(group) - 1)), True
+        )
+
     def test_all_reduce_product(self):
         group, group_id, rank = self._init_global_test()
         self._test_all_reduce_helper(
@@ -348,20 +387,18 @@ class _DistTestBase(object):
         for dest in group:
             tensor = _build_tensor(dest + 1, -1)
             expected_tensor = _build_tensor(dest + 1, rank)
-            if rank == dest:
-                tensors = [_build_tensor(dest + 1, i) for i in group]
-                dist.scatter_send(tensors, tensor, group_id)
-                self.assertEqual(tensor, expected_tensor)
-            else:
-                dist.scatter_recv(tensor, dest, group_id)
-                self.assertEqual(tensor, expected_tensor)
+            tensors = [_build_tensor(dest + 1, i) for i in group] if rank == dest else []
+            dist.scatter(tensor, src=dest, scatter_list=tensors, group=group_id)
+            self.assertEqual(tensor, expected_tensor)
 
         self._barrier()
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support scatter")
     def test_scatter(self):
         group, group_id, rank = self._init_global_test()
         self._test_scatter_helper(group, group_id, rank)
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support scatter")
     def test_scatter_group(self):
         group, group_id, rank = self._init_group_test()
         self._test_scatter_helper(group, group_id, rank)
@@ -370,22 +407,21 @@ class _DistTestBase(object):
     def _test_gather_helper(self, group, group_id, rank):
         for dest in group:
             tensor = _build_tensor(dest + 1, rank)
+            tensors = [_build_tensor(dest + 1, -1) for i in group] if rank == dest else []
+            dist.gather(tensor, dst=dest, gather_list=tensors, group=group_id)
             if rank == dest:
-                tensors = [_build_tensor(dest + 1, -1) for i in group]
-                dist.gather_recv(tensors, tensor, group_id)
-
                 expected_tensors = [_build_tensor(dest + 1, i) for i in group]
                 for t1, t2 in zip(tensors, expected_tensors):
                     self.assertEqual(t1, t2)
-            else:
-                dist.gather_send(tensor, dest, group_id)
 
         self._barrier()
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support gather")
     def test_gather(self):
         group, group_id, rank = self._init_global_test()
         self._test_gather_helper(group, group_id, rank)
 
+    @unittest.skipIf(BACKEND == 'gloo', "Gloo does not support gather")
     def test_gather_group(self):
         group, group_id, rank = self._init_group_test()
         self._test_gather_helper(group, group_id, rank)
@@ -437,13 +473,13 @@ class _DistTestBase(object):
         group, group_id, rank = self._init_group_test()
         self._test_barrier_helper(group, group_id, rank)
 
-if BACKEND == 'tcp':
+if BACKEND == 'tcp' or BACKEND == 'gloo':
     WORLD_SIZE = os.environ['WORLD_SIZE']
 
-    class TestTCP(TestCase, _DistTestBase):
+    class TestTCPOrGloo(TestCase, _DistTestBase):
 
         MANAGER_PROCESS_RANK = -1
-        JOIN_TIMEOUT = 5
+        JOIN_TIMEOUT = 10
 
         @staticmethod
         def manager_join(fn):
@@ -486,7 +522,11 @@ if BACKEND == 'tcp':
 
         def _run(self, rank):
             self.rank = rank
-            dist.init_process_group(backend=BACKEND)
+            try:
+                dist.init_process_group(backend=BACKEND)
+            except RuntimeError as e:
+                if 'recompile' in e.args[0]:
+                    sys.exit(0)
             # self.id() == e.g. '__main__.TestDistributed.test_get_rank'
             # We're retreiving a corresponding test and executing it.
             getattr(self, self.id().split(".")[2])()

test/test_legacy_nn.py

@@ -184,16 +184,16 @@ tests = [
     OldModuleTest(nn.Sum,
                   (1,),
                   input_size=(2, 4, 5),
-                  reference_fn=lambda i, _: i.sum(1).squeeze(1)),
+                  reference_fn=lambda i, _: i.sum(1, keepdim=False)),
     OldModuleTest(nn.Sum,
                   (1, True),
                   input_size=(2, 4, 5),
-                  reference_fn=lambda i, _: i.sum(1).div(i.size(1)).squeeze(1),
+                  reference_fn=lambda i, _: i.sum(1, keepdim=False).div(i.size(1)),
                   desc='sizeAverage'),
     OldModuleTest(nn.Mean,
                   (1,),
                   input_size=(2, 4, 5),
-                  reference_fn=lambda i, _: torch.mean(i, 1).squeeze(1)),
+                  reference_fn=lambda i, _: torch.mean(i, 1, keepdim=False)),
     OldModuleTest(lambda: nn.Sequential().add(nn.GradientReversal()).add(nn.GradientReversal()),
                   input_size=(4, 3, 2, 2),
                   fullname='GradientReversal'),
@@ -233,19 +233,19 @@ tests = [
                   reference_fn=lambda i, _: torch.bmm(i[0], i[1].view(i[1].size(0), i[1].size(1), 1)).squeeze()),
     OldModuleTest(nn.Max,
                   input_size=(4, 5, 3),
-                  reference_fn=lambda i, _: torch.max(i, 0)[0].squeeze()),
+                  reference_fn=lambda i, _: torch.max(i, 0, False)[0]),
     OldModuleTest(nn.Max,
                   (1,),
                   input_size=(4, 5, 3),
-                  reference_fn=lambda i, _: torch.max(i, 1)[0].squeeze(),
+                  reference_fn=lambda i, _: torch.max(i, 1, False)[0],
                   desc='with_dimension'),
     OldModuleTest(nn.Min,
                   input_size=(4, 5, 3),
-                  reference_fn=lambda i, _: torch.min(i, 0)[0].squeeze()),
+                  reference_fn=lambda i, _: torch.min(i, 0, False)[0]),
     OldModuleTest(nn.Min,
                   (1,),
                   input_size=(4, 5, 3),
-                  reference_fn=lambda i, _: torch.min(i, 1)[0].squeeze(),
+                  reference_fn=lambda i, _: torch.min(i, 1, False)[0],
                   desc='with_dimension'),
     OldModuleTest(nn.MixtureTable,
                   tuple(),
@@ -483,14 +483,14 @@ tests = [
                   input_size=(1, 2, 4, 4, 4)),
     OldModuleTest(nn.VolumetricMaxPooling,
                   (2, 2, 2),
-                  input_size=(2, 3, 5, 5, 5)),
+                  input=(torch.randn(2, 3, 5, 5, 5) * 1000)),
     OldModuleTest(nn.VolumetricMaxPooling,
                   (2, 2, 2, 2, 2, 2),
-                  input_size=(2, 3, 5, 5, 5),
+                  input=(torch.randn(2, 3, 5, 5, 5) * 1000),
                   desc='stride'),
     OldModuleTest(nn.VolumetricMaxPooling,
                   (2, 2, 2, 2, 2, 2, 1, 1, 1),
-                  input_size=(2, 3, 5, 5, 5),
+                  input=(torch.randn(2, 3, 5, 5, 5) * 1000),
                   desc='stride_padding'),
     OldModuleTest(nn.VolumetricReplicationPadding,
                   (1, 2, 3, 4, 5, 6),
@@ -532,7 +532,7 @@ for p in (1, 2, 1.5):
                       (p,),
                       input_size=(4, 5),
                       # Eh, we need to use p as a default, so it's passed by value
-                      reference_fn=lambda i, _, p=p: i.div(i.norm(p, 1).expand_as(i)),
+                      reference_fn=lambda i, _, p=p: i.div(i.norm(p, 1, True).expand_as(i)),
                       desc=str(p)),
     )
 for p in range(1, 4 + 1):
@@ -807,14 +807,14 @@ class TestNN(NNTestCase):
         str(m)
 
         output = m.forward(input)
-        output2 = input.sum(1).expand(4, 5).repeat(num_modules, 1)
+        output2 = input.sum(1, True).expand(4, 5).repeat(num_modules, 1)
         self.assertEqual(output2, output)
 
         gradInput = m.backward(input, torch.ones(output2.size()))
         gradInput2 = torch.ones(4, 2).fill_(num_modules * 5)
         self.assertEqual(gradInput, gradInput2)
 
-        gradWeight = input.sum(0).expand(5, 2)
+        gradWeight = input.sum(0, keepdim=True).expand(5, 2)
         for l in linears:
             self.assertEqual(gradWeight, l.gradWeight)
 
@@ -884,8 +884,8 @@ class TestNN(NNTestCase):
         output2 = [input, input, input]
         self.assertEqual(output2, output)
         gradInput = module.backward(input, gradOutput)
-        gradInput2 = [_gradOutput[0].sum(0).squeeze(0), _gradOutput[1].sum(
-            0).squeeze(0), [_gradOutput[2].sum(0).squeeze(0)]]
+        gradInput2 = [_gradOutput[0].sum(0, keepdim=False), _gradOutput[1].sum(
+            0, keepdim=False), [_gradOutput[2].sum(0, keepdim=False)]]
         self.assertTrue(isinstance(gradInput, list))
         self.assertFalse(isinstance(gradInput[0], list))
         self.assertFalse(isinstance(gradInput[1], list))
@@ -1251,6 +1251,8 @@ class TestNN(NNTestCase):
                 self.assertIsInstance(module, type(reference))
 
 
+prepare_tests()
+
+
 if __name__ == '__main__':
-    prepare_tests()
     run_tests()

test/test_multiprocessing.py

@@ -75,13 +75,13 @@ def autograd_sharing(queue, ready, master_modified):
     ready.set()
     master_modified.wait()
 
-    expected_var = torch.range(1, 25).view(5, 5)
+    expected_var = torch.arange(1, 26).view(5, 5)
     expected_var[0, 0] = 1000
     is_ok = var.data.equal(expected_var)
     var.data[:] = torch.ones(5, 5)
 
-    is_ok &= var.grad.data.equal(torch.zeros(5, 5))
-    var.grad.data[:] = torch.ones(5, 5)
+    is_ok &= var.grad is None
+    var._grad = Variable(torch.ones(5, 5), requires_grad=False)
 
     queue.put(is_ok)
 
@@ -112,9 +112,10 @@ class leak_checker(object):
             # test is no more than 4 higher than the 10th available at the
             # start. This attempts to catch file descriptor leaks, but allows
             # one-off initialization that may use up a file descriptor
-            available_fds = self._get_next_fds(10)
-            self.test_case.assertLessEqual(
-                available_fds[-1] - self.next_fds[-1], 5)
+            # TODO: Disabled because this check is too flaky
+            # available_fds = self._get_next_fds(10)
+            # self.test_case.assertLessEqual(
+            #     available_fds[-1] - self.next_fds[-1], 5)
             self.test_case.assertFalse(self.has_shm_files())
         return False
 
@@ -149,9 +150,6 @@ class leak_checker(object):
 
 class TestMultiprocessing(TestCase):
 
-    def __init__(self, *args, **kwargs):
-        super(TestMultiprocessing, self).__init__(*args, **kwargs)
-
     def _test_sharing(self, ctx=mp, type=torch.FloatTensor, repeat=1):
         def test_fill():
             x = torch.zeros(5, 5).type(type)
@@ -160,9 +158,11 @@ class TestMultiprocessing(TestCase):
             data = [x, x[:, 1]]
             q.put(data)
             p = ctx.Process(target=simple_fill, args=(q, e))
+            p.daemon = True
             lc.check_pid(p.pid)
             p.start()
-            e.wait()
+            e.wait(10)
+            self.assertTrue(e.is_set())
             self.assertTrue(data[0].eq(4).all())
             self.assertTrue(data[1].eq(4).all())
             p.join(1)
@@ -172,6 +172,7 @@ class TestMultiprocessing(TestCase):
             q = ctx.Queue()
             e = ctx.Event()
             p = ctx.Process(target=send_tensor, args=(q, e, type))
+            p.daemon = True
             lc.check_pid(p.pid)
             p.start()
             t1 = q.get()
@@ -183,7 +184,7 @@ class TestMultiprocessing(TestCase):
             self.assertFalse(p.is_alive())
 
         with leak_checker(self) as lc:
-            for i in range(repeat):
+            for _ in range(repeat):
                 test_fill()
                 test_receive()
 
@@ -193,7 +194,7 @@ class TestMultiprocessing(TestCase):
             data = [x.storage(), x.storage()[1:4], x, x[2], x[:, 1]]
             q = ctx.Queue()
             q.put(data)
-            new_data = q.get()
+            new_data = q.get(timeout=1)
             self.assertEqual(new_data, data, 0)
             storage_cdata = data[0]._cdata
             self.assertEqual(new_data[0]._cdata, storage_cdata)
@@ -264,15 +265,15 @@ class TestMultiprocessing(TestCase):
             q.get()
 
         with fs_sharing(), leak_checker(self) as lc:
-            for i in range(TEST_REPEATS):
+            for _ in range(TEST_REPEATS):
                 queue_put()
 
     def test_inherit_tensor(self):
         class SubProcess(mp.Process):
-
             def __init__(self, tensor):
                 super(SubProcess, self).__init__()
                 self.tensor = tensor
+                self.daemon = True
 
             def run(self):
                 self.tensor.add_(3)
@@ -280,7 +281,7 @@ class TestMultiprocessing(TestCase):
         t = torch.zeros(5, 5)
         p = SubProcess(t.share_memory_())
         p.start()
-        p.join()
+        p.join(1)
         self.assertEqual(t, torch.ones(5, 5) * 3, 0)
 
     @unittest.skipIf(not TEST_CUDA_IPC, 'CUDA IPC not available')
@@ -296,7 +297,8 @@ class TestMultiprocessing(TestCase):
         ctx = mp.get_context('spawn')
         tensors = []
         for i in range(5):
-            tensors += [torch.range(i * 5, (i * 5) + 4).cuda()]
+            device = i % 2
+            tensors += [torch.arange(i * 5, (i + 1) * 5).cuda(device)]
 
         inq = ctx.Queue()
         outq = ctx.Queue()
@@ -311,8 +313,8 @@ class TestMultiprocessing(TestCase):
 
         for i, tensor in enumerate(tensors):
             v, device, tensor_size, storage_size = results[i]
-            self.assertEqual(v, torch.range(i * 5, (i * 5) + 4).sum())
-            self.assertEqual(device, 0)
+            self.assertEqual(v, torch.arange(i * 5, (i + 1) * 5).sum())
+            self.assertEqual(device, i % 2)
             self.assertEqual(tensor_size, 5)
             self.assertEqual(storage_size, 5)
 
@@ -357,8 +359,9 @@ class TestMultiprocessing(TestCase):
         master_modified = mp.Event()
         queue = mp.Queue()
         p = mp.Process(target=autograd_sharing, args=(queue, ready, master_modified))
+        p.daemon = True
         p.start()
-        var.grad.data.zero_()
+        var._grad = Variable(torch.zeros(5, 5), requires_grad=False)
         queue.put(var)
 
         ready.wait()
@@ -371,7 +374,8 @@ class TestMultiprocessing(TestCase):
 
         self.assertEqual(var.data, torch.ones(5, 5))
         self.assertEqual(var.grad.data, torch.ones(5, 5) * 4)
-        p.join()
+        p.join(1)
+        self.assertFalse(p.is_alive())
 
     def test_variable_sharing(self):
         configs = [
@@ -380,15 +384,19 @@ class TestMultiprocessing(TestCase):
             (False, True),
         ]
         for requires_grad, volatile in configs:
-            var = Variable(torch.range(1, 25).view(5, 5),
+            var = Variable(torch.arange(1, 26).view(5, 5),
                            requires_grad=requires_grad,
                            volatile=volatile)
             self._test_autograd_sharing(var)
 
     def test_parameter_sharing(self):
-        param = Parameter(torch.range(1, 25).view(5, 5))
+        param = Parameter(torch.arange(1, 26).view(5, 5))
         self._test_autograd_sharing(param)
 
+    def test_empty_shared(self):
+        t = torch.Tensor()
+        t.share_memory_()
+
     def _test_is_shared(self):
         t = torch.randn(5, 5)
         self.assertFalse(t.is_shared())

test/test_nccl.py

@@ -6,12 +6,10 @@ import torch.cuda
 
 from common import TestCase, run_tests
 
-if not torch.cuda.is_available():
-    print('CUDA not available, skipping tests')
-    import sys
-    sys.exit()
-
 nGPUs = torch.cuda.device_count()
+if nGPUs == 0:
+    print('CUDA not available, skipping tests')
+    TestCase = object  # noqa: F811
 
 
 class TestNCCL(TestCase):

test/test_optim.py

@@ -4,8 +4,11 @@ from copy import deepcopy
 import torch
 import torch.optim as optim
 import torch.legacy.optim as old_optim
+import torch.nn.functional as F
+from torch.optim import SGD
 from torch.autograd import Variable
-
+from torch import sparse
+from torch.optim.lr_scheduler import LambdaLR, StepLR, MultiStepLR, ExponentialLR, ReduceLROnPlateau
 from common import TestCase, run_tests
 
 
@@ -58,6 +61,49 @@ class TestOptim(TestCase):
 
         self.assertLessEqual(params.data.dist(solution), initial_dist)
 
+    def _test_rosenbrock_sparse(self, constructor):
+        params_t = torch.Tensor([1.5, 1.5])
+
+        params = Variable(torch.Tensor([1.5, 1.5]), requires_grad=True)
+        params_c = Variable(torch.Tensor([1.5, 1.5]), requires_grad=True)
+        optimizer = constructor([params])
+        optimizer_c = constructor([params_c])
+
+        solution = torch.Tensor([1, 1])
+        initial_dist = params.data.dist(solution)
+
+        def eval(params, sparse_grad, w):
+            # Depending on w, provide only the x or y gradient
+            optimizer.zero_grad()
+            loss = rosenbrock(params)
+            loss.backward()
+            grad = drosenbrock(params.data)
+            # NB: We torture test the optimizer by returning an
+            # uncoalesced sparse tensor
+            if w:
+                i = torch.LongTensor([[0, 0]])
+                x = grad[0]
+                v = torch.DoubleTensor([x / 4., x - x / 4.])
+            else:
+                i = torch.LongTensor([[1, 1]])
+                y = grad[1]
+                v = torch.DoubleTensor([y - y / 4., y / 4.])
+            x = sparse.DoubleTensor(i, v, torch.Size([2]))
+            if sparse_grad:
+                params.grad.data = x
+            else:
+                params.grad.data = x.to_dense()
+            return loss
+
+        for i in range(2000):
+            # Do cyclic coordinate descent
+            w = i % 2
+            optimizer.step(functools.partial(eval, params, True, w))
+            optimizer_c.step(functools.partial(eval, params_c, False, w))
+            self.assertEqual(params.data, params_c.data)
+
+        self.assertLessEqual(params.data.dist(solution), initial_dist)
+
     def _test_basic_cases_template(self, weight, bias, input, constructor):
         weight = Variable(weight, requires_grad=True)
         bias = Variable(bias, requires_grad=True)
@@ -155,6 +201,9 @@ class TestOptim(TestCase):
     def _build_params_dict(self, weight, bias, **kwargs):
         return [dict(params=[weight]), dict(params=[bias], **kwargs)]
 
+    def _build_params_dict_single(self, weight, bias, **kwargs):
+        return [dict(params=bias, **kwargs)]
+
     def test_sgd(self):
         self._test_rosenbrock(
             lambda params: optim.SGD(params, lr=1e-3),
@@ -174,6 +223,11 @@ class TestOptim(TestCase):
                 self._build_params_dict(weight, bias, lr=1e-2),
                 lr=1e-3)
         )
+        self._test_basic_cases(
+            lambda weight, bias: optim.SGD(
+                self._build_params_dict_single(weight, bias, lr=1e-2),
+                lr=1e-3)
+        )
 
     def test_adam(self):
         self._test_rosenbrock(
@@ -236,6 +290,11 @@ class TestOptim(TestCase):
                 lr=1e-1)
         )
 
+    def test_adagrad_sparse(self):
+        self._test_rosenbrock_sparse(
+            lambda params: optim.Adagrad(params, lr=1e-1)
+        )
+
     def test_adamax(self):
         self._test_rosenbrock(
             lambda params: optim.Adamax(params, lr=1e-1),
@@ -343,5 +402,157 @@ class TestOptim(TestCase):
             optim.SGD(Variable(torch.randn(5, 5)), lr=3)
 
 
+class SchedulerTestNet(torch.nn.Module):
+    def __init__(self):
+        super(SchedulerTestNet, self).__init__()
+        self.conv1 = torch.nn.Conv2d(1, 1, 1)
+        self.conv2 = torch.nn.Conv2d(1, 1, 1)
+
+    def forward(self, x):
+        return self.conv2(F.relu(self.conv1(x)))
+
+
+class TestLRScheduler(TestCase):
+    def setUp(self):
+        self.net = SchedulerTestNet()
+        self.opt = SGD(
+            [{'params': self.net.conv1.parameters()}, {'params': self.net.conv2.parameters(), 'lr': 0.5}],
+            lr=0.05)
+
+    def test_step_lr(self):
+        # lr = 0.05     if epoch < 3
+        # lr = 0.005    if 30 <= epoch < 6
+        # lr = 0.0005   if epoch >= 9
+        single_targets = [0.05] * 3 + [0.005] * 3 + [0.0005] * 3 + [0.00005] * 3
+        targets = [single_targets, list(map(lambda x: x * 10, single_targets))]
+        scheduler = StepLR(self.opt, gamma=0.1, step_size=3)
+        epochs = 10
+        self._test(scheduler, targets, epochs)
+
+    def test_multi_step_lr(self):
+        # lr = 0.05     if epoch < 2
+        # lr = 0.005    if 2 <= epoch < 5
+        # lr = 0.0005   if epoch < 9
+        # lr = 0.00005   if epoch >= 9
+        single_targets = [0.05] * 2 + [0.005] * 3 + [0.0005] * 4 + [0.00005] * 3
+        targets = [single_targets, list(map(lambda x: x * 10, single_targets))]
+        scheduler = MultiStepLR(self.opt, gamma=0.1, milestones=[2, 5, 9])
+        epochs = 10
+        self._test(scheduler, targets, epochs)
+
+    def test_exp_lr(self):
+        single_targets = [0.05 * (0.9 ** x) for x in range(10)]
+        targets = [single_targets, list(map(lambda x: x * 10, single_targets))]
+        scheduler = ExponentialLR(self.opt, gamma=0.9)
+        epochs = 10
+        self._test(scheduler, targets, epochs)
+
+    def test_reduce_lr_on_plateau1(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 20]
+        metrics = [10 - i * 0.0167 for i in range(20)]
+        scheduler = ReduceLROnPlateau(self.opt, threshold_mode='abs', mode='min',
+                                      threshold=0.01, patience=5, cooldown=5)
+        epochs = 10
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau2(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 6 + [0.05] * 7 + [0.005] * 7 + [0.0005] * 2]
+        metrics = [10 - i * 0.0165 for i in range(22)]
+        scheduler = ReduceLROnPlateau(self.opt, patience=5, cooldown=0, threshold_mode='abs',
+                                      mode='min', threshold=0.1)
+        epochs = 22
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau3(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * (2 + 6) + [0.05] * (5 + 6) + [0.005] * 4]
+        metrics = [-0.8] * 2 + [-0.234] * 20
+        scheduler = ReduceLROnPlateau(self.opt, mode='max', patience=5, cooldown=5,
+                                      threshold_mode='abs')
+        epochs = 22
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau4(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 20]
+        metrics = [1.5 * (1.025 ** i) for i in range(20)]  # 1.025 > 1.1**0.25
+        scheduler = ReduceLROnPlateau(self.opt, mode='max', patience=3,
+                                      threshold_mode='rel', threshold=0.1)
+        epochs = 20
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau5(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 6 + [0.05] * (5 + 6) + [0.005] * 4]
+        metrics = [1.5 * (1.005 ** i) for i in range(20)]
+        scheduler = ReduceLROnPlateau(self.opt, mode='max', threshold_mode='rel',
+                                      threshold=0.1, patience=5, cooldown=5)
+        epochs = 20
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau6(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 20]
+        metrics = [1.5 * (0.85 ** i) for i in range(20)]
+        scheduler = ReduceLROnPlateau(self.opt, mode='min', threshold_mode='rel',
+                                      threshold=0.1)
+        epochs = 20
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau7(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 6 + [0.05] * (5 + 6) + [0.005] * 4]
+        metrics = [1] * 7 + [0.6] + [0.5] * 12
+        scheduler = ReduceLROnPlateau(self.opt, mode='min', threshold_mode='rel',
+                                      threshold=0.1, patience=5, cooldown=5)
+        epochs = 20
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_reduce_lr_on_plateau8(self):
+        for param_group in self.opt.param_groups:
+            param_group['lr'] = 0.5
+        targets = [[0.5] * 6 + [0.4] * 14, [0.5] * 6 + [0.3] * 14]
+        metrics = [1.5 * (1.005 ** i) for i in range(20)]
+        scheduler = ReduceLROnPlateau(self.opt, mode='max', threshold_mode='rel', min_lr=[0.4, 0.3],
+                                      threshold=0.1, patience=5, cooldown=5)
+        epochs = 20
+        self._test_reduce_lr_on_plateau(scheduler, targets, metrics, epochs)
+
+    def test_lambda_lr(self):
+        self.opt.param_groups[0]['lr'] = 0.05
+        self.opt.param_groups[1]['lr'] = 0.4
+        targets = [[0.05 * (0.9 ** x) for x in range(10)], [0.4 * (0.8 ** x) for x in range(10)]]
+        scheduler = LambdaLR(self.opt,
+                             lr_lambda=[lambda x1: 0.9 ** x1, lambda x2: 0.8 ** x2])
+        epochs = 10
+        self._test(scheduler, targets, epochs)
+
+    def _test(self, scheduler, targets, epochs=10):
+        for epoch in range(epochs):
+            scheduler.step(epoch)
+            for param_group, target in zip(self.opt.param_groups, targets):
+                self.assertAlmostEqual(target[epoch], param_group['lr'],
+                                       msg='LR is wrong in epoch {}: expected {}, got {}'.format(
+                                           epoch, target[epoch], param_group['lr']), delta=1e-5)
+
+    def _test_reduce_lr_on_plateau(self, scheduler, targets, metrics, epochs=10, verbose=False):
+        for epoch in range(epochs):
+            scheduler.step(metrics[epoch])
+            if verbose:
+                print('epoch{}:\tlr={}'.format(epoch, self.opt.param_groups[0]['lr']))
+            for param_group, target in zip(self.opt.param_groups, targets):
+                self.assertAlmostEqual(target[epoch], param_group['lr'],
+                                       msg='LR is wrong in epoch {}: expected {}, got {}'.format(
+                                           epoch, target[epoch], param_group['lr']), delta=1e-5)
+
 if __name__ == '__main__':
     run_tests()

test/test_sparse.py

@@ -5,55 +5,129 @@ import itertools
 import random
 import unittest
 from common import TestCase, run_tests
+from common_nn import TEST_CUDA
 from numbers import Number
 
-SparseTensor = sparse.DoubleTensor
+
+def cpu_only(inner):
+    def outer(self, *args, **kwargs):
+        if self.is_cuda:
+            raise unittest.SkipTest("Test is CPU-only")
+        inner(self, *args, **kwargs)
+    return outer
+
+
+def cuda_only(inner):
+    def outer(self, *args, **kwargs):
+        if not self.is_cuda:
+            raise unittest.SkipTest("Test is GPU-only")
+        inner(self, *args, **kwargs)
+    return outer
 
 
 class TestSparse(TestCase):
 
-    @staticmethod
-    def _gen_sparse(d, nnz, with_size):
-        v = torch.randn(nnz)
-        if isinstance(with_size, Number):
-            i = (torch.rand(d, nnz) * with_size).type(torch.LongTensor)
-            x = SparseTensor(i, v)
-        else:
-            i = torch.rand(d, nnz) * \
-                torch.Tensor(with_size).repeat(nnz, 1).transpose(0, 1)
-            i = i.type(torch.LongTensor)
-            x = SparseTensor(i, v, torch.Size(with_size))
+    def setUp(self):
+        # These parameters control the various ways we can run the test.
+        # We will subclass and override this method to implement CUDA
+        # tests
+        self.is_cuda = False
+        self.is_uncoalesced = False
+        self.IndexTensor = torch.LongTensor
+        self.ValueTensor = torch.DoubleTensor
+        self.SparseTensor = torch.sparse.DoubleTensor
 
-        return x, i, v
+    def _gen_sparse(self, d, nnz, with_size):
+        # TODO: Consider implementing this in the CUDA case by directly
+        # performing the operations on the GPU.  You won't be able to
+        # use torch.rand/torch.randn in this case because they are
+        # CPU-only.  If you do this, you can remove the is_cuda branch
+        # at the end.
+        #
+        # If you do this, be sure to update assert_uncoalesced too
+
+        if isinstance(with_size, Number):
+            with_size = [with_size] * d
+
+        if self.is_uncoalesced:
+            # We want to generate a tensor with a lot of uncoalesced
+            # entries to stress test whether or not we handle this
+            # (subtle) case correctly
+            v_size = [nnz * 2] + list(with_size[d:])
+            v = torch.randn(*v_size)
+            r = torch.rand(d, nnz)
+            # Repeat the indexes, so every position shows up twice
+            i = torch.cat([r, r], dim=1) * \
+                torch.Tensor(with_size[:d]).repeat(nnz * 2, 1).transpose(0, 1)
+            i = i.type(torch.LongTensor)
+            x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
+            self.assert_uncoalesced(x)
+        else:
+            # Generate a sparse tensor with d sparse dimensions; the
+            # rest the dimensions with_size[d:] are dense.
+            v_size = [nnz] + list(with_size[d:])
+            v = torch.randn(*v_size)
+            i = torch.rand(d, nnz) * \
+                torch.Tensor(with_size[:d]).repeat(nnz, 1).transpose(0, 1)
+            i = i.type(torch.LongTensor)
+            x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
+
+        if self.is_cuda:
+            return x.cuda(), i.cuda(), v.cuda()
+        else:
+            return x, i.clone(), v.clone()
+
+    def assert_uncoalesced(self, x):
+        """
+        Test if a CPU tensor is uncoalesced.  This is used to ensure
+        correctness of the uncoalesced tensor generation algorithm.
+        """
+        assert not x.is_coalesced()
+        # Strategy: construct a new sparse tensor with the raw value
+        # field overwritten to a tensor of ones, coalesce it, and then
+        # check if any value entries are > 1 (which indicates that the
+        # original was uncoalesced.)
+        i = x._indices().clone()
+        v = x._values().clone().fill_(1)
+        y = torch.sparse.DoubleTensor(i, v, x.size())
+        z = self.safeCoalesce(y)
+        assert (z._values() > 1).sum() > 0
+
+    def randn(self, *args, **kwargs):
+        """
+        Variant of torch.randn that also works in the TEST_CUDA case.
+        """
+        # TODO: Put this in torch.cuda.randn
+        return self.ValueTensor(*args, **kwargs).normal_()
 
     def test_basic(self):
         x, i, v = self._gen_sparse(3, 10, 100)
 
-        self.assertEqual(i, x.indices())
-        self.assertEqual(v, x.values())
+        self.assertEqual(i, x._indices())
+        self.assertEqual(v, x._values())
 
         x, i, v = self._gen_sparse(3, 10, [100, 100, 100])
-        self.assertEqual(i, x.indices())
-        self.assertEqual(v, x.values())
+        self.assertEqual(i, x._indices())
+        self.assertEqual(v, x._values())
         self.assertEqual(x.ndimension(), 3)
-        self.assertEqual(x.nnz(), 10)
+        self.assertEqual(x.coalesce()._nnz(), 10)
         for i in range(3):
             self.assertEqual(x.size(i), 100)
 
         # Make sure we can access empty indices / values
-        x = SparseTensor()
-        self.assertEqual(x.indices().numel(), 0)
-        self.assertEqual(x.values().numel(), 0)
+        x = self.SparseTensor()
+        self.assertEqual(x._indices().numel(), 0)
+        self.assertEqual(x._values().numel(), 0)
 
     def test_to_dense(self):
-        i = torch.LongTensor([
+        i = self.IndexTensor([
             [0, 1, 2, 2],
             [0, 0, 0, 3],
             [0, 0, 1, 4],
         ])
-        v = torch.Tensor([2, 1, 3, 4])
-        x = SparseTensor(i, v, torch.Size([3, 4, 5]))
-        res = torch.Tensor([
+        v = self.ValueTensor([2, 1, 3, 4])
+        x = self.SparseTensor(i, v, torch.Size([3, 4, 5]))
+        res = self.ValueTensor([
             [[2, 0, 0, 0, 0],
              [0, 0, 0, 0, 0],
              [0, 0, 0, 0, 0],
@@ -73,58 +147,164 @@ class TestSparse(TestCase):
         x.to_dense()
         self.assertEqual(res, x.to_dense())
 
+    def test_shared(self):
+        i = self.IndexTensor([[2]])
+        v = self.ValueTensor([5])
+        x = self.SparseTensor(i, v, torch.Size([3]))
+        v[0] = 6
+        self.assertEqual(self.ValueTensor([0, 0, 6]), x.to_dense())
+        i[0][0] = 0
+        self.assertEqual(self.ValueTensor([6, 0, 0]), x.to_dense())
+
+    def test_to_dense_hybrid(self):
+        i = self.IndexTensor([
+            [0, 1, 2, 2],
+            [0, 0, 0, 3],
+        ])
+        v = self.ValueTensor([[2, 3], [1, 2], [3, 4], [4, 5]])
+        x = self.SparseTensor(i, v, torch.Size([3, 4, 2]))
+        res = self.ValueTensor([
+            [[2, 3],
+             [0, 0],
+             [0, 0],
+             [0, 0]],
+            [[1, 2],
+             [0, 0],
+             [0, 0],
+             [0, 0]],
+            [[3, 4],
+             [0, 0],
+             [0, 0],
+             [4, 5]],
+        ])
+
+        x.to_dense()  # Tests double to_dense for memory corruption
+        x.to_dense()
+        x.to_dense()
+        self.assertEqual(res, x.to_dense())
+
     def test_contig(self):
-        i = torch.LongTensor([
+        i = self.IndexTensor([
             [1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
             [92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
         ])
-        v = torch.Tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
-        x = SparseTensor(i, v, torch.Size([100, 100]))
-        exp_i = torch.LongTensor([
+        v = self.ValueTensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
+        x = self.SparseTensor(i, v, torch.Size([100, 100]))
+        exp_i = self.IndexTensor([
             [0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
             [31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
         ])
-        exp_v = torch.Tensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7])
-        x.contiguous()
-        self.assertEqual(exp_i, x.indices())
-        self.assertEqual(exp_v, x.values())
+        exp_v = self.ValueTensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7])
+        x = self.safeCoalesce(x)
+        self.assertEqual(exp_i, x._indices())
+        self.assertEqual(exp_v, x._values())
 
-        i = torch.LongTensor([
+        i = self.IndexTensor([
             [2, 0, 2, 1],
             [0, 0, 3, 0],
             [1, 0, 4, 0],
         ])
-        v = torch.Tensor([3, 2, 4, 1])
-        x = SparseTensor(i, v, torch.Size([3, 4, 5]))
-        exp_i = torch.LongTensor([
+        v = self.ValueTensor([3, 2, 4, 1])
+        x = self.SparseTensor(i, v, torch.Size([3, 4, 5]))
+        exp_i = self.IndexTensor([
             [0, 1, 2, 2],
             [0, 0, 0, 3],
             [0, 0, 1, 4],
         ])
-        exp_v = torch.Tensor([2, 1, 3, 4])
+        exp_v = self.ValueTensor([2, 1, 3, 4])
 
-        x.contiguous()
-        self.assertEqual(exp_i, x.indices())
-        self.assertEqual(exp_v, x.values())
+        x = self.safeCoalesce(x)
+        self.assertEqual(exp_i, x._indices())
+        self.assertEqual(exp_v, x._values())
 
         # Duplicate indices
-        i = torch.LongTensor([
+        i = self.IndexTensor([
             [0, 0, 2, 0],
             [0, 0, 3, 0],
             [0, 0, 4, 0],
         ])
-        v = torch.Tensor([3, 2, 4, 1])
-        x = SparseTensor(i, v, torch.Size([3, 4, 5]))
-        exp_i = torch.LongTensor([
+        v = self.ValueTensor([3, 2, 4, 1])
+        x = self.SparseTensor(i, v, torch.Size([3, 4, 5]))
+        exp_i = self.IndexTensor([
             [0, 2],
             [0, 3],
             [0, 4],
         ])
-        exp_v = torch.Tensor([6, 4])
+        exp_v = self.ValueTensor([6, 4])
 
-        x.contiguous()
-        self.assertEqual(exp_i, x.indices())
-        self.assertEqual(exp_v, x.values())
+        x = self.safeCoalesce(x)
+        self.assertEqual(exp_i, x._indices())
+        self.assertEqual(exp_v, x._values())
+
+    def test_contig_hybrid(self):
+        i = self.IndexTensor([
+            [1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
+            [92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
+        ])
+        v = self.ValueTensor([
+            [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
+            [6, 7], [7, 8], [8, 9], [9, 10], [10, 11],
+        ])
+        x = self.SparseTensor(i, v, torch.Size([100, 100, 2]))
+        exp_i = self.IndexTensor([
+            [0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
+            [31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
+        ])
+        exp_v = self.ValueTensor([
+            [2, 3], [1, 2], [6, 7], [4, 5], [10, 11],
+            [3, 4], [5, 6], [9, 10], [8, 9], [7, 8],
+        ])
+        x = self.safeCoalesce(x)
+        self.assertEqual(exp_i, x._indices())
+        self.assertEqual(exp_v, x._values())
+
+        i = self.IndexTensor([
+            [2, 0, 2, 1],
+            [0, 0, 3, 0],
+            [1, 0, 4, 0],
+        ])
+        v = self.ValueTensor([[3, 3, 3], [2, 2, 2], [4, 4, 4], [1, 1, 1]])
+        x = self.SparseTensor(i, v, torch.Size([3, 4, 5, 3]))
+        exp_i = self.IndexTensor([
+            [0, 1, 2, 2],
+            [0, 0, 0, 3],
+            [0, 0, 1, 4],
+        ])
+        exp_v = self.ValueTensor([[2, 2, 2], [1, 1, 1], [3, 3, 3], [4, 4, 4]])
+
+        x = self.safeCoalesce(x)
+        self.assertEqual(exp_i, x._indices())
+        self.assertEqual(exp_v, x._values())
+
+        # Duplicate indices
+        i = self.IndexTensor([
+            [0, 0, 2, 0],
+            [0, 0, 3, 0],
+            [0, 0, 4, 0],
+        ])
+        v = self.ValueTensor([[3, 2, 3], [2, 1, 1], [4, 3, 4], [1, 1, 1]])
+        x = self.SparseTensor(i, v, torch.Size([3, 4, 5, 3]))
+        exp_i = self.IndexTensor([
+            [0, 2],
+            [0, 3],
+            [0, 4],
+        ])
+        exp_v = self.ValueTensor([[6, 4, 5], [4, 3, 4]])
+
+        x = self.safeCoalesce(x)
+        self.assertEqual(exp_i, x._indices())
+        self.assertEqual(exp_v, x._values())
+
+    def test_clone(self):
+        x, _, _ = self._gen_sparse(4, 20, 5)
+        if self.is_uncoalesced:
+            self.assertFalse(x.is_coalesced())
+            y = x.clone()
+            self.assertFalse(y.is_coalesced())
+        x = x.coalesce()
+        self.assertTrue(x.is_coalesced())
+        y = x.clone()
+        self.assertTrue(y.is_coalesced())
 
     def test_transpose(self):
         x = self._gen_sparse(4, 20, 5)[0]
@@ -139,6 +319,7 @@ class TestSparse(TestCase):
             y = y.transpose(i, j)
             self.assertEqual(x.to_dense(), y)
 
+    @cpu_only
     def test_mm(self):
         def test_shape(di, dj, dk):
             x, _, _ = self._gen_sparse(2, 20, [di, dj])
@@ -147,22 +328,23 @@ class TestSparse(TestCase):
             alpha = random.random()
             beta = random.random()
 
-            expected = torch.addmm(alpha, t, beta, x.to_dense(), y)
             res = torch.addmm(alpha, t, beta, x, y)
+            expected = torch.addmm(alpha, t, beta, x.to_dense(), y)
             self.assertEqual(res, expected)
 
-            expected = torch.addmm(t, x.to_dense(), y)
             res = torch.addmm(t, x, y)
+            expected = torch.addmm(t, x.to_dense(), y)
             self.assertEqual(res, expected)
 
-            expected = torch.mm(x.to_dense(), y)
             res = torch.mm(x, y)
+            expected = torch.mm(x.to_dense(), y)
             self.assertEqual(res, expected)
 
         test_shape(10, 100, 100)
         test_shape(100, 1000, 200)
         test_shape(64, 10000, 300)
 
+    @cpu_only
     def test_saddmm(self):
         def test_shape(di, dj, dk):
             x = self._gen_sparse(2, 20, [di, dj])[0]
@@ -171,50 +353,273 @@ class TestSparse(TestCase):
             alpha = random.random()
             beta = random.random()
 
-            expected = torch.addmm(alpha, t.to_dense(), beta, x.to_dense(), y)
             res = torch.saddmm(alpha, t, beta, x, y)
+            expected = torch.addmm(alpha, t.to_dense(), beta, x.to_dense(), y)
             self.assertEqual(res.to_dense(), expected)
 
-            expected = torch.addmm(t.to_dense(), x.to_dense(), y)
             res = torch.saddmm(t, x, y)
+            expected = torch.addmm(t.to_dense(), x.to_dense(), y)
             self.assertEqual(res.to_dense(), expected)
 
-            expected = torch.mm(x.to_dense(), y)
             res = torch.smm(x, y)
+            expected = torch.mm(x.to_dense(), y)
             self.assertEqual(res.to_dense(), expected)
 
         test_shape(7, 5, 3)
         test_shape(1000, 100, 100)
         test_shape(3000, 64, 300)
 
+    def test_dsmm(self):
+        def test_shape(di, dj, dk):
+            x = self._gen_sparse(2, 20, [di, dj])[0]
+            y = self.randn(dj, dk)
+
+            res = torch.dsmm(x, y)
+            expected = torch.mm(x.to_dense(), y)
+            self.assertEqual(res, expected)
+
+        test_shape(7, 5, 3)
+        test_shape(1000, 100, 100)
+        test_shape(3000, 64, 300)
+
+    def test_hsmm(self):
+        def test_shape(di, dj, dk):
+            x = self._gen_sparse(2, 20, [di, dj])[0]
+            y = self.randn(dj, dk)
+
+            res = torch.hsmm(x, y)
+            expected = torch.mm(x.to_dense(), y)
+            self.assertEqual(res.to_dense(), expected)
+
+        test_shape(7, 5, 3)
+        test_shape(1000, 100, 100)
+        test_shape(3000, 64, 300)
+
+    def _test_spadd_shape(self, shape_i, shape_v=None):
+        shape = shape_i + (shape_v or [])
+        x, _, _ = self._gen_sparse(len(shape_i), 10, shape)
+        y = self.randn(*shape)
+        r = random.random()
+
+        res = torch.add(y, r, x)
+        expected = y + r * x.to_dense()
+
+        self.assertEqual(res, expected)
+
+        # Non contiguous dense tensor
+        s = list(shape)
+        s[0] = shape[-1]
+        s[-1] = shape[0]
+        y = self.randn(*s)
+        y.transpose_(0, len(s) - 1)
+        r = random.random()
+
+        res = torch.add(y, r, x)
+        expected = y + r * x.to_dense()
+
+        self.assertEqual(res, expected)
+
     def test_spadd(self):
-        def test_shape(*shape):
-            x, _, _ = self._gen_sparse(len(shape), 10, shape)
-            y = torch.randn(*shape)
-            r = random.random()
+        self._test_spadd_shape([5, 6])
+        self._test_spadd_shape([10, 10, 10])
+        self._test_spadd_shape([50, 30, 20])
+        self._test_spadd_shape([5, 5, 5, 5, 5, 5])
 
-            expected = y + r * x.to_dense()
-            res = torch.add(y, r, x)
+    def test_spadd_hybrid(self):
+        self._test_spadd_shape([5, 6], [2, 3])
+        self._test_spadd_shape([10, 10, 10], [3])
+        self._test_spadd_shape([50, 30, 20], [2])
+        self._test_spadd_shape([5, 5, 5, 5, 5, 5], [2])
 
-            self.assertEqual(res, expected)
+    def _test_basic_ops_shape(self, shape_i, shape_v=None):
+        shape = shape_i + (shape_v or [])
+        x1, _, _ = self._gen_sparse(len(shape_i), 9, shape)
+        x2, _, _ = self._gen_sparse(len(shape_i), 12, shape)
 
-            # Non contiguous dense tensor
-            s = list(shape)
-            s[0] = shape[-1]
-            s[-1] = shape[0]
-            y = torch.randn(*s).transpose_(0, len(s) - 1)
-            r = random.random()
+        y1 = x1 + x2
+        y2 = x1.clone()
+        y2.add_(x2)
+        expected = x1.to_dense() + x2.to_dense()
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
 
-            expected = y + r * x.to_dense()
-            res = torch.add(y, r, x)
+        y1 = x1 - x2
+        y2 = x1.clone()
+        y2.sub_(x2)
+        expected = x1.to_dense() - x2.to_dense()
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
 
-            self.assertEqual(res, expected)
+        y1 = x1 * x2
+        y2 = x1.clone()
+        y2.mul_(x2)
+        expected = x1.to_dense() * x2.to_dense()
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
 
-        test_shape(5, 6)
-        test_shape(10, 10, 10)
-        test_shape(50, 30, 20)
-        test_shape(5, 5, 5, 5, 5, 5)
+        y1 = x1 * 37.5
+        y2 = x1.clone()
+        y2.mul_(37.5)
+        expected = x1.to_dense() * 37.5
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
 
+        y1 = x1 / 37.5
+        y2 = x1.clone()
+        y2.div_(37.5)
+        expected = x1.to_dense() / 37.5
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
+
+        # TODO: add back inplace support
+        y1 = x1 ** 2
+        y2 = x1.clone()
+        y2 = y2.pow(2)
+        expected = x1.to_dense() ** 2
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
+
+        y = x1.clone()
+        y.zero_()
+        expected = torch.zeros(x1.size())
+        self.assertEqual(y.to_dense(), expected)
+
+        self.assertFalse(x1.is_coalesced())
+        y = x1.coalesce()
+        z = x1.coalesce()
+        self.assertFalse(x1.is_coalesced())
+        self.assertTrue(y.is_coalesced())
+        self.assertEqual(x1, y)
+        # check that coalesce is out of place
+        y._values().add_(1)
+        self.assertEqual(z._values() + 1, y._values())
+
+    def test_basic_ops(self):
+        self._test_basic_ops_shape([5, 6])
+        self._test_basic_ops_shape([10, 10, 10])
+        self._test_basic_ops_shape([50, 30, 20])
+        self._test_basic_ops_shape([5, 5, 5, 5, 5, 5])
+
+    def test_basic_ops_hybrid(self):
+        self._test_basic_ops_shape([5, 6], [2, 3])
+        self._test_basic_ops_shape([10, 10, 10], [3])
+        self._test_basic_ops_shape([50, 30, 20], [2])
+        self._test_basic_ops_shape([5, 5, 5, 5, 5, 5], [2])
+
+    def _test_sparse_mask_shape(self, shape_i, shape_v=None):
+        shape = shape_i + (shape_v or [])
+        x1, _, _ = self._gen_sparse(len(shape_i), 9, shape)
+        x2, _, _ = self._gen_sparse(len(shape_i), 12, shape)
+
+        y1 = x1 + x2
+        y2 = x1.clone()
+        y2.add_(x2)
+        expected = x1.to_dense() + x2.to_dense()
+        self.assertEqual(y1.to_dense(), expected)
+        self.assertEqual(y2.to_dense(), expected)
+
+    def _test_sparse_mask_fixed(self):
+        i = self.IndexTensor([
+            [1, 3, 0, 4],
+            [2, 1, 2, 3],
+        ])
+        v = self.ValueTensor([1, 2, 3, 4])
+        x = self.SparseTensor(i, v, torch.Size([5, 4])).coalesce()
+        dense = self.ValueTensor([
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+            [13, 14, 15, 16],
+            [17, 18, 19, 20],
+        ])
+        exp_v = self.ValueTensor([7, 14, 3, 20])
+        res = dense._sparse_mask(x)
+        expected = self.SparseTensor(i, exp_v, torch.Size([5, 4]))
+        self.assertEqual(res, expected)
+
+    def test_sparse_mask(self):
+        self._test_sparse_mask_fixed()
+
+        self._test_sparse_mask_shape([5, 6])
+        self._test_sparse_mask_shape([10, 10, 10])
+        self._test_sparse_mask_shape([50, 30, 20])
+        self._test_sparse_mask_shape([5, 5, 5, 5, 5, 5])
+
+    def _test_sparse_mask_hybrid_fixed(self):
+        i = self.IndexTensor([
+            [1, 3, 0, 4],
+            [2, 1, 2, 3],
+        ])
+        v = self.ValueTensor([[1, 2], [2, 3], [3, 4], [4, 5]])
+        # TODO: This is also testing that, if coalesce is a no-op,
+        # the indices don't get permuted. I don't know if we actually
+        # want to give this invariant.
+        x = self.SparseTensor(i, v, torch.Size([5, 4, 2])).coalesce()
+        dense = self.ValueTensor([
+            [[1, 3], [2, 2], [3, 3], [4, 2]],
+            [[5, 7], [6, 7], [7, 9], [8, 9]],
+            [[9, 2], [10, 4], [11, 1], [12, 3]],
+            [[13, 5], [14, 1], [15, 1], [16, 6]],
+            [[17, 7], [18, 2], [19, 7], [20, 1]],
+        ])
+        res = dense._sparse_mask(x)
+        exp_v = self.ValueTensor([[7, 9], [14, 1], [3, 3], [20, 1]])
+        expected = self.SparseTensor(i, exp_v, torch.Size([5, 4, 2]))
+        self.assertEqual(res, expected)
+
+    def test_sparse_mask_hybrid(self):
+        self._test_sparse_mask_hybrid_fixed()
+
+        self._test_sparse_mask_shape([5, 6], [2, 3])
+        self._test_sparse_mask_shape([10, 10, 10], [3])
+        self._test_sparse_mask_shape([50, 30, 20], [2])
+        self._test_sparse_mask_shape([5, 5, 5, 5, 5, 5], [2])
+
+    @cuda_only
+    def test_storage_not_null(self):
+        x = torch.cuda.sparse.FloatTensor(2)
+        self.assertNotEqual(x.get_device(), -1)
+
+    @cuda_only
+    @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
+    def test_same_gpu(self):
+        i = self.IndexTensor([[2]]).cuda(1)
+        v = self.ValueTensor([5]).cuda(1)
+        x = self.SparseTensor(i, v, torch.Size([3]), device=1)
+        self.assertEqual(x.get_device(), 1)
+        self.assertEqual(x._values().get_device(), 1)
+        self.assertEqual(x._indices().get_device(), 1)
+
+        x = self.SparseTensor(3, device=1)
+        self.assertEqual(x.get_device(), 1)
+        self.assertEqual(x._values().get_device(), 1)
+        self.assertEqual(x._indices().get_device(), 1)
+
+        v = self.ValueTensor([5]).cuda(0)
+        self.assertRaises(RuntimeError, lambda: self.SparseTensor(i, v, torch.Size([3])))
+
+
+class TestUncoalescedSparse(TestSparse):
+    def setUp(self):
+        super(TestUncoalescedSparse, self).setUp()
+        self.is_uncoalesced = True
+
+
+@unittest.skipIf(not TEST_CUDA, 'CUDA not available')
+class TestCudaSparse(TestSparse):
+    def setUp(self):
+        super(TestCudaSparse, self).setUp()
+        self.is_cuda = True
+        self.IndexTensor = torch.cuda.LongTensor
+        self.ValueTensor = torch.cuda.DoubleTensor
+        self.SparseTensor = torch.cuda.sparse.DoubleTensor
+
+
+@unittest.skipIf(not TEST_CUDA, 'CUDA not available')
+class TestCudaUncoalescedSparse(TestCudaSparse):
+    def setUp(self):
+        super(TestCudaUncoalescedSparse, self).setUp()
+        self.is_uncoalesced = True
 
 if __name__ == '__main__':
     run_tests()

test/test_utils.py

@@ -6,9 +6,9 @@ import shutil
 import random
 import tempfile
 import unittest
-import sys
 import traceback
 import torch
+import torch.utils.data
 import torch.cuda
 import warnings
 from torch.autograd import Variable
@@ -19,7 +19,7 @@ from torch.utils.serialization import load_lua
 
 HAS_CUDA = torch.cuda.is_available()
 
-from common import TestCase, run_tests
+from common import TestCase, run_tests, download_file
 
 try:
     import cffi
@@ -108,6 +108,44 @@ class DatasetMock(object):
         return 10
 
 
+class TestDataLoader(TestCase):
+    def setUp(self):
+        self.dataset = torch.randn(5, 3, 3, 2)
+        self.batch_size = 3
+
+    def test_single_keep(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=0,
+                                                 drop_last=False)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 2)
+
+    def test_single_drop(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=0,
+                                                 drop_last=True)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 1)
+
+    def test_multi_keep(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=2,
+                                                 drop_last=False)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 2)
+
+    def test_multi_drop(self):
+        dataloader = torch.utils.data.DataLoader(self.dataset,
+                                                 batch_size=self.batch_size,
+                                                 num_workers=2,
+                                                 drop_last=True)
+        dataiter = iter(dataloader)
+        self.assertEqual(len(list(dataiter)), 1)
+
+
 class TestTrainer(TestCase):
 
     intervals = [
@@ -296,40 +334,13 @@ class TestLuaReader(TestCase):
             self.assertEqual(grad_input, test['grad_input'])
         return do_test
 
-    @classmethod
-    def _download_data(cls, test_file_path):
-        if os.path.exists(test_file_path):
-            return
-        print('Downloading test file for TestLuaReader.')
-        DATA_URL = 'https://s3.amazonaws.com/pytorch/legacy_modules.t7'
-        urllib = cls._get_urllib('request')
-        data = urllib.urlopen(DATA_URL, timeout=15).read()
-        with open(test_file_path, 'wb') as f:
-            f.write(data)
-
-    @staticmethod
-    def _get_urllib(submodule):
-        if sys.version_info < (3,):
-            import urllib2
-            return urllib2
-        else:
-            import urllib.error
-            import urllib.request
-            return getattr(urllib, submodule)
-
     @classmethod
     def init(cls):
-        data_dir = os.path.join(os.path.dirname(__file__), 'data')
-        test_file_path = os.path.join(data_dir, 'legacy_modules.t7')
-        urllib = cls._get_urllib('error')
         try:
-            cls._download_data(test_file_path)
-        except urllib.URLError as e:
-            warnings.warn(("Couldn't download the test file for TestLuaReader! "
-                           "Tests will be incomplete!"), RuntimeWarning)
+            path = download_file('https://download.pytorch.org/test_data/legacy_modules.t7')
+        except unittest.SkipTest:
             return
-
-        tests = load_lua(test_file_path)
+        tests = load_lua(path)
         for name, test in tests['modules'].items():
             test_name = 'test_' + name.replace('nn.', '')
             setattr(cls, test_name, cls._module_test(name, test))

tools/cwrap/cwrap.py

@@ -4,9 +4,12 @@ from string import Template
 from copy import deepcopy
 from .plugins import ArgcountChecker, OptionalArguments, ArgumentReferences, \
     BeforeAfterCall, ConstantArguments, ReturnArguments, GILRelease
+from ..shared import cwrap_common
 
 
 class cwrap(object):
+    BASE_INDENT_SIZE = 6
+
     RETURN_WRAPPERS = {
         'void': Template('Py_RETURN_NONE;'),
         'long': Template('return PyLong_FromLong($result);'),
@@ -16,24 +19,28 @@ class cwrap(object):
 
     OPTION_TEMPLATE = Template("""
     ${els}if ($arg_check) {
+      $pre_arg_assign
+      $arg_assign
       $code
     """)
 
+    ARG_ASSIGN_TEMPLATE = Template("""${type} ${name} = ${unpack};""")
+
     OPTION_CODE_TEMPLATE = [
         '$call',
         '$return_result',
     ]
 
-    FUNCTION_CALL_TEMPLATE = Template("$capture_result$cname($arg_unpack);")
+    FUNCTION_CALL_TEMPLATE = Template("$capture_result$cname($call_arg);")
 
     DEFAULT_PLUGIN_CLASSES = [ArgcountChecker, ConstantArguments, OptionalArguments,
                               ArgumentReferences, BeforeAfterCall, ReturnArguments, GILRelease]
 
-    def __init__(self, source, destination=None, plugins=[], default_plugins=True):
+    def __init__(self, source, destination=None, plugins=None, default_plugins=True):
         if destination is None:
             destination = source.replace('.cwrap', '.cpp')
 
-        self.plugins = plugins
+        self.plugins = [] if plugins is None else plugins
         if default_plugins:
             defaults = [cls() for cls in self.DEFAULT_PLUGIN_CLASSES]
             self.plugins = defaults + self.plugins
@@ -45,7 +52,10 @@ class cwrap(object):
         with open(source, 'r') as f:
             declarations = f.read()
 
+        # wrap all the declarations in the source .cwrap file
         wrapper = self.wrap_declarations(declarations)
+
+        # let each plugin do any post-processing of the wrapped file
         for plugin in self.plugins:
             wrapper = plugin.process_full_file(wrapper)
 
@@ -67,7 +77,7 @@ class cwrap(object):
             elif line == ']]':
                 in_declaration = False
                 declaration = yaml.load('\n'.join(declaration_lines))
-                self.set_declaration_defaults(declaration)
+                cwrap_common.set_declaration_defaults(declaration)
 
                 # Pass declaration in a list - maybe some plugins want to add
                 # multiple wrappers
@@ -95,24 +105,6 @@ class cwrap(object):
 
         return '\n'.join(output)
 
-    def set_declaration_defaults(self, declaration):
-        declaration.setdefault('arguments', [])
-        declaration.setdefault('return', 'void')
-        if 'cname' not in declaration:
-            declaration['cname'] = declaration['name']
-        # Simulate multiple dispatch, even if it's not necessary
-        if 'options' not in declaration:
-            declaration['options'] = [{'arguments': declaration['arguments']}]
-            del declaration['arguments']
-        # Parse arguments (some of them can be strings)
-        for option in declaration['options']:
-            option['arguments'] = self.parse_arguments(option['arguments'])
-        # Propagate defaults from declaration to options
-        for option in declaration['options']:
-            for k, v in declaration.items():
-                if k != 'name' and k != 'options':
-                    option.setdefault(k, v)
-
     def parse_arguments(self, args):
         new_args = []
         for arg in args:
@@ -130,6 +122,10 @@ class cwrap(object):
         return new_args
 
     def search_plugins(self, fnname, args, fallback):
+        """Search plugins for the given function to call with args.
+
+        If not found, call fallback with args.
+        """
         for plugin in self.plugins:
             wrapper = getattr(plugin, fnname)(*args)
             if wrapper is not None:
@@ -148,6 +144,9 @@ class cwrap(object):
     def get_wrapper_template(self, declaration):
         return self.search_plugins('get_wrapper_template', (declaration,), lambda _: None)
 
+    def get_assign_args(self, arguments):
+        return self.search_plugins('get_assign_args', (arguments,), lambda _: arguments)
+
     def get_arg_accessor(self, arg, option):
         def wrap_accessor(arg, _):
             if arg.get('idx') is None:
@@ -178,9 +177,44 @@ class cwrap(object):
             res = tmpl.substitute(arg=accessor, idx=arg.get('idx'))
             for plugin in self.plugins:
                 res = getattr(plugin, plugin_fn_name)(res, arg, accessor)
+
             result.append(res)
         return result
 
+    def build_option_args(self, arguments, arg_unpack):
+        assignement = []
+        call_arg = []
+        # If types or names needs to be changed
+        arguments = self.get_assign_args(arguments)
+        for arg, unpack in zip(arguments, arg_unpack):
+            if arg['type'] == 'CONSTANT':
+                call_arg.append(unpack)
+            else:
+                var_name = "arg_" + str(arg.get('assign_name', arg['name']))
+                res = self.ARG_ASSIGN_TEMPLATE.substitute(
+                    type=arg['type'],
+                    name=var_name,
+                    unpack=unpack)
+
+                if var_name not in call_arg:
+                    assignement.append(res)
+                call_arg.append(var_name)
+        return assignement, call_arg
+
+    def indent_code(self, code):
+        if code == '':
+            return code
+        code_lines = map(lambda s: s.strip(), code.split('\n'))
+        code = '\n'
+        depth = self.BASE_INDENT_SIZE
+        for line in code_lines:
+            depth -= line.count('}') * 2
+            code += ' ' * depth + line + '\n'
+            depth += line.count('{') * 2
+            depth += line.count('(') * 4
+            depth -= line.count(')') * 4
+        return code[:-1]
+
     def generate_option(self, option, is_first):
         checked_args = list(filter(
             lambda arg: 'ignore_check' not in arg or not arg['ignore_check'],
@@ -199,22 +233,29 @@ class cwrap(object):
         for plugin in self.plugins:
             arg_checks = plugin.process_all_checks(arg_checks, option)
 
-        # Generate unpacks
+        # Generate pre_arg assign
+        pre_arg_assign = []
+        for plugin in self.plugins:
+            pre_arg_assign = plugin.process_pre_arg_assign(pre_arg_assign, option)
+
+        # Generate arg assignment and call arguments
         arg_unpack = self.map_selected_arguments('get_type_unpack',
                                                  'process_single_unpack', option, option['arguments'])
-        arg_unpack = ', '.join(arg_unpack)
+        arg_assign, call_arg = self.build_option_args(option['arguments'], arg_unpack)
+
+        call_arg = ', '.join(call_arg)
         for plugin in self.plugins:
-            arg_unpack = plugin.process_all_unpacks(arg_unpack, option)
+            call_arg = plugin.process_all_call_arg(call_arg, option)
 
         # Generate call
         try:
             return_result = self.get_return_wrapper(option).substitute()
             call = self.FUNCTION_CALL_TEMPLATE.substitute(capture_result='',
-                                                          cname=option['cname'], arg_unpack=arg_unpack)
+                                                          cname=option['cname'], call_arg=call_arg)
         except KeyError:
             return_result = self.get_return_wrapper(option).substitute(result='__result')
             call = self.FUNCTION_CALL_TEMPLATE.substitute(capture_result=(option['return'] + ' __result = '),
-                                                          cname=option['cname'], arg_unpack=arg_unpack)
+                                                          cname=option['cname'], call_arg=call_arg)
 
         code_template = deepcopy(self.OPTION_CODE_TEMPLATE)
         for plugin in self.plugins:
@@ -222,19 +263,15 @@ class cwrap(object):
                                                                 option)
         code_template = Template('\n'.join(code_template))
         code = code_template.substitute(call=call, return_result=return_result)
-        code_lines = map(lambda s: s.strip(), code.split('\n'))
-        code = '\n'
-        depth = 6
-        for line in code_lines:
-            depth -= line.count('}') * 2
-            code += ' ' * depth + line + '\n'
-            depth += line.count('{') * 2
-            depth += line.count('(') * 4
-            depth -= line.count(')') * 4
+        code = self.indent_code(code)
+        pre_arg_assign = self.indent_code('\n'.join(pre_arg_assign))
+        arg_assign = self.indent_code('\n'.join(arg_assign))
 
         # Put everything together
         return self.OPTION_TEMPLATE.substitute(
             els=('} else ' if not is_first else ''),
             arg_check=arg_checks,
+            pre_arg_assign=pre_arg_assign,
+            arg_assign=arg_assign,
             code=code,
         )

tools/cwrap/plugins/ArgcountSortPlugin.py

@@ -1,4 +1,6 @@
+import os
 from . import CWrapPlugin
+from ...shared import cwrap_common
 
 
 class ArgcountSortPlugin(CWrapPlugin):
@@ -7,8 +9,7 @@ class ArgcountSortPlugin(CWrapPlugin):
         self.descending = descending
 
     def process_declarations(self, declarations):
-        def num_checked_args(option):
-            return sum(map(lambda a: not a.get('ignore_check', False), option['arguments']))
         for declaration in declarations:
-            declaration['options'].sort(key=num_checked_args, reverse=self.descending)
+            cwrap_common.sort_by_number_of_options(declaration,
+                                                   self.descending)
         return declarations

tools/cwrap/plugins/AssertNDim.py

@@ -0,0 +1,29 @@
+from . import CWrapPlugin
+from string import Template
+
+
+class AssertNDim(CWrapPlugin):
+
+    PRE_CODE_TEMPLATE = Template(
+        """if(THTensor_(nDimension)(LIBRARY_STATE ${arg_op}) != ${dim_value}) {
+             THError("Expected argument %s to have %d dimension(s), but has %d",
+                     "${op}", ${dim_value}, THTensor_(nDimension)(LIBRARY_STATE ${arg_op}));
+           }
+        """)
+
+    def process_option_code_template(self, template, option):
+        new_code_pre = []
+
+        for _, arg in enumerate(option['arguments']):
+            if 'assert_ndim' not in arg:
+                continue
+
+            dim_value = arg.get('assert_ndim')
+            op = arg.get('assign_name', arg['name'])
+            arg_op = "arg_" + op
+            new_code_pre.append(self.PRE_CODE_TEMPLATE.substitute(op=op,
+                                                                  arg_op=arg_op,
+                                                                  dim_value=dim_value))
+            template = new_code_pre + template
+
+        return template

tools/cwrap/plugins/AutoGPU.py

@@ -15,7 +15,9 @@ class AutoGPU(CWrapPlugin):
 #endif
 """
 
-    def process_option_code_template(self, template, option):
+    def process_pre_arg_assign(self, template, option):
+        if not option.get('auto_gpu', True):
+            return template
         call = 'THCPAutoGPU __autogpu_guard = THCPAutoGPU(args{});'.format(
             ', (PyObject*)self' if self.has_self else '')
 

tools/cwrap/plugins/BeforeAfterCall.py

@@ -18,6 +18,11 @@ class BeforeAfterCall(CWrapPlugin):
             prepend_str = before_call_template.substitute(args)
         template.insert(offset, prepend_str)
 
+    def process_pre_arg_assign(self, template, option):
+        if option.get('before_arg_assign'):
+            self.insert_snippet(template, option, 0, 'before_arg_assign')
+        return template
+
     def process_option_code_template(self, template, option):
         if option.get('before_call') or option.get('after_call'):
             call_idx = template.index('$call')

tools/cwrap/plugins/BoolOption.py

@@ -1,6 +1,12 @@
 from . import CWrapPlugin
 from string import Template
 
+import sys
+if sys.version_info[0] == 3:
+    string_type = str
+else:
+    string_type = basestring
+
 
 class BoolOption(CWrapPlugin):
 
@@ -9,11 +15,21 @@ class BoolOption(CWrapPlugin):
     def is_bool_option(self, arg):
         return arg['type'] == 'bool' and 'if_true' in arg and 'if_false' in arg
 
+    def process_declarations(self, declarations):
+        for declaration in declarations:
+            for option in declaration['options']:
+                for arg in option['arguments']:
+                    if self.is_bool_option(arg):
+                        arg['is_bool_option'] = True
+                        if isinstance(arg['if_true'], string_type):
+                            arg['type'] = 'const char*'
+        return declarations
+
     def get_type_check(self, arg, option):
-        if self.is_bool_option(arg):
+        if arg.get('is_bool_option', False):
             return Template('PyBool_Check($arg)')
 
     def get_type_unpack(self, arg, option):
-        if self.is_bool_option(arg):
+        if arg.get('is_bool_option', False):
             return Template(self.UNPACK_TEMPLATE.safe_substitute(
                 if_true=arg['if_true'], if_false=arg['if_false']))

tools/cwrap/plugins/Broadcast.py

@@ -0,0 +1,318 @@
+from . import CWrapPlugin
+from string import Template
+
+# Arguments to the Broadcast Plugin:
+# broadcast: args_to_broadcast_against [inplace] [fallback]
+# [args_to_broadcast_against]: either a single argument (e.g. "arg1") or a comma-seperated
+#                              list of two arguments (e.g. "tensor1,tensor2") indicating
+#                              arguments to broadcast specified argument (usually "self") against
+# [inplace] will generate code for in-place function, which doesn't allow the in-place
+#           argument to be broadcast
+# [fallback] if tensors aren't broadcastable, preserves "element number" pointwise behavior,
+#            where only number of elements need to match, and tensors are viewed as 1-dimensional.
+# [dims] specify if the tensors shouldn't be broadcast to a specific tensor or tensors, but a combination
+#        of individual dimension sizes of a set of tensors.  For example: addbmm(C,A,B) a.k.a. [C + A @ B]
+#        broadcasts C to the first dimension of A and the second dimension of B.  Each dimension is specified as
+#        [arg].dim[#] and dimensions are comma-separated.  So, to specify that the tensor should be
+#        broadcast to 3-dimensions with sizes:
+#        tensor0->size[0] x tensor1->size[1] x tensor2->size[2]
+#        you would write:
+#        dims:tensor0.dim0,tensor1.dim1,tensor2.dim2
+# [types] if the tensors should be of different types than THTensor, specify as X where
+#         the actual type to use is THXTensor (i.e. Byte for THByteTensor).  If the type
+#         should be THTensor, use 'Real'
+
+# For out of place:
+# Two args: expand the two args together
+# Three args (fused kernels): (e.g. addcmul) expand all three args together
+# Sketch of proof that this is the same:
+# consider addcmul, under expansion we want: a + (b * c) = (a + b * c) [all expanded together]
+# Let e(i, j) be the expansion of i with j, e(i, j, k) be the expansion of i with j,k
+#
+# Then a + (b * c) = e(a, e(b,c) * e(c,b)) + e(e(b,c)   * e(c,b), a)
+#                  = e(a, e(b,c))          + e(e(b,c)   * e(c,b), a)    (only size matters for second param)
+#                  = e(a,b,c)              + e(e(b,c)   * e(c,b), a)    (by associativity of max in expand)
+#                  = e(a,b,c)              + e(b,c,a)   * e(c,b,a)      (see L1)
+# which is a + b * c all expanded together
+#
+# L1: Show e(i * j, a) = e(i,a) * e(j,a) where i,j have same size
+# Consider any index _{ s_0, ..., s_n}
+# e(i * j, a) = (i*j)_{f(s_0), ...,f(s_n)} where f is the expansion of that dimension with a
+#             = i_{f(s_0), ..., f(s_n)} * j_{f(s_0), ..., f(s_n)} by definition of pointwise operator
+#             = e(i,a) * e(j,a)
+
+
+class Broadcast(CWrapPlugin):
+
+    # Save and restore passed in arguments in case later plugins use
+    POST_TEMPLATE = Template(
+        """${arg_op_other} = ${arg_op_other}_save;\n""")
+
+    def getPreArgStringTemplate(self, type=None):
+        if type is None:
+            ret = """THTensor *${arg_op_other}_save = ${arg_op_other};
+                     THTensorPtr ${arg_op_other}_guard(THTensor_(new)(LIBRARY_STATE_NOARGS));\n"""
+        else:
+            cpu_t = "TH" + type + "Tensor"
+            gpu_t = "THCuda" + type + "Tensor"
+            ret = ("#if !IS_CUDA\n" +
+                   cpu_t + " *${arg_op_other}_save = ${arg_op_other};\n" +
+                   cpu_t + "Ptr ${arg_op_other}_guard(" + cpu_t + "_new(LIBRARY_STATE_NOARGS));\n" +
+                   "#else\n" +
+                   gpu_t + " *${arg_op_other}_save = ${arg_op_other};\n" +
+                   "THPPointer<" + gpu_t + "> ${arg_op_other}_guard(\n" + gpu_t + "_new(LIBRARY_STATE_NOARGS));\n" +
+                   "#endif\n")
+        return Template(ret)
+
+    def getExpandTemplate(self, expand_call, success_code, raise_errors):
+        if not raise_errors:
+            return Template(
+                "bool expand_success = false;\n" +
+                "try {\n" +
+                expand_call +
+                "\nexpand_success = true;\n" +
+                "}\n"
+                "catch (std::exception &e) {}\n" +
+                "if(expand_success) {\n" +
+                success_code +
+                "\n}\n")
+        else:
+            return Template(
+                expand_call + "\n" +
+                success_code + "\n")
+
+    def getOutPlacePreExpand2Template(self, raise_errors):
+        expand_code = """expand_outplace2(LIBRARY_STATE ${arg_op_a}_guard.get(), ${arg_op_other}_guard.get(),
+                                          ${arg_op_a}, ${arg_op_other},
+                                          \"${op_a}\", \"${op_other}\", !${raise_errors});"""
+        success_code = """${arg_op_a} = ${arg_op_a}_guard.get();
+                          ${arg_op_other} = ${arg_op_other}_guard.get();"""
+        return self.getExpandTemplate(expand_code, success_code, raise_errors)
+
+    def getOutPlacePreExpand3Template(self, raise_errors):
+        expand_code = """expand_outplace3(LIBRARY_STATE ${arg_op_a}_guard.get(),
+                                          ${arg_op_other1}_guard.get(), ${arg_op_other2}_guard.get(),
+                                          ${arg_op_a}, ${arg_op_other1}, ${arg_op_other2},
+                                          \"${op_a}\", \"${op_other1}\", \"${op_other2}\", !${raise_errors});"""
+        success_code = """${arg_op_a} = ${arg_op_a}_guard.get();
+                          ${arg_op_other1} = ${arg_op_other1}_guard.get();
+                          ${arg_op_other2} = ${arg_op_other2}_guard.get();"""
+        return self.getExpandTemplate(expand_code, success_code, raise_errors)
+
+    OUT_PLACE_PRE_EXPAND_PRE_DIM_TEMPLATE = Template(
+        """if(THTensor_(nDimension)(LIBRARY_STATE ${arg_op_dim}) <= ${arg_op_dim_value}) {
+             THError("Argument %s requires at least %d dimensions, but only has %d",
+                     "${op_dim}", ${arg_op_dim_value} + 1, THTensor_(nDimension)(LIBRARY_STATE ${arg_op_dim}));
+           }
+           long ${arg_op_a}_dim${idx}_size = THTensor_(size)(LIBRARY_STATE ${arg_op_dim}, ${arg_op_dim_value});\n""")
+
+    OUT_PLACE_PRE_EXPAND1_DIM_TEMPLATE = Template(
+        """THLongStoragePtr ${arg_op_a}_storage(THLongStorage_newWithSize1(${arg_op_a}_dim0_size));\n""")
+
+    OUT_PLACE_PRE_EXPAND2_DIM_TEMPLATE = Template(
+        """THLongStoragePtr ${arg_op_a}_storage(
+               THLongStorage_newWithSize2(${arg_op_a}_dim0_size, ${arg_op_a}_dim1_size));\n""")
+
+    OUT_PLACE_PRE_EXPAND3_DIM_TEMPLATE = Template(
+        """THLongStoragePtr ${arg_op_a}_storage(
+               THLongStorage_newWithSize3(${arg_op_a}_dim0_size, ${arg_op_a}_dim1_size, ${arg_op_a}_dim2_size));\n""")
+
+    def getOutPlacePreExpandPostDimTemplate(self, raise_errors):
+        expand_code = """expand(LIBRARY_STATE ${arg_op_a}_guard.get(), ${arg_op_a}, ${arg_op_a}_storage);"""
+        success_code = """${arg_op_a} = ${arg_op_a}_guard.get();"""
+        return self.getExpandTemplate(expand_code, success_code, raise_errors)
+
+    OUT_PLACE_PRE_TEMPLATE = Template(
+        """${code_arg_op_a}${code_arg_op_other1}${code_arg_op_other2}
+           ${expand_code}""")
+
+    def getInPlacePreExpand1Template(self, raise_errors):
+        expand_code = """expand_inplace1(LIBRARY_STATE ${arg_op_other}_guard.get(), ${arg_op_other}, ${arg_op_a},
+                                         \"${op_other}\", \"${op_a}\", !${raise_errors});"""
+        success_code = """${arg_op_other} = ${arg_op_other}_guard.get();"""
+        return self.getExpandTemplate(expand_code, success_code, raise_errors)
+
+    def getInPlacePreExpand2Template(self, raise_errors):
+        expand_code = """expand_inplace2(LIBRARY_STATE ${arg_op_other1}_guard.get(), ${arg_op_other2}_guard.get(),
+                                         ${arg_op_other1}, ${arg_op_other2}, ${arg_op_a},
+                                         \"${op_other1}\", \"${op_other2}\", \"${op_a}\", !${raise_errors});"""
+        success_code = """${arg_op_other1} = ${arg_op_other1}_guard.get();
+                          ${arg_op_other2} = ${arg_op_other2}_guard.get();"""
+        return self.getExpandTemplate(expand_code, success_code, raise_errors)
+
+    IN_PLACE_PRE_TEMPLATE = Template(
+        """${code_arg_op_other1}${code_arg_op_other2}
+           ${expand_code}""")
+
+    def initialize(self, cwrap):
+        self.cwrap = cwrap
+
+    # Arguments:
+    # [0]: name of tensor to broadcast with (possibly two comma separated)
+    # [1] inplace (optional).  In place operations only broadcast on second tensor argument
+    # [2] fallback (optional).  Will fallback to applying to tensor of equal nElem if broadcast fails
+    def process_option_code_template(self, template, option):
+        new_code_pre = []
+        new_code_post = []
+        for _, arg in enumerate(option['arguments']):
+            if 'broadcast' not in arg:
+                continue
+
+            params = arg.get('broadcast').split(" ")
+            op_a = arg.get('assign_name', arg['name'])
+            in_place = "inplace" in params
+            raise_errors = "false" if "fallback" in params else "true"
+
+            param_others = params[0].split(",")
+            if len(param_others) > 2:
+                raise ValueError('Broadcast only supports up to 2 secondary parameters')
+            op_b = param_others[0]
+            op_c = param_others[1] if len(param_others) == 2 else None
+            arg_op_b = "arg_" + op_b
+            arg_op_a = "arg_" + op_a
+            arg_op_c = ("arg_" + op_c) if op_c else None
+
+            dims_kvs = []
+            for p in params:
+                if p.startswith("dims:"):
+                    assert(raise_errors == "true")
+                    if len(dims_kvs) != 0:
+                        raise ValueError("multiple specifications of dims")
+                    dims = p[len("dims:"):].split(",")
+                    for dim in dims:
+                        batchdim = dim.split(".")
+                        assert len(batchdim) == 2
+                        assert batchdim[1].startswith("dim")
+                        dim_val = batchdim[1][len("dim"):]
+                        dims_kvs.append({"op": batchdim[0], "arg_op": "arg_" + batchdim[0], "val": dim_val})
+
+            assert len(dims_kvs) <= 3
+            for p in params[1:]:
+                if p != "inplace" and p != "fallback" and not p.startswith("dims:") and not p.startswith("types:"):
+                    raise ValueError("invalid parameter {}".format(p))
+
+            type_op_b = None
+            type_op_c = None
+            for p in params:
+                if p.startswith("types:"):
+                    if not in_place and len(dims_kvs) > 0:
+                        raise ValueError("type specification not supported yet for out-of-place functions "
+                                         "that specify explicit dimensions")
+                    types = p[len("types:"):].split(",")
+                    assert(len(types) == (2 if op_c else 1))
+                    type_op_b = None if types[0] == "Real" else types[0]
+                    if op_c:
+                        type_op_c = None if types[1] == "Real" else types[1]
+
+            op_b_mapping = {
+                "op_a": op_a,
+                "op_other": op_b,
+                "arg_op_a": arg_op_a,
+                "arg_op_other": arg_op_b,
+                "raise_errors": raise_errors
+            }
+            op_c_mapping = {
+                "op_a": op_a,
+                "op_other": op_c,
+                "arg_op_a": arg_op_a,
+                "arg_op_other": arg_op_c,
+                "raise_errors": raise_errors
+            }
+
+            if in_place:
+                code_arg_op_other1 = self.getPreArgStringTemplate(type=type_op_b).substitute(op_b_mapping)
+                code_arg_op_other2 = (
+                    self.getPreArgStringTemplate(type=type_op_c).substitute(op_c_mapping) if op_c else "")
+
+                if op_c:
+                    expand_code = self.getInPlacePreExpand2Template(raise_errors == "true").substitute(
+                        op_b_mapping,
+                        op_other1=op_b,
+                        op_other2=op_c,
+                        arg_op_other1=arg_op_b,
+                        arg_op_other2=arg_op_c)
+                else:
+                    expand_code = self.getInPlacePreExpand1Template(raise_errors == "true").substitute(op_b_mapping)
+
+                new_code_pre.append(self.IN_PLACE_PRE_TEMPLATE.substitute(
+                    arg_op_a=arg_op_a,
+                    code_arg_op_other1=code_arg_op_other1,
+                    code_arg_op_other2=code_arg_op_other2,
+                    expand_code=expand_code,
+                    raise_errors=raise_errors))
+                new_code_pre.append("")
+
+                post_code = self.POST_TEMPLATE.substitute(op_b_mapping)
+                if op_c:
+                    post_code += self.POST_TEMPLATE.substitute(op_c_mapping)
+
+                new_code_post.append(post_code)
+                new_code_post.append("")
+            else:
+                if len(dims_kvs) != 0:
+                    code_arg_op_a = self.getPreArgStringTemplate().substitute(arg_op_other=arg_op_a)
+                    code_arg_op_other1 = ""
+                    code_arg_op_other2 = ""
+                    expand_code = ""
+                    for idx, kv in enumerate(dims_kvs):
+                        expand_code += self.OUT_PLACE_PRE_EXPAND_PRE_DIM_TEMPLATE.substitute(
+                            arg_op_a=arg_op_a,
+                            op_dim=kv["op"],
+                            arg_op_dim=kv["arg_op"],
+                            arg_op_dim_value=kv["val"],
+                            idx=idx)
+
+                    if len(dims_kvs) == 1:
+                        expand_code += self.OUT_PLACE_PRE_EXPAND1_DIM_TEMPLATE.substitute(
+                            arg_op_a=arg_op_a,
+                            arg_op_dim0=dims_kvs[0]["arg_op"])
+                    elif len(dims_kvs) == 2:
+                        expand_code += self.OUT_PLACE_PRE_EXPAND2_DIM_TEMPLATE.substitute(
+                            arg_op_a=arg_op_a,
+                            arg_op_dim0=dims_kvs[0]["arg_op"],
+                            arg_op_dim1=dims_kvs[1]["arg_op"])
+                    else:
+                        expand_code += self.OUT_PLACE_PRE_EXPAND3_DIM_TEMPLATE.substitute(
+                            arg_op_a=arg_op_a,
+                            arg_op_dim0=dims_kvs[0]["arg_op"],
+                            arg_op_dim1=dims_kvs[1]["arg_op"],
+                            arg_op_dim2=dims_kvs[2]["arg_op"])
+                    expand_code += self.getOutPlacePreExpandPostDimTemplate(raise_errors == "true").substitute(
+                        arg_op_a=arg_op_a,
+                        raise_errors=raise_errors)
+                    post_code = self.POST_TEMPLATE.substitute(arg_op_other=arg_op_a)
+
+                else:
+                    code_arg_op_a = self.getPreArgStringTemplate().substitute(arg_op_other=arg_op_a)
+                    code_arg_op_other1 = self.getPreArgStringTemplate(type=type_op_b).substitute(op_b_mapping)
+                    code_arg_op_other2 = (self.getPreArgStringTemplate(type=type_op_c).substitute(op_c_mapping)
+                                          if op_c else "")
+
+                    if op_c:
+                        expand_code = self.getOutPlacePreExpand3Template(raise_errors == "true").substitute(
+                            op_b_mapping,
+                            op_other1=op_b,
+                            op_other2=op_c,
+                            arg_op_other1=arg_op_b,
+                            arg_op_other2=arg_op_c)
+
+                    else:
+                        expand_code = self.getOutPlacePreExpand2Template(
+                            raise_errors == "true").substitute(op_b_mapping)
+
+                    post_code = self.POST_TEMPLATE.substitute(arg_op_other=arg_op_a)
+                    post_code += self.POST_TEMPLATE.substitute(op_b_mapping)
+                    post_code += self.POST_TEMPLATE.substitute(op_c_mapping) if op_c else ""
+
+                new_code_pre.append(self.OUT_PLACE_PRE_TEMPLATE.substitute(
+                    code_arg_op_a=code_arg_op_a,
+                    code_arg_op_other1=code_arg_op_other1,
+                    code_arg_op_other2=code_arg_op_other2,
+                    expand_code=expand_code))
+                new_code_pre.append("")
+
+                new_code_post.append(post_code)
+                new_code_post.append("")
+
+        template = new_code_pre + template + new_code_post
+        return template

tools/cwrap/plugins/CuDNNPlugin.py

@@ -1,4 +1,5 @@
 from string import Template
+import copy
 from copy import deepcopy
 from . import CWrapPlugin
 from itertools import product
@@ -17,6 +18,10 @@ class CuDNNPlugin(CWrapPlugin):
         'double': Template('THPDoubleUtils_unpackReal($arg)'),
     }
 
+    INPUT_ARGUMENT_MAP = {
+        'THTensor*': 'THVoidTensor*',
+    }
+
     TYPE_CHECK = {
         'Convolution*': Template('THPWrapper_check($arg)'),
         'THTensor*': Template('(PyObject*)Py_TYPE($arg) == tensorClass'),
@@ -79,6 +84,16 @@ static PyObject * $name(PyObject *self, PyObject *args, PyObject *kwargs)
     def get_type_check(self, arg, option):
         return self.TYPE_CHECK.get(arg['type'], None)
 
+    def get_assign_args(self, arguments):
+        assign_args = []
+        for arg in arguments:
+            arg = copy.copy(arg)
+            new_type = self.INPUT_ARGUMENT_MAP.get(arg['type'])
+            if new_type is not None:
+                arg['type'] = new_type
+            assign_args.append(arg)
+        return assign_args
+
     def get_wrapper_template(self, declaration):
         arg_desc = []
         for option in declaration['options']:
@@ -120,7 +135,7 @@ static PyObject * $name(PyObject *self, PyObject *args, PyObject *kwargs)
                     if arg['name'] in ['self', 'state', 'dataType', 'handle']:
                         arg['ignore_check'] = True
             declaration['options'] = self.filter_unique_options(declaration['options'])
-        return declarations
+        return [d for d in declarations if not d.get('only_register', False)]
 
     def filter_unique_options(self, options):
         def signature(option):
@@ -143,7 +158,7 @@ static PyObject * $name(PyObject *self, PyObject *args, PyObject *kwargs)
             return self.preprocessor_guard(code, declaration['defined_if'])
         return code
 
-    def process_all_unpacks(self, code, option):
+    def process_all_call_arg(self, code, option):
         return 'state, ' + code
 
     def declare_methods(self):

tools/cwrap/plugins/GILRelease.py

@@ -23,6 +23,8 @@ class GILRelease(CWrapPlugin):
     ]
 
     def process_option_code_template(self, template, option):
+        if option.get('with_gil', False):
+            return template
         call_idx = template.index('$call')
         template.insert(call_idx, self.BEFORE_CALL)
         template.insert(call_idx + 2, self.AFTER_CALL)

tools/cwrap/plugins/GenericNN.py

@@ -64,8 +64,9 @@ void $name($args)
         'THTensor*': 'thpp::Tensor*',
         'THCTensor*': 'thpp::Tensor*',
         'THIndexTensor*': 'thpp::Tensor*',
+        'THCIndexTensor*': 'thpp::Tensor*',
         'THIndex_t': 'long',
-        'real': 'double',
+        'accreal': 'double',
     }
 
     def __init__(self, header=False):
@@ -89,8 +90,8 @@ void $name($args)
             base_args = declaration['options'][0]['arguments']
             for option in declaration['options']:
                 for idx, arg in enumerate(option['arguments']):
-                    arg['formal_name'] = base_args[idx]['name']
-                    arg['formal_type'] = base_args[idx]['type']
+                    arg['assign_name'] = base_args[idx]['name']
+                    arg['assign_type'] = base_args[idx]['type']
                     if idx != 1:
                         arg['ignore_check'] = True
         return declarations
@@ -98,11 +99,19 @@ void $name($args)
     def get_arg_accessor(self, arg, option):
         return self.get_type_unpack(arg, option)
 
+    def process_pre_arg_assign(self, pre_arg_assign, option):
+        if option['backend'] == 'cunn':
+            # Enclose arg_assign with CUDA guard
+            pre_arg_assign.append('#ifdef WITH_CUDA')
+        return pre_arg_assign
+
     def process_option_code_template(self, template, option):
-        code = '// fill me in'
+        template = []
+        if option['backend'] == 'cunn':
+            template.append('#endif')
 
         def base_cast(arg, CReal, real):
-            name = arg['formal_name']
+            name = 'arg_' + arg['assign_name']
             type = arg['type']
             if type in self.REAL_TENSOR_TYPES:
                 return ('(TH{CReal}Tensor*){name}->cdata()'
@@ -120,7 +129,7 @@ void $name($args)
         def cast(arg, CReal, real):
             expr = base_cast(arg, CReal, real)
             if arg.get('optional', False):
-                name = arg['formal_name']
+                name = 'arg_' + arg['assign_name']
                 return '{name} ? {expr} : NULL'.format(name=name, expr=expr)
             return expr
 
@@ -135,6 +144,7 @@ void $name($args)
                 name=option['cname'],
                 float_args=',\n'.join(float_args),
                 double_args=',\n'.join(double_args))
+            template.append(code)
 
         elif option['backend'] == 'cunn':
             float_args = []
@@ -150,11 +160,13 @@ void $name($args)
                 float_args=',\n'.join(float_args),
                 double_args=',\n'.join(double_args),
                 half_args=',\n'.join(half_args))
+            template.append(code)
 
-        return [code, '']
+        template.append('')
+        return template
 
     def get_type_unpack(self, arg, option):
-        return Template(arg['name'])
+        return Template(arg.get('assign_name', arg['name']))
 
     def get_type_check(self, arg, option):
         if option['backend'] == 'cunn':
@@ -162,20 +174,20 @@ void $name($args)
         else:
             return Template('!is_cuda')
 
-    def get_formal_args(self, arguments):
-        formal_args = []
+    def get_assign_args(self, arguments):
+        assign_args = []
         for arg in arguments:
             arg = copy.copy(arg)
             new_type = self.INPUT_ARGUMENT_MAP.get(arg['type'])
             if new_type is not None:
                 arg['type'] = new_type
-            formal_args.append(arg)
-        return formal_args
+            assign_args.append(arg)
+        return assign_args
 
     def get_wrapper_template(self, declaration):
-        # get formal arguments string
+        # get assign arguments string
         base_arguments = declaration['options'][0]['arguments']
-        args = self.get_formal_args(base_arguments)
+        args = self.get_assign_args(base_arguments)
         arg_str = ', '.join([arg['type'] + ' ' + arg['name'] for arg in args])
 
         if self.header:
@@ -185,7 +197,7 @@ void $name($args)
             checked_args = []
             for arg in base_arguments:
                 if arg['type'] in tensor_types:
-                    name = arg.get('formal_name', arg['name'])
+                    name = arg.get('assign_name', arg['name'])
                     name_str = name
                     if arg.get('optional', False):
                         name_str = '?' + name_str

tools/cwrap/plugins/KwargsPlugin.py

@@ -24,6 +24,14 @@ class KwargsPlugin(CWrapPlugin):
                 for option in declaration['options']:
                     for arg in option['arguments']:
                         arg['no_kwargs'] = True
+        # we need to use offsets for arg position in *arg if kwarg_only args
+        # are not at the end
+        for declaration in declarations:
+            for option in declaration['options']:
+                offset = 0
+                for arg in option['arguments']:
+                    if arg.get('kwarg_only'):
+                        arg['no_idx'] = True
         return declarations
 
     def get_arg_accessor(self, arg, option):
@@ -53,9 +61,9 @@ class KwargsPlugin(CWrapPlugin):
                         name not in seen_args):
                     seen_args.add(name)
                     args.append(name)
-        declarations = '\n    '.join(['PyObject *__kw_{} = NULL;'.format(name) for name in args])
+        declarations = '\n    '.join(['PyObject *__kw_{} = NULL;'.format(a) for a in args])
         lookups = '\n      '.join(
-            ['__kw_{name} = PyDict_GetItemString(kwargs, "{name}");'.format(name=name) for name in args])
+            ['__kw_{name} = PyDict_GetItemString(kwargs, "{name}");'.format(name=a) for a in args])
         start_idx = code.find('{') + 1
         new_code = self.WRAPPER_TEMPLATE.substitute(declarations=declarations, lookups=lookups)
         return code[:start_idx] + new_code + code[start_idx:]

tools/cwrap/plugins/OptionalArguments.py

@@ -1,58 +1,18 @@
+import os
 from copy import deepcopy
 from . import CWrapPlugin
 from itertools import product
+from ...shared import cwrap_common
 
 
 class OptionalArguments(CWrapPlugin):
 
     def process_declarations(self, declarations):
-        new_options = []
         for declaration in declarations:
-            for option in declaration['options']:
-                optional_args = []
-                for i, arg in enumerate(option['arguments']):
-                    if 'default' in arg:
-                        optional_args.append(i)
-                for permutation in product((True, False), repeat=len(optional_args)):
-                    option_copy = deepcopy(option)
-                    for i, bit in zip(optional_args, permutation):
-                        arg = option_copy['arguments'][i]
-                        if not bit:
-                            arg['type'] = 'CONSTANT'
-                            arg['ignore_check'] = True
-                            # PyYAML interprets NULL as None...
-                            arg['name'] = 'NULL' if arg['default'] is None else arg['default']
-                    new_options.append(option_copy)
-            declaration['options'] = self.filter_unique_options(new_options)
-        return declarations
+            cwrap_common.enumerate_options_due_to_default(
+                declaration,
+                allow_kwarg=True,
+                type_to_signature={},
+                remove_self=False)
 
-    def filter_unique_options(self, options):
-        def signature(option, kwarg_only_count):
-            if kwarg_only_count == 0:
-                kwarg_only_count = None
-            else:
-                kwarg_only_count = -kwarg_only_count
-            arg_signature = '#'.join(
-                arg['type']
-                for arg in option['arguments'][:kwarg_only_count]
-                if not arg.get('ignore_check'))
-            if kwarg_only_count is None:
-                return arg_signature
-            kwarg_only_signature = '#'.join(
-                arg['name'] + '#' + arg['type']
-                for arg in option['arguments'][kwarg_only_count:]
-                if not arg.get('ignore_check'))
-            return arg_signature + "#-#" + kwarg_only_signature
-        seen_signatures = set()
-        unique = []
-        for option in options:
-            for num_kwarg_only in range(0, len(option['arguments']) + 1):
-                sig = signature(option, num_kwarg_only)
-                if sig not in seen_signatures:
-                    if num_kwarg_only > 0:
-                        for arg in option['arguments'][-num_kwarg_only:]:
-                            arg['kwarg_only'] = True
-                    unique.append(option)
-                    seen_signatures.add(sig)
-                    break
-        return unique
+        return declarations

tools/cwrap/plugins/ProcessorSpecificPlugin.py

@@ -0,0 +1,90 @@
+from copy import deepcopy
+from . import CWrapPlugin
+import yaml
+
+
+class ProcessorSpecificPlugin(CWrapPlugin):
+
+    def process_declarations(self, declarations):
+        # In order to move Torch's random functions into the same cwrap
+        # declaration, we need to be able to handle the fact that on the CPU
+        # these functions take a generator argument, while on the GPU, they
+        # do not. As such, we would like to split those declarations at cwrap
+        # runtime into two separate declarations, one for the CPU (unchanged),
+        # and one for the GPU (with the generator argument removed).
+        #
+        # For example, the declaration arguments:
+        # arguments:
+        #   - THTensor* self
+        #   - arg: THGenerator* generator
+        #     default: THPDefaultGenerator->cdata
+        #     kwarg_only: True
+        #
+        # Would have the generator argument removed when generating for the GPU
+        # backend.
+
+        def arg_contains_generator(arg):
+            return (arg['type'] == 'THGenerator*' or (arg.get('default', None)
+                    is not None and 'THPDefaultGenerator' in
+                    str(arg.get('default', ""))))
+
+        def split_candidate(declaration):
+            # First, check and see if it is a declaration for both CPU/GPU
+            if all([proc in declaration['backends'] for
+                    proc in ['CPU', 'CUDA']]):
+                for option in declaration['options']:
+                    for argument in option['arguments']:
+                        if arg_contains_generator(argument):
+                            return True
+
+            return False
+
+        def can_we_handle_the_split(declaration):
+            # hook into here if the split cannot happen for some reason
+            return True
+
+        def generator_split(declaration):
+            # the split must make two changes: 1. remove the generator argument
+            # for the GPU, and 2. assign the correct backends/types to the
+            # split declaration
+            dec_cpu = declaration
+            dec_gpu = deepcopy(declaration)
+
+            # Remove GPU backend and types from dec_cpu
+            dec_cpu['backends'].remove('CUDA')
+            if dec_cpu.get('backend_type_pairs', False):
+                dec_cpu['backend_type_pairs'] = (
+                    [pair for pair in dec_cpu['backend_type_pairs'] if
+                     pair[1] == 'CPU'])
+            # also need to reach into options
+            for option in dec_cpu['options']:
+                option['backends'].remove('CUDA')
+
+            # Remove CPU backend and types from dec_gpu
+            dec_gpu['backends'].remove('CPU')
+            if dec_gpu.get('backend_type_pairs', False):
+                dec_gpu['backend_type_pairs'] = (
+                    [pair for pair in dec_gpu['backend_type_pairs'] if
+                     pair[1] == 'CUDA'])
+            # also need to reach into options
+            for option in dec_gpu['options']:
+                option['backends'].remove('CPU')
+
+            # Remove generator arguments from dec_gpu options
+            for option in dec_gpu['options']:
+                option['arguments'] = (
+                    [arg for arg in option['arguments'] if
+                     not arg_contains_generator(arg)])
+
+            return [dec_cpu, dec_gpu]
+
+        decs = []
+        for declaration in declarations:
+            if split_candidate(declaration):
+                assert(can_we_handle_the_split(declaration))
+                newdecs = generator_split(declaration)
+                decs.extend(newdecs)
+            else:
+                decs.append(declaration)
+
+        return decs

tools/cwrap/plugins/THPPlugin.py

@@ -15,9 +15,12 @@ class THPPlugin(CWrapPlugin):
         'THTensor*': Template('((THPTensor*)$arg)->cdata'),
         'THBoolTensor*': Template('((THPBoolTensor*)$arg)->cdata'),
         'THIndexTensor*': Template('((THPIndexTensor*)$arg)->cdata'),
+        'THIntegerTensor*': Template('((THPIntegerTensor*)$arg)->cdata'),
 
         'THCudaTensor*': Template('((THCPFloatTensor*)$arg)->cdata'),
         'THCudaDoubleTensor*': Template('((THCPDoubleTensor*)$arg)->cdata'),
+        'THCudaIntTensor*': Template('((THCPIntTensor*)$arg)->cdata'),
+        'THCudaLongTensor*': Template('((THCPLongTensor*)$arg)->cdata'),
 
         'THSFloatTensor*': Template('((THSPFloatTensor*)$arg)->cdata'),
         'THSDoubleTensor*': Template('((THSPDoubleTensor*)$arg)->cdata'),
@@ -50,9 +53,12 @@ class THPPlugin(CWrapPlugin):
         'THTensor*': Template('(PyObject*)Py_TYPE($arg) == THPTensorClass'),
         'THBoolTensor*': Template('(PyObject*)Py_TYPE($arg) == THPBoolTensorClass'),
         'THIndexTensor*': Template('(PyObject*)Py_TYPE($arg) == THPIndexTensorClass'),
+        'THIntegerTensor*': Template('(PyObject*)Py_TYPE($arg) == THPIntegerTensorClass'),
 
         'THCudaTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPFloatTensorClass'),
         'THCudaDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPDoubleTensorClass'),
+        'THCudaIntTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPIntTensorClass'),
+        'THCudaLongTensor*': Template('(PyObject*)Py_TYPE($arg) == THCPLongTensorClass'),
 
         'THSDoubleTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPDoubleTensorClass'),
         'THSFloatTensor*': Template('(PyObject*)Py_TYPE($arg) == THSPFloatTensorClass'),
@@ -82,8 +88,11 @@ class THPPlugin(CWrapPlugin):
     RETURN_WRAPPER = {
         'THTensor*': Template('return THPTensor_(New)($result);'),
         'THSTensor*': Template('return THSPTensor_(New)($result);'),
+        'THIndexTensor*': Template('return THPIndexTensor_(New)($result);'),
         'THLongTensor*': Template('return THPLongTensor_New($result);'),
         'THLongStorage*': Template('return THPLongStorage_New($result);'),
+        'THCudaIntTensor*': Template('return THCPIntTensor_New($result);'),
+        'THCudaLongTensor*': Template('return THCPLongTensor_New($result);'),
         # TODO: make it smarter - it should return python long if result doesn't fit into an int
         'long': Template('return PyInt_FromLong($result);'),
         'accreal': Template('return THPUtils_(newAccreal)($result);'),
@@ -118,7 +127,7 @@ PyObject * $name(PyObject *self, PyObject *args, PyObject *kwargs)
     """)
 
     ALLOCATE_TMPL = Template("""\
-THP${type}TensorPtr _${name}_guard = (THP${type}Tensor*) THP${type}Tensor_NewEmpty();
+THP${type}TensorPtr _${name}_guard((THP${type}Tensor*) THP${type}Tensor_NewEmpty());
 if (!_${name}_guard.get()) return NULL;
 THP${type}Tensor* $name = _${name}_guard.get();
 """)
@@ -149,6 +158,7 @@ ${cpu}
         'THIntTensor*': _allocate('Int', ALLOCATE_TMPL),
         'THBoolTensor*': _allocate('Byte', ALLOCATE_TMPL, ALLOCATE_CUDA),
         'THIndexTensor*': _allocate('Long', ALLOCATE_TMPL, ALLOCATE_CUDA),
+        'THIntegerTensor*': _allocate('Int', ALLOCATE_TMPL, ALLOCATE_CUDA),
 
         'THSTensor*': _allocate('', ALLOCATE_TMPL, sparse=True),
     }
@@ -163,10 +173,13 @@ ${cpu}
         'THIntTensor*': '" THPModuleStr "IntTensor',
         'THBoolTensor*': '" THPModuleStr "ByteTensor',
         'THIndexTensor*': '" THPModuleStr "LongTensor',
+        'THIntegerTensor*': '" THPModuleStr "IntTensor',
         'THFloatTensor*': '" THPModuleStr "FloatTensor',
         'THDoubleTensor*': '" THPModuleStr "DoubleTensor',
         'THCudaTensor*': 'torch.cuda.FloatTensor',
         'THCudaDoubleTensor*': 'torch.cuda.DoubleTensor',
+        'THCudaIntTensor*': 'torch.cuda.IntTensor',
+        'THCudaLongTensor*': 'torch.cuda.LongTensor',
         'THSize*': 'torch.Size',
         'THStride*': 'tuple',
         'long': 'int',
@@ -174,10 +187,12 @@ ${cpu}
         'double': 'float',
         'accreal': '" RealStr "',
         'bool': 'bool',
+        'const char*': 'bool',  # Can come only from bool option.
     }
 
     OUT_INIT = """
     __out = kwargs ? PyDict_GetItemString(kwargs, "out") : NULL;
+    if (__out == Py_None) { __out = NULL; __dictcount--; __argcount--; }
     """
 
     def __init__(self):
@@ -303,8 +318,6 @@ ${cpu}
 
     def process_declarations(self, declarations):
         new_declarations = []
-        register_only = [d for d in declarations if d.get('only_register', False)]
-        declarations = [d for d in declarations if not d.get('only_register', False)]
 
         def has_arg_type(declaration, type_name):
             return any(arg['type'] == type_name
@@ -321,9 +334,101 @@ ${cpu}
                        for option in declaration['options']
                        for arg in option['arguments'])
 
+        def backends_types_to_defined_if_string(declaration):
+            # A declaration has two fields: 'backend', which stores a list of
+            # backends (currently 'cpu' and 'cuda') the declaration applies
+            # to, and 'types', which stores a list of real types the
+            # declaration applies to. In PyTorch, when a function is only
+            # supported by a subset of types, we wrap it in macro definition
+            # checks.
+            #
+            # Previously, we manually required the cwrap declaration to
+            # specify for which backend/type combinations a function was
+            # defined for. Now, we explicitly list the types and backends for
+            # a declaration, if it should only be supported for a specific
+            # subset of types, backends, or type-backend pairs.
+
+            types = declaration.get('types', [])
+            backends = declaration['backends']
+            all_backends = ['CPU', 'CUDA']
+
+            def get_defined_string(backend, real):
+                if backend == 'CUDA':
+                    if real == 'all':
+                        return "IS_CUDA"
+                    else:
+                        return 'CUDA_{0}'.format(real.upper())
+                else:
+                    if real == 'all':
+                        return "!IS_CUDA"
+                    else:
+                        return 'defined(TH_REAL_IS_{0})'.format(real.upper())
+
+            def expand_composite_type(p, t):
+                if t == 'floating_point':
+                    result = ['double', 'float']
+                    if p == 'CUDA':
+                        result.append('half')
+                elif t == 'integral':
+                    result = ['byte', 'char', 'short', 'int', 'long']
+                else:
+                    result = [t]
+                return result
+
+            defineds = []
+
+            # The logic below does not handle corner cases well. We allow the
+            # declaration to have a field 'backend_type_pairs' that stores a
+            # dictionary from type --> backend representing allowed
+            # combinations. Let's use these first.
+            for pair in declaration.get('backend_type_pairs', []):
+                p, t = pair
+                defineds.extend([get_defined_string(p, et) for et in
+                                expand_composite_type(p, t)])
+
+            # In the base case, types is empty and backends contains both
+            # 'CPU' and 'CUDA' --> this means we support all types, and our
+            # string should be empty, or simply the list of explict type
+            # backend pairs
+            if (len(types) == 0 and all([proc in backends for proc in
+                                         all_backends])):
+                return " || ".join(defineds)
+
+            # Case 2: types is empty, but only one backend type is specified
+            if len(types) == 0 and len(backends) == 1:
+                defineds.append('IS_CUDA' if backends[0] == 'CUDA' else
+                                "!IS_CUDA")
+                return " || ".join(defineds)
+
+            # Else, we loop overall all of the backend, type pairs and add
+            # them
+            for p in backends:
+                for t in types:
+                    defineds.extend([get_defined_string(p, et) for et in
+                                    expand_composite_type(p, t)])
+
+            return " || ".join(defineds)
+
         for declaration in declarations:
+            # Disable all methods for THHalfTensor, unless cpu_half is True
+
+            dfstr = backends_types_to_defined_if_string(declaration)
+            if len(dfstr) > 0:
+                # for now, need to check for distributed defined if as well
+                if 'defined_if' in declaration:
+                    declaration['defined_if'] += ' && (' + dfstr + ')'
+                else:
+                    declaration['defined_if'] = dfstr
+
+            if not declaration.get('cpu_half', False):
+                defined_if = '!defined(TH_REAL_IS_HALF)'
+                if 'defined_if' in declaration:
+                    defined_if += ' && (' + declaration['defined_if'] + ')'
+                declaration['defined_if'] = defined_if
+
             if declaration.get('only_register', False):
                 continue
+
             declaration.setdefault('python_name', declaration['name'])
             declaration.setdefault('variables', [])
             if has_arg_type(declaration, 'THSize*'):
@@ -334,15 +439,23 @@ ${cpu}
                 declaration['variables'] += ['PyObject *__out;']
                 self.generate_out_options(declaration)
             if has_long_args(declaration):
-                declaration['no_kwargs'] = True
+                for option in declaration['options']:
+                    for arg in option['arguments']:
+                        if arg.get('long_args', False):
+                            arg['no_kwargs'] = True
             for option in declaration['options']:
                 option['cname'] = 'TH{}Tensor_({})'.format(
                     'S' if option.get('sparse', False) else '', option['cname'])
-            if declaration.get('with_stateless', False) or declaration.get('only_stateless', False):
+                if option.get('sparse', False):
+                    defined_if = option.get('defined_if', '')
+                    option['defined_if'] = '!IS_DISTRIBUTED' + (' && ' if defined_if else '') + defined_if
+
+            variants = declaration.get('variants', ['method'])
+            if 'function' in variants:
                 stateless_declaration = self.make_stateless(declaration)
                 new_declarations.append(stateless_declaration)
                 self.stateless_declarations.append(stateless_declaration)
-            if declaration.get('only_stateless', False):
+            if 'method' not in variants:
                 continue
 
             self.declarations.append(declaration)
@@ -353,9 +466,15 @@ ${cpu}
                     if arg['name'] == 'self':
                         arg['ignore_check'] = True
 
-        declarations = [d for d in declarations if not d.get('only_stateless', False)]
-        self.declarations.extend(filter(lambda x: not x.get('only_stateless', False), register_only))
-        self.stateless_declarations.extend(filter(lambda x: x.get('only_stateless', False), register_only))
+        register_only = [d for d in declarations if d.get('only_register', False)]
+        declarations = [d for d in declarations
+                        if (('method' in d.get('variants', ['method'])) and
+                            (not d.get('only_register', False)))]
+        self.declarations.extend(filter(lambda x: 'method' in x.get('variants',
+                                 ['method']), register_only))
+        self.stateless_declarations.extend(filter(lambda x: 'method' not in
+                                           x.get('variants', ['method']),
+                                           register_only))
 
         self.process_docstrings()
 
@@ -369,6 +488,7 @@ ${cpu}
         for option in declaration['options']:
             for arg in option['arguments']:
                 if arg['name'] == 'self':
+                    arg['assign_name'] = 'self'
                     arg['name'] = 'source'
         return declaration
 
@@ -390,11 +510,14 @@ ${cpu}
             if 'defined_if' in declaration:
                 entry = self.preprocessor_guard(entry, declaration['defined_if'])
             tensor_methods += entry
-        return self.TENSOR_METHODS_DECLARATION.substitute(
+        generated = self.TENSOR_METHODS_DECLARATION.substitute(
             methods=tensor_methods,
             stateless=('' if not stateless else 'stateless_'),
             sparse=('' if not sparse else 'S'),
         )
+        if sparse:
+            generated = '#if !defined(TH_REAL_IS_HALF) && !IS_DISTRIBUTED\n' + generated + '\n#endif\n\n'
+        return generated
 
     def process_full_file(self, code):
         # We have to find a place before all undefs
@@ -415,7 +538,7 @@ ${cpu}
             return self.preprocessor_guard(code, declaration['defined_if'])
         return code
 
-    def process_all_unpacks(self, code, option):
+    def process_all_call_arg(self, code, option):
         return 'LIBRARY_STATE ' + code
 
     def process_all_checks(self, code, option):
@@ -434,12 +557,25 @@ ${cpu}
 
         if any(arg.get('long_args', False) for arg in option['arguments']):
             code = code.replace('__argcount ==', '__argcount >=')
-            expected = str(int(option.get('output_provided', False)))
+            expected = str(int(option.get('output_provided', False)) +
+                           sum(not arg.get('no_kwargs', False) and not arg.get('ignore_check', False)
+                               for arg in option['arguments']))
             code = '__dictcount == ' + expected + ' &&\n          ' + code
 
         return code
 
-    def process_option_code_template(self, template, option):
+    def process_option_code(self, code, option):
+        if option.get('defined_if', ''):
+            defined_if = option['defined_if']
+            placeholder = ''
+            # This means that it's a first option, so we need a dummy if,
+            # so the next option can be an else if.
+            if 'else if' not in code:
+                placeholder = '\n    #else\n    if (false) {'
+            return '#if ' + defined_if + '\n          ' + code + placeholder + '\n    #endif\n'
+        return code
+
+    def process_pre_arg_assign(self, template, option):
         new_args = []
         for arg in option['arguments']:
             if not option.get('output_provided', True) and arg.get('output'):

tools/cwrap/plugins/WrapDim.py

@@ -0,0 +1,40 @@
+from . import CWrapPlugin
+from string import Template
+
+
+class WrapDim(CWrapPlugin):
+
+    NDIM_TEMPLATE = Template(
+        """${arg_tensor}->nDimension""")
+
+    CODE_TEMPLATE = Template(
+        """THPUtils_assert(${arg_dim} >= -(${ndim}) && ${arg_dim} < (${ndim}),
+         "dimension out of range (expected to be in range of [%d, %d], but got %d)",
+         -(${ndim}), (${ndim})-1, ${arg_dim});
+         if (${arg_dim} < 0) ${arg_dim} += (${ndim});""")
+
+    def initialize(self, cwrap):
+        self.cwrap = cwrap
+
+    def process_option_code_template(self, template, option):
+        new_code = []
+        for i, arg in enumerate(option['arguments']):
+            if 'wrap_dim' not in arg:
+                continue
+
+            params = arg.get('wrap_dim').split("+")
+            arg_tensor = params[0]
+
+            arg_tensor = "arg_" + arg_tensor
+            arg_dim = "arg_" + arg.get('assign_name', arg['name'])
+
+            params[0] = self.NDIM_TEMPLATE.substitute(arg_tensor=arg_tensor)
+            ndim = "+".join(params)
+
+            new_code.append(self.CODE_TEMPLATE.substitute(
+                arg_dim=arg_dim,
+                ndim=ndim))
+            new_code.append("")
+
+        template = new_code + template
+        return template

tools/cwrap/plugins/__init__.py

@@ -1,49 +1,422 @@
 
 class CWrapPlugin(object):
+    """Base class from which all cwrap plugins should inherit.
+
+    Override any of the following methods to implement the desired wrapping
+    behavior.
+    """
 
     def initialize(self, cwrap):
+        """Initialize the Plugin class prior to calling any other functions.
+
+        It is used to give the Plugin access to the cwrap object's helper
+        functions and state.
+
+        Args:
+            cwrap: the cwrap object performing the wrapping.
+
+        """
         pass
 
     def get_type_check(self, arg, option):
+        """Used to generate code for runtime checks of object types.
+
+        The type can be found in arg['type']. For example, it could be
+        THTensor*. If this Plugin recognizes the type in arg, it should
+        return a Template string containing code that checks whether a
+        Python object is of this type. For example, the return type in
+        this case would be:
+
+        Template('(PyObject*)Py_TYPE($arg) == THPTensorClass')
+
+        As a simpler example, if the type == 'bool' then we would return:
+
+        Template('PyBool_Check($arg)')
+
+        Note that the name of the identifier that will be subsituted must be
+        $arg.
+
+        Args:
+            arg: a Python object with a 'type' field representing the type
+            to generate a check string for.
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            A Template string as described above, or None if this Plugin does
+            not have a corresponding type check for the passed type.
+
+        """
         pass
 
     def get_type_unpack(self, arg, option):
+        """Used to generate code unpacking of Python objects into C types.
+
+        Similar to get_type_check, but for unpacking Python objects into their
+        corresponding C types. The type is once again accessible via
+        arg['type']. This time we return a Template string that unpacks an
+        object. For a THTensor*, we know that the corresponding PyTorch type is
+        a THPTensor*, so we need to get the cdata from the object. So we would
+        return:
+
+        Template('((THPTensor*)$arg)->cdata')
+
+        For a simpler type, such as a long, we could do:
+
+        Template('PyLong_AsLong($arg)')
+
+        though in practice we will use our own custom unpacking code. Once
+        again, $arg must be used as the identifier.
+
+        Args:
+            arg: a Python object with a 'type' field representing the type
+            to generate a unpack string for.
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            A Template string as described above, or None if this Plugin does
+            not have a corresponding type unpack for the passed type.
+
+        """
         pass
 
     def get_return_wrapper(self, option):
+        """Used to generate code wrapping a function's return value.
+
+        Wrapped functions should always return a PyObject *. However,
+        internally, the code will be working with C objects or primitives.
+        Therefore, if a function has a return value we need to convert it back
+        to a PyObject * before the function returns. Plugins can override this
+        function to generate wrapper code for returning specific C types. The
+        type is accessible via option['return'].
+
+        Continuing on with our THTensor* example, we might do something like:
+
+        Template('return THPTensor_(New)($result);')
+
+        In general, you want to do return <statement>; In this case, we call
+        into THP's library routine that takes a THTensor* (the $result
+        identifier) and returns a PyObject *.
+
+        For a bool, we could do Template('return PyBool_FromLong($result);').
+
+        Note that in other cases, our logic might be more complicated. For
+        example, if our return value is also an argument to the function call,
+        we could need to increase the reference count prior to returning.
+
+        Args:
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            A Template string as described above, or None if this Plugin does
+            not have a corresponding return wrapper for the functions return
+            type or specifier.
+
+        """
         pass
 
     def get_wrapper_template(self, declaration):
+        """Used to create a code template to wrap the options.
+
+        This function returns a Template string that contains the function call
+        for the overall declaration, including the method definition, opening
+        and closing brackets, and any additional code within the method body.
+        Look through the examples to get a sense of what this might look like.
+        The only requirements are that it contains unsubstituted template
+        identifiers for anything the cwrap engine expects.
+
+        Note that for any declaration only one Plugin can generate the wrapper
+        template.
+
+        Args:
+            declaration: the declaration for the wrapped method.
+
+        Returns:
+            A template string representing the entire function declaration,
+            with identifiers as necessary.
+
+        """
+        pass
+
+    def get_assign_args(self, arguments):
+        """Used to modify argument metadata prior to assignment.
+
+        We have already setup argument checking, and how to unpack arguments.
+        This function allows you to modify the metadata of an argument prior to
+        actually performing the assignment. For example, you might want to
+        check that an argument is of a specific type, but when unpacking it you
+        might want to treat it as a different type. This function will allow
+        you to do stuff like that --> e.g. you could set the 'type' field for a
+        particular argument to be something else.
+
+        Args:
+            arguments: a list of argument metadata dictionaries.
+
+        Returns:
+            The same list of arguments, with any modifications as you see fit.
+
+        """
         pass
 
     def get_arg_accessor(self, arg, option):
+        """Used to generate a string for accessing the passed arg.
+
+        One of the key components of the YAML definition for a method to be
+        wrapped are the arguments to that method. Override this function to
+        show how to access that specific arg in the code. For example, you
+        might do something different if the argument is a keyword argument, or
+        a constant, or self. The base cwrap plugin has a fallback arg accessor
+        for loading elements from the args PyObject * tuple passed to the
+        function.
+
+        Its best to look at some of the existing Plugins to get a sense of what
+        one might do.
+
+        Args:
+            arg: a dictionary specifying attributes of the arg to be accessed
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            A a string (note: not a Template string!) of code that can be used
+            to access the given arg. If the plugin does not know how to access
+            the arg, return None.
+        """
         pass
 
     def process_full_file(self, code):
+        """Used to modify the code for the entire output file.
+
+        The last thing any plugin can do. Code contains the results of wrapping
+        all the declarations. The plugin can do things like adding header
+        guards, include statements, etc.
+
+        Args:
+            code: a string source code for the wrapped declarations.
+
+        Returns:
+            The same code, modified as the plugin sees fit.
+
+        """
         return code
 
     def process_single_check(self, code, arg, arg_accessor):
+        """Used to postprocess a type check.
+
+        Above we defined a function get_type_check that returns a Template
+        string that allows for type checking a PyObject * for a specific type.
+        In this function, the passed "code" is a combination of that type check
+        along with a specific arg_accessor pasted in. For example:
+
+        '(PyObject*)Py_TYPE(PyTuple_GET_ITEM(args, 1)) == THPTensorClass'
+
+        This function can be overriden to support modifying this check string.
+        For example, if an argument can be null, we might want to check and see
+        if the type is Py_None, as well.
+
+        Args:
+            code: The string code representing a type check for a specific
+            argument being accessed.
+            arg: dictionary containing properties of that specific argument
+            arg_accessor: the arg_accessor string for that specific argument.
+            Note that this is likely also embedded in code, but if you want to
+            be able to access this arg and throw away the other code, you can
+            do so.
+
+        Returns:
+            A string representing the processed check/access string for this
+            arg. If the plugin does not know how to modify a specific input, it
+            should return the original code.
+
+        """
         return code
 
     def process_all_checks(self, code, option):
+        """Used to generate additional checks based on all the individual ones.
+
+        After individually processing each argument with get_type_check,
+        get_arg_accessor, process_single_check, this function allows you to
+        inspect the combined checks and do any additional checking/modify that
+        string as you see fit. In particular, given code is a string like:
+
+        CHECK_TYPE(GET_ARG(0)) && CHECK_TYPE(GET_ARG(1)) && ..
+
+        We can process it as we see fit. For example, we may want to add a
+        check at the beginning that we have the specified number of arguments.
+
+        Args:
+            code: A string representing each argument check separated by an
+            '&&'. code can be None if there are no arguments to be checked.
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            The modified code string with any additional checks, or just the
+            existing code if no modifications are to be made.
+
+        """
         return code
 
     def process_single_unpack(self, code, arg, arg_accessor):
+        """Used to postprocess a type unpack.
+
+        Same as process_single_check above, but for type unpacking. E.g. an
+        example code could be:
+
+        PyLong_FromLong(PyTuple_GET_ITEM(args, 0))
+
+        And this code could modify that as it sees fit. For example, if the
+        result of accessing the argument is None, we would not want to call the
+        unpacking code.
+
+        Args:
+            code: The string code representing a type unpack for a specific
+            argument being accessed.
+            arg: dictionary containing properties of that specific argument
+            arg_accessor: the arg_accessor string for that specific argument.
+            Note that this is likely also embedded in code, but if you want to
+            be able to access this arg and throw away the other code, you can
+            do so.
+
+        Returns:
+            A string representing the processed unpack/access string for this
+            arg. If the plugin does not know how to modify a specific input, it
+            should return the original code.
+
+        """
         return code
 
-    def process_all_unpacks(self, code, option):
+    def process_all_call_arg(self, code, option):
+        """Used to modify the arguments to the underlying C function call.
+
+        Code is the string of comma-separated arguments that will be passed to
+        the wrapped C function. You can use this function to modify that string
+        as you see fit. For example, THP prepends the LIBRARY_STATE definition
+        so that the generated code will follow the conventions it uses for
+        writing one function for both TH/THC calls.
+
+        Args:
+            code: A string as described above.
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            The same code, modified as the plugin sees fit.
+
+        """
         return code
 
     def process_option_code(self, code, option):
+        """Used to modify the entire code body for an option.
+
+        Code in this case is a string containing the entire generated code for
+        a specific option. Note that this body includes the checks for each
+        option, i.e. if (type checks for one permutation) { ... } else if (type
+        checks for another permutation) { ... } etc.
+
+        Args:
+            code: string representing the generated code for the option
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            The same code, modified as the plugin sees fit.
+
+        """
         return code
 
     def process_wrapper(self, code, declaration):
+        """Used to modify the entire code body for a declaration.
+
+        Code in this case is a string containing the entire generated code for
+        a specific declaration. This code can be modified as the plugin sees
+        fit. For example, we might want to wrap the function in preprocessor
+        guards if it is only enabled for floats.
+
+        Args:
+            code: string representing the generated code for the declaration
+            declaration: the declaration metadata.
+
+        Returns:
+            The same code, modified as the plugin sees fit.
+
+        """
         return code
 
     def process_declarations(self, declarations):
+        """Used to process/modify the function's declaration.
+
+        Cwrap loads the YAML of a function to be cwrap'd into a dictionary.
+        This is known as the declaration. The cwrap code sets some defaults as
+        necessary, and then passes this dictionary to process_declarations.
+        Overriding this code allows the plugin to modify this declaration as it
+        sees fit prior to any code generation. The plugin may add, remove or
+        modify the fields of the declaration dictionary. It can also save state
+        to the Plugin for use in subsequent function overrides.
+
+        Its best to look at some of the existing Plugins to get a sense of what
+        one might do.
+
+        Args:
+            declarations: a list of declarations, i.e. dictionaries that define
+            the function(s) being wrapped. Note that this can be plural, so the
+            function must take care to modify each input declaration.
+
+        Returns:
+            Those same declarations, modified as the Plugin sees fit. Note that
+            you could insert a declaration, if you wanted to take an input
+            declaration and e.g. wrap it multiple times.
+
+        """
         return declarations
 
     def process_option_code_template(self, template, option):
+        """Used to modify the code template for the option.
+
+        The "code template" can be thought of the actual body implementing the
+        wrapped function call --> i.e. it is not the argument check,
+        assignment, etc. but the actual logic of the function. The template is
+        a list containing two operations: the $call, and the $return_result.
+        These represent the "locations" where the function call will happen,
+        and the function will return.
+
+        This function can modify the list to insert arbitrary code around the
+        $call and $return_result. For example, one might want to wrap the code
+        in a try/catch, or post-process the result in some way. This allows a
+        plugin to do that.
+
+        Args:
+            template: a list containing $call and $return_result, in addition
+            to any arbitrary code inserted by other plugins.
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            The same "code template", possibly modified by this plugin.
+
+        """
+        return template
+
+    def process_pre_arg_assign(self, template, option):
+        """Used to include any code before argument assignment.
+
+        This function can be used to insert any code that will be part of the
+        resulting function. The code is inserted after argument checks occur,
+        but before argument assignment.
+
+        Args:
+            template: String representing the code to be inserted. If other
+            plugins have included code for pre_arg_assign, it will be included
+            here.
+            option: dictionary containing the information for this specific
+            option.
+
+        Returns:
+            template, with any additional code if needed.
+
+        """
         return template
 
 
@@ -59,3 +432,5 @@ from .GILRelease import GILRelease
 from .AutoGPU import AutoGPU
 from .CuDNNPlugin import CuDNNPlugin
 from .GenericNN import GenericNN
+from .WrapDim import WrapDim
+from .Broadcast import Broadcast

tools/docker/Dockerfile-v6

@@ -1,40 +0,0 @@
-FROM nvidia/cuda:8.0-devel-ubuntu14.04 
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-         build-essential \
-         cmake \
-         git \
-         curl \
-         ca-certificates \
-         libjpeg-dev \
-         libpng-dev &&\
-     rm -rf /var/lib/apt/lists/*
-
-
-RUN curl -fsSL http://developer.download.nvidia.com/compute/redist/cudnn/v6.0/cudnn-8.0-linux-x64-v6.0-rc.tgz -O && \
-tar -xzf cudnn-8.0-linux-x64-v6.0-rc.tgz -C /usr/local && \
-    rm cudnn-8.0-linux-x64-v6.0-rc.tgz && \
-    ldconfig
-RUN ln -s /usr/local/cuda/lib64/libcudnn.so.6.0.5 /usr/lib/x86_64-linux-gnu/libcudnn.so.6.0.5
-
-RUN curl -o ~/miniconda.sh -O  https://repo.continuum.io/miniconda/Miniconda3-4.2.12-Linux-x86_64.sh  && \
-     chmod +x ~/miniconda.sh && \
-     ~/miniconda.sh -b -p /opt/conda && \     
-     rm ~/miniconda.sh && \
-     /opt/conda/bin/conda install conda-build && \
-     /opt/conda/bin/conda create -y --name pytorch-py35 python=3.5.2 numpy scipy ipython mkl&& \
-     /opt/conda/bin/conda clean -ya 
-ENV PATH /opt/conda/envs/pytorch-py35/bin:$PATH
-RUN conda install --name pytorch-py35 -c soumith magma-cuda80
-# This must be done before pip so that requirements.txt is available
-WORKDIR /opt/pytorch
-COPY . .
-
-RUN cat requirements.txt | xargs -n1 pip install --no-cache-dir && \
-    TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1+PTX" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
-    CMAKE_LIBRARY_PATH=/opt/conda/envs/pytorch-py35/lib \
-    CMAKE_INCLUDE_PATH=/opt/conda/envs/pytorch-py35/include \
-    pip install -v .
-
-WORKDIR /workspace
-RUN chmod -R a+w /workspace

tools/docker/Dockerfile_runtime

@@ -0,0 +1,27 @@
+FROM ubuntu:16.04 
+
+LABEL com.nvidia.volumes.needed="nvidia_driver"
+RUN apt-get update && apt-get install -y --no-install-recommends \
+         build-essential \ 
+         git \
+         curl \
+         ca-certificates \
+         libjpeg-dev \
+         libpng-dev && \
+     rm -rf /var/lib/apt/lists/*
+
+RUN curl -o ~/miniconda.sh -O  https://repo.continuum.io/miniconda/Miniconda3-4.2.12-Linux-x86_64.sh  && \
+     chmod +x ~/miniconda.sh && \
+     ~/miniconda.sh -b -p /opt/conda && \     
+     rm ~/miniconda.sh && \
+     /opt/conda/bin/conda install conda-build && \
+     /opt/conda/bin/conda create -y --name pytorch-py35 python=3.5.2 numpy pyyaml scipy ipython mkl&& \
+     /opt/conda/bin/conda clean -ya 
+ENV PATH /opt/conda/envs/pytorch-py35/bin:$PATH
+RUN conda install --name pytorch-py35 -c soumith magma-cuda80 && /opt/conda/bin/conda clean -ya
+RUN conda install --name pytorch-py35 pytorch torchvision cuda80 -c soumith && /opt/conda/bin/conda clean -ya
+
+ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
+
+WORKDIR /workspace
+RUN chmod -R a+w /workspace

tools/nnwrap/__init__.py

@@ -1 +1,2 @@
-from .generate_wrappers import generate_wrappers, wrap_function, import_module
+from .generate_wrappers import generate_wrappers, wrap_function, \
+    import_module, wrap_generic_function

tools/nnwrap/generate_wrappers.py

@@ -3,26 +3,13 @@ import sys
 from string import Template, ascii_lowercase
 from ..cwrap import cwrap
 from ..cwrap.plugins import StandaloneExtension, GenericNN, NullableArguments, AutoGPU
+from ..shared import import_module
 
 BASE_PATH = os.path.realpath(os.path.join(__file__, '..', '..', '..'))
 WRAPPER_PATH = os.path.join(BASE_PATH, 'torch', 'csrc', 'nn')
 THNN_UTILS_PATH = os.path.join(BASE_PATH, 'torch', '_thnn', 'utils.py')
 
 
-def import_module(name, path):
-    if sys.version_info >= (3, 5):
-        import importlib.util
-        spec = importlib.util.spec_from_file_location(name, path)
-        module = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(module)
-        return module
-    elif sys.version_info >= (3, 0):
-        from importlib.machinery import SourceFileLoader
-        return SourceFileLoader(name, path).load_module()
-    else:
-        import imp
-        return imp.load_source(name, path)
-
 thnn_utils = import_module('torch._thnn.utils', THNN_UTILS_PATH)
 
 FUNCTION_TEMPLATE = Template("""\
@@ -52,22 +39,27 @@ TYPE_TRANSFORMS = {
     'Float': {
         'THTensor*': 'THFloatTensor*',
         'real': 'float',
+        'accreal': 'double',
     },
     'Double': {
         'THTensor*': 'THDoubleTensor*',
         'real': 'double',
+        'accreal': 'double',
     },
     'CudaHalf': {
         'THCTensor*': 'THCudaHalfTensor*',
         'real': 'half',
+        'accreal': 'float',
     },
     'Cuda': {
         'THCTensor*': 'THCudaTensor*',
         'real': 'float',
+        'accreal': 'float',
     },
     'CudaDouble': {
         'THCTensor*': 'THCudaDoubleTensor*',
         'real': 'double',
+        'accreal': 'double',
     },
 }
 for t, transforms in TYPE_TRANSFORMS.items():
@@ -83,14 +75,17 @@ def wrap_function(name, type, arguments):
     cname = 'THNN_' + type + name
     declaration = ''
     declaration += 'extern "C" void ' + cname + \
-        '(' + ', '.join(TYPE_TRANSFORMS[type].get(arg.type, arg.type) for arg in arguments) + ');\n'
+        '(' + ', '.join(TYPE_TRANSFORMS[type].get(arg.type, arg.type)
+                        for arg in arguments) + ');\n'
     declaration += FUNCTION_TEMPLATE.substitute(name=type + name, cname=cname)
     indent = ' ' * 4
     dict_indent = ' ' * 6
     prefix = indent + '- '
     for arg in arguments:
         if not arg.is_optional:
-            declaration += prefix + TYPE_TRANSFORMS[type].get(arg.type, arg.type) + ' ' + arg.name + '\n'
+            declaration += prefix + \
+                TYPE_TRANSFORMS[type].get(
+                    arg.type, arg.type) + ' ' + arg.name + '\n'
         else:
             t = TYPE_TRANSFORMS[type].get(arg.type, arg.type)
             declaration += prefix + 'type: ' + t + '\n' + \
@@ -135,6 +130,7 @@ def wrap_cunn():
         AutoGPU(has_self=False),
     ])
 
+
 GENERIC_FUNCTION_TEMPLATE = Template("""\
 [[
   name: $name
@@ -163,7 +159,7 @@ def wrap_generic():
     defs = OrderedDict()
 
     def should_wrap_function(name):
-        if name.startswith('LookupTable'):
+        if name.startswith('LookupTable_'):
             return False
         return (name.endswith('updateOutput') or
                 name.endswith('updateGradInput') or

tools/pytorch.version

@@ -0,0 +1,12 @@
+{
+     global:
+         _TH*;
+         TH*;
+         *THP*;
+         *THCP*;
+         PyInit*;
+         init*;
+         state;
+     local:
+         *;
+ };

tools/setup_helpers/cuda.py

@@ -1,17 +1,39 @@
-import ctypes.util
 import os
+import platform
+import ctypes.util
+from subprocess import Popen, PIPE
 
 from .env import check_env_flag
 
+
+def find_nvcc():
+    proc = Popen(['which', 'nvcc'], stdout=PIPE, stderr=PIPE)
+    out, err = proc.communicate()
+    out = out.decode().strip()
+    if len(out) > 0:
+        return os.path.dirname(out)
+    else:
+        return None
+
+
 if check_env_flag('NO_CUDA'):
     WITH_CUDA = False
     CUDA_HOME = None
 else:
     CUDA_HOME = os.getenv('CUDA_HOME', '/usr/local/cuda')
     if not os.path.exists(CUDA_HOME):
-        cudart_path = ctypes.util.find_library('cudart')
-        if cudart_path is not None:
-            CUDA_HOME = os.path.dirname(cudart_path)
+        # We use nvcc path on Linux and cudart path on macOS
+        osname = platform.system()
+        if osname == 'Linux':
+            cuda_path = find_nvcc()
+        else:
+            cudart_path = ctypes.util.find_library('cudart')
+            if cudart_path is not None:
+                cuda_path = os.path.dirname(cudart_path)
+            else:
+                cuda_path = None
+        if cuda_path is not None:
+            CUDA_HOME = os.path.dirname(cuda_path)
         else:
             CUDA_HOME = None
     WITH_CUDA = CUDA_HOME is not None

tools/setup_helpers/cudnn.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import glob
 from itertools import chain
 
@@ -9,6 +10,8 @@ from .cuda import WITH_CUDA, CUDA_HOME
 def gather_paths(env_vars):
     return list(chain(*(os.getenv(v, '').split(':') for v in env_vars)))
 
+is_conda = 'conda' in sys.version or 'Continuum' in sys.version
+conda_dir = os.path.join(os.path.dirname(sys.executable), '..')
 
 WITH_CUDNN = False
 CUDNN_LIB_DIR = None
@@ -19,6 +22,7 @@ if WITH_CUDA and not check_env_flag('NO_CUDNN'):
         os.path.join(CUDA_HOME, 'lib'),
         os.path.join(CUDA_HOME, 'lib64'),
         '/usr/lib/x86_64-linux-gnu/',
+        '/usr/lib/powerpc64le-linux-gnu/',
     ] + gather_paths([
         'LIBRARY_PATH',
     ])))
@@ -31,6 +35,9 @@ if WITH_CUDA and not check_env_flag('NO_CUDNN'):
         'C_INCLUDE_PATH',
         'CPLUS_INCLUDE_PATH',
     ])))
+    if is_conda:
+        lib_paths.append(os.path.join(conda_dir, 'lib'))
+        include_paths.append(os.path.join(conda_dir, 'include'))
     for path in lib_paths:
         if path is None or not os.path.exists(path):
             continue

tools/setup_helpers/split_types.py

@@ -0,0 +1,58 @@
+import os
+
+this_file = os.path.dirname(os.path.abspath(__file__))
+generated_dir = os.path.abspath(os.path.join(this_file, '..', '..', 'torch', 'csrc', 'generated'))
+
+line_start = '//generic_include '
+
+types = [
+    'Double',
+    'Float',
+    'Half',
+    'Long',
+    'Int',
+    'Short',
+    'Char',
+    'Byte'
+]
+
+generic_include = '#define {lib}_GENERIC_FILE "{path}"'
+generate_include = '#include "{lib}/{lib}Generate{type}Type.h"'
+
+
+def split_types(file_name):
+    assert file_name.startswith('torch/csrc/')
+    if not os.path.exists(generated_dir):
+        os.makedirs(generated_dir)
+
+    with open(file_name, 'r') as f:
+        lines = f.read().split('\n')
+
+    # Find //generic_include
+    for i, l in enumerate(lines):
+        if l.startswith(line_start):
+            args = l[len(line_start):]
+            lib_prefix, generic_file = filter(bool, args.split())
+            break
+    else:
+        raise RuntimeError("generic include not found")
+
+    gen_name_prefix = file_name[len('torch/csrc/'):].replace('/', '_').replace('.cpp', '')
+    gen_path_prefix = os.path.join(generated_dir, gen_name_prefix)
+
+    prefix = '\n'.join(lines[:i])
+    suffix = '\n'.join(lines[i + 1:])
+
+    to_build = []
+
+    g_include = generic_include.format(lib=lib_prefix, path=generic_file)
+    for t in types:
+        t_include = generate_include.format(lib=lib_prefix, type=t)
+        gen_path = gen_path_prefix + t + '.cpp'
+        to_build.append(gen_path)
+        with open(gen_path, 'w') as f:
+            f.write(prefix + '\n' +
+                    g_include + '\n' +
+                    t_include + '\n' +
+                    suffix)
+    return to_build

tools/shared/__init__.py

@@ -0,0 +1,3 @@
+from .module_loader import import_module
+from .cwrap_common import set_declaration_defaults, \
+    sort_by_number_of_options, enumerate_options_due_to_default

tools/shared/cwrap_common.py

@@ -0,0 +1 @@
+../../torch/lib/ATen/common_with_cwrap.py

tools/shared/module_loader.py

@@ -0,0 +1,16 @@
+import sys
+
+
+def import_module(name, path):
+    if sys.version_info >= (3, 5):
+        import importlib.util
+        spec = importlib.util.spec_from_file_location(name, path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+    elif sys.version_info >= (3, 0):
+        from importlib.machinery import SourceFileLoader
+        return SourceFileLoader(name, path).load_module()
+    else:
+        import imp
+        return imp.load_source(name, path)

torch/__init__.py

@@ -5,7 +5,7 @@ Additionally, it provides many utilities for efficient serializing of
 Tensors and arbitrary types, and other useful utilities.
 
 It has a CUDA counterpart, that enables you to run your tensor computations
-on an NVIDIA GPU with compute capability >= 2.0.
+on an NVIDIA GPU with compute capability >= 3.0.
 """
 
 import sys
@@ -15,7 +15,7 @@ from .version import __version__
 __all__ = [
     'typename', 'is_tensor', 'is_storage', 'set_default_tensor_type',
     'set_rng_state', 'get_rng_state', 'manual_seed', 'initial_seed',
-    'save', 'load', 'set_printoptions', 'chunk', 'split', 'stack',
+    'save', 'load', 'set_printoptions', 'chunk', 'split', 'stack', 'matmul',
     'DoubleStorage', 'FloatStorage', 'LongStorage', 'IntStorage',
     'ShortStorage', 'CharStorage', 'ByteStorage',
     'DoubleTensor', 'FloatTensor', 'LongTensor', 'IntTensor',
@@ -31,6 +31,13 @@ __all__ = [
 # automatically filled by the dynamic loader.
 import os as _dl_flags
 
+# if we have numpy, it *must* be imported before the call to setdlopenflags()
+# or there is risk that later c modules will segfault when importing numpy
+try:
+    import numpy as np
+except:
+    pass
+
 # first check if the os package has the required flags
 if not hasattr(_dl_flags, 'RTLD_GLOBAL') or not hasattr(_dl_flags, 'RTLD_NOW'):
     try:
@@ -81,7 +88,7 @@ def is_tensor(obj):
     Args:
         obj (Object): Object to test
     """
-    return obj.__class__ in _tensor_classes
+    return type(obj) in _tensor_classes
 
 
 def is_storage(obj):
@@ -90,7 +97,7 @@ def is_storage(obj):
     Args:
         obj (Object): Object to test
     """
-    return obj.__class__ in _storage_classes
+    return type(obj) in _storage_classes
 
 
 def set_default_tensor_type(t):
@@ -122,6 +129,9 @@ def manual_seed(seed):
     Args:
         seed (int or long): The desired seed.
     """
+    if torch.cuda.is_available() and not torch.cuda._in_bad_fork:
+        torch.cuda.manual_seed_all(seed)
+
     return default_generator.manual_seed(seed)
 
 
@@ -151,6 +161,10 @@ class FloatStorage(_C.FloatStorageBase, _StorageBase):
     pass
 
 
+class HalfStorage(_C.HalfStorageBase, _StorageBase):
+    pass
+
+
 class LongStorage(_C.LongStorageBase, _StorageBase):
     pass
 
@@ -191,6 +205,16 @@ class FloatTensor(_C.FloatTensorBase, _TensorBase):
         return FloatStorage
 
 
+class HalfTensor(_C.HalfTensorBase, _TensorBase):
+
+    def is_signed(self):
+        return True
+
+    @classmethod
+    def storage_type(cls):
+        return HalfStorage
+
+
 class LongTensor(_C.LongTensorBase, _TensorBase):
 
     def is_signed(self):
@@ -244,12 +268,12 @@ class ByteTensor(_C.ByteTensorBase, _TensorBase):
 
 _storage_classes = {
     DoubleStorage, FloatStorage, LongStorage, IntStorage, ShortStorage,
-    CharStorage, ByteStorage,
+    CharStorage, ByteStorage, HalfStorage
 }
 
 _tensor_classes = {
     DoubleTensor, FloatTensor, LongTensor, IntTensor, ShortTensor,
-    CharTensor, ByteTensor,
+    CharTensor, ByteTensor, HalfTensor
 }
 
 
@@ -261,19 +285,21 @@ set_default_tensor_type('torch.FloatTensor')
 
 from .functional import *
 
+
 ################################################################################
 # Initialize extension
 ################################################################################
 
+def manager_path():
+    import os
+    path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'lib', 'torch_shm_manager')
+    if not os.path.exists(path):
+        raise RuntimeError("Unable to find torch_shm_manager at " + path)
+    return path.encode('utf-8')
+
+
 # Shared memory manager needs to know the exact location of manager executable
-import os
-manager_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'lib', 'torch_shm_manager')
-if sys.version_info[0] >= 3:
-    manager_path = bytes(manager_path, 'ascii')
-
-_C._initExtension(manager_path)
-
-del os
+_C._initExtension(manager_path())
 del manager_path
 
 ################################################################################
@@ -312,7 +338,10 @@ import torch.autograd
 import torch.nn
 import torch.optim
 import torch.multiprocessing
+import torch.sparse
+import torch.utils.backcompat
+_C._init_names(list(torch._tensor_classes) + list(torch._storage_classes))
 
 # attach docstrings to torch and tensor functions
-from . import _torch_docs, _tensor_docs
-del _torch_docs, _tensor_docs
+from . import _torch_docs, _tensor_docs, _storage_docs
+del _torch_docs, _tensor_docs, _storage_docs

torch/_six.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2010-2017 Benjamin Peterson
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+def with_metaclass(meta, *bases):
+    """Create a base class with a metaclass."""
+    # This requires a bit of explanation: the basic idea is to make a dummy
+    # metaclass for one level of class instantiation that replaces itself with
+    # the actual metaclass.
+    class metaclass(meta):
+
+        def __new__(cls, name, this_bases, d):
+            return meta(name, bases, d)
+    return type.__new__(metaclass, 'temporary_class', (), {})

torch/_storage_docs.py

@@ -0,0 +1,43 @@
+"""Adds docstrings to Storage functions"""
+
+import torch._C
+from torch._C import _add_docstr as add_docstr
+
+
+storage_classes = [
+    'DoubleStorageBase',
+    'FloatStorageBase',
+    'LongStorageBase',
+    'IntStorageBase',
+    'ShortStorageBase',
+    'CharStorageBase',
+    'ByteStorageBase',
+]
+
+
+def add_docstr_all(method, docstr):
+    for cls_name in storage_classes:
+        cls = getattr(torch._C, cls_name)
+        try:
+            add_docstr(getattr(cls, method), docstr)
+        except AttributeError:
+            pass
+
+
+add_docstr_all('from_file',
+               """
+from_file(filename, shared=False, size=0) -> Storage
+
+If shared is True then memory is shared between all processes. All changes are
+written to the file. If shared is False then the changes on the storage do not
+affect the file.
+
+Size is the number of elements in the storage. If shared is False then the file
+must contain at least `size * sizeof(Type)` bytes (`Type` is the type of
+storage). If shared is True the file will be created if needed.
+
+Args:
+    filename (str): file name to map
+    shared (bool): whether to share memory
+    size (int): number of elements in the storage
+""")

torch/_tensor_str.py

@@ -67,7 +67,7 @@ def set_printoptions(
 
 def _number_format(tensor, min_sz=-1):
     min_sz = max(min_sz, 2)
-    tensor = torch.DoubleTensor(tensor.nelement()).copy_(tensor).abs_()
+    tensor = torch.DoubleTensor(tensor.size()).copy_(tensor).abs_().view(tensor.nelement())
 
     pos_inf_mask = tensor.eq(float('inf'))
     neg_inf_mask = tensor.eq(float('-inf'))

torch/_utils.py

@@ -1,8 +1,10 @@
 import torch
+import importlib
 
 
 def _type(self, new_type=None, async=False):
-    """Casts this object to the specified type.
+    """Returns the type if `new_type` is not provided, else casts this object to
+    the specified type.
 
     If this is already of the correct type, no copy is performed and the
     original object is returned.
@@ -21,6 +23,15 @@ def _type(self, new_type=None, async=False):
         new_type = _import_dotted_name(new_type)
     if new_type == type(self):
         return self
+    if self.is_sparse:
+        if not new_type.is_sparse:
+            raise RuntimeError("Cannot cast sparse tensor to dense tensor")
+        new_type_name = new_type.__module__ + '.' + new_type.__name__
+        new_values_type_name = new_type_name.replace('.sparse', '')
+        new_values = self._values().type(new_values_type_name, async)
+        return new_type(self._indices(), new_values, self.size())
+    if new_type.is_sparse:
+        raise RuntimeError("Cannot cast dense tensor to sparse tensor")
     return new_type(self.size()).copy_(self, async)
 
 
@@ -39,16 +50,27 @@ def _cuda(self, device=None, async=False):
     if self.is_cuda:
         if device is None:
             device = torch.cuda.current_device()
-        if self.get_device() != device:
-            with torch.cuda.device(device):
-                return type(self)(self.size()).copy_(self, async)
-        else:
+        if self.get_device() == device:
             return self
     else:
         if device is None:
             device = -1
-        with torch.cuda.device(device):
-            return self.type(getattr(torch.cuda, self.__class__.__name__), async)
+    with torch.cuda.device(device):
+        if self.is_sparse:
+            new_type = getattr(torch.cuda.sparse, self.__class__.__name__)
+            indices = self._indices().cuda(device, async)
+            values = self._values().cuda(device, async)
+            return new_type(indices, values, self.size())
+        else:
+            new_type = getattr(torch.cuda, self.__class__.__name__)
+            return new_type(self.size()).copy_(self, async)
+
+
+def _rebuild_tensor(storage, storage_offset, size, stride):
+    class_name = storage.__class__.__name__.replace('Storage', 'Tensor')
+    module = importlib.import_module(storage.__module__)
+    tensor_class = getattr(module, class_name)
+    return tensor_class().set_(storage, storage_offset, size, stride)
 
 
 def _range(*args, **kwargs):
@@ -77,3 +99,47 @@ def _accumulate(iterable, fn=lambda x, y: x + y):
     for element in it:
         total = fn(total, element)
         yield total
+
+
+def _flatten_tensors(tensors):
+    """Flatten tensors into a single contiguous 1D buffer"""
+    if len(tensors) == 1:
+        return tensors[0].contiguous().view(-1)
+    numels = [tensor.numel() for tensor in tensors]
+    size = sum(numels)
+    offset = 0
+    flat = tensors[0].new(size)
+    for tensor, numel in zip(tensors, numels):
+        flat.narrow(0, offset, numel).copy_(tensor, broadcast=False)
+        offset += numel
+    return flat
+
+
+def _unflatten_tensors(flat, tensors):
+    """View a flat buffer using the sizes of tensors"""
+    outputs = []
+    offset = 0
+    for tensor in tensors:
+        numel = tensor.numel()
+        outputs.append(flat.narrow(0, offset, numel).view_as(tensor))
+        offset += numel
+    return tuple(outputs)
+
+
+def _take_tensors(tensors, size_limit):
+    """Groups tensors into lists of up to size_limit bytes"""
+    buf = []
+    size = 0
+    last_type = type(tensors[0]) if len(tensors) > 0 else None
+    for tensor in tensors:
+        t = type(tensor)
+        param_size = tensor.numel() * tensor.element_size()
+        if t is not last_type or (size + param_size > size_limit and size > 0):
+            yield buf
+            last_type = t
+            size = 0
+            buf = []
+        buf.append(tensor)
+        size += param_size
+    if len(buf) > 0:
+        yield buf

torch/autograd/__init__.py

@@ -5,21 +5,51 @@ changes to the existing code - you only need to wrap all tensors in
 :class:`.Variable` objects.
 """
 import torch
+import warnings
 
 from .variable import Variable
 from .function import Function, NestedIOFunction
 from .stochastic_function import StochasticFunction
+from .gradcheck import gradcheck
 
 __all__ = ['Variable', 'Function', 'StochasticFunction', 'backward']
 
 
-def backward(variables, grad_variables, retain_variables=False):
+def _make_grads(outputs, grads, user_create_graph):
+    if user_create_graph is not None:
+        create_graph = user_create_graph
+    else:
+        create_graph = any(isinstance(grad, Variable) and not grad.volatile
+                           for grad in grads)
+
+    new_grads = []
+    for out, grad in zip(outputs, grads):
+        if isinstance(grad, Variable):
+            new_grads.append(grad)
+        elif torch.is_tensor(grad):
+            new_grads.append(Variable(grad, volatile=not create_graph))
+        elif grad is None:
+            if out.requires_grad:
+                if out.numel() != 1:
+                    raise RuntimeError("grad can be implicitly created only for scalar outputs")
+                data = out.data
+                new_grads.append(
+                    Variable(data.new().resize_as_(data).fill_(1), volatile=not create_graph))
+            else:
+                new_grads.append(None)
+        else:
+            raise TypeError("gradients can be either Tensors, Variables or None, but got " +
+                            type(grad).__name__)
+    return tuple(new_grads), create_graph
+
+
+def backward(variables, grad_variables=None, retain_graph=None, create_graph=None, retain_variables=None):
     """Computes the sum of gradients of given variables w.r.t. graph leaves.
 
     The graph is differentiated using the chain rule. If any of ``variables``
     are non-scalar (i.e. their data has more than one element) and require
     gradient, the function additionaly requires specifying ``grad_variables``.
-    It should be a sequence of matching length, that containins gradient of
+    It should be a sequence of matching length, that contains gradient of
     the differentiated function w.r.t. corresponding variables (``None`` is an
     acceptable value for all variables that don't need gradient tensors).
 
@@ -29,15 +59,98 @@ def backward(variables, grad_variables, retain_variables=False):
     Arguments:
         variables (sequence of Variable): Variables of which the derivative will be
             computed.
-        grad_variables (sequence of Tensor): Gradients w.r.t. each element of
-            corresponding variables. Required only for non-scalar variables that
-            require gradient.
-        retain_variables (bool): If ``True``, buffers necessary for computing
-            gradients won't be freed after use. It is only necessary to
-            specify ``True`` if you want to differentiate some subgraph multiple
-            times.
+        grad_variables (sequence of (Tensor, Variable or None)): Gradients w.r.t.
+            each element of corresponding variables.  Any tensors will be
+            automatically converted to Variables that are volatile unless
+            ``create_graph`` is True.  None values can be specified for scalar
+            Variables or ones that don't require grad. If a None value would
+            be acceptable for all grad_variables, then this argument is optional.
+        retain_graph (bool, optional): If False, the graph used to compute the grad
+            will be freed. Note that in nearly all cases setting this option to True
+            is not needed and often can be worked around in a much more efficient
+            way. Defaults to the value of ``create_graph``.
+        create_graph (bool, optional): If true, graph of the derivative will
+            be constructed, allowing to compute higher order derivative products.
+            Defaults to False, unless ``grad_variables`` contains at least one
+            non-volatile Variable.
     """
-    Variable._execution_engine.run_backward(
-        tuple(variables), tuple(grad_variables), retain_variables)
+    variables = (variables,) if isinstance(variables, Variable) else tuple(variables)
 
-assert torch._C._autograd_init()
+    if grad_variables is None:
+        grad_variables = [None] * len(variables)
+    elif isinstance(grad_variables, Variable) or torch.is_tensor(grad_variables):
+        grad_variables = [grad_variables]
+    else:
+        grad_variables = list(grad_variables)
+
+    grad_variables, create_graph = _make_grads(variables, grad_variables, create_graph)
+
+    if retain_variables is not None:
+        if retain_graph is not None:
+            raise ValueError("only one of retain_graph and retain_variables can be specified")
+        retain_graph = retain_variables
+        warnings.warn("retain_variables option is deprecated and will be removed in 0.3. "
+                      "Use retain_graph instead.")
+    elif retain_graph is None:
+        retain_graph = create_graph
+
+    Variable._execution_engine.run_backward(
+        variables, grad_variables, retain_graph)
+
+
+def grad(outputs, inputs, grad_outputs=None, retain_graph=None, create_graph=None, only_inputs=True):
+    """Computes and returns the sum of gradients of outputs w.r.t. the inputs.
+
+    ``grad_outputs`` should be a sequence of length matching ``output``
+    containing the pre-computed gradients w.r.t. each of the outputs. If an
+    output doesn't require_grad, then the gradient can be ``None``).
+    Gradients can be given as Tensors when one doesn't need the graph of the
+    derivative, or as Variables, in which case the graph will be created.
+
+    If ``only_inputs`` is True, the function will only return a list of gradients
+    w.r.t the specified inputs. If it's False, then gradient w.r.t. all remaining
+    leaves will still be computed, and will be accumulated into their ``.grad``
+    attribute.
+
+    Arguments:
+        outputs (sequence of Variable): outputs of the differentiated function.
+        inputs (sequence of Variable): Inputs w.r.t. which the gradient will be
+            returned (and not accumulated into ``.grad``).
+        grad_outputs (sequence of Tensor or Variable): Gradients w.r.t. each output.
+            Any tensors will be automatically converted to Variables that are
+            volatile unless ``create_graph`` is True.  None values can be
+            specified for scalar Variables or ones that don't require grad.
+            If a None value would be acceptable for all grad_variables, then
+            this argument is optional.
+        retain_graph (bool, optional): If False, the graph used to compute the grad
+            will be freed. Note that in nearly all cases setting this option to True
+            is not needed and often can be worked around in a much more efficient
+            way. Defaults to the value of ``create_graph``.
+        create_graph (bool, optional): If True, graph of the derivative will
+            be constructed, allowing to compute higher order derivative products.
+            Defaults to False, unless ``grad_variables`` contains at least one
+            non-volatile Variable.
+        only_inputs (bool, optional): If True, gradient w.r.t. leaves that are
+            part of the graph, but don't appear in ``inputs`` won't be computed
+            and accumulated. Defaults to True.
+    """
+
+    outputs = (outputs,) if isinstance(outputs, Variable) else tuple(outputs)
+    inputs = (inputs,) if isinstance(inputs, Variable) else tuple(inputs)
+    if grad_outputs is None:
+        grad_outputs = [None] * len(outputs)
+    elif isinstance(grad_outputs, Variable) or torch.is_tensor(grad_outputs):
+        grad_outputs = [grad_outputs]
+    else:
+        grad_outputs = list(grad_outputs)
+
+    grad_outputs, create_graph = _make_grads(outputs, grad_outputs, create_graph)
+    if retain_graph is None:
+        retain_graph = create_graph
+
+    return Variable._execution_engine.run_backward(
+        outputs, grad_outputs, retain_graph,
+        inputs, only_inputs)
+
+if not torch._C._autograd_init():
+    raise RuntimeError("autograd initialization failed")

torch/autograd/_functions/basic_ops.py

@@ -1,189 +1,228 @@
 import torch
 from ..function import Function, InplaceFunction
+from .utils import maybe_unexpand, maybe_unexpand_or_view
 import math
 
 
 class Add(InplaceFunction):
 
-    def forward(self, a, b):
-        if self.inplace:
-            self.mark_dirty(a)
+    @staticmethod
+    def forward(ctx, a, b, inplace=False):
+        ctx.a_size = a.size()
+        ctx.b_size = b.size()
+        if inplace:
+            ctx.mark_dirty(a)
             return a.add_(b)
         else:
             return a.add(b)
 
-    def backward(self, grad_output):
-        return grad_output, grad_output
+    @staticmethod
+    def backward(ctx, grad_output):
+        return maybe_unexpand(grad_output, ctx.a_size), maybe_unexpand_or_view(grad_output, ctx.b_size), None
 
 
 class Sub(InplaceFunction):
 
-    def forward(self, a, b):
-        if self.inplace:
-            self.mark_dirty(a)
+    @staticmethod
+    def forward(ctx, a, b, inplace=False):
+        ctx.a_size = a.size()
+        ctx.b_size = b.size()
+        if inplace:
+            ctx.mark_dirty(a)
             return a.sub_(b)
         else:
             return a.sub(b)
 
-    def backward(self, grad_output):
-        return grad_output, grad_output.neg()
+    @staticmethod
+    def backward(ctx, grad_output):
+        return maybe_unexpand(grad_output, ctx.a_size), maybe_unexpand_or_view(grad_output.neg(), ctx.b_size), None
 
 
 class Mul(Function):
 
-    def forward(self, a, b):
-        self.save_for_backward(a, b)
+    @staticmethod
+    def forward(ctx, a, b):
+        ctx.a_size = a.size()
+        ctx.b_size = b.size()
+        ctx.save_for_backward(a, b)
         return a.mul(b)
 
-    def backward(self, grad_output):
-        a, b = self.saved_tensors
-        return grad_output.mul(b), grad_output.mul(a)
+    @staticmethod
+    def backward(ctx, grad_output):
+        a, b = ctx.saved_variables
+        return maybe_unexpand(grad_output.mul(b), ctx.a_size), maybe_unexpand_or_view(grad_output.mul(a), ctx.b_size)
 
 
 class Div(Function):
 
-    def forward(self, a, b):
-        self.save_for_backward(a, b)
+    @staticmethod
+    def forward(ctx, a, b):
+        ctx.a_size = a.size()
+        ctx.b_size = b.size()
+        ctx.save_for_backward(a, b)
         return a.div(b)
 
-    def backward(self, grad_output):
-        a, b = self.saved_tensors
-        return grad_output.div(b), grad_output.neg().mul(a).div_(b).div_(b)
+    @staticmethod
+    def backward(ctx, grad_output):
+        a, b = ctx.saved_variables
+        b_rec = b.reciprocal()
+        grad_a = grad_output.mul(b_rec)
+        grad_b = grad_output.neg().mul(a).mul(b_rec).mul(b_rec)
+        return maybe_unexpand(grad_a, ctx.a_size), maybe_unexpand_or_view(grad_b, ctx.b_size)
 
 
 class Pow(Function):
 
-    def forward(self, a, b):
-        self.save_for_backward(a, b)
+    @staticmethod
+    def forward(ctx, a, b):
+        ctx.a_size = a.size()
+        ctx.b_size = b.size()
+        ctx.save_for_backward(a, b)
         return a.pow(b)
 
-    def backward(self, grad_output):
-        a, b = self.saved_tensors
-        return grad_output.mul(b).mul_(a.pow(b - 1)), grad_output.mul(a.pow(b)).mul_(a.log())
+    @staticmethod
+    def backward(ctx, grad_output):
+        a, b = ctx.saved_variables
+        grad_a = grad_output.mul(b).mul(a.pow(b - 1))
+        grad_b = grad_output.mul(a.pow(b)).mul(a.log())
+        return maybe_unexpand(grad_a, ctx.a_size), maybe_unexpand_or_view(grad_b, ctx.b_size)
+
+
+def sort_args(a, b):
+    return (a, b, True) if torch.is_tensor(a) else (b, a, False)
 
 
 class AddConstant(InplaceFunction):
 
-    def __init__(self, constant, inplace=False):
-        super(AddConstant, self).__init__(inplace)
-        self.constant = constant
-
-    def forward(self, a):
-        if self.inplace:
-            self.mark_dirty(a)
-            return a.add_(self.constant)
+    @staticmethod
+    def forward(ctx, a, b, inplace=False):
+        tensor, constant, ctx.tensor_first = sort_args(a, b)
+        if inplace:
+            ctx.mark_dirty(tensor)
+            return tensor.add_(constant)
         else:
-            return a.add(self.constant)
+            return tensor.add(constant)
 
-    def backward(self, grad_output):
-        return grad_output
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.tensor_first:
+            return grad_output, None, None
+        else:
+            return None, grad_output, None
 
 
 class SubConstant(InplaceFunction):
 
-    def __init__(self, constant, sub_tensor=False, inplace=False):
-        super(SubConstant, self).__init__(inplace)
-        self.constant = constant
-        self.sub_tensor = sub_tensor
-
-    def forward(self, a):
-        if self.sub_tensor:
-            if a.is_signed() and self.inplace:
-                self.mark_dirty(a)
-                return a.neg_().add_(self.constant)
+    @staticmethod
+    def forward(ctx, a, b, inplace=False):
+        tensor, constant, ctx.tensor_first = sort_args(a, b)
+        if ctx.tensor_first:
+            if inplace:
+                ctx.mark_dirty(tensor)
+                return tensor.sub_(constant)
             else:
-                assert not self.inplace, "can't perform (constant - tensor) " \
-                    "subtraction in-place on an unsigned type"
-                return a.new().resize_as_(a).fill_(self.constant).sub_(a)
+                return tensor.sub(constant)
         else:
-            if self.inplace:
-                self.mark_dirty(a)
-                return a.sub_(self.constant)
+            if inplace:
+                ctx.mark_dirty(tensor)
+                return tensor.neg_().add_(constant)
             else:
-                return a.sub(self.constant)
+                return tensor.neg().add_(constant)
 
-    def backward(self, grad_output):
-        if self.sub_tensor:
-            return grad_output.neg()
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.tensor_first:
+            return grad_output, None, None
         else:
-            return grad_output
+            return None, grad_output.neg(), None
 
 
 class MulConstant(InplaceFunction):
 
-    def __init__(self, constant, inplace=False):
-        super(MulConstant, self).__init__(inplace)
-        self.constant = constant
-
-    def forward(self, a):
-        if self.inplace:
-            self.mark_dirty(a)
-            return a.mul_(self.constant)
+    @staticmethod
+    def forward(ctx, a, b, inplace=False):
+        tensor, ctx.constant, ctx.tensor_first = sort_args(a, b)
+        if inplace:
+            ctx.mark_dirty(tensor)
+            return tensor.mul_(ctx.constant)
         else:
-            return a.mul(self.constant)
+            return tensor.mul(ctx.constant)
 
-    def backward(self, grad_output):
-        return grad_output.mul(self.constant)
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.mul(ctx.constant)
+        if ctx.tensor_first:
+            return grad_input, None, None
+        else:
+            return None, grad_input, None
 
 
 class DivConstant(InplaceFunction):
 
-    def __init__(self, constant, div_by_tensor=False, inplace=False):
-        super(DivConstant, self).__init__(inplace)
-        self.constant = constant
-        self.div_by_tensor = div_by_tensor
-        if self.inplace and self.div_by_tensor:
-            # TODO: actually, as long as the type is floating point, we can
-            raise RuntimeError("can't perform (constant / tensor) division in-place")
-
-    def forward(self, a):
-        if self.div_by_tensor:
-            self.save_for_backward(a)
-            return a.new().resize_as_(a).fill_(self.constant).div_(a)
-        else:
-            if self.inplace:
-                return a.div_(self.constant)
+    @staticmethod
+    def forward(ctx, a, b, inplace=False):
+        tensor, ctx.constant, ctx.tensor_first = sort_args(a, b)
+        ctx.inplace = inplace
+        if ctx.tensor_first:
+            if inplace:
+                ctx.mark_dirty(tensor)
+                return tensor.div_(ctx.constant)
             else:
-                return a.div(self.constant)
-
-    def backward(self, grad_output):
-        if self.div_by_tensor:
-            a = self.saved_tensors[0]
-            return grad_output.neg().mul_(self.constant).div_(a).div_(a)
+                return tensor.div(ctx.constant)
         else:
-            return grad_output.div(self.constant)
+            ctx.save_for_backward(tensor)
+            if inplace:
+                ctx.mark_dirty(tensor)
+                return tensor.reciprocal_().mul_(ctx.constant)
+            else:
+                return tensor.reciprocal().mul_(ctx.constant)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.tensor_first:
+            return grad_output.div(ctx.constant), None, None
+        else:
+            v, = ctx.saved_variables
+            if ctx.inplace:
+                return None, grad_output.mul(v).mul(v).div_(-ctx.constant), None
+            else:
+                v_rep = v.reciprocal()
+                return None, grad_output.mul(v_rep).mul(v_rep).mul_(-ctx.constant), None
 
 
 class PowConstant(Function):
 
-    def __init__(self, constant, tensor_power=False):
-        super(PowConstant, self).__init__()
-        self.constant = constant
-        self.tensor_power = tensor_power
-
-    def forward(self, a):
-        if self.tensor_power:
-            self.fw_result = torch.pow(self.constant, a)
-            return self.fw_result
+    @staticmethod
+    def forward(ctx, a, b):
+        tensor, ctx.constant, ctx.tensor_first = sort_args(a, b)
+        if ctx.tensor_first:
+            ctx.save_for_backward(tensor)
+            return tensor.pow(ctx.constant)
         else:
-            self.save_for_backward(a)
-            return a.pow(self.constant)
+            result = torch.pow(ctx.constant, tensor)
+            ctx.save_for_backward(result)
+            return result
 
-    def backward(self, grad_output):
-        if self.tensor_power:
-            return grad_output.mul(self.fw_result).mul_(math.log(self.constant))
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.tensor_first:
+            var, = ctx.saved_variables
+            return grad_output.mul(ctx.constant).mul(var.pow(ctx.constant - 1)), None
         else:
-            a = self.saved_tensors[0]
-            return grad_output.mul(self.constant).mul_(a.pow(self.constant - 1))
+            var_result, = ctx.saved_variables
+            return None, grad_output.mul(var_result).mul_(math.log(ctx.constant))
 
 
 class Negate(InplaceFunction):
 
-    def forward(self, i):
-        if self.inplace:
+    @staticmethod
+    def forward(ctx, i, inplace=False):
+        if inplace:
+            ctx.mark_dirty(i)
             return i.neg_()
         else:
             return i.neg()
 
-    def backward(self, grad_output):
-        return grad_output.neg()
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output.neg(), None

torch/autograd/_functions/blas.py

@@ -1,201 +1,224 @@
 import torch
 
 from ..function import Function, InplaceFunction
+from .utils import maybe_unexpand
 
 
 # TODO: no need to save all args if the grad w.r.t. some of them is not needed
-class _BlasBase(InplaceFunction):
-
-    def __init__(self, alpha=1, beta=1, inplace=False):
-        super(_BlasBase, self).__init__(inplace)
-        self.alpha = alpha
-        self.beta = beta
-
-    def _get_output(self, arg):
-        if self.inplace:
-            self.mark_dirty(arg)
-            return arg
-        else:
-            return arg.new().resize_as_(arg)
+def _get_output(ctx, arg, inplace=False):
+    if inplace:
+        ctx.mark_dirty(arg)
+        return arg
+    else:
+        return arg.new().resize_as_(arg)
 
 
-class Addmm(_BlasBase):
+class Addmm(InplaceFunction):
 
-    def forward(self, add_matrix, matrix1, matrix2):
-        self.save_for_backward(matrix1, matrix2)
-        output = self._get_output(add_matrix)
-        return torch.addmm(self.alpha, add_matrix, self.beta,
+    @staticmethod
+    def forward(ctx, add_matrix, matrix1, matrix2, alpha=1, beta=1, inplace=False):
+        ctx.alpha = alpha
+        ctx.beta = beta
+        ctx.add_matrix_size = add_matrix.size()
+        ctx.save_for_backward(matrix1, matrix2)
+        output = _get_output(ctx, add_matrix, inplace=inplace)
+        return torch.addmm(alpha, add_matrix, beta,
                            matrix1, matrix2, out=output)
 
-    def backward(self, grad_output):
-        matrix1, matrix2 = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        matrix1, matrix2 = ctx.saved_variables
         grad_add_matrix = grad_matrix1 = grad_matrix2 = None
 
-        if self.needs_input_grad[0]:
-            grad_add_matrix = grad_output
-            if self.alpha != 1:
-                grad_add_matrix = grad_add_matrix.mul(self.alpha)
+        if ctx.needs_input_grad[0]:
+            grad_add_matrix = maybe_unexpand(grad_output, ctx.add_matrix_size)
+            if ctx.alpha != 1:
+                grad_add_matrix = grad_add_matrix.mul(ctx.alpha)
 
-        if self.needs_input_grad[1]:
-            grad_matrix1 = torch.mm(grad_output, matrix2.t())
-            if self.beta != 1:
-                grad_matrix1 *= self.beta
+        if ctx.needs_input_grad[1]:
+            if matrix1.stride() == (1, matrix1.size(0)):
+                # column major gradient if input is column major
+                grad_matrix1 = torch.mm(matrix2, grad_output.t()).t()
+            else:
+                grad_matrix1 = torch.mm(grad_output, matrix2.t())
+            if ctx.beta != 1:
+                grad_matrix1 *= ctx.beta
 
-        if self.needs_input_grad[2]:
-            grad_matrix2 = torch.mm(matrix1.t(), grad_output)
-            if self.beta != 1:
-                grad_matrix2 *= self.beta
+        if ctx.needs_input_grad[2]:
+            if matrix2.stride() == (1, matrix2.size(0)):
+                # column major gradient if input is column major
+                grad_matrix2 = torch.mm(grad_output.t(), matrix1).t()
+            else:
+                grad_matrix2 = torch.mm(matrix1.t(), grad_output)
+            if ctx.beta != 1:
+                grad_matrix2 *= ctx.beta
 
-        return grad_add_matrix, grad_matrix1, grad_matrix2
+        return grad_add_matrix, grad_matrix1, grad_matrix2, None, None, None
 
 
-class Addbmm(_BlasBase):
+class Addbmm(InplaceFunction):
 
-    def forward(self, add_matrix, batch1, batch2):
-        self.save_for_backward(batch1, batch2)
-        output = self._get_output(add_matrix)
-        return torch.addbmm(self.alpha, add_matrix, self.beta,
+    @staticmethod
+    def forward(ctx, add_matrix, batch1, batch2, alpha=1, beta=1, inplace=False):
+        ctx.alpha = alpha
+        ctx.beta = beta
+        ctx.add_matrix_size = add_matrix.size()
+        ctx.save_for_backward(batch1, batch2)
+        output = _get_output(ctx, add_matrix, inplace=inplace)
+        return torch.addbmm(alpha, add_matrix, beta,
                             batch1, batch2, out=output)
 
-    def backward(self, grad_output):
-        batch1, batch2 = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        batch1, batch2 = ctx.saved_variables
         grad_add_matrix = grad_batch1 = grad_batch2 = None
 
-        if self.needs_input_grad[0]:
-            grad_add_matrix = grad_output
-            if self.alpha != 1:
-                grad_add_matrix = grad_add_matrix.mul(self.alpha)
+        if ctx.needs_input_grad[0]:
+            grad_add_matrix = maybe_unexpand(grad_output, ctx.add_matrix_size)
+            if ctx.alpha != 1:
+                grad_add_matrix = grad_add_matrix.mul(ctx.alpha)
 
-        if any(self.needs_input_grad[1:]):
+        if any(ctx.needs_input_grad[1:]):
             batch_grad_output = (grad_output
                                  .unsqueeze(0)
                                  .expand(batch1.size(0), batch1.size(1), batch2.size(2)))
 
-        if self.needs_input_grad[1]:
+        if ctx.needs_input_grad[1]:
             grad_batch1 = torch.bmm(batch_grad_output, batch2.transpose(1, 2))
-            if self.beta != 1:
-                grad_batch1 *= self.beta
+            if ctx.beta != 1:
+                grad_batch1 *= ctx.beta
 
-        if self.needs_input_grad[2]:
+        if ctx.needs_input_grad[2]:
             grad_batch2 = torch.bmm(batch1.transpose(1, 2), batch_grad_output)
-            if self.beta != 1:
-                grad_batch2 *= self.beta
+            if ctx.beta != 1:
+                grad_batch2 *= ctx.beta
 
-        return grad_add_matrix, grad_batch1, grad_batch2
+        return grad_add_matrix, grad_batch1, grad_batch2, None, None, None
 
 
-class Baddbmm(_BlasBase):
+class Baddbmm(InplaceFunction):
 
-    def forward(self, add_batch, batch1, batch2):
-        self.save_for_backward(batch1, batch2)
-        output = self._get_output(add_batch)
-        return torch.baddbmm(self.alpha, add_batch, self.beta,
+    @staticmethod
+    def forward(ctx, add_batch, batch1, batch2, alpha=1, beta=1, inplace=False):
+        ctx.alpha = alpha
+        ctx.beta = beta
+        ctx.add_batch_size = add_batch.size()
+        ctx.save_for_backward(batch1, batch2)
+        output = _get_output(ctx, add_batch, inplace=inplace)
+        return torch.baddbmm(alpha, add_batch, beta,
                              batch1, batch2, out=output)
 
-    def backward(self, grad_output):
-        batch1, batch2 = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        batch1, batch2 = ctx.saved_variables
         grad_add_batch = grad_batch1 = grad_batch2 = None
 
-        if self.needs_input_grad[0]:
-            grad_add_batch = grad_output
-            if self.alpha != 1:
-                grad_add_batch = grad_add_batch.mul(self.alpha)
+        if ctx.needs_input_grad[0]:
+            grad_add_batch = maybe_unexpand(grad_output, ctx.add_batch_size)
+            if ctx.alpha != 1:
+                grad_add_batch = grad_add_batch.mul(ctx.alpha)
 
-        if self.needs_input_grad[1]:
+        if ctx.needs_input_grad[1]:
             grad_batch1 = torch.bmm(grad_output, batch2.transpose(1, 2))
-            if self.beta != 1:
-                grad_batch1 *= self.beta
+            if ctx.beta != 1:
+                grad_batch1 *= ctx.beta
 
-        if self.needs_input_grad[2]:
+        if ctx.needs_input_grad[2]:
             grad_batch2 = torch.bmm(batch1.transpose(1, 2), grad_output)
-            if self.beta != 1:
-                grad_batch2 *= self.beta
+            if ctx.beta != 1:
+                grad_batch2 *= ctx.beta
 
-        return grad_add_batch, grad_batch1, grad_batch2
+        return grad_add_batch, grad_batch1, grad_batch2, None, None, None
 
 
-class Addmv(_BlasBase):
+class Addmv(InplaceFunction):
 
-    def forward(self, add_vector, matrix, vector):
-        self.save_for_backward(matrix, vector)
-        output = self._get_output(add_vector)
-        return torch.addmv(self.alpha, add_vector, self.beta,
+    @staticmethod
+    def forward(ctx, add_vector, matrix, vector, alpha=1, beta=1, inplace=False):
+        ctx.alpha = alpha
+        ctx.beta = beta
+        ctx.add_vector_size = add_vector.size()
+        ctx.save_for_backward(matrix, vector)
+        output = _get_output(ctx, add_vector, inplace=inplace)
+        return torch.addmv(alpha, add_vector, beta,
                            matrix, vector, out=output)
 
-    def backward(self, grad_output):
-        matrix, vector = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        matrix, vector = ctx.saved_variables
         grad_add_vector = grad_matrix = grad_vector = None
 
-        if self.needs_input_grad[0]:
-            grad_add_vector = grad_output
-            if self.alpha != 1:
-                grad_add_vector = grad_add_vector.mul(self.alpha)
+        if ctx.needs_input_grad[0]:
+            grad_add_vector = maybe_unexpand(grad_output, ctx.add_vector_size)
+            if ctx.alpha != 1:
+                grad_add_vector = grad_add_vector.mul(ctx.alpha)
 
-        if self.needs_input_grad[1]:
+        if ctx.needs_input_grad[1]:
             grad_matrix = torch.ger(grad_output, vector)
-            if self.beta != 1:
-                grad_matrix *= self.beta
+            if ctx.beta != 1:
+                grad_matrix *= ctx.beta
 
-        if self.needs_input_grad[2]:
+        if ctx.needs_input_grad[2]:
             grad_vector = torch.mv(matrix.t(), grad_output)
-            if self.beta != 1:
-                grad_vector *= self.beta
+            if ctx.beta != 1:
+                grad_vector *= ctx.beta
 
-        return grad_add_vector, grad_matrix, grad_vector
+        return grad_add_vector, grad_matrix, grad_vector, None, None, None
 
 
-class Addr(_BlasBase):
+class Addr(InplaceFunction):
 
-    def forward(self, add_matrix, vector1, vector2):
-        self.save_for_backward(vector1, vector2)
-        output = self._get_output(add_matrix)
-        return torch.addr(self.alpha, add_matrix, self.beta,
+    @staticmethod
+    def forward(ctx, add_matrix, vector1, vector2, alpha=1, beta=1, inplace=False):
+        ctx.alpha = alpha
+        ctx.beta = beta
+        ctx.add_matrix_size = add_matrix.size()
+        ctx.save_for_backward(vector1, vector2)
+        output = _get_output(ctx, add_matrix, inplace=inplace)
+        return torch.addr(alpha, add_matrix, beta,
                           vector1, vector2, out=output)
 
-    def backward(self, grad_output):
-        vector1, vector2 = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        vector1, vector2 = ctx.saved_variables
         grad_add_matrix = grad_vector1 = grad_vector2 = None
 
-        if self.needs_input_grad[0]:
-            grad_add_matrix = grad_output
-            if self.alpha != 1:
-                grad_add_matrix = grad_add_matrix.mul(self.alpha)
+        if ctx.needs_input_grad[0]:
+            grad_add_matrix = maybe_unexpand(grad_output, ctx.add_matrix_size)
+            if ctx.alpha != 1:
+                grad_add_matrix = grad_add_matrix.mul(ctx.alpha)
 
-        if self.needs_input_grad[1]:
+        if ctx.needs_input_grad[1]:
             grad_vector1 = torch.mv(grad_output, vector2)
-            if self.beta != 1:
-                grad_vector1 *= self.beta
+            if ctx.beta != 1:
+                grad_vector1 *= ctx.beta
 
-        if self.needs_input_grad[2]:
+        if ctx.needs_input_grad[2]:
             # TODO: maybe it's better to do transpose + mv + transpose
-            grad_vector2 = torch.mm(vector1.unsqueeze(0), grad_output)
-            if self.beta != 1:
-                grad_vector2 *= self.beta
+            grad_vector2 = torch.mm(vector1.unsqueeze(0), grad_output).squeeze(0)
+            if ctx.beta != 1:
+                grad_vector2 *= ctx.beta
 
-        return grad_add_matrix, grad_vector1, grad_vector2
+        return grad_add_matrix, grad_vector1, grad_vector2, None, None, None
 
 
 class Dot(Function):
 
-    def forward(self, vector1, vector2):
-        self.save_for_backward(vector1, vector2)
+    @staticmethod
+    def forward(ctx, vector1, vector2):
+        ctx.save_for_backward(vector1, vector2)
+        ctx.sizes = (vector1.size(), vector2.size())
         return vector1.new((vector1.dot(vector2),))
 
-    def backward(self, grad_output):
-        vector1, vector2 = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        vector1, vector2 = ctx.saved_variables
         grad_vector1 = grad_vector2 = None
 
-        if self.needs_input_grad[0]:
-            grad_vector1 = vector2.mul(grad_output[0])
+        if ctx.needs_input_grad[0]:
+            grad_vector1 = vector2.mul(grad_output.expand(ctx.sizes[1])).view(ctx.sizes[0])
 
-        if self.needs_input_grad[1]:
-            grad_vector2 = vector1.mul(grad_output[0])
+        if ctx.needs_input_grad[1]:
+            grad_vector2 = vector1.mul(grad_output.expand(ctx.sizes[0])).view(ctx.sizes[1])
 
         return grad_vector1, grad_vector2
-
-
-# TODO: cross
-# TODO: diag
-# TODO: trace
-# TODO: tril
-# TODO: triu

torch/autograd/_functions/compare.py

@@ -1,20 +1,29 @@
 import torch
 
 from ..function import Function
+from .utils import maybe_unexpand, maybe_unexpand_or_view
 
 
+# TODO: once Cpp-style functions are implemented we can detach a and b
+# before calling forward.
 class _CompareOp(Function):
 
-    def __init__(self, scalar=None):
-        super(_CompareOp, self).__init__()
-        self.scalar = scalar
-
-    def forward(self, tensor1, tensor2=None):
-        other = tensor2 if tensor2 is not None else self.scalar
-        mask = getattr(tensor1, self.fn_name)(other)
-        self.mark_non_differentiable(mask)
+    @classmethod
+    def forward(cls, ctx, a, b):
+        ctx.a_size = a.size()
+        ctx.b_tensor = torch.is_tensor(b)
+        ctx.b_size = b.size() if ctx.b_tensor else None
+        ctx.input_type = type(a)
+        mask = getattr(a, cls.fn_name)(b)
+        ctx.mark_non_differentiable(mask)
         return mask
 
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = (grad_output * 0).type(ctx.input_type)
+        return (maybe_unexpand(grad_input, ctx.a_size),
+                maybe_unexpand_or_view(grad_input, ctx.b_size) if ctx.b_tensor else None)
+
 
 class Eq(_CompareOp):
     fn_name = 'eq'

torch/autograd/_functions/linalg.py

@@ -1,44 +1,104 @@
 import torch
 
 from ..function import Function
+from ..variable import Variable
 
 
 class Diag(Function):
 
-    def __init__(self, diagonal_idx=0):
-        super(Diag, self).__init__()
-        self.diagonal_idx = diagonal_idx
+    @staticmethod
+    def forward(ctx, input, diagonal_idx=0):
+        ctx.diagonal_idx = diagonal_idx
+        return input.diag(ctx.diagonal_idx)
 
-    def forward(self, input):
-        return input.diag()
-
-    def backward(self, grad_output):
-        return grad_output.diag()
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output.diag(ctx.diagonal_idx), None
 
 
 class Tril(Function):
 
-    def __init__(self, diagonal_idx=0):
-        super(Tril, self).__init__()
-        self.diagonal_idx = diagonal_idx
+    @staticmethod
+    def forward(ctx, input, diagonal_idx=0):
+        ctx.diagonal_idx = diagonal_idx
+        return input.tril(ctx.diagonal_idx)
 
-    def forward(self, input):
-        return input.tril(self.diagonal_idx)
-
-    def backward(self, grad_output):
-        return grad_output.tril(self.diagonal_idx)
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output.tril(ctx.diagonal_idx), None
 
 
 class Triu(Function):
 
-    def __init__(self, diagonal_idx=0):
-        super(Triu, self).__init__()
-        self.diagonal_idx = diagonal_idx
+    @staticmethod
+    def forward(ctx, input, diagnoal_idx=0):
+        ctx.diagonal_idx = diagnoal_idx
+        return input.triu(ctx.diagonal_idx)
 
-    def forward(self, input):
-        return input.triu(self.diagonal_idx)
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output.triu(ctx.diagonal_idx), None
 
-    def backward(self, grad_output):
-        return grad_output.triu(self.diagonal_idx)
 
-# TODO: trace
+class Trace(Function):
+
+    @staticmethod
+    def forward(ctx, input):
+        ctx.isize = input.size()
+        return input.new((input.trace(), ))
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        isize = ctx.isize
+        min_size = min(isize)
+        grad_input = Variable(grad_output.data.new(isize).zero_()).view(-1)
+        grad_input[::(isize[1] + 1)] = grad_output.expand(min_size)
+        return grad_input.view(isize)
+
+
+class Cross(Function):
+
+    @staticmethod
+    def forward(ctx, input, other, dim=-1):
+        ctx.dim = dim
+        ctx.save_for_backward(input, other)
+        return torch.cross(input, other, ctx.dim)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, other = ctx.saved_variables
+        grad_input = other.cross(grad_output, ctx.dim)
+        grad_other = grad_output.cross(input, ctx.dim)
+        return grad_input, grad_other, None
+
+
+class Inverse(Function):
+
+    @staticmethod
+    def forward(ctx, input):
+        inverse = torch.inverse(input)
+        ctx.save_for_backward(inverse)
+        return inverse
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        inverse, = ctx.saved_variables
+        return -torch.mm(inverse.t(), torch.mm(grad_output, inverse.t()))
+
+
+class Gesv(Function):
+
+    @staticmethod
+    def forward(ctx, b, a):
+        # TODO see if one can backprop through LU
+        X, LU = torch.gesv(b, a)
+        ctx.save_for_backward(X, a)
+        ctx.mark_non_differentiable(LU)
+        return X, LU
+
+    @staticmethod
+    def backward(ctx, grad_output, grad_LU=None):
+        X, a = ctx.saved_variables
+        grad_b, _ = torch.gesv(grad_output, a.t())
+        grad_a = -torch.mm(grad_b, X.t())
+        return grad_b, grad_a

torch/autograd/_functions/pointwise.py

@@ -1,282 +1,351 @@
 from itertools import repeat
 
+from ..._thnn import type2backend
 from ..function import Function, InplaceFunction
+from ..variable import Variable
+from .utils import maybe_unexpand, maybe_unexpand_or_view
 
 
 class Exp(InplaceFunction):
 
-    def forward(self, i):
-        if self.inplace:
-            self.mark_dirty(i)
+    @staticmethod
+    def forward(ctx, i, inplace=False):
+        if inplace:
+            ctx.mark_dirty(i)
             result = i.exp_()
         else:
             result = i.exp()
-        self.save_for_backward(result)
+        ctx.save_for_backward(result)
         return result
 
-    def backward(self, grad_output):
-        return self.saved_tensors[0] * grad_output
+    @staticmethod
+    def backward(ctx, grad_output):
+        result, = ctx.saved_variables
+        return grad_output * result, None
 
 
 class Log(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.log()
 
-    def backward(self, grad_output):
-        return grad_output.div(self.saved_tensors[0])
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
+        return grad_output.div(i)
 
 
 class Log1p(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.log1p()
 
-    def backward(self, grad_output):
-        return grad_output.div(self.saved_tensors[0].add(1))
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
+        return grad_output.div(i.add(1))
 
 
 class Tanh(InplaceFunction):
 
-    def forward(self, i):
-        if self.inplace:
-            self.mark_dirty(i)
+    @staticmethod
+    def forward(ctx, i, inplace=False):
+        if inplace:
+            ctx.mark_dirty(i)
             result = i.tanh_()
         else:
             result = i.tanh()
-        self.save_for_backward(result)
+        ctx.save_for_backward(result)
         return result
 
-    def backward(self, grad_output):
-        result, = self.saved_tensors
-        return grad_output * (1 - result * result)
+    @staticmethod
+    def backward(ctx, grad_output):
+        result, = ctx.saved_variables
+        if grad_output.volatile:
+            grad_input = Variable(grad_output.data.new(grad_output.size()), volatile=True)
+            backend = type2backend[type(result.data)]
+            backend.Tanh_updateGradInput(backend.library_state, None, grad_output.data,
+                                         grad_input.data, result.data)
+        else:
+            grad_input = grad_output * (1 - result * result)
+        return grad_input, None
 
 
 class Sigmoid(InplaceFunction):
 
-    def forward(self, i):
-        if self.inplace:
-            self.mark_dirty(i)
+    @staticmethod
+    def forward(ctx, i, inplace=False):
+        if inplace:
+            ctx.mark_dirty(i)
             result = i.sigmoid_()
         else:
             result = i.sigmoid()
-        self.save_for_backward(result)
+        ctx.save_for_backward(result)
         return result
 
-    def backward(self, grad_output):
-        result, = self.saved_tensors
-        return grad_output * ((1 - result) * result)
+    @staticmethod
+    def backward(ctx, grad_output):
+        result, = ctx.saved_variables
+        if grad_output.volatile:
+            grad_input = Variable(grad_output.data.new(grad_output.size()), volatile=True)
+            backend = type2backend[type(result.data)]
+            backend.Sigmoid_updateGradInput(backend.library_state, None, grad_output.data,
+                                            grad_input.data, result.data)
+        else:
+            grad_input = grad_output * ((1 - result) * result)
+        return grad_input, None
 
 
 class Sinh(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.sinh()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
         return grad_output * i.cosh()
 
 
 class Cosh(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.cosh()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
         return grad_output * i.sinh()
 
 
 class Abs(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.abs()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
         return grad_output * i.sign()
 
 
 class Clamp(Function):
 
-    def __init__(self, min_val, max_val):
-        super(Clamp, self).__init__()
-        self.min_val = min_val
-        self.max_val = max_val
+    @staticmethod
+    def forward(ctx, i, min_val, max_val):
+        ctx._mask = (i.ge(min_val) * i.le(max_val))
+        return i.clamp(min_val, max_val)
 
-    def forward(self, i):
-        self.save_for_backward(i)
-        return i.clamp(self.min_val, self.max_val)
-
-    def backward(self, grad_output):
-        i, = self.saved_tensors
-        mask = i.ge(self.min_val) * i.le(self.max_val)
-        return grad_output * mask.type_as(grad_output)
+    @staticmethod
+    def backward(ctx, grad_output):
+        mask = Variable(ctx._mask.type_as(grad_output.data))
+        return grad_output * mask, None, None
 
 
 class Sqrt(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.sqrt()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
-        return grad_output.mul(i.pow(-0.5)).div(2)
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
+        return grad_output.mul(i.pow(-0.5)).div_(2)
 
 
 class Sin(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.sin()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
         return grad_output * i.cos()
 
 
 class Cos(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.cos()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
         return grad_output.mul(i.sin()).neg_()
 
 
 class Tan(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.tan()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
         return grad_output.div(i.cos().pow(2))
 
 
 class Asin(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.asin()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
-        return grad_output * (1 - i.mul(i)).sqrt_().reciprocal_()
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
+        return grad_output * (1 - i.mul(i)).sqrt().reciprocal()
 
 
 class Acos(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.acos()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
-        return grad_output.mul((1 - i.mul(i)).sqrt_().reciprocal_()).neg_()
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
+        return grad_output.mul((1 - i.mul(i)).sqrt().reciprocal()).neg_()
 
 
 class Atan(Function):
 
-    def forward(self, i):
-        self.save_for_backward(i)
+    @staticmethod
+    def forward(ctx, i):
+        ctx.save_for_backward(i)
         return i.atan()
 
-    def backward(self, grad_output):
-        i, = self.saved_tensors
-        return grad_output * i.mul(i).add_(1).reciprocal_()
+    @staticmethod
+    def backward(ctx, grad_output):
+        i, = ctx.saved_variables
+        return grad_output * i.mul(i).add_(1).reciprocal()
 
 
+class Atan2(Function):
+
+    @staticmethod
+    def forward(ctx, y, x):
+        ctx.save_for_backward(y, x)
+        return y.atan2(x)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        y, x, = ctx.saved_variables
+        denominator = y.mul(y).add(x.mul(x)).reciprocal()
+        return grad_output * x.mul(denominator), grad_output * y.neg().mul(denominator)
+
+
+# TODO: make inplace and update grad formulas
 class Reciprocal(Function):
 
-    def forward(self, i):
+    @staticmethod
+    def forward(ctx, i):
         result = i.reciprocal()
-        self.save_for_backward(result)
+        ctx.save_for_backward(result)
         return result
 
-    def backward(self, grad_output):
-        result, = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        result, = ctx.saved_variables
         return grad_output * result.mul(result).neg_()
 
 
 class Cmax(Function):
 
-    def forward(self, a, b):
-        self._max_buffer = a.gt(b).type_as(a)
+    @staticmethod
+    def forward(ctx, a, b):
+        ctx._a_size = a.size()
+        ctx._b_size = b.size()
+        ctx._mask = a.gt(b)
         return a.max(b)
 
-    def backward(self, grad_output):
+    @staticmethod
+    def backward(ctx, grad_output):
+        mask = Variable(ctx._mask.type_as(grad_output.data))
         return (
-            grad_output * self._max_buffer,
-            grad_output * self._max_buffer.eq(0).type_as(grad_output)
+            maybe_unexpand(grad_output * mask, ctx._a_size),
+            maybe_unexpand_or_view(grad_output * Variable(ctx._mask.eq(0).type_as(grad_output.data)), ctx._b_size)
         )
 
 
 class CmaxConstant(Function):
 
-    def __init__(self, constant):
-        super(CmaxConstant, self).__init__()
-        self.constant = constant
+    @staticmethod
+    def forward(ctx, i, constant):
+        ctx._mask = i.gt(constant)
+        return i.clamp(min=constant)
 
-    def forward(self, i):
-        self._max_buffer = i.gt(self.constant).type_as(i)
-        return i.clamp(min=self.constant)
-
-    def backward(self, grad_output):
-        return grad_output * self._max_buffer
+    @staticmethod
+    def backward(ctx, grad_output):
+        mask = Variable(ctx._mask.type_as(grad_output.data))
+        return grad_output * mask, None
 
 
 class Cmin(Function):
 
-    def forward(self, a, b):
-        self._min_buffer = a.lt(b).type_as(a)
+    @staticmethod
+    def forward(ctx, a, b):
+        ctx._a_size = a.size()
+        ctx._b_size = b.size()
+        ctx._mask = a.lt(b).type_as(a)
         return a.min(b)
 
-    def backward(self, grad_output):
+    @staticmethod
+    def backward(ctx, grad_output):
+        mask = Variable(ctx._mask.type_as(grad_output.data))
         return (
-            grad_output * self._min_buffer,
-            grad_output * self._min_buffer.eq(0).type_as(grad_output)
+            maybe_unexpand(grad_output * mask, ctx._a_size),
+            maybe_unexpand_or_view(grad_output * Variable(ctx._mask.eq(0).type_as(grad_output.data)), ctx._b_size)
         )
 
 
 class CminConstant(Function):
 
-    def __init__(self, constant):
-        super(CminConstant, self).__init__()
-        self.constant = constant
+    @staticmethod
+    def forward(ctx, i, constant):
+        ctx._mask = i.lt(constant)
+        return i.clamp(max=constant)
 
-    def forward(self, i):
-        self._min_buffer = i.lt(self.constant).type_as(i)
-        return i.clamp(max=self.constant)
-
-    def backward(self, grad_output):
-        return grad_output * self._min_buffer
+    @staticmethod
+    def backward(ctx, grad_output):
+        mask = Variable(ctx._mask.type_as(grad_output.data))
+        return grad_output * mask, None
 
 
 class _ConstantGrad(Function):
     grad_value = 0
 
-    def __init__(self, *args):
-        super(_ConstantGrad, self).__init__()
-        self.args = args
+    @classmethod
+    def forward(cls, ctx, *args):
+        ctx._num_args = len(args)
+        ctx._args0_size = args[0].size()
+        return getattr(args[0], cls.__name__.lower())(*args[1:])
 
-    def forward(self, i):
-        return getattr(i, type(self).__name__.lower())(*self.args)
-
-    def backward(self, grad_output):
-        grad_input = grad_output.new(*repeat(1, grad_output.dim()))
-        grad_input = grad_input.fill_(self.grad_value).expand_as(grad_output)
-        return grad_input.mul(grad_output)
+    @classmethod
+    def backward(cls, ctx, grad_output):
+        return (maybe_unexpand(grad_output.mul(cls.grad_value), ctx._args0_size),) + (ctx._num_args - 1) * (None,)
 
 
 class Floor(_ConstantGrad):
@@ -313,91 +382,96 @@ class Remainder(_ConstantGrad):
 
 class Lerp(Function):
 
-    def __init__(self, weight):
-        super(Lerp, self).__init__()
-        self.weight = float(weight)
+    @staticmethod
+    def forward(ctx, a, b, weight):
+        ctx._a_size = a.size()
+        ctx._b_size = b.size()
+        ctx._weight = float(weight)
+        return a.lerp(b, ctx._weight)
 
-    def forward(self, a, b):
-        return a.lerp(b, self.weight)
-
-    def backward(self, grad_output):
-        return grad_output.mul(1 - self.weight), grad_output.mul(self.weight)
+    @staticmethod
+    def backward(ctx, grad_output):
+        return (maybe_unexpand(grad_output.mul(1 - ctx._weight), ctx._a_size),
+                maybe_unexpand_or_view(grad_output.mul(ctx._weight), ctx._b_size), None)
 
 
 class Rsqrt(InplaceFunction):
 
-    def forward(self, input):
-        if self.inplace:
-            self.mark_dirty(input)
-            result = input.rsqrt_()
+    @staticmethod
+    def forward(ctx, i, inplace=False):
+        if inplace:
+            ctx.mark_dirty(i)
+            result = i.rsqrt_()
         else:
-            result = input.rsqrt()
-        self.save_for_backward(result)
+            result = i.rsqrt()
+        ctx.save_for_backward(result)
         return result
 
-    def backward(self, grad_output):
-        result, = self.saved_tensors
-        return result.pow(3).div_(-2).mul_(grad_output)
+    @staticmethod
+    def backward(ctx, grad_output):
+        result, = ctx.saved_variables
+        return result.pow(3).div_(-2).mul(grad_output), None
 
 
 class Addcmul(InplaceFunction):
 
-    def __init__(self, scale=1, inplace=False):
-        super(Addcmul, self).__init__(inplace)
-        self.scale = scale
-
-    def forward(self, add_tensor, mul_tensor1, mul_tensor2):
-        self.save_for_backward(mul_tensor1, mul_tensor2)
-        if self.inplace:
-            return add_tensor.addcmul_(self.scale, mul_tensor1, mul_tensor2)
+    @staticmethod
+    def forward(ctx, add_tensor, mul_tensor1, mul_tensor2, scale=1.0, inplace=False):
+        ctx._scale = scale
+        ctx._add_tensor_size = add_tensor.size()
+        ctx.save_for_backward(mul_tensor1, mul_tensor2)
+        if inplace:
+            ctx.mark_dirty(add_tensor)
+            return add_tensor.addcmul_(scale, mul_tensor1, mul_tensor2)
         else:
-            return add_tensor.addcmul(self.scale, mul_tensor1, mul_tensor2)
+            return add_tensor.addcmul(scale, mul_tensor1, mul_tensor2)
 
-    def backward(self, grad_output):
+    @staticmethod
+    def backward(ctx, grad_output):
         grad_add = grad_mul1 = grad_mul2 = None
-        mul_tensor1, mul_tensor2 = self.saved_tensors
+        mul_tensor1, mul_tensor2 = ctx.saved_variables
 
-        if self.needs_input_grad[0]:
-            grad_add = grad_output
+        if ctx.needs_input_grad[0]:
+            grad_add = maybe_unexpand(grad_output, ctx._add_tensor_size)
 
-        if self.needs_input_grad[1]:
-            grad_mul1 = grad_output.mul(mul_tensor2).mul(self.scale)
+        if ctx.needs_input_grad[1]:
+            grad_mul1 = maybe_unexpand_or_view(grad_output.mul(mul_tensor2).mul_(ctx._scale), mul_tensor1.size())
 
-        if self.needs_input_grad[2]:
-            grad_mul2 = grad_output.mul(mul_tensor1).mul(self.scale)
+        if ctx.needs_input_grad[2]:
+            grad_mul2 = maybe_unexpand_or_view(grad_output.mul(mul_tensor1).mul_(ctx._scale), mul_tensor2.size())
 
-        return grad_add, grad_mul1, grad_mul2
+        return grad_add, grad_mul1, grad_mul2, None, None
 
 
 class Addcdiv(InplaceFunction):
 
-    def __init__(self, scale=1, inplace=False):
-        super(Addcdiv, self).__init__(inplace)
-        self.scale = scale
-
-    def forward(self, add_tensor, div_tensor1, div_tensor2):
-        self.save_for_backward(div_tensor1, div_tensor2)
-        if self.inplace:
-            return add_tensor.addcdiv_(self.scale, div_tensor1, div_tensor2)
+    @staticmethod
+    def forward(ctx, add_tensor, div_tensor1, div_tensor2, scale=1.0, inplace=False):
+        ctx._scale = scale
+        ctx._add_tensor_size = add_tensor.size()
+        ctx.save_for_backward(div_tensor1, div_tensor2)
+        if inplace:
+            ctx.mark_dirty(add_tensor)
+            return add_tensor.addcdiv_(ctx._scale, div_tensor1, div_tensor2)
         else:
-            return add_tensor.addcdiv(self.scale, div_tensor1, div_tensor2)
+            return add_tensor.addcdiv(ctx._scale, div_tensor1, div_tensor2)
 
-    def backward(self, grad_output):
+    @staticmethod
+    def backward(ctx, grad_output):
         grad_add = grad_div1 = grad_div2 = None
-        div_tensor1, div_tensor2 = self.saved_tensors
+        div_tensor1, div_tensor2 = ctx.saved_variables
 
-        if self.needs_input_grad[0]:
-            grad_add = grad_output
+        if ctx.needs_input_grad[0]:
+            grad_add = maybe_unexpand(grad_output, ctx._add_tensor_size)
 
-        if self.needs_input_grad[1]:
-            grad_div1 = grad_output.div(div_tensor2).mul(self.scale)
+        if ctx.needs_input_grad[1]:
+            grad_div1 = maybe_unexpand_or_view(grad_output.div(div_tensor2).mul_(ctx._scale), div_tensor1.size())
 
-        if self.needs_input_grad[2]:
+        if ctx.needs_input_grad[2]:
             div_tensor2_sq = div_tensor2.mul(div_tensor2)
-            grad_div2 = grad_output.mul(div_tensor1).div_(div_tensor2_sq)
-            grad_div2.neg_().mul_(self.scale)
-
-        return grad_add, grad_div1, grad_div2
+            grad_div2 = maybe_unexpand_or_view(grad_output.mul(div_tensor1).div(div_tensor2_sq).mul(-ctx._scale),
+                                               div_tensor2.size())
 
+        return grad_add, grad_div1, grad_div2, None, None
 
 # TODO: atan2 + inplace

torch/autograd/_functions/reduce.py

@@ -1,73 +1,141 @@
 from functools import reduce
 
 from ..function import Function
+from ..variable import Variable
+import torch
 
 
-class _DimReduceFunction(Function):
+class Sum(Function):
 
-    def __init__(self, dim=None):
-        super(_DimReduceFunction, self).__init__()
-        self.dim = dim
-
-    def forward(self, input):
-        self.input_size = input.size()
-        fn = getattr(input, self.fn_name)
-        if self.dim is None:
-            return input.new((fn(),))
+    @staticmethod
+    def forward(ctx, input, dim=None, keepdim=None):
+        ctx.dim = dim
+        ctx.keepdim = False if keepdim is None else keepdim
+        ctx.input_size = input.size()
+        if dim is None:
+            return input.new((input.sum(),))
         else:
-            return fn(self.dim)
+            if keepdim is not None:
+                return input.sum(dim, keepdim=keepdim)
+            else:
+                return input.sum(dim)
 
-
-class Sum(_DimReduceFunction):
-    fn_name = 'sum'
-
-    def backward(self, grad_output):
-        if self.dim is None:
-            return grad_output.new(self.input_size).fill_(grad_output[0])
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.dim is None:
+            return grad_output.expand(ctx.input_size), None, None
         else:
-            repeats = [1 for _ in self.input_size]
-            repeats[self.dim] = self.input_size[self.dim]
-            return grad_output.repeat(*repeats),
+            if ctx.keepdim is False and len(ctx.input_size) != 1:
+                grad_output = grad_output.unsqueeze(ctx.dim)
+
+            repeats = [1 for _ in ctx.input_size]
+            repeats[ctx.dim] = ctx.input_size[ctx.dim]
+            return grad_output.repeat(*repeats), None, None
 
 
-class Prod(_DimReduceFunction):
+class Prod(Function):
 
-    def forward(self, input):
-        self.input_size = input.size()
-        if self.dim is None:
-            self.result = input.prod()
-            self.save_for_backward(input)
-            return input.new((self.result,))
+    @staticmethod
+    def forward(ctx, input, dim=None, keepdim=None):
+        ctx.dim = dim
+        ctx.keepdim = False if keepdim is None else keepdim
+        ctx.input_size = input.size()
+        if dim is None:
+            ctx.result = input.prod()
+            ctx.save_for_backward(input)
+            return input.new((ctx.result,))
         else:
-            output = input.prod(self.dim)
-            self.save_for_backward(input, output)
+            if keepdim is not None:
+                output = input.prod(dim, keepdim=keepdim)
+            else:
+                output = input.prod(dim)
+            ctx.save_for_backward(input, output)
             return output
 
-    def backward(self, grad_output):
-        if self.dim is None:
-            input, = self.saved_tensors
-            grad_input = grad_output.new(self.input_size).fill_(self.result)
-            return grad_input.div(input)
+    @staticmethod
+    def backward(ctx, grad_output):
+        def safe_zeros_backward(inp, dim):
+            # note that the gradient is equivalent to:
+            # cumprod(exclusive, normal) * cumprod(exclusive, reverse), e.g.:
+            # input:                        [    a,     b,     c]
+            # cumprod(exclusive, normal):   [1    ,     a, a * b]
+            # cumprod(exclusive, reverse):  [b * c,     c,     1]
+            # product:                      [b * c, a * c, a * b]
+            # and this is safe under input with 0s.
+            if inp.size(dim) == 1:
+                return grad_output
+
+            ones_size = torch.Size((inp.size()[:dim] + (1,) + inp.size()[dim + 1:]))
+            ones = Variable(grad_output.data.new(ones_size).fill_(1))
+            exclusive_normal_nocp = torch.cat((ones, inp.narrow(dim, 0, inp.size(dim) - 1)), dim)
+            exclusive_normal = exclusive_normal_nocp.cumprod(dim)
+
+            def reverse_dim(var, dim):
+                return var.index_select(dim, Variable(torch.arange(var.size(dim) - 1, -1, -1)).long())
+
+            narrow_reverse = reverse_dim(inp.narrow(dim, 1, inp.size(dim) - 1), dim)
+            exclusive_reverse_nocp = torch.cat((ones, narrow_reverse), dim)
+            exclusive_reverse = reverse_dim(exclusive_reverse_nocp.cumprod(dim), dim)
+
+            grad_input = grad_output.expand_as(exclusive_normal).mul(exclusive_normal.mul(exclusive_reverse))
+            return grad_input
+
+        if ctx.dim is None:
+            input, = ctx.saved_variables
+            zero_idx = (input.data == 0).nonzero()
+            if zero_idx.dim() == 0:
+                return grad_output.mul(ctx.result).expand_as(input).div(input), None, None
+            elif zero_idx.size(0) > 1:
+                return (grad_output * 0).expand_as(input), None, None
+            else:
+                return safe_zeros_backward(input.contiguous().view(-1), 0).view_as(input), None, None
+
         else:
-            input, output = self.saved_tensors
-            repeats = [1 for _ in self.input_size]
-            repeats[self.dim] = self.input_size[self.dim]
-            return output.mul(grad_output).repeat(*repeats).div_(input)
+            input, output = ctx.saved_variables
+            dim = ctx.dim if ctx.dim >= 0 else ctx.dim + input.dim()
+            if ctx.keepdim is False and len(ctx.input_size) != 1:
+                grad_output = grad_output.unsqueeze(dim)
+                output = output.unsqueeze(dim)
+
+            zero_mask = input == 0
+            slice_zero_count = zero_mask.sum(dim, True)
+            total_zeros = slice_zero_count.data.sum()
+            if total_zeros == 0:
+                grad_input = grad_output.mul(output).expand_as(input).div(input)
+            else:
+                grad_input = safe_zeros_backward(input, dim)
+
+            return grad_input, None, None
 
 
-class Mean(_DimReduceFunction):
-    fn_name = 'mean'
+class Mean(Function):
 
-    def backward(self, grad_output):
-        if self.dim is None:
-            grad_input_val = grad_output[0]
-            grad_input_val /= reduce(lambda x, y: x * y, self.input_size, 1)
-            return grad_output.new(*self.input_size).fill_(grad_input_val)
+    @staticmethod
+    def forward(ctx, input, dim=None, keepdim=None):
+        ctx.dim = dim
+        ctx.keepdim = False if keepdim is None else keepdim
+        ctx.input_size = input.size()
+        if dim is None:
+            return input.new((input.mean(),))
         else:
-            repeats = [1 for _ in self.input_size]
-            dim_size = self.input_size[self.dim]
-            repeats[self.dim] = dim_size
-            return grad_output.repeat(*repeats).div_(dim_size)
+            if keepdim is not None:
+                return input.mean(dim, keepdim=keepdim)
+            else:
+                return input.mean(dim)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.dim is None:
+            grad_input_val = grad_output / reduce(lambda x, y: x * y, ctx.input_size, 1)
+            return grad_input_val.expand(ctx.input_size), None, None
+        else:
+            if ctx.keepdim is False and len(ctx.input_size) != 1:
+                grad_output = grad_output.unsqueeze(ctx.dim)
+
+            repeats = [1 for _ in ctx.input_size]
+            dim_size = ctx.input_size[ctx.dim]
+            repeats[ctx.dim] = dim_size
+            return grad_output.repeat(*repeats).div_(dim_size), None, None
 
 
 class _SelectionFunction(Function):
@@ -75,44 +143,53 @@ class _SelectionFunction(Function):
     # additional_args is prepended before dim when calling the tensor
     # function. It's a no-op for subclasses other than kthvalue.
     # kthvalue not only requires us to pass a dim, but also preceed it with k.
-    additional_args = tuple()
 
-    def __init__(self, dim=None):
-        super(_SelectionFunction, self).__init__()
-        self.dim = dim
-
-    def forward(self, input):
-        fn = getattr(input, type(self).__name__.lower())
-        self.input_size = input.size()
-        if self.dim is None and self.has_all_reduce:
-            value = fn(*self.additional_args)
-            self.indices = tuple(input.eq(value).nonzero()[0])
+    @classmethod
+    def forward(cls, ctx, input, dim=None, keepdim=None, additional_args=tuple()):
+        fn = getattr(input, cls.__name__.lower())
+        ctx.dim = dim
+        ctx.keepdim = False if keepdim is None else keepdim
+        ctx.additional_args = additional_args
+        ctx.input_size = input.size()
+        if ctx.dim is None and cls.has_all_reduce:
+            value = fn(*additional_args)
+            ctx.indices_tuple = tuple(input.eq(value).nonzero()[0])
             return input.new((value,))
         else:
-            if self.dim is None:
+            if ctx.dim is None:
                 dim = input.dim() - 1
             else:
-                dim = self.dim
+                dim = ctx.dim
             args = (dim,)
-            if self.additional_args:
-                args = self.additional_args + args
-            output, indices = fn(*args)
-            self.save_for_backward(indices)
-            self.mark_non_differentiable(indices)
+            if additional_args:
+                args = additional_args + args
+            if keepdim is not None:
+                output, indices = fn(*args, keepdim=keepdim)
+            else:
+                output, indices = fn(*args)
+            ctx.save_for_backward(indices)
+            ctx.mark_non_differentiable(indices)
             return output, indices
 
-    def backward(self, grad_output, grad_indices=None):
-        grad_input = grad_output.new(*self.input_size).zero_()
-        if self.dim is None and self.has_all_reduce:
-            grad_input[self.indices] = grad_output[0]
+    @classmethod
+    def backward(cls, ctx, grad_output, grad_indices=None):
+        grad_input = Variable(grad_output.data.new(*ctx.input_size).zero_())
+        if ctx.dim is None and cls.has_all_reduce:
+            grad_input[ctx.indices_tuple] = grad_output
         else:
-            if self.dim is None:
-                dim = input.dim() - 1
+            if ctx.dim is None:
+                dim = len(ctx.input_size) - 1
             else:
-                dim = self.dim
-            indices, = self.saved_tensors
+                dim = ctx.dim
+
+            indices, = ctx.saved_variables
+            if ctx.keepdim is False and len(ctx.input_size) != 1:
+                grad_output = grad_output.unsqueeze(dim)
+                grad_indices = grad_indices.unsqueeze(dim)
+                indices = indices.unsqueeze(dim)
+
             grad_input.scatter_(dim, indices, grad_output)
-        return grad_input
+        return grad_input, None, None, None
 
 
 class Max(_SelectionFunction):
@@ -128,53 +205,63 @@ class Mode(_SelectionFunction):
 
 
 class Median(_SelectionFunction):
-    has_all_reduce = False
+    pass
 
 
 class Kthvalue(_SelectionFunction):
     has_all_reduce = False
 
-    def __init__(self, k, dim=None):
-        super(Kthvalue, self).__init__(dim)
-        self.additional_args = (k,)
+    @classmethod
+    def forward(cls, ctx, input, k, dim=None, keepdim=None):
+        return super(Kthvalue, cls).forward(ctx, input, dim, keepdim, (k,))
 
 
 class Norm(Function):
 
-    def __init__(self, norm_type=2, dim=None):
-        super(Norm, self).__init__()
-        self.norm_type = norm_type
-        self.dim = dim
+    @staticmethod
+    def forward(ctx, input, p=2, dim=None, keepdim=None):
+        ctx.p = p
+        ctx.dim = dim
+        ctx.keepdim = False if keepdim is None else keepdim
 
-    def forward(self, input):
-        if self.dim is None:
-            self.norm = input.norm(self.norm_type)
-            self.save_for_backward(input)
-            return input.new((self.norm,))
+        if dim is None:
+            ctx.norm = input.norm(p)
+            ctx.save_for_backward(input)
+            return input.new((ctx.norm,))
         else:
-            output = input.norm(self.norm_type, self.dim)
-            self.save_for_backward(input, output)
+            if keepdim is not None:
+                output = input.norm(p, dim, keepdim=keepdim)
+            else:
+                output = input.norm(p, dim)
+            ctx.save_for_backward(input, output)
             return output
 
-    def backward(self, grad_output):
-        if self.dim is None:
-            input, = self.saved_tensors
-            if self.norm_type == 2:
-                return input.mul(grad_output[0] / self.norm)
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.dim is None:
+            input, = ctx.saved_variables
+            if ctx.p == 2:
+                scale_v = (grad_output / ctx.norm).expand_as(input)
+                return input.mul(scale_v), None, None, None
             else:
-                pow = input.abs().pow(self.norm_type - 2)
-                scale = grad_output[0] / self.norm ** (self.norm_type - 1)
-                return input.mul(pow).mul(scale)
+                pow = input.abs().pow(ctx.p - 2)
+                scale_v = (grad_output / ctx.norm ** (ctx.p - 1)).expand_as(input)
+                return input.mul(pow).mul(scale_v), None, None, None
         else:
-            input, output = self.saved_tensors
+            input, output = ctx.saved_variables
+
+            if ctx.keepdim is False and input.dim() != 1:
+                grad_output = grad_output.unsqueeze(ctx.dim)
+                output = output.unsqueeze(ctx.dim)
+
             big_grad_output = grad_output.expand_as(input)
-            if self.norm_type == 2:
+            if ctx.p == 2:
                 big_output = output.expand_as(input)
-                return input.mul(big_grad_output).div(big_output)
+                return input.mul(big_grad_output).div(big_output), None, None, None
             else:
-                pow = input.abs().pow(self.norm_type - 2)
-                big_output = output.pow(self.norm_type - 1).expand_as(input)
-                return input.mul(pow).mul(big_grad_output).div(big_output)
+                pow = input.abs().pow(ctx.p - 2)
+                big_output = output.pow(ctx.p - 1).expand_as(input)
+                return input.mul(pow).mul(big_grad_output).div(big_output), None, None, None
 
 
 # TODO: renorm

torch/autograd/_functions/replace.vim

@@ -0,0 +1,3 @@
+%s/self/ctx/g
+%s/\s\+def forward/    @staticmethod\r    def forward/g
+%s/\s\+def backward/    @staticmethod\r    @once_differentiable\r    def backward/g

torch/autograd/_functions/stochastic.py

@@ -23,8 +23,9 @@ class Multinomial(StochasticFunction):
         if probs.dim() == 1:
             probs = probs.unsqueeze(0)
             samples = samples.unsqueeze(0)
+            reward = reward.unsqueeze(0)
         # normalize probs (multinomial accepts weights)
-        probs /= probs.sum(1).expand_as(probs)
+        probs /= probs.sum(1, True).expand_as(probs)
         grad_probs = probs.new().resize_as_(probs).zero_()
         output_probs = probs.gather(1, samples)
         output_probs.add_(1e-6).reciprocal_()
@@ -83,8 +84,9 @@ class Normal(StochasticFunction):
             stddevs_cb = stddevs_sq * stddevs
             stddevs_sq += 1e-6
             stddevs_cb += 1e-6
-            grad_stddevs = (grad_means * grad_means) / stddevs_cb
-            grad_stddevs = (stddevs - grad_stddevs) * reward
+            grad_stddevs = (stddevs_sq - (grad_means * grad_means))
+            grad_stddevs /= stddevs_cb
+            grad_stddevs *= reward
             grad_means /= stddevs_sq
         grad_means *= reward
         return grad_means, grad_stddevs

torch/autograd/_functions/utils.py

@@ -0,0 +1,39 @@
+import torch
+
+
+def maybe_view(variable, size):
+    if variable.size() == size:
+        return variable
+    return variable.contiguous().view(size)
+
+
+def maybe_unexpand(variable, old_size):
+    num_unsqueezed = variable.dim() - len(old_size)
+    expanded_dims = [dim for dim, (expanded, original)
+                     in enumerate(zip(variable.size()[num_unsqueezed:], old_size))
+                     if expanded != original]
+
+    for _ in range(num_unsqueezed):
+        variable = variable.sum(0, keepdim=False)
+    for dim in expanded_dims:
+        variable = variable.sum(dim, keepdim=True)
+    return variable
+
+
+def variable_expandable(variable, old_size):
+    try:
+        torch._C._infer_size(variable.size(), old_size)
+    except RuntimeError:
+        return False
+    return True
+
+
+def maybe_unexpand_or_view(variable, old_size):
+    var_expanded = True
+    if maybe_view:
+        var_expanded = variable_expandable(variable, old_size)
+
+    if var_expanded:
+        return maybe_unexpand(variable, old_size)
+    else:
+        return maybe_view(variable, old_size)

torch/autograd/engine.py

@@ -1,85 +0,0 @@
-from collections import deque, defaultdict
-from torch._C import _ImperativeEngine as ImperativeEngine
-from .variable import Variable
-
-
-class BasicEngine(object):
-
-    def _compute_dependencies(self, function):
-        dependencies = defaultdict(int)
-        seen = {function}
-        queue = [function]
-        while len(queue) > 0:
-            fn = queue.pop()
-            for prev_fn, output_nr in fn.previous_functions:
-                if not prev_fn.requires_grad or isinstance(prev_fn, Variable):
-                    continue
-                dependencies[prev_fn] += 1
-                if prev_fn not in seen:
-                    queue.append(prev_fn)
-                    seen.add(prev_fn)
-        return dependencies
-
-    def _free_backward_dependency(self, dependencies, prev_fn):
-        dependencies[prev_fn] -= 1
-        if dependencies[prev_fn] == 0:
-            del dependencies[prev_fn]
-            return True
-        return False
-
-    def _add_grad(self, need_copy, prev_grad, output_nr, d_prev_fn):
-        copy_id = (id(prev_grad), output_nr)
-        if not prev_grad[output_nr]:
-            prev_grad[output_nr] = d_prev_fn
-            need_copy.add(copy_id)
-        else:
-            grad_tensor = prev_grad[output_nr]
-            if copy_id in need_copy:
-                need_copy.remove(copy_id)
-                grad_tensor = grad_tensor.clone()
-                prev_grad[output_nr] = grad_tensor
-            grad_tensor.add_(d_prev_fn)
-
-    def run_backward(self, variable, grad, retain_variables):
-        if variable.creator is None:
-            variable._do_backward((grad,), retain_variables)
-            return
-
-        initial_grad = [None for _ in range(variable.creator.num_outputs)]
-        initial_grad[variable.output_nr] = grad
-        ready = deque([(variable.creator, initial_grad)])
-        not_ready = {}
-        need_copy = set()
-
-        dependencies = self._compute_dependencies(variable.creator)
-
-        while len(ready) > 0:
-            fn, grad = ready.pop()
-            grad_input = fn._do_backward(tuple(grad), retain_variables)
-            for (prev_fn, output_nr), d_prev_fn in zip(fn.previous_functions, grad_input):
-                if not prev_fn.requires_grad:
-                    # TODO: check that d_prev_fn is None and warn otherwise
-                    continue
-                if isinstance(prev_fn, Variable):
-                    prev_fn._do_backward((d_prev_fn,), retain_variables)
-                    continue
-                is_ready = self._free_backward_dependency(dependencies, prev_fn)
-                if is_ready:
-                    if prev_fn in not_ready:
-                        prev_grad = not_ready[prev_fn]
-                        self._add_grad(need_copy, prev_grad, output_nr, d_prev_fn)
-                    else:
-                        if prev_fn.num_outputs != 1:
-                            raise RuntimeError("one of the function outputs "
-                                               "wasn't used - this is an error not, but "
-                                               "it's going to be fixed soon")
-                        prev_grad = (d_prev_fn,)
-                    ready.appendleft((prev_fn, prev_grad))
-                else:
-                    if prev_fn in not_ready:
-                        prev_grad = not_ready[prev_fn]
-                    else:
-                        prev_grad = [None for _ in range(prev_fn.num_outputs)]
-
-                    self._add_grad(need_copy, prev_grad, output_nr, d_prev_fn)
-                    not_ready[prev_fn] = prev_grad

torch/autograd/function.py

@@ -1,47 +1,12 @@
 import torch
 import torch._C as _C
 import torch.utils.hooks as hooks
+from torch._six import with_metaclass
+import functools
 from collections import OrderedDict
 
 
-class Function(_C._FunctionBase):
-    """Records operation history and defines formulas for differentiating ops.
-
-    Every operation performed on :class:`Variable` s creates a new function
-    object, that performs the computation, and records that it happened.
-    The history is retained in the form of a DAG of functions, with edges
-    denoting data dependencies (``input <- output``). Then, when backward is
-    called, the graph is processed in the topological ordering, by calling
-    :func:`backward` methods of each :class:`Function` object, and passing
-    returned gradients on to next :class:`Function` s.
-
-    Normally, the only way users interact with functions is by creating
-    subclasses and defining new operations. This is a recommended way of
-    extending torch.autograd.
-
-    Since Function logic is a hotspot in most scripts, almost all of it
-    was moved to our C backend, to ensure that the framework overhead is
-    minimal.
-
-    Each function is meant to be used only once (in the forward pass).
-
-    Attributes:
-        saved_tensors: Tuple of Tensors that were saved in the call to
-            :func:`forward`.
-        needs_input_grad: Tuple of booleans of length :attr:`num_inputs`,
-            indicating whether a given input requires gradient. This can be
-            used to optimize buffers saved for backward, and ignoring gradient
-            computation in :func:`~Function.backward`.
-        num_inputs: Number of inputs given to :func:`forward`.
-        num_outputs: Number of tensors returned by :func:`forward`.
-        requires_grad: Boolean indicating whether the :func:`backward` will
-            ever need to be called.
-        previous_functions: Tuple of (int, Function) pairs of length
-            :attr:`num_inputs`. Each entry contains a reference to a
-            :class:`Function` that created corresponding input, and an index
-            of the previous function output that's been used.
-    """
-    __call__ = _C._FunctionBase._do_forward
+class _ContextMethodMixin(object):
 
     def save_for_backward(self, *tensors):
         """Saves given tensors for a future call to :func:`~Function.backward`.
@@ -50,9 +15,10 @@ class Function(_C._FunctionBase):
         :func:`forward` **method.**
 
         Later, saved tensors can be accessed through the :attr:`saved_tensors`
-        attribute. Before returning them to the user, a check is made, to
-        ensure they weren't used in any in-place operation that modified
-        their content.
+        attribute; or, if the corresponding Variable is needed (e.g. for double
+        backwards), those can be accessed through the :attr:`saved_variables`
+        attribute.  Before returning them to the user, a check is made, to ensure
+        they weren't used in any in-place operation that modified their content.
 
         Arguments can also be ``None``.
         """
@@ -65,7 +31,7 @@ class Function(_C._FunctionBase):
         :func:`forward` **method, and all arguments should be inputs.**
 
         Every tensor that's been modified in-place in a call to :func:`forward`
-        should be given to this function, to ensure correcness of our checks.
+        should be given to this function, to ensure correctness of our checks.
         It doesn't matter wheter the function is called before or after
         modification.
         """
@@ -106,14 +72,95 @@ class Function(_C._FunctionBase):
         """
         self.non_differentiable = args
 
-    def register_hook(self, hook):
-        if self._backward_hooks is None:
-            self._backward_hooks = OrderedDict()
-        handle = hooks.RemovableHandle(self._backward_hooks)
-        self._backward_hooks[id(handle)] = hook
-        return handle
 
-    def forward(self, *input):
+class _HookMixin(object):
+
+    @staticmethod
+    def _register_hook(backward_hooks, hook):
+        if backward_hooks is None:
+            backward_hooks = OrderedDict()
+        handle = hooks.RemovableHandle(backward_hooks)
+        backward_hooks[handle.id] = hook
+        return backward_hooks, handle
+
+
+class BackwardCFunction(_C._FunctionBase, _ContextMethodMixin, _HookMixin):
+    _is_legacy = False
+
+    def apply(self, *args):
+        return self._forward_cls.backward(self, *args)
+
+
+class FunctionMeta(type):
+    """Function metaclass.
+
+    This metaclass sets up the following properties:
+        _is_legacy: True if forward is not defined as a static method.
+        _backward_cls: The Function class corresponding to the differentiated
+            version of this function (which is generated on the fly by this
+            metaclass).
+    """
+
+    def __init__(cls, name, bases, attrs):
+        for super_cls in cls.mro():
+            forward = super_cls.__dict__.get('forward')
+            if forward is not None:
+                has_static_forward = isinstance(forward, staticmethod) or isinstance(forward, classmethod)
+                break
+
+        setattr(cls, '_is_legacy', not has_static_forward)
+
+        # old-style functions
+        if not has_static_forward:
+            return super(FunctionMeta, cls).__init__(name, bases, attrs)
+
+        backward_fn = type(name + 'Backward', (BackwardCFunction,), {'_forward_cls': cls})
+        setattr(cls, '_backward_cls', backward_fn)
+
+        return super(FunctionMeta, cls).__init__(name, bases, attrs)
+
+
+class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin)):
+    """Records operation history and defines formulas for differentiating ops.
+
+    Every operation performed on :class:`Variable` s creates a new function
+    object, that performs the computation, and records that it happened.
+    The history is retained in the form of a DAG of functions, with edges
+    denoting data dependencies (``input <- output``). Then, when backward is
+    called, the graph is processed in the topological ordering, by calling
+    :func:`backward` methods of each :class:`Function` object, and passing
+    returned gradients on to next :class:`Function` s.
+
+    Normally, the only way users interact with functions is by creating
+    subclasses and defining new operations. This is a recommended way of
+    extending torch.autograd.
+
+    Since Function logic is a hotspot in most scripts, almost all of it
+    was moved to our C backend, to ensure that the framework overhead is
+    minimal.
+
+    Each function is meant to be used only once (in the forward pass).
+
+    Attributes:
+        saved_tensors: Tuple of Tensors that were saved in the call to
+            :func:`forward`.
+        saved_variables: Tuple of Variables that correspond to the tensors
+            saved in the call to :func:`forward`.
+        needs_input_grad: Tuple of booleans of length :attr:`num_inputs`,
+            indicating whether a given input requires gradient. This can be
+            used to optimize buffers saved for backward, and ignoring gradient
+            computation in :func:`~Function.backward`.
+        num_inputs: Number of inputs given to :func:`forward`.
+        num_outputs: Number of tensors returned by :func:`forward`.
+        requires_grad: Boolean indicating whether the :func:`backward` will
+            ever need to be called.
+    """
+
+    # only for backward compatibility
+    __call__ = _C._FunctionBase._do_forward
+
+    @staticmethod
+    def forward(*args, **kwargs):
         """Performs the operation.
 
         This function is to be overriden by all subclasses.
@@ -122,7 +169,8 @@ class Function(_C._FunctionBase):
         """
         raise NotImplementedError
 
-    def backward(self, *grad_output):
+    @staticmethod
+    def backward(*grad_outputs):
         """Defines a formula for differentiating the operation.
 
         This function is to be overriden by all subclasses.
@@ -136,6 +184,41 @@ class Function(_C._FunctionBase):
         raise NotImplementedError
 
 
+def once_differentiable(fn):
+    from .variable import Variable
+
+    @functools.wraps(fn)
+    def wrapper(ctx, *args):
+        tensor_args = [arg.data if isinstance(arg, Variable) else arg
+                       for arg in args]
+        outputs = fn(ctx, *tensor_args)
+        # XXX: this is only an approximation of these flags - there's no way
+        # to figure out if fn didn't use ctx.saved_variables and as a result
+        # some Variables might require grad, even if no args do.
+        # Unfortunately, this leads to unexpected error messages ("no nodes
+        # require computing gradients"), but I don't have a better idea.
+        # These functions would raise an error in backward anyway.
+        volatile = any(arg.volatile if isinstance(arg, Variable) else False
+                       for arg in args)
+        requires_grad = any(arg.requires_grad if isinstance(arg, Variable) else False
+                            for arg in args)
+        if volatile:
+            def err_fn(*args):
+                return args
+            kwargs = {'volatile': True}
+        else:
+            err_fn = torch._C._functions.DelayedError(
+                b"trying to differentiate twice a function that was marked"
+                b"with @once_differentiable")
+            kwargs = {'requires_grad': requires_grad}
+        if not isinstance(outputs, tuple):
+            var = Variable(outputs, **kwargs) if outputs is not None else None
+            return err_fn(var)
+        return err_fn(*[Variable(o, **kwargs) if o is not None else None
+                      for o in outputs])
+    return wrapper
+
+
 class InplaceFunction(Function):
 
     def __init__(self, inplace=False):

torch/autograd/gradcheck.py

@@ -0,0 +1,225 @@
+import torch
+from torch.autograd import Variable
+from collections import Iterable
+
+
+def iter_variables(x):
+    if isinstance(x, Variable):
+        if x.requires_grad:
+            yield (x.grad.data, x.data) if x.grad is not None else (None, None)
+    elif isinstance(x, Iterable):
+        for elem in x:
+            for result in iter_variables(elem):
+                yield result
+
+
+def zero_gradients(x):
+    if isinstance(x, Variable):
+        if x.grad is not None:
+            x.grad.detach_()
+            x.grad.data.zero_()
+    elif isinstance(x, Iterable):
+        for elem in x:
+            zero_gradients(elem)
+
+
+def make_jacobian(input, num_out):
+    if isinstance(input, Variable) and not input.requires_grad:
+        return None
+    elif torch.is_tensor(input) or isinstance(input, Variable):
+        return torch.zeros(input.nelement(), num_out)
+    elif isinstance(input, Iterable):
+        jacobians = list(filter(
+            lambda x: x is not None, (make_jacobian(elem, num_out) for elem in input)))
+        if not jacobians:
+            return None
+        return type(input)(jacobians)
+    else:
+        return None
+
+
+def iter_tensors(x, only_requiring_grad=False):
+    if torch.is_tensor(x):
+        yield x
+    elif isinstance(x, Variable):
+        if x.requires_grad or not only_requiring_grad:
+            yield x.data
+    elif isinstance(x, Iterable):
+        for elem in x:
+            for result in iter_tensors(elem, only_requiring_grad):
+                yield result
+
+
+def contiguous(input):
+    if torch.is_tensor(input):
+        return input.contiguous()
+    elif isinstance(input, Variable):
+        return input.contiguous()
+    elif isinstance(input, Iterable):
+        return type(input)(contiguous(e) for e in input)
+    return input
+
+
+def get_numerical_jacobian(fn, input, target, eps=1e-3):
+    # To be able to use .view(-1) input must be contiguous
+    input = contiguous(input)
+    output_size = fn(input).numel()
+    jacobian = make_jacobian(target, output_size)
+
+    # It's much easier to iterate over flattened lists of tensors.
+    # These are reference to the same objects in jacobian, so any changes
+    # will be reflected in it as well.
+    x_tensors = [t for t in iter_tensors(target, True)]
+    j_tensors = [t for t in iter_tensors(jacobian)]
+
+    outa = torch.DoubleTensor(output_size)
+    outb = torch.DoubleTensor(output_size)
+
+    # TODO: compare structure
+    for x_tensor, d_tensor in zip(x_tensors, j_tensors):
+        flat_tensor = x_tensor.view(-1)
+        for i in range(flat_tensor.nelement()):
+            orig = flat_tensor[i]
+            flat_tensor[i] = orig - eps
+            outa.copy_(fn(input), broadcast=False)
+            flat_tensor[i] = orig + eps
+            outb.copy_(fn(input), broadcast=False)
+            flat_tensor[i] = orig
+
+            outb.add_(-1, outa).div_(2 * eps)
+            d_tensor[i] = outb
+
+    return jacobian
+
+
+def get_analytical_jacobian(input, output):
+    jacobian = make_jacobian(input, output.numel())
+    jacobian_reentrant = make_jacobian(input, output.numel())
+    grad_output = output.data.clone().zero_()
+    flat_grad_output = grad_output.view(-1)
+    reentrant = True
+    correct_grad_sizes = True
+
+    for i in range(flat_grad_output.numel()):
+        flat_grad_output.zero_()
+        flat_grad_output[i] = 1
+        for jacobian_c in (jacobian, jacobian_reentrant):
+            zero_gradients(input)
+            output.backward(grad_output, create_graph=True)
+            for jacobian_x, (d_x, x) in zip(jacobian_c, iter_variables(input)):
+                if d_x is None:
+                    jacobian_x[:, i].zero_()
+                else:
+                    if d_x.size() != x.size():
+                        correct_grad_sizes = False
+                    jacobian_x[:, i] = d_x.to_dense() if d_x.is_sparse else d_x
+
+    for jacobian_x, jacobian_reentrant_x in zip(jacobian, jacobian_reentrant):
+        if (jacobian_x - jacobian_reentrant_x).abs().max() != 0:
+            reentrant = False
+
+    return jacobian, reentrant, correct_grad_sizes
+
+
+def _as_tuple(x):
+    if isinstance(x, tuple):
+        return x
+    elif isinstance(x, list):
+        return tuple(x)
+    else:
+        return x,
+
+
+def gradcheck(func, inputs, eps=1e-6, atol=1e-5, rtol=1e-3):
+    """Check gradients computed via small finite differences
+       against analytical gradients
+
+    The check between numerical and analytical has the same behaviour as
+    numpy.allclose https://docs.scipy.org/doc/numpy/reference/generated/numpy.allclose.html
+    meaning it check that
+        absolute(a - n) <= (atol + rtol * absolute(n))
+    is true for all elements of analytical jacobian a and numerical jacobian n.
+
+    Args:
+        func: Python function that takes Variable inputs and returns
+            a tuple of Variables
+        inputs: tuple of Variables
+        eps: perturbation for finite differences
+        atol: absolute tolerance
+        rtol: relative tolerance
+
+    Returns:
+        True if all differences satisfy allclose condition
+    """
+    output = func(*inputs)
+    output = _as_tuple(output)
+
+    for i, o in enumerate(output):
+        if not o.requires_grad:
+            continue
+
+        def fn(input):
+            return _as_tuple(func(*input))[i].data
+
+        analytical, reentrant, correct_grad_sizes = get_analytical_jacobian(_as_tuple(inputs), o)
+        numerical = get_numerical_jacobian(fn, inputs, inputs, eps)
+
+        for a, n in zip(analytical, numerical):
+            if not ((a - n).abs() <= (atol + rtol * n.abs())).all():
+                return False
+
+        if not reentrant:
+            return False
+
+        if not correct_grad_sizes:
+            return False
+
+    # check if the backward multiplies by grad_output
+    zero_gradients(inputs)
+    output = _as_tuple(func(*inputs))
+    torch.autograd.backward(output, [o.data.new(o.size()).zero_() for o in output])
+    var_inputs = list(filter(lambda i: isinstance(i, Variable), inputs))
+    if not var_inputs:
+        raise RuntimeError("no Variables found in input")
+    for i in var_inputs:
+        if i.grad is None:
+            continue
+        if not i.grad.data.eq(0).all():
+            return False
+
+    return True
+
+
+def gradgradcheck(func, inputs, grad_outputs, eps=1e-6, atol=1e-5, rtol=1e-3):
+    """Check gradients of gradients computed via small finite differences
+       against analytical gradients
+    This function checks that backpropagating through the gradients computed
+    to the given grad_outputs are correct.
+
+    The check between numerical and analytical has the same behaviour as
+    numpy.allclose https://docs.scipy.org/doc/numpy/reference/generated/numpy.allclose.html
+    meaning it check that
+        absolute(a - n) <= (atol + rtol * absolute(n))
+    is true for all elements of analytical gradient a and numerical gradient n.
+
+    Args:
+        func: Python function that takes Variable inputs and returns
+            a tuple of Variables
+        inputs: tuple of Variables
+        grad_outputs: tuple of Variables
+        eps: perturbation for finite differences
+        atol: absolute tolerance
+        rtol: relative tolerance
+
+    Returns:
+        True if all differences satisfy allclose condition
+    """
+    def new_func(*input_args):
+        input_args = input_args[:-len(grad_outputs)]
+        outputs = func(*input_args)
+        outputs = _as_tuple(outputs)
+        input_args = tuple(x for x in input_args if isinstance(x, Variable) and x.requires_grad)
+        grad_inputs = torch.autograd.grad(outputs, input_args, grad_outputs)
+        return grad_inputs
+
+    return gradcheck(new_func, inputs + grad_outputs, eps, atol, rtol)

torch/autograd/stochastic_function.py

@@ -1,3 +1,5 @@
+import torch
+from numbers import Number
 from .function import Function
 
 _NOT_PROVIDED = object()
@@ -17,5 +19,26 @@ class StochasticFunction(Function):
             self.reward = None
         return result
 
+    def _do_forward(self, *inputs):
+        result = super(StochasticFunction, self)._do_forward(*inputs)
+        # save output type and size, to check the type of reward
+        assert isinstance(result, torch.autograd.Variable), \
+            "stochastic functions support only a single output at the moment"
+        self.reward_info = (type(inputs[0].data), result.size())
+        return result
+
+    __call__ = _do_forward
+
     def _reinforce(self, reward):
+        is_number = isinstance(reward, Number)
+        if not is_number and type(reward) != self.reward_info[0]:
+            raise TypeError("mismatch between reward and output type: got {}, "
+                            "but expected {}".format(torch.typename(reward),
+                                                     torch.typename(self.reward_info[0])))
+        if not is_number and reward.size() != self.reward_info[1]:
+            raise ValueError("got reward of size {}, but expected a tensor of size {}".format(
+                             'x'.join(map(str, reward.size())),
+                             'x'.join(map(str, self.reward_info[1]))))
+        if self.reward is not _NOT_PROVIDED:
+            raise RuntimeError("you can only reinforce a stochastic Function once")
         self.reward = reward

torch/autograd/variable.py

@@ -1,9 +1,11 @@
 import sys
+import torch
 import torch._C as _C
 from collections import OrderedDict
+import torch.sparse as sparse
 import torch.utils.hooks as hooks
-
-from ._functions import *
+import warnings
+import weakref
 
 
 class Variable(_C._VariableBase):
@@ -12,7 +14,7 @@ class Variable(_C._VariableBase):
     Variable is a thin wrapper around a Tensor object, that also holds
     the gradient w.r.t. to it, and a reference to a function that created it.
     This reference allows retracing the whole chain of operations that
-    created the data. If the Variable has been created by the user, its creator
+    created the data. If the Variable has been created by the user, its grad_fn
     will be ``None`` and we call such objects *leaf* Variables.
 
     Since autograd only supports scalar valued function differentiation, grad
@@ -32,8 +34,9 @@ class Variable(_C._VariableBase):
             inference mode, i.e. don't save the history. See
             :ref:`excluding-subgraphs` for more details.
             Can be changed only on leaf Variables.
-        creator: Function of which the variable was an output. For leaf
-            (user created) variables it's ``None``. Read-only attribute.
+        is_leaf: Boolean indicating if the Variable is a graph leaf (i.e
+            if it was created by the user).
+        grad_fn: Gradient function graph trace.
 
     Parameters:
         data (any tensor class): Tensor to wrap.
@@ -59,29 +62,30 @@ class Variable(_C._VariableBase):
     def __getattr__(self, name):
         if name in self._fallthrough_methods:
             return getattr(self.data, name)
-        raise AttributeError(name)
+        return object.__getattribute__(self, name)
 
     def __getitem__(self, key):
-        if (isinstance(key, Variable) and
-                type(key.data).__name__ == 'ByteTensor'):
-            return MaskedSelect()(self, key)
-        return Index(key)(self)
+        if torch.is_tensor(key):
+            key = Variable(key)  # auto-wrap tensors
+        if isinstance(key, Variable):
+            if type(key.data).__name__ == 'ByteTensor':
+                return MaskedSelect.apply(self, key)
+            elif type(key.data).__name__ == 'LongTensor':
+                return IndexSelect.apply(self, 0, key)
+            # else fall through and raise an error in Index
+        return Index.apply(self, key)
 
     def __setitem__(self, key, value):
-        if (isinstance(key, Variable) and
-                type(key.data).__name__ == 'ByteTensor'):
+        if isinstance(key, Variable) and type(key.data).__name__ == 'ByteTensor':
             if isinstance(value, Variable):
-                return MaskedCopy(inplace=True)(self, key, value)
+                return MaskedScatter.apply(self, key, value, True)
             else:
-                return MaskedFill(value, inplace=True)(self, key)
+                return MaskedFill.apply(self, key, value, True)
         else:
-            if isinstance(value, Variable):
-                return SetItem(key)(self, value)
-            else:
-                return SetItem(key, value)(self)
+            return SetItem.apply(self, key, value)
 
     def __deepcopy__(self, memo):
-        if self.creator is not None:
+        if not self.is_leaf:
             raise RuntimeError("Only Variables created explicitly by the user "
                                "(graph leaves) support the deepcopy protocol at the moment")
         result = type(self)(self.data.clone())
@@ -105,44 +109,51 @@ class Variable(_C._VariableBase):
             # legacy serialization of Variable
             self.data = state[0]
             state = (state[3], state[4], state[2])
-        if self.creator is not None:
+        if not self.is_leaf:
             raise RuntimeError('__setstate__ can be only called on leaf variables')
         self.requires_grad, self.volatile, self._backward_hooks = state
 
     def __repr__(self):
         return 'Variable containing:' + self.data.__repr__()
 
-    def backward(self, gradient=None, retain_variables=False):
+    def __bool__(self):
+        if self.data.numel() == 0:
+            return False
+        raise RuntimeError("bool value of Variable objects containing non-empty " +
+                           torch.typename(self.data) + " is ambiguous")
+
+    __nonzero__ = __bool__
+
+    def backward(self, gradient=None, retain_graph=None, create_graph=None, retain_variables=None):
         """Computes the gradient of current variable w.r.t. graph leaves.
 
         The graph is differentiated using the chain rule. If the variable is
         non-scalar (i.e. its data has more than one element) and requires
         gradient, the function additionaly requires specifying ``gradient``.
-        It should be a tensor of matching type and location, that containins
+        It should be a tensor of matching type and location, that contains
         the gradient of the differentiated function w.r.t. ``self``.
 
-        This function accumulates gradients in the leaves - you might need to zero
-        them before calling it.
+        This function accumulates gradients in the leaves - you might need to
+        zero them before calling it.
 
         Arguments:
-            gradient (Tensor): Gradient of the differentiated function
-                w.r.t. the data. Required only if the data has more than one
-                element. Type and location should match these of ``self.data``.
-            retain_variables (bool): If ``True``, buffers necessary for computing
-                gradients won't be freed after use. It is only necessary to
-                specify ``True`` if you want to differentiate some subgraph multiple
-                times (in some cases it will be much more efficient to use
-                `autograd.backward`).
+            grad_variables (Tensor, Variable or None): Gradient w.r.t. the
+                variable. If it is a tensor, it will be automatically converted
+                to a Variable that is volatile unless ``create_graph`` is True.
+                None values can be specified for scalar Variables or ones that
+                don't require grad. If a None value would be acceptable then
+                this argument is optional.
+            retain_graph (bool, optional): If False, the graph used to compute
+                the grads will be freed. Note that in nearly all cases setting
+                this option to True is not needed and often can be worked around
+                in a much more efficient way. Defaults to the value of
+                ``create_graph``.
+            create_graph (bool, optional): If true, graph of the derivative will
+                be constructed, allowing to compute higher order derivative
+                products. Defaults to False, unless ``gradient`` is a volatile
+                Variable.
         """
-        if self.volatile:
-            raise RuntimeError('calling backward on a volatile variable')
-        if gradient is None and self.requires_grad:
-            if self.data.numel() != 1:
-                raise RuntimeError(
-                    'backward should be called only on a scalar (i.e. 1-element tensor) '
-                    'or with gradient w.r.t. the variable')
-            gradient = self.data.new().resize_as_(self.data).fill_(1)
-        self._execution_engine.run_backward((self,), (gradient,), retain_variables)
+        torch.autograd.backward(self, gradient, retain_graph, create_graph, retain_variables)
 
     def register_hook(self, hook):
         """Registers a backward hook.
@@ -150,7 +161,7 @@ class Variable(_C._VariableBase):
         The hook will be called every time a gradient with respect to the
         variable is computed. The hook should have the following signature::
 
-            hook(grad) -> Tensor or None
+            hook(grad) -> Variable or None
 
         The hook should not modify its argument, but it can optionally return
         a new gradient which will be used in place of :attr:`grad`.
@@ -176,25 +187,12 @@ class Variable(_C._VariableBase):
                                "doesn't require gradient")
         if self._backward_hooks is None:
             self._backward_hooks = OrderedDict()
-            if self.creator is not None:
-                self.creator._register_hook_dict(self)
+            if self.grad_fn is not None:
+                self.grad_fn._register_hook_dict(self)
         handle = hooks.RemovableHandle(self._backward_hooks)
-        self._backward_hooks[id(handle)] = hook
+        self._backward_hooks[handle.id] = hook
         return handle
 
-    def _do_backward(self, grad_output, retain_variables):
-        assert len(grad_output) == 1
-        assert self._version == 0 and self.creator is None, \
-            "leaf variable was used in an inplace operation"
-        unpacked_grad = grad_output[0]
-        if self._backward_hooks:
-            for hook in self._backward_hooks.values():
-                result = hook(unpacked_grad)
-                if result is not None:
-                    unpacked_grad = result
-        self.grad.data.add_(unpacked_grad)
-        return tuple()
-
     def reinforce(self, reward):
         """Registers a reward obtained as a result of a stochastic process.
 
@@ -206,10 +204,10 @@ class Variable(_C._VariableBase):
             reward(Tensor): Tensor with per-element rewards. It has to match
                 the device location and shape of Variable's data.
         """
-        if not isinstance(self.creator, StochasticFunction):
+        if not isinstance(self.grad_fn, StochasticFunction):
             raise RuntimeError("reinforce() can be only called on outputs "
                                "of stochastic functions")
-        self.creator._reinforce(reward)
+        self.grad_fn._reinforce(reward)
 
     def detach(self):
         """Returns a new Variable, detached from the current graph.
@@ -224,32 +222,61 @@ class Variable(_C._VariableBase):
           errors in correctness checks.
         """
         result = NoGrad()(self)  # this is needed, because it merges version counters
-        result._creator = None
+        result._grad_fn = None
         return result
 
     def detach_(self):
-        """Detaches the Variable from the graph that created it, making it a leaf."""
-        self._creator = None
+        """Detaches the Variable from the graph that created it, making it a
+        leaf.
+        """
+        self._grad_fn = None
         self.requires_grad = False
 
+    def retain_grad(self):
+        """Enables .grad attribute for non-leaf Variables."""
+        if self.grad_fn is None:  # no-op for leaves
+            return
+        if not self.requires_grad:
+            raise RuntimeError("can't retain_grad on Variable that has requires_grad=False")
+        if hasattr(self, 'retains_grad'):
+            return
+        weak_self = weakref.ref(self)
+
+        def retain_grad_hook(grad):
+            var = weak_self()
+            if var is None:
+                return
+            if var._grad is None:
+                var._grad = grad.clone()
+            else:
+                var._grad = var._grad + grad
+
+        self.register_hook(retain_grad_hook)
+        self.retains_grad = True
+
     def contiguous(self):
         self.data = self.data.contiguous()
         return self
 
     def clone(self):
-        return Clone()(self)
+        return Clone.apply(self)
 
     def type(self, t):
         if t != type(self.data):
-            return Type(t)(self)
+            return Type.apply(self, t)
         return self
 
+    def type_as(self, t):
+        if isinstance(t, Variable):
+            t = t.data
+        return self.type(type(t))
+
     def _get_type(self, name):
         module = torch._import_dotted_name(self.data.__module__)
         return getattr(module, name)
 
     def cuda(self, device_id=None, async=False):
-        return CudaTransfer(device_id, async)(self)
+        return CudaTransfer.apply(self, device_id, async)
 
     def cpu(self):
         return self.type(getattr(torch, type(self.data).__name__))
@@ -283,10 +310,10 @@ class Variable(_C._VariableBase):
 
     def _add(self, other, inplace):
         if isinstance(other, Variable):
-            return Add(inplace)(self, other)
+            return Add.apply(self, other, inplace)
         else:
             assert not torch.is_tensor(other)
-            return AddConstant(other, inplace)(self)
+            return AddConstant.apply(self, other, inplace)
 
     def add(self, other):
         return self._add(other, False)
@@ -296,10 +323,10 @@ class Variable(_C._VariableBase):
 
     def _sub(self, other, inplace):
         if isinstance(other, Variable):
-            return Sub(inplace=inplace)(self, other)
+            return Sub.apply(self, other, inplace)
         else:
             assert not torch.is_tensor(other)
-            return SubConstant(other, inplace=inplace)(self)
+            return SubConstant.apply(self, other, inplace)
 
     def sub(self, other):
         return self._sub(other, False)
@@ -309,219 +336,229 @@ class Variable(_C._VariableBase):
 
     def mul(self, other):
         if isinstance(other, Variable):
-            return Mul()(self, other)
+            return Mul.apply(self, other)
         else:
             assert not torch.is_tensor(other)
-            return MulConstant(other)(self)
+            return MulConstant.apply(self, other)
 
     def mul_(self, other):
         if not isinstance(other, Variable) and not torch.is_tensor(other):
-            return MulConstant(other, inplace=True)(self)
+            return MulConstant.apply(self, other, True)
         raise RuntimeError("mul_ only supports scalar multiplication")
 
     def div(self, other):
         if isinstance(other, Variable):
-            return Div()(self, other)
+            return Div.apply(self, other)
         else:
             assert not torch.is_tensor(other)
-            return DivConstant(other)(self)
+            return DivConstant.apply(self, other)
 
     def div_(self, other):
         if not isinstance(other, Variable) and not torch.is_tensor(other):
-            return DivConstant(other, inplace=True)(self)
+            return DivConstant.apply(self, other, True)
         raise RuntimeError("div_ only supports scalar multiplication")
 
     def pow(self, other):
         if isinstance(other, Variable):
-            return Pow()(self, other)
+            return Pow.apply(self, other)
         else:
             assert not torch.is_tensor(other)
-            return PowConstant(other)(self)
+            return PowConstant.apply(self, other)
 
     def exp(self):
-        return Exp()(self)
+        return Exp.apply(self)
 
     def exp_(self):
-        return Exp(inplace=True)(self)
+        return Exp.apply(self, True)
 
     def log(self):
-        return Log()(self)
+        return Log.apply(self)
 
     def log1p(self):
-        return Log1p()(self)
+        return Log1p.apply(self)
 
     def neg(self):
-        return Negate()(self)
+        return Negate.apply(self)
 
     def neg_(self):
-        return Negate(inplace=True)(self)
+        return Negate.apply(self, True)
 
     def tanh(self):
-        return Tanh()(self)
+        return Tanh.apply(self)
 
     def tanh_(self):
-        return Tanh(True)(self)
+        return Tanh.apply(self, True)
 
     def sigmoid(self):
-        return Sigmoid()(self)
+        return Sigmoid.apply(self)
 
     def sigmoid_(self):
-        return Sigmoid(True)(self)
+        return Sigmoid.apply(self, True)
 
     def sin(self):
-        return Sin()(self)
+        return Sin.apply(self)
 
     def cos(self):
-        return Cos()(self)
+        return Cos.apply(self)
 
     def tan(self):
-        return Tan()(self)
+        return Tan.apply(self)
 
     def asin(self):
-        return Asin()(self)
+        return Asin.apply(self)
 
     def acos(self):
-        return Acos()(self)
+        return Acos.apply(self)
 
     def atan(self):
-        return Atan()(self)
+        return Atan.apply(self)
+
+    def atan2(self, x):
+        return Atan2.apply(self, x)
 
     def sinh(self):
-        return Sinh()(self)
+        return Sinh.apply(self)
 
     def cosh(self):
-        return Cosh()(self)
+        return Cosh.apply(self)
 
     def abs(self):
-        return Abs()(self)
+        return Abs.apply(self)
 
     def clamp(self, min=None, max=None):
         if min is None and max is None:
             raise ValueError("clamp requires specifying at least one of "
                              "min and max arguments")
         elif min is None and max is not None:
-            return CminConstant(max)(self)
+            return CminConstant.apply(self, max)
         elif min is not None and max is None:
-            return CmaxConstant(min)(self)
+            return CmaxConstant.apply(self, min)
         else:
-            return Clamp(min, max)(self)
+            return Clamp.apply(self, min, max)
 
     def reciprocal(self):
-        return Reciprocal()(self)
+        return Reciprocal.apply(self)
 
     def floor(self):
-        return Floor()(self)
+        return Floor.apply(self)
 
     def ceil(self):
-        return Ceil()(self)
+        return Ceil.apply(self)
 
     def frac(self):
-        return Frac()(self)
+        return Frac.apply(self)
 
     def sqrt(self):
-        return Sqrt()(self)
+        return Sqrt.apply(self)
 
     def round(self):
-        return Round()(self)
+        return Round.apply(self)
 
     def sign(self):
-        return Sign()(self)
+        return Sign.apply(self)
 
     def trunc(self):
-        return Trunc()(self)
-
-    def floor(self):
-        return Floor()(self)
-
-    def ceil(self):
-        return Ceil()(self)
+        return Trunc.apply(self)
 
     def fmod(self, value):
-        return Fmod(value)(self)
+        return Fmod.apply(self, value)
 
     def remainder(self, value):
-        return Remainder(value)(self)
+        return Remainder.apply(self, value)
 
     def lerp(self, tensor, weight):
-        return Lerp(weight)(self, tensor)
+        return Lerp.apply(self, tensor, weight)
 
     def rsqrt(self):
-        return Rsqrt()(self)
+        return Rsqrt.apply(self)
 
-    def sum(self, dim=None):
-        return Sum(dim)(self)
+    def sum(self, dim=None, keepdim=None):
+        return Sum.apply(self, dim, keepdim)
 
-    def prod(self, dim=None):
-        return Prod(dim)(self)
+    def prod(self, dim=None, keepdim=None):
+        return Prod.apply(self, dim, keepdim)
 
-    def mean(self, dim=None):
-        return Mean(dim)(self)
+    def mean(self, dim=None, keepdim=None):
+        return Mean.apply(self, dim, keepdim)
 
-    def max(self, dim=None):
+    def max(self, dim=None, keepdim=None):
         if isinstance(dim, Variable):
-            return Cmax()(self, dim)
-        return Max(dim)(self)
+            return Cmax.apply(self, dim)
+        return Max.apply(self, dim, keepdim)
 
-    def min(self, dim=None):
+    def min(self, dim=None, keepdim=None):
         if isinstance(dim, Variable):
-            return Cmin()(self, dim)
-        return Min(dim)(self)
+            return Cmin.apply(self, dim)
+        return Min.apply(self, dim, keepdim)
 
-    def mode(self, dim):
-        return Mode(dim)(self)
+    def mode(self, dim=None, keepdim=None):
+        return Mode.apply(self, dim, keepdim)
 
-    def median(self, dim):
-        return Median(dim)(self)
+    def median(self, dim=None, keepdim=None):
+        return Median.apply(self, dim, keepdim)
 
-    def kthvalue(self, dim):
-        return Kthvalue(dim)(self)
+    def kthvalue(self, k, dim=None, keepdim=None):
+        return Kthvalue.apply(self, k, dim, keepdim)
 
     def sort(self, dim=None, descending=False):
-        return Sort(dim, descending)(self)
+        return Sort.apply(self, dim, descending, True)
 
     def topk(self, k, dim=None, largest=True, sorted=True):
-        return Topk(k, dim, largest, sorted)(self)
+        return Topk.apply(self, k, dim, largest, sorted, True)
 
     def view(self, *sizes):
-        return View(*sizes)(self)
+        return View.apply(self, sizes)
 
     def view_as(self, tensor):
-        return View(*tensor.size())(self)
+        return View.apply(self, tensor.size())
 
     def split(self, split_size, dim=0):
         return torch.split(self, split_size, dim)
 
-    def chunk(self, n_chunks, dim=0):
-        return torch.chunk(self, n_chunks, dim)
-
     def repeat(self, *repeats):
         if len(repeats) == 1 and isinstance(repeats[0], torch.Size):
             repeats = repeats[0]
         else:
             repeats = torch.Size(repeats)
-        return Repeat(repeats)(self)
+        return Repeat.apply(self, repeats)
 
-    def var(self, dim=None, unbiased=True):
-        mean = self.mean(dim)
+    def cumsum(self, dim):
+        return Cumsum.apply(self, dim)
+
+    def cumprod(self, dim):
+        return Cumprod.apply(self, dim)
+
+    def unfold(self, dim, size, step):
+        return Unfold.apply(self, dim, size, step)
+
+    def var(self, dim=None, keepdim=None, unbiased=True):
+        keepdim_ = False if keepdim is None else keepdim
+        mean = self.mean(dim, keepdim)
         if dim is None:
             mean = mean.view(*(1 for s in self.size()))
+        # we could just set keepdim to True, but this preserves some fidelity
+        elif keepdim_ is False and self.dim() != 1:
+            mean = mean.unsqueeze(dim)
         mean_expanded = mean.expand_as(self)
         zero_centered = self.sub(mean_expanded)
-        var = zero_centered.mul(zero_centered).sum(dim)
+        var = zero_centered.mul(zero_centered).sum(dim, keepdim=keepdim_)
         numel = self.numel() if dim is None else self.size(dim)
         return var.div(numel - int(unbiased))
 
-    def std(self, dim=None, unbiased=True):
-        return self.var(dim, unbiased).sqrt()
+    def std(self, dim=None, keepdim=None, unbiased=True):
+        return self.var(dim, keepdim, unbiased).sqrt()
 
-    def renorm(self, norm_type, dim, maxnorm):
+    def renorm(self, p, dim, maxnorm):
         t = self.transpose(dim, 0)
         flat = t.contiguous().view(self.size(0), -1)
-        norms = flat.norm(norm_type, 1)
+        norms = flat.norm(p, 1, True)
         norms = norms.clamp(max=maxnorm).div(norms.add(1e-7))
         flat_out = flat.mul(norms.expand_as(flat))
         return flat_out.view(t.size()).transpose(dim, 0)
 
+    def matmul(self, other):
+        return torch.matmul(self, other)
+
     @staticmethod
     def _static_blas(cls, args, inplace):
         num_args = len(args)
@@ -532,14 +569,14 @@ class Variable(_C._VariableBase):
             alpha, beta = args[1:3]
         if num_args == 4:
             alpha = args[1]
-        return cls(alpha, beta, inplace)(*(args[:1] + args[-2:]))
+        return cls.apply(*(args[:1] + args[-2:] + (alpha, beta, inplace)))
 
     def _blas(self, cls, args, inplace):
         return self._static_blas(cls, (self,) + args, inplace)
 
     def mm(self, matrix):
         output = Variable(self.data.new(self.data.size(0), matrix.data.size(1)))
-        return self._static_blas(Addmm, (output, 0, 1, self, matrix), False)
+        return Addmm.apply(output, self, matrix, 0, 1, True)
 
     def bmm(self, batch):
         output = Variable(self.data.new(self.data.size(0), self.data.size(1),
@@ -555,10 +592,10 @@ class Variable(_C._VariableBase):
         return self._static_blas(Addr, (output, 0, 1, self, vector), False)
 
     def resize(self, *sizes):
-        return Resize(*sizes)(self)
+        return Resize.apply(self, sizes)
 
     def resize_as(self, variable):
-        return Resize(*variable.size())(self)
+        return Resize.apply(self, variable.size())
 
     def addmm(self, *args):
         return self._blas(Addmm, args, False)
@@ -591,162 +628,186 @@ class Variable(_C._VariableBase):
         return self._blas(Addr, args, True)
 
     def dot(self, other):
-        return Dot()(self, other)
+        return Dot.apply(self, other)
 
-    def _addcop(self, op, args):
+    def _addcop(self, op, args, inplace):
         if len(args) == 3:
-            # scale, tensor1, tensor2
-            return op(args[0])(self, *args[1:])
+            # args == [scale, tensor1, tensor2]
+            return op.apply(self, args[1], args[2], args[0], inplace)
         else:
-            # tensor1, tensor2
-            return op()(self, *args)
+            # args == [tensor1, tensor2]
+            return op.apply(self, args[0], args[1], 1.0, inplace)
 
     def addcmul(self, *args):
-        return self._addcop(Addcmul, args)
+        return self._addcop(Addcmul, args, False)
 
     def addcdiv(self, *args):
-        return self._addcop(Addcdiv, args)
+        return self._addcop(Addcdiv, args, False)
 
-    def norm(self, norm_type=2, dim=None):
-        return Norm(norm_type, dim)(self)
+    def addcmul_(self, *args):
+        return self._addcop(Addcmul, args, True)
 
-    def dist(self, tensor, norm_type=2):
-        return Norm(norm_type)(self - tensor)
+    def addcdiv_(self, *args):
+        return self._addcop(Addcdiv, args, True)
+
+    def norm(self, p=2, dim=None, keepdim=None):
+        return Norm.apply(self, p, dim, keepdim)
+
+    def dist(self, tensor, p=2):
+        return Norm.apply(self - tensor, p)
 
     def index_add(self, dim, index, tensor):
-        return IndexAdd(dim)(self, index, tensor)
+        return IndexAdd.apply(self, dim, index, tensor)
+
+    def _advanced_index_add(self, index, tensor):
+        return AdvancedIndexAdd.apply(self, index, tensor)
 
     def index_add_(self, dim, index, tensor):
-        return IndexAdd(dim, True)(self, index, tensor)
+        return IndexAdd.apply(self, dim, index, tensor, True)
 
     def index_copy(self, dim, index, tensor):
-        return IndexCopy(dim)(self, index, tensor)
+        return IndexCopy.apply(self, dim, index, tensor)
 
     def index_copy_(self, dim, index, tensor):
-        return IndexCopy(dim, True)(self, index, tensor)
+        return IndexCopy.apply(self, dim, index, tensor, True)
 
     def index_fill(self, dim, index, value):
-        return IndexFill(dim, value)(self, index)
+        return IndexFill.apply(self, dim, index, value)
 
     def index_fill_(self, dim, index, value):
-        return IndexFill(dim, value, True)(self, index)
+        return IndexFill.apply(self, dim, index, value, True)
 
     def index_select(self, dim, index):
-        return IndexSelect(dim)(self, index)
+        return IndexSelect.apply(self, dim, index)
 
     def gather(self, dim, index):
-        return Gather(dim)(self, index)
+        return Gather.apply(self, dim, index)
 
     def scatter(self, dim, index, source):
-        return Scatter(dim)(self, index, source)
+        return Scatter.apply(self, dim, index, source)
 
     def scatter_(self, dim, index, source):
-        return Scatter(dim, True)(self, index, source)
+        return Scatter.apply(self, dim, index, source, True)
+
+    def scatter_add(self, dim, index, source):
+        return ScatterAdd.apply(self, dim, index, source)
+
+    def scatter_add_(self, dim, index, source):
+        return ScatterAdd.apply(self, dim, index, source, True)
 
     def masked_copy(self, mask, variable):
-        return MaskedCopy()(self, mask, variable)
+        warnings.warn("masked_copy is deprecated and renamed to masked_scatter, and will be removed in v0.3")
+        return MaskedScatter.apply(self, mask, variable)
 
     def masked_copy_(self, mask, variable):
-        return MaskedCopy(True)(self, mask, variable)
+        warnings.warn("masked_copy_ is deprecated and renamed to masked_scatter_, and will be removed in v0.3")
+        return MaskedScatter.apply(self, mask, variable, True)
+
+    def masked_scatter(self, mask, variable):
+        return MaskedScatter.apply(self, mask, variable)
+
+    def masked_scatter_(self, mask, variable):
+        return MaskedScatter.apply(self, mask, variable, True)
 
     def masked_fill(self, mask, value):
-        return MaskedFill(value)(self, mask)
+        return MaskedFill.apply(self, mask, value)
 
     def masked_fill_(self, mask, value):
-        return MaskedFill(value, True)(self, mask)
+        return MaskedFill.apply(self, mask, value, True)
 
     def masked_select(self, mask):
-        return MaskedSelect()(self, mask)
+        return MaskedSelect.apply(self, mask)
 
     def expand(self, *sizes):
-        if isinstance(sizes[0], torch.Size):
-            if len(sizes) > 1:
-                raise ValueError("expand expects a several ints or a single "
-                                 "torch.Size argument")
-            sizes = sizes[0]
-        return Expand(sizes)(self)
+        return Expand.apply(self, sizes)
 
     def expand_as(self, tensor):
-        return Expand(tensor.size())(self)
+        return Expand.apply(self, (tensor.size(),))
 
     def t(self):
-        return Transpose(0, 1)(self)
+        if self.dim() != 2:
+            raise RuntimeError("t() expects a 2D Variable, but self is {}D".format(self.dim()))
+        return Transpose.apply(self, 0, 1)
 
     def transpose(self, dim1, dim2):
-        return Transpose(dim1, dim2)(self)
+        return Transpose.apply(self, dim1, dim2)
 
     def select(self, dim, _index):
+        dim = dim if dim >= 0 else dim + self.dim()
         index = tuple(slice(None, None) for _ in range(dim)) + (_index,)
-        return Index(index)(self)
+        return Index.apply(self, index)
 
     def narrow(self, dim, start_index, length):
+        dim = dim if dim >= 0 else dim + self.dim()
         index = tuple(slice(None, None) for _ in range(dim)) + \
             (slice(start_index, start_index + length),)
-
-        return Index(index)(self)
+        return Index.apply(self, index)
 
     def chunk(self, num_chunks, dim=0):
-        return Chunk(num_chunks, dim)(self)
+        return Chunk.apply(self, num_chunks, dim)
 
     def squeeze(self, dim=None):
-        return Squeeze(dim)(self)
+        return Squeeze.apply(self, dim)
+
+    def squeeze_(self, dim=None):
+        return Squeeze.apply(self, dim, True)
 
     def unsqueeze(self, dim):
-        return Unsqueeze(dim)(self)
+        return Unsqueeze.apply(self, dim)
 
     def permute(self, *permutation):
-        return Permute(permutation)(self)
+        return Permute.apply(self, permutation)
 
-    def diag(self, diagonal_idx=0):
-        return Diag(diagonal_idx)(self)
+    def diag(self, diagonal=0):
+        return Diag.apply(self, diagonal)
 
-    def tril(self, diagonal_idx=0):
-        return Tril(diagonal_idx)(self)
+    def tril(self, diagonal=0):
+        return Tril.apply(self, diagonal)
 
-    def triu(self, diagonal_idx=0):
-        return Triu(diagonal_idx)(self)
+    def triu(self, diagonal=0):
+        return Triu.apply(self, diagonal)
 
-    def multinomial(self, num_samples=1, with_replacement=False):
-        return Multinomial(num_samples, with_replacement)(self)
+    def trace(self):
+        return Trace.apply(self)
+
+    def cross(self, other, dim=-1):
+        return Cross.apply(self, other)
+
+    def inverse(self):
+        return Inverse.apply(self)
+
+    def gesv(self, a):
+        return Gesv.apply(self, a)
+
+    def multinomial(self, num_samples=1, replacement=False):
+        return Multinomial(num_samples, replacement)(self)
 
     def bernoulli(self):
         return Bernoulli()(self)
 
     def eq(self, other):
-        if isinstance(other, Variable):
-            return Eq()(self, other)
         assert not torch.is_tensor(other), "can't compare Variable and tensor"
-        return Eq(other)(self)
+        return Eq.apply(self, other)
 
     def ne(self, other):
-        if isinstance(other, Variable):
-            return Ne()(self, other)
         assert not torch.is_tensor(other), "can't compare Variable and tensor"
-        return Ne(other)(self)
+        return Ne.apply(self, other)
 
     def gt(self, other):
-        if isinstance(other, Variable):
-            return Gt()(self, other)
         assert not torch.is_tensor(other), "can't compare Variable and tensor"
-        return Gt(other)(self)
+        return Gt.apply(self, other)
 
     def ge(self, other):
-        if isinstance(other, Variable):
-            return Ge()(self, other)
         assert not torch.is_tensor(other), "can't compare Variable and tensor"
-        return Ge(other)(self)
+        return Ge.apply(self, other)
 
     def lt(self, other):
-        if isinstance(other, Variable):
-            return Lt()(self, other)
         assert not torch.is_tensor(other), "can't compare Variable and tensor"
-        return Lt(other)(self)
+        return Lt.apply(self, other)
 
     def le(self, other):
-        if isinstance(other, Variable):
-            return Le()(self, other)
         assert not torch.is_tensor(other), "can't compare Variable and tensor"
-        return Le(other)(self)
+        return Le.apply(self, other)
 
     def __add__(self, other):
         return self.add(other)
@@ -762,7 +823,7 @@ class Variable(_C._VariableBase):
         return self.sub_(other)
 
     def __rsub__(self, other):
-        return SubConstant(other, sub_tensor=True)(self)
+        return SubConstant.apply(other, self)
 
     def __mul__(self, other):
         return self.mul(other)
@@ -772,28 +833,16 @@ class Variable(_C._VariableBase):
         return self.mul_(other)
 
     def __matmul__(self, other):
-        dim_self = self.dim()
-        try:
-            dim_other = other.dim()
-        except AttributeError:  # not a Variable
+        if not isinstance(other, Variable):
             return NotImplemented
-        if dim_self == 1 and dim_other == 1:
-            return self.dot(other)
-        if dim_self == 2 and dim_other == 1:
-            return self.mv(other)
-        if dim_self == 1 and dim_other == 2:
-            return self.unsqueeze(0).mm(other).squeeze(0)
-        elif dim_self == 2 and dim_other == 2:
-            return self.mm(other)
-        raise ValueError("both arguments to __matmul__ need to be 1D or 2D, "
-                         "but they are {}D and {}D".format(dim_self, dim_other))
+        return self.matmul(other)
 
     def __div__(self, other):
         return self.div(other)
     __truediv__ = __div__
 
     def __rdiv__(self, other):
-        return DivConstant(other, div_by_tensor=True)(self)
+        return DivConstant.apply(other, self)
     __rtruediv__ = __rdiv__
 
     def __idiv__(self, other):
@@ -806,10 +855,10 @@ class Variable(_C._VariableBase):
         raise NotImplementedError("in-place pow not implemented")
 
     def __rpow__(self, other):
-        return PowConstant(other, tensor_power=True)(self)
+        return PowConstant.apply(other, self)
 
     def __neg__(self):
-        return Negate()(self)
+        return Negate.apply(self)
 
     def __len__(self):
         return len(self.data)
@@ -845,7 +894,7 @@ class Variable(_C._VariableBase):
 
         @staticmethod
         def cat(iterable, dim=0):
-            return Concat(dim)(*iterable)
+            return Concat.apply(dim, *iterable)
 
         @staticmethod
         def normal(means, std=1):
@@ -868,7 +917,7 @@ class Variable(_C._VariableBase):
                 tensors = args[1:]
             else:
                 tensors = args
-            return cls(alpha, beta, inplace)(*tensors)
+            return cls.apply(*(tensors + (alpha, beta, inplace)))
 
         @classmethod
         def addmm(cls, *args):
@@ -902,5 +951,6 @@ for method in dir(Variable):
     setattr(Variable._torch, method, as_static)
 
 
-from .engine import ImperativeEngine
+from ._functions import *
+from torch._C import _ImperativeEngine as ImperativeEngine
 Variable._execution_engine = ImperativeEngine()

torch/backends/cudnn/__init__.py

@@ -1,45 +1,37 @@
-import torch._C as _C
 import ctypes
-import warnings
-import torch.cuda
 import sys
-import os.path as path
+import torch
+import warnings
 
 enabled = True  # set to False to globally disable cuDNN
 
 lib = None
-# TODO: fix libname for Windows
+__cudnn_version = None
 # TODO: dynamic version checks via cudnnGetVersion
-# TODO: load 5.1.3 if using CUDA 7.5 and 5.1.5 if using CUDA 8.0
-thisdir = path.dirname(__file__)
-libpaths = ['', path.join(thisdir, '../../lib')]
-if sys.platform.startswith('linux'):
-    libnames = ['libcudnn.so.6.0.5', 'libcudnn.so.6.0.10', 'libcudnn.so.5.1.5', 'libcudnn.so.5.1.3',
-                'libcudnn.so.5.0.5', 'libcudnn.so.5.1.10']
-elif sys.platform == 'darwin':
-    libnames = ['libcudnn.6.dylib', 'libcudnn.5.dylib']
-else:
-    libnames = []
 
 
-def _loadlib():
-    global lib
-    loaded = False
-    for libpath in libpaths:
-        for libname in libnames:
-            try:
-                lib = ctypes.cdll.LoadLibrary(path.join(libpath, libname))
-                loaded = True
-                break
-            except OSError:
-                continue
-        if loaded:
-            break
-    if loaded:
-        lib.cudnnGetErrorString.restype = ctypes.c_char_p
-    else:
-        lib = None
-        raise OSError("Could not load cuDNN")
+def _libcudnn():
+    global lib, __cudnn_version
+    if lib is None:
+        lib = ctypes.cdll.LoadLibrary(None)
+        if hasattr(lib, 'cudnnGetErrorString'):
+            lib.cudnnGetErrorString.restype = ctypes.c_char_p
+            __cudnn_version = lib.cudnnGetVersion()
+            compile_version = torch._C._cudnn_version()
+            # Check that cuDNN major and minor versions match
+            if (__cudnn_version // 100) != (compile_version // 100):
+                raise RuntimeError(
+                    'cuDNN version mismatch: PyTorch was compiled against {} '
+                    'but linked against {}'.format(compile_version, __cudnn_version))
+        else:
+            lib = None
+    return lib
+
+
+def version():
+    if _libcudnn() is None:
+        return None
+    return __cudnn_version
 
 
 def is_acceptable(tensor):
@@ -49,59 +41,30 @@ def is_acceptable(tensor):
             isinstance(tensor, torch.cuda.FloatTensor) or
             isinstance(tensor, torch.cuda.DoubleTensor)):
         return False
-    if lib is None:
-        try:
-            _loadlib()
-        except Exception:
-            warnings.warn('cuDNN library not found. Check your {libpath}'.format(
-                libpath={
-                    'darwin': 'DYLD_LIBRARY_PATH',
-                    'win32': 'PATH'
-                }.get(sys.platform, 'LD_LIBRARY_PATH')))
-            return False
-    if not _C.has_cudnn:
-        warnings.warn("cuDNN library has been detected, but your pytorch "
-                      "installation was compiled without support for it. You "
-                      "might want to rebuild pytorch, making sure the library "
-                      "is visible to the build system.")
+    if not torch._C.has_cudnn:
+        warnings.warn(
+            "PyTorch was compiled without cuDNN support. To use cuDNN, rebuild "
+            "PyTorch making sure the library is visible to the build system.")
+        return False
+    if _libcudnn() is None:
+        warnings.warn('cuDNN library not found. Check your {libpath}'.format(
+            libpath={
+                'darwin': 'DYLD_LIBRARY_PATH',
+                'win32': 'PATH'
+            }.get(sys.platform, 'LD_LIBRARY_PATH')))
         return False
     return True
 
-__cudnn_version = []
-
-
-def version():
-    if not lib:
-        raise RuntimeError("cuDNN not initialized")
-    if len(__cudnn_version) == 0:
-        __cudnn_version.append(lib.cudnnGetVersion())
-    return __cudnn_version[0]
 
 _handles = {}
 
 benchmark = False
 verbose = False
-workspace_limit = None
 
 CUDNN_DATA_FLOAT = 0
 CUDNN_DATA_DOUBLE = 1
 CUDNN_DATA_HALF = 2
 
-CUDNN_CONVOLUTION = 0
-CUDNN_CROSS_CORRELATION = 1
-
-CUDNN_CONVOLUTION_FWD_NO_WORKSPACE = 0
-CUDNN_CONVOLUTION_FWD_PREFER_FASTEST = 1
-CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT = 2
-
-CUDNN_CONVOLUTION_BWD_FILTER_NO_WORKSPACE = 0
-CUDNN_CONVOLUTION_BWD_FILTER_PREFER_FASTEST = 1
-CUDNN_CONVOLUTION_BWD_FILTER_SPECIFY_WORKSPACE_LIMIT = 2
-
-CUDNN_CONVOLUTION_BWD_DATA_NO_WORKSPACE = 0
-CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST = 1
-CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT = 2
-
 CUDNN_TENSOR_NCHW = 0
 CUDNN_TENSOR_NHWC = 1
 
@@ -113,16 +76,12 @@ CUDNN_GRU = 3
 CUDNN_LINEAR_INPUT = 0
 CUDNN_SKIP_INPUT = 1
 
-CUDNN_NON_DETERMINISTIC = 0
-CUDNN_DETERMINISTIC = 1
-
 CUDNN_RNN_ALGO_STANDARD = 0
 CUDNN_RNN_ALGO_PERSIST_STATIC = 1
 CUDNN_RNN_ALGO_PERSIST_DYNAMIC = 2
 
 
 class CuDNNHandle:
-
     def __init__(self):
         ptr = ctypes.c_void_p()
         check_error(lib.cudnnCreate(ctypes.byref(ptr)))
@@ -133,7 +92,6 @@ class CuDNNHandle:
 
 
 class CuDNNError(RuntimeError):
-
     def __init__(self, status):
         self.status = status
         msg = '{}: {}'.format(status, get_error_string(status))
@@ -141,7 +99,6 @@ class CuDNNError(RuntimeError):
 
 
 class TensorDescriptor(object):
-
     def __init__(self):
         ptr = ctypes.c_void_p()
         check_error(lib.cudnnCreateTensorDescriptor(ctypes.byref(ptr)))
@@ -164,7 +121,6 @@ class TensorDescriptor(object):
 
 
 class TensorDescriptorArray(object):
-
     def __init__(self, N):
         self.ptrs = (ctypes.c_void_p * N)()
         for i in range(N):
@@ -179,44 +135,22 @@ class TensorDescriptorArray(object):
     def __getitem__(self, key):
         return ctypes.c_void_p(self.ptrs[key])
 
-    def set(self, tensor):
-        self._type = tensor.type()
-        self._size = tensor.size()
-        self._stride = tensor.stride()
+    def set_all(self, tensor):
+        _type = _typemap[tensor.type()]
+        _ndim = tensor.dim()
+        _size = int_array(tensor.size())
+        _stride = int_array(tensor.stride())
         for ptr in self.ptrs:
             check_error(lib.cudnnSetTensorNdDescriptor(
-                ctypes.c_void_p(ptr), _typemap[tensor.type()], tensor.dim(),
-                int_array(tensor.size()), int_array(tensor.stride())))
+                ctypes.c_void_p(ptr), _type, _ndim, _size, _stride))
 
-    def as_tuple(self):
-        return (self._type, tuple(self._size), tuple(self._stride))
-
-
-class ConvolutionDescriptor(object):
-
-    def __init__(self):
-        ptr = ctypes.c_void_p()
-        check_error(lib.cudnnCreateConvolutionDescriptor(ctypes.byref(ptr)))
-        self._as_parameter_ = ptr
-
-    def __del__(self):
-        check_error(lib.cudnnDestroyConvolutionDescriptor(self._as_parameter_))
-        del self._as_parameter_
-
-    def set(self, typename, pad, stride):
-        self._pad = pad
-        self._stride = stride
-        upscale = int_array([1, 1])
-        check_error(lib.cudnnSetConvolutionNdDescriptor(
-            self, 2, int_array(pad), int_array(stride), upscale,
-            CUDNN_CROSS_CORRELATION, _typemap[typename]))
-
-    def as_tuple(self):
-        return (self._pad, self._stride)
+    def set_raw(self, i, _type, _ndim, _size, _stride):
+        ptr = self.ptrs[i]
+        check_error(lib.cudnnSetTensorNdDescriptor(
+            ctypes.c_void_p(ptr), _type, _ndim, _size, _stride))
 
 
 class FilterDescriptor(object):
-
     def __init__(self):
         ptr = ctypes.c_void_p()
         check_error(lib.cudnnCreateFilterDescriptor(ctypes.byref(ptr)))
@@ -230,41 +164,58 @@ class FilterDescriptor(object):
         self._size = weight.size()
         datatype = _typemap[weight.type()]
         check_error(lib.cudnnSetFilterNdDescriptor(
-            self, datatype, CUDNN_TENSOR_NCHW, weight.ndimension(), int_array(weight.size())))
+            self, datatype, CUDNN_TENSOR_NCHW, weight.ndimension(),
+            int_array(weight.size())))
 
     def as_tuple(self):
         return tuple(self._size)
 
 
 class DropoutDescriptor(object):
-
     def __init__(self, handle, dropout, seed):
         ptr = ctypes.c_void_p()
         check_error(lib.cudnnCreateDropoutDescriptor(ctypes.byref(ptr)))
+
         self._as_parameter_ = ptr
+        self.state = None
+        self.dropout = dropout
+        self.handle = handle
 
-        dropout_states_size = ctypes.c_long()
-        check_error(lib.cudnnDropoutGetStatesSize(
-            handle,
-            ctypes.byref(dropout_states_size)))
+        self._set(dropout, seed)
 
-        self.state = torch.cuda.ByteTensor(dropout_states_size.value)
+    def set_dropout(self, dropout, seed):
+        if dropout != self.dropout:
+            self._set(dropout, seed)
+
+    def _set(self, dropout, seed):
+        if self.state is None and dropout > 0:
+            dropout_states_size = ctypes.c_long()
+            check_error(lib.cudnnDropoutGetStatesSize(
+                self.handle,
+                ctypes.byref(dropout_states_size)))
+            self.state = torch.cuda.ByteTensor(dropout_states_size.value)
+            state_ptr = self.state.data_ptr()
+            state_size = self.state.size(0)
+        else:
+            state_ptr = None
+            state_size = 0
 
         check_error(lib.cudnnSetDropoutDescriptor(
             self,
-            handle,
+            self.handle,
             ctypes.c_float(dropout),
-            ctypes.c_void_p(self.state.data_ptr()),
-            ctypes.c_size_t(self.state.size(0)),
+            ctypes.c_void_p(state_ptr),
+            ctypes.c_size_t(state_size),
             ctypes.c_ulonglong(seed),
         ))
 
+        self.dropout = dropout
+
     def __del__(self):
         check_error(lib.cudnnDestroyDropoutDescriptor(self))
 
 
 class RNNDescriptor(object):
-
     def __init__(self, handle, hidden_size, num_layers, dropout_desc, input_mode,
                  bidirectional, mode, datatype):
         ptr = ctypes.c_void_p()
@@ -299,26 +250,6 @@ class RNNDescriptor(object):
         check_error(lib.cudnnDestroyRNNDescriptor(self))
 
 
-class ConvolutionAlgoPerf_v5(ctypes.Structure):
-    _fields_ = [
-        ("algo", ctypes.c_int),
-        ("status", ctypes.c_int),
-        ("time", ctypes.c_float),
-        ("memory", ctypes.c_size_t),
-    ]
-
-
-class ConvolutionAlgoPerf_v6(ctypes.Structure):
-    _fields_ = [
-        ("algo", ctypes.c_int),
-        ("status", ctypes.c_int),
-        ("time", ctypes.c_float),
-        ("memory", ctypes.c_size_t),
-        ("determinism", ctypes.c_int),
-        ("reserved", ctypes.c_int * 4)
-    ]
-
-
 def check_error(status):
     if status is not 0:
         raise CuDNNError(status)
@@ -329,8 +260,8 @@ def get_error_string(status):
 
 
 def get_handle():
-    if lib is None:
-        _loadlib()
+    if _libcudnn() is None:
+        raise RuntimeError('cuDNN not available')
     current_device = torch.cuda.current_device()
     handle = _handles.get(current_device, None)
     if handle is None:
@@ -338,6 +269,7 @@ def get_handle():
         _handles[current_device] = handle
     return handle
 
+
 _typemap = {
     'torch.cuda.HalfTensor': CUDNN_DATA_HALF,
     'torch.cuda.FloatTensor': CUDNN_DATA_FLOAT,
@@ -368,136 +300,28 @@ def int_array(itr):
 
 
 def descriptor(tensor, N=None):
+    padded_size = tensor.size() + ((1,) * (5 - tensor.dim()))
+    tensor = tensor.view(padded_size)
     if N is not None:
         descriptor = TensorDescriptorArray(N)
+        descriptor.set_all(tensor)
     else:
         descriptor = TensorDescriptor()
-    if tensor.dim() == 2:
-        tensor = tensor.view(tensor.size(0), tensor.size(1), 1, 1)
-    elif tensor.dim() == 3:
-        tensor = tensor.view(tensor.size(0), tensor.size(1), tensor.size(2), 1)
-    descriptor.set(tensor)
+        descriptor.set(tensor)
     return descriptor
 
-_autotuner_forward = {}
-_autotuner_backward_data = {}
-_autotuner_backward_filter = {}
 
-
-def convolution_autotuner_key(idesc, weight_desc, conv_desc):
-    return (idesc.as_tuple(), weight_desc.as_tuple(), conv_desc.as_tuple())
-
-
-def convolution_forward_algorithm(idesc, weight_desc, conv_desc, odesc):
-    k = convolution_autotuner_key(idesc, weight_desc, conv_desc)
-    if k in _autotuner_forward:
-        return _autotuner_forward[k]
-
-    if benchmark:
-        if version() < 6000:
-            perf_results = ConvolutionAlgoPerf_v5()
-        else:
-            perf_results = ConvolutionAlgoPerf_v6()
-        algo_count = ctypes.c_int()
-        check_error(lib.cudnnFindConvolutionForwardAlgorithm(
-            get_handle(), idesc, weight_desc, conv_desc, odesc, 1,
-            ctypes.byref(algo_count), ctypes.byref(perf_results)))
-        _autotuner_forward[k] = perf_results.algo
-        return perf_results.algo
-
-    search_mode = CUDNN_CONVOLUTION_FWD_PREFER_FASTEST
-    wlimit = 0
-    if workspace_limit is not None:
-        wlimit = workspace_limit
-        search_mode = CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
-
-    fwd_alg = ctypes.c_int()
-    check_error(lib.cudnnGetConvolutionForwardAlgorithm(
-        get_handle(), idesc, weight_desc, conv_desc, odesc, search_mode,
-        wlimit, ctypes.byref(fwd_alg)))
-    return fwd_alg
-
-
-def convolution_forward_workspace_size(*args):
-    check_error(lib.cudnnGetConvolutionForwardWorkspaceSize(*args))
-
-
-def convolution_forward(*args):
-    check_error(lib.cudnnConvolutionForward(*args))
-
-
-def convolution_backward_data(*args):
-    return check_error(lib.cudnnConvolutionBackwardData(*args))
-
-
-def convolution_backward_data_algorithm(weight_desc, odesc, conv_desc, idesc):
-    k = convolution_autotuner_key(idesc, weight_desc, conv_desc)
-    if k in _autotuner_backward_data:
-        return _autotuner_backward_data[k]
-
-    if benchmark:
-        perf_results = ConvolutionAlgoPerf()
-        algo_count = ctypes.c_int()
-        check_error(lib.cudnnFindConvolutionBackwardDataAlgorithm(
-            get_handle(), weight_desc, odesc, conv_desc, idesc, 1,
-            ctypes.byref(algo_count), ctypes.byref(perf_results)))
-        _autotuner_backward_data[k] = perf_results.algo
-        return perf_results.algo
-
-    search_mode = CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST
-    wlimit = 0
-    if workspace_limit is not None:
-        wlimit = workspace_limit
-        search_mode = CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT
-
-    bwd_data_alg = ctypes.c_int()
-    check_error(lib.cudnnGetConvolutionBackwardDataAlgorithm(
-        get_handle(), weight_desc, odesc, conv_desc, idesc, search_mode,
-        wlimit, ctypes.byref(bwd_data_alg)))
-    return bwd_data_alg
-
-
-def convolution_backward_data_workspace_size(*args):
-    return check_error(lib.cudnnGetConvolutionBackwardDataWorkspaceSize(*args))
-
-
-def convolution_backward_filter(*args):
-    return check_error(lib.cudnnConvolutionBackwardFilter(*args))
-
-
-def convolution_backward_filter_algorithm(idesc, odesc, conv_desc, weight_desc):
-    k = convolution_autotuner_key(idesc, weight_desc, conv_desc)
-    if k in _autotuner_backward_filter:
-        return _autotuner_backward_filter[k]
-
-    if benchmark:
-        perf_results = ConvolutionAlgoPerf()
-        algo_count = ctypes.c_int()
-        check_error(lib.cudnnFindConvolutionBackwardFilterAlgorithm(
-            get_handle(), idesc, odesc, conv_desc, weight_desc, 1,
-            ctypes.byref(algo_count), ctypes.byref(perf_results)))
-        _autotuner_backward_filter[k] = perf_results.algo
-        return perf_results.algo
-
-    search_mode = CUDNN_CONVOLUTION_BWD_FILTER_PREFER_FASTEST
-    wlimit = 0
-    if workspace_limit is not None:
-        wlimit = workspace_limit
-        search_mode = CUDNN_CONVOLUTION_BWD_FILTER_SPECIFY_WORKSPACE_LIMIT
-
-    bwd_filter_alg = ctypes.c_int()
-    check_error(lib.cudnnGetConvolutionBackwardFilterAlgorithm(
-        get_handle(), idesc, odesc, conv_desc, weight_desc, search_mode,
-        wlimit, ctypes.byref(bwd_filter_alg)))
-    return bwd_filter_alg
-
-
-def convolution_backward_filter_workspace_size(*args):
-    return check_error(lib.cudnnGetConvolutionBackwardFilterWorkspaceSize(*args))
-
-
-def convolution_backward_bias(*args):
-    check_error(lib.cudnnConvolutionBackwardBias(*args))
+def descriptor_sequence(tensor, batch_sizes):
+    descriptors = TensorDescriptorArray(len(batch_sizes))
+    _type = _typemap[tensor.type()]
+    _ndim = 5
+    dim_pad = (1,) * (5 - tensor.dim())
+    _size = int_array(tensor.size() + dim_pad)
+    _stride = int_array(tensor.stride() + dim_pad)
+    for i, batch_size in enumerate(batch_sizes):
+        _size[0] = batch_size
+        descriptors.set_raw(i, _type, _ndim, _size, _stride)
+    return descriptors
 
 
 def add_tensor(*args):

torch/backends/cudnn/rnn.py

@@ -34,20 +34,20 @@ class Unserializable(object):
         self.inner = None
 
 
-def init_dropout_descriptor(fn, handle):
-    return cudnn.DropoutDescriptor(
-        handle,
-        fn.dropout,
-        fn.dropout_seed
-    )
-
-
 def init_rnn_descriptor(fn, handle):
+    dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
+    dropout_p = fn.dropout if fn.train else 0
+    if (dropout_desc_name not in fn.dropout_state) or (fn.dropout_state[dropout_desc_name].get() is None):
+        fn.dropout_state[dropout_desc_name] = Unserializable(
+            cudnn.DropoutDescriptor(handle, dropout_p, fn.dropout_seed)
+        )
+    dropout_desc = fn.dropout_state[dropout_desc_name].get()
+    dropout_desc.set_dropout(dropout_p, fn.dropout_seed)
     return cudnn.RNNDescriptor(
         handle,
         fn.hidden_size,
         fn.num_layers,
-        fn.dropout_state['desc'].get(),
+        dropout_desc,
         fn.input_mode,
         fn.bidirectional,
         fn.mode,
@@ -62,16 +62,22 @@ def init_weight_descriptor(fn, weight):
     return w_desc
 
 
-def _input_size(fn):
-    return (fn.seq_length, fn.mini_batch, fn.input_size)
+def _input_size(fn, input):
+    if fn.batch_sizes is not None:
+        return (input.size(0), fn.input_size)
+    else:
+        return (fn.seq_length, fn.mini_batch, fn.input_size)
 
 
 def _hidden_size(fn):
     return (fn.num_layers * fn.num_directions, fn.mini_batch, fn.hidden_size)
 
 
-def _output_size(fn):
-    return (fn.seq_length, fn.mini_batch, fn.hidden_size * fn.num_directions)
+def _output_size(fn, input):
+    if fn.batch_sizes is not None:
+        return (input.size(0), fn.hidden_size * fn.num_directions)
+    else:
+        return (fn.seq_length, fn.mini_batch, fn.hidden_size * fn.num_directions)
 
 
 def get_num_weights(handle, rnn_desc, x_desc, datatype):
@@ -157,9 +163,9 @@ def get_parameters(fn, handle, weight_buf):
                 # might as well merge the CUDNN ones into a single tensor as well
                 if linear_id == 0 or linear_id == num_linear_layers / 2:
                     assert filter_dim_a.prod() == filter_dim_a[0]
+                    size = (filter_dim_a[0] * num_linear_layers // 2, filter_dim_a[2])
                     param = fn.weight_buf.new().set_(
-                        weight_buf.storage(), offset,
-                        filter_dim_a[0] * num_linear_layers // 2, filter_dim_a[2])
+                        weight_buf.storage(), offset, size)
                     layer_params.append(param)
                 else:
                     assert cur_offset == offset
@@ -172,10 +178,13 @@ def get_parameters(fn, handle, weight_buf):
 
 
 def _copyParams(params_from, params_to):
+    assert len(params_from) == len(params_to)
     for layer_params_from, layer_params_to in zip(params_from, params_to):
+        # NOTE: these lists have all weights before all biases, so if the layer doesn't
+        # use biases, zip will terminate once layer_params_from ends and ignore them.
         for param_from, param_to in zip(layer_params_from, layer_params_to):
             assert param_from.type() == param_to.type()
-            param_to.copy_(param_from)
+            param_to.copy_(param_from, broadcast=False)
 
 
 def forward(fn, input, hx, weight, output, hy):
@@ -183,6 +192,7 @@ def forward(fn, input, hx, weight, output, hy):
         lib = cudnn.lib
         handle = cudnn.get_handle()
         fn.datatype = cudnn._typemap[input.type()]
+        is_input_packed = fn.batch_sizes is not None
 
         if fn.mode == cudnn.CUDNN_LSTM:
             hx, cx = hx
@@ -190,22 +200,27 @@ def forward(fn, input, hx, weight, output, hy):
         else:
             cx, cy = None, None
 
-        if fn.batch_first:
+        if fn.batch_first and not is_input_packed:
             input = input.transpose(0, 1)
 
-        if input.dim() != 3:
+        if (not is_input_packed and input.dim() != 3) or (is_input_packed and input.dim() != 2):
             raise RuntimeError(
                 'input must have 3 dimensions, got {}'.format(input.dim()))
-        if fn.input_size != input.size(2):
-            raise RuntimeError('input.size(2) must be equal to input_size. Expected {}, got {}'.format(
-                fn.input_size, input.size(2)
+        if fn.input_size != input.size(-1):
+            raise RuntimeError('input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
+                fn.input_size, input.size(-1)
             ))
         if fn.dropout != 0 and cudnn.version() < 5103:
             raise RuntimeError('dropout supported only in cudnn v5.1 and above')
 
-        fn.seq_length, fn.mini_batch, fn.input_size = input.size()
+        if is_input_packed:
+            fn.seq_length = len(fn.batch_sizes)
+            fn.mini_batch = fn.batch_sizes[0]
+            fn.input_size = input.size(-1)
+        else:
+            fn.seq_length, fn.mini_batch, fn.input_size = input.size()
         hidden_size = _hidden_size(fn)
-        output_size = _output_size(fn)
+        output_size = _output_size(fn, input)
 
         assert hx.is_contiguous()
         assert cx is None or cx.is_contiguous()
@@ -217,30 +232,34 @@ def forward(fn, input, hx, weight, output, hy):
         y = output
 
         # init descriptors
-        if ('desc' not in fn.dropout_state) or (fn.dropout_state['desc'].get() is None):
-            fn.dropout_state['desc'] = Unserializable(
-                init_dropout_descriptor(fn, handle)
-            )
         fn.rnn_desc = init_rnn_descriptor(fn, handle)
-        fn.x_descs = cudnn.descriptor(x[0], fn.seq_length)
-        fn.y_descs = cudnn.descriptor(y[0], fn.seq_length)
+        if is_input_packed:
+            fn.x_descs = cudnn.descriptor_sequence(x, fn.batch_sizes)
+            fn.y_descs = cudnn.descriptor_sequence(y, fn.batch_sizes)
+        else:
+            fn.x_descs = cudnn.descriptor(x[0], fn.seq_length)
+            fn.y_descs = cudnn.descriptor(y[0], fn.seq_length)
         fn.hx_desc = cudnn.descriptor(hx)
         fn.hy_desc = cudnn.descriptor(hx)
         fn.cx_desc = cudnn.descriptor(cx) if cx is not None else None
         fn.cy_desc = cudnn.descriptor(cx) if cx is not None else None
 
         # create the weight buffer and copy the weights into it
-        num_weights = get_num_weights(
-            handle, fn.rnn_desc, fn.x_descs[0], fn.datatype)
-        fn.weight_buf = input.new(num_weights)
-        fn.w_desc = init_weight_descriptor(fn, fn.weight_buf)
-        w = fn.weight_buf
-        # this zero might not seem necessary, but it is in the case
-        # where biases are disabled; then they won't be copied and must be zero'd.
-        # Alternatively, _copyParams could be written more carefully.
-        w.zero_()
-        params = get_parameters(fn, handle, w)
-        _copyParams(weight, params)
+        if fn.weight_buf is None:
+            num_weights = get_num_weights(
+                handle, fn.rnn_desc, fn.x_descs[0], fn.datatype)
+            fn.weight_buf = x.new(num_weights)
+            fn.w_desc = init_weight_descriptor(fn, fn.weight_buf)
+            w = fn.weight_buf
+            # this zero might not seem necessary, but it is in the case
+            # where biases are disabled; then they won't be copied and must be zero'd.
+            # Alternatively, _copyParams could be written more carefully.
+            w.zero_()
+            params = get_parameters(fn, handle, w)
+            _copyParams(weight, params)
+        else:
+            fn.w_desc = init_weight_descriptor(fn, fn.weight_buf)
+            w = fn.weight_buf
 
         if tuple(hx.size()) != hidden_size:
             raise RuntimeError('Expected hidden size {}, got {}'.format(
@@ -257,8 +276,10 @@ def forward(fn, input, hx, weight, output, hy):
             fn.x_descs,
             ctypes.byref(workspace_size)
         ))
-        fn.workspace = torch.cuda.ByteTensor(workspace_size.value)
-        if fn.train:
+        fn.workspace_size = workspace_size.value
+        with torch.cuda.device_of(input):
+            workspace = torch.cuda.ByteTensor(fn.workspace_size)
+        if fn.requires_grad:
             reserve_size = ctypes.c_long()
             check_error(lib.cudnnGetRNNTrainingReserveSize(
                 handle,
@@ -280,7 +301,7 @@ def forward(fn, input, hx, weight, output, hy):
                 fn.y_descs, ctypes.c_void_p(y.data_ptr()),
                 fn.hy_desc, ctypes.c_void_p(hy.data_ptr()),
                 fn.cy_desc, ctypes.c_void_p(cy.data_ptr()) if cx is not None else None,
-                ctypes.c_void_p(fn.workspace.data_ptr()), fn.workspace.size(0),
+                ctypes.c_void_p(workspace.data_ptr()), workspace.size(0),
                 ctypes.c_void_p(fn.reserve.data_ptr()), fn.reserve.size(0)
             ))
         else:  # inference
@@ -295,15 +316,16 @@ def forward(fn, input, hx, weight, output, hy):
                 fn.y_descs, ctypes.c_void_p(y.data_ptr()),
                 fn.hy_desc, ctypes.c_void_p(hy.data_ptr()),
                 fn.cy_desc, ctypes.c_void_p(cy.data_ptr()) if cx is not None else None,
-                ctypes.c_void_p(fn.workspace.data_ptr()), fn.workspace.size(0)
+                ctypes.c_void_p(workspace.data_ptr()), workspace.size(0)
             ))
 
-        if fn.batch_first:
-            output = output.transpose_(0, 1)
+        if fn.batch_first and not is_input_packed:
+            output.transpose_(0, 1)
 
 
 def backward_grad(fn, input, hx, weight, output, grad_output, grad_hy, grad_input, grad_hx):
     with torch.cuda.device_of(input):
+        is_input_packed = fn.batch_sizes is not None
         handle = cudnn.get_handle()
 
         if fn.mode == cudnn.CUDNN_LSTM:
@@ -313,14 +335,14 @@ def backward_grad(fn, input, hx, weight, output, grad_output, grad_hy, grad_inpu
         else:
             cx, grad_cx, grad_cy = None, None, None
 
-        if fn.batch_first:
+        if fn.batch_first and not is_input_packed:
             input = input.transpose(0, 1)
             grad_output = grad_output.transpose(0, 1)
             output = output.transpose(0, 1)
 
-        input_size = _input_size(fn)
+        input_size = _input_size(fn, input)
         hidden_size = _hidden_size(fn)
-        output_size = _output_size(fn)
+        output_size = _output_size(fn, input)
 
         assert hx.is_contiguous()
         assert cx is None or cx.is_contiguous()
@@ -336,12 +358,12 @@ def backward_grad(fn, input, hx, weight, output, grad_output, grad_hy, grad_inpu
 
         if fn.dropout != 0 and cudnn.version() < 5103:
             raise RuntimeError('dropout supported only in cudnn v 5.1 and above')
-        if not fn.train:
-            raise RuntimeError('backward_grad can only be called when training!')
+        if not fn.requires_grad:
+            raise RuntimeError('backward_grad can only be called when the function requires grad!')
         if tuple(input.size()) != input_size:
             raise RuntimeError('Expected input size {}, got {}'.format(
                 input_size, tuple(input.size())))
-        if tuple(output.size()) != _output_size(fn):
+        if tuple(output.size()) != output_size:
             raise RuntimeError('Expected output size {}, got {}'.format(
                 output_size, output.size()))
         if hx is not None and tuple(hx.size()) != hidden_size:
@@ -359,6 +381,8 @@ def backward_grad(fn, input, hx, weight, output, grad_output, grad_hy, grad_inpu
         if not dhy.is_cuda or not dy.is_cuda or (dcy is not None and not dcy.is_cuda):
             raise RuntimeError('Gradients aren\'t CUDA tensors')
 
+        with torch.cuda.device_of(input):
+            workspace = torch.cuda.ByteTensor(fn.workspace_size)
         check_error(cudnn.lib.cudnnRNNBackwardData(
             handle,
             fn.rnn_desc,
@@ -373,11 +397,11 @@ def backward_grad(fn, input, hx, weight, output, grad_output, grad_hy, grad_inpu
             fn.x_descs, ctypes.c_void_p(dx.data_ptr()),
             fn.hx_desc, ctypes.c_void_p(dhx.data_ptr()),
             fn.cx_desc, ctypes.c_void_p(dcx.data_ptr()) if cx is not None else None,
-            ctypes.c_void_p(fn.workspace.data_ptr()), fn.workspace.size(0),
+            ctypes.c_void_p(workspace.data_ptr()), workspace.size(0),
             ctypes.c_void_p(fn.reserve.data_ptr()), fn.reserve.size(0)
         ))
 
-        if fn.batch_first:
+        if fn.batch_first and not is_input_packed:
             grad_input = grad_input.transpose_(0, 1)
 
 
@@ -396,6 +420,7 @@ def _num_linear_layers(fn):
 
 def backward_weight(fn, input, hx, output, weight, grad_weight):
     with torch.cuda.device_of(input):
+        is_input_packed = fn.batch_sizes is not None
         handle = cudnn.get_handle()
 
         if fn.mode == cudnn.CUDNN_LSTM:
@@ -403,13 +428,13 @@ def backward_weight(fn, input, hx, output, weight, grad_weight):
         else:
             cx = None
 
-        if fn.batch_first:
+        if fn.batch_first and not is_input_packed:
             input = input.transpose(0, 1)
             output = output.transpose(0, 1)
-        input_size = _input_size(fn)
+        input_size = _input_size(fn, input)
         hidden_size = _hidden_size(fn)
-        if not fn.train:
-            raise RuntimeError('backward_weight can only be called when training!')
+        if not fn.requires_grad:
+            raise RuntimeError('backward_weight can only be called when the function requires grad!')
         if fn.dropout != 0 and cudnn.version() < 5103:
             raise RuntimeError('dropout supported only in cudnn v 5.1 and above')
         if tuple(input.size()) != input_size:
@@ -425,6 +450,8 @@ def backward_weight(fn, input, hx, output, weight, grad_weight):
         y = output
         dw = fn.weight_buf.new().resize_as_(fn.weight_buf).zero_()
 
+        with torch.cuda.device_of(input):
+            workspace = torch.cuda.ByteTensor(fn.workspace_size)
         check_error(cudnn.lib.cudnnRNNBackwardWeights(
             handle,
             fn.rnn_desc,
@@ -432,7 +459,7 @@ def backward_weight(fn, input, hx, output, weight, grad_weight):
             fn.x_descs, ctypes.c_void_p(x.data_ptr()),
             fn.hx_desc, ctypes.c_void_p(hx.data_ptr()),
             fn.y_descs, ctypes.c_void_p(y.data_ptr()),
-            ctypes.c_void_p(fn.workspace.data_ptr()), fn.workspace.size(0),
+            ctypes.c_void_p(workspace.data_ptr()), workspace.size(0),
             fn.w_desc, ctypes.c_void_p(dw.data_ptr()),
             ctypes.c_void_p(fn.reserve.data_ptr()), fn.reserve.size(0)
         ))