Convolutions

For basic convolution operations, JAX provides NumPy- and SciPy-style interfaces: jax.numpy.convolve() (with jax.numpy.correlate()) and jax.scipy.signal.convolve() (with its 2D and correlation variants). Those work like their NumPy and SciPy counterparts, and their API docs describe them.

This page is about the general case: the batched, N-dimensional convolutions used in deep neural networks, computed by jax.lax.conv_general_dilated(). It’s a very general function, and it’s not very obvious how to use it, so we’ll work through the common use-cases.

A survey of the family of convolutional operators, a guide to convolutional arithmetic, is highly recommended reading.

Setup: a kernel and an image

Let’s define a simple diagonal edge kernel, and a synthetic image of colored squares to convolve it with:

import matplotlib.pyplot as plt
import numpy as np

import jax.numpy as jnp
from jax import lax

# 2D kernel - HWIO layout
kernel = jnp.zeros((3, 3, 3, 3), dtype=jnp.float32)
kernel += jnp.array([[1, 1, 0],
                     [1, 0,-1],
                     [0,-1,-1]])[:, :, jnp.newaxis, jnp.newaxis]

# NHWC layout
img = jnp.zeros((1, 200, 198, 3), dtype=jnp.float32)
for k in range(3):
  x = 30 + 60*k
  y = 20 + 60*k
  img = img.at[0, x:x+10, y:y+10, k].set(1.0)

fig, ax = plt.subplots(1, 2, figsize=(10, 5))
ax[0].imshow(kernel[:, :, 0, 0])
ax[0].set_title('edge conv kernel')
ax[1].imshow(img[0])
ax[1].set_title('original image');

Dimension numbers define dimensional layout

The important argument to conv_general_dilated() is the 3-tuple of axis layout arguments: (Input Layout, Kernel Layout, Output Layout), where

• N - batch dimension

• H - spatial height

• W - spatial width

• C - channel dimension

• I - kernel input channel dimension

• O - kernel output channel dimension

Here we choose a NHWC image and HWIO kernel convention, matching the arrays above:

dn = lax.conv_dimension_numbers(img.shape,     # only ndim matters, not shape
                                kernel.shape,  # only ndim matters, not shape
                                ('NHWC', 'HWIO', 'NHWC'))  # the important bit
print(dn)

ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))

With dimension numbers in hand, here’s a basic convolution. 'SAME' padding preserves the spatial shape:

out = lax.conv_general_dilated(img,    # lhs = image tensor
                               kernel, # rhs = conv kernel tensor
                               (1,1),  # window strides
                               'SAME', # padding mode
                               (1,1),  # lhs/image dilation
                               (1,1),  # rhs/kernel dilation
                               dn)     # dimension_numbers
print("out shape: ", out.shape)
plt.imshow(np.array(out)[0,:,:,0]);

out shape:  (1, 200, 198, 3)

Strides downsample the output. With (2,2) window strides, the output has half the spatial size (and 'VALID' padding, which applies no padding at all, would shrink it further):

out = lax.conv_general_dilated(img, kernel,
                               (2,2),  # window strides
                               'SAME', (1,1), (1,1), dn)
print("out shape: ", out.shape, " <-- half the size")
plt.imshow(np.array(out)[0,:,:,0]);

out shape:  (1, 100, 99, 3)  <-- half the size

Kernel (rhs) dilation spreads the kernel taps apart, giving atrous or dilated convolution. Here we use an excessive dilation to make the effect obvious:

out = lax.conv_general_dilated(img, kernel,
                               (1,1), 'VALID',
                               (1,1),
                               (12,12),  # rhs/kernel dilation
                               dn)
print("out shape: ", out.shape)
plt.imshow(np.array(out)[0,:,:,0]);

out shape:  (1, 176, 174, 3)

Note

JAX also provides the convenience functions lax.conv and lax.conv_with_general_padding, which fix the layout convention instead of taking dimension numbers: they assume NCHW images and OIHW kernels.

Transposed convolutions

Transposed convolution (sometimes called deconvolution, or fractionally strided convolution) upsamples rather than downsamples: it’s the transpose of a strided convolution, and it’s what the backward pass of a strided convolution computes. JAX provides it directly as jax.lax.conv_transpose():

out = lax.conv_transpose(img, kernel,
                         (2,2),   # strides of the conv being transposed
                         'SAME',
                         dimension_numbers=dn,
                         transpose_kernel=True)
print("out shape: ", out.shape, "<-- larger than the input")
plt.imshow(np.array(out)[0,:,:,0]);

out shape:  (1, 400, 396, 3) <-- larger than the input

Internally, this is another conv_general_dilated: a transposed convolution is a 180° kernel rotation plus lhs (input) dilation, which spreads the input values apart:

# transposed conv = 180deg kernel rotation plus LHS dilation
kernel_rot = jnp.rot90(jnp.rot90(kernel, axes=(0,1)), axes=(0,1))
manual = lax.conv_general_dilated(img, kernel_rot,
                                  (1,1),
                                  ((2, 1), (2, 1)),  # custom output padding
                                  (2,2),             # lhs/image dilation
                                  (1,1), dn)
print("same as conv_transpose:", jnp.allclose(out, manual))

same as conv_transpose: True

1D convolutions

You aren’t limited to 2D convolutions; the same dimension-number machinery handles any spatial rank. A simple 1D demo:

# 1D kernel - WIO layout
kernel = jnp.array([[[1, 0, -1], [-1,  0,  1]],
                    [[1, 1,  1], [-1, -1, -1]]],
                    dtype=jnp.float32).transpose([2,1,0])
# 1D data - NWC layout
data = np.zeros((1, 200, 2), dtype=jnp.float32)
for i in range(2):
  for k in range(2):
      x = 35*i + 30 + 60*k
      data[0, x:x+30, k] = 1.0

print("in shapes:", data.shape, kernel.shape)

plt.figure(figsize=(10,5))
plt.plot(data[0]);
dn = lax.conv_dimension_numbers(data.shape, kernel.shape,
                                ('NWC', 'WIO', 'NWC'))
print(dn)

out = lax.conv_general_dilated(data, kernel,
                               (1,),   # 1D window strides
                               'SAME', (1,), (1,), dn)
print("out shape: ", out.shape)
plt.figure(figsize=(10,5))
plt.plot(out[0]);

in shapes: (1, 200, 2) (3, 2, 2)
ConvDimensionNumbers(lhs_spec=(0, 2, 1), rhs_spec=(2, 1, 0), out_spec=(0, 2, 1))
out shape:  (1, 200, 2)

3D convolutions

import matplotlib as mpl

# Random 3D kernel - HWDIO layout
kernel = jnp.array([
  [[0, 0,  0], [0,  1,  0], [0,  0,   0]],
  [[0, -1, 0], [-1, 0, -1], [0,  -1,  0]],
  [[0, 0,  0], [0,  1,  0], [0,  0,   0]]],
  dtype=jnp.float32)[:, :, :, jnp.newaxis, jnp.newaxis]

# 3D data - NHWDC layout
data = jnp.zeros((1, 30, 30, 30, 1), dtype=jnp.float32)
x, y, z = np.mgrid[0:1:30j, 0:1:30j, 0:1:30j]
data += (jnp.sin(2*x*jnp.pi)*jnp.cos(2*y*jnp.pi)*jnp.cos(2*z*jnp.pi))[None,:,:,:,None]

print("in shapes:", data.shape, kernel.shape)
dn = lax.conv_dimension_numbers(data.shape, kernel.shape,
                                ('NHWDC', 'HWDIO', 'NHWDC'))
print(dn)

out = lax.conv_general_dilated(data, kernel,
                               (1,1,1), 'SAME', (1,1,1), (1,1,1), dn)
print("out shape: ", out.shape)

# Make some simple 3d density plots:
def make_alpha(cmap):
  my_cmap = cmap(jnp.arange(cmap.N))
  my_cmap[:,-1] = jnp.linspace(0, 1, cmap.N)**3
  return mpl.colors.ListedColormap(my_cmap)
my_cmap = make_alpha(plt.cm.viridis)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=data.ravel(), cmap=my_cmap)
ax.axis('off')
ax.set_title('input')
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=out.ravel(), cmap=my_cmap)
ax.axis('off')
ax.set_title('3D conv output');

in shapes: (1, 30, 30, 30, 1) (3, 3, 3, 1, 1)
ConvDimensionNumbers(lhs_spec=(0, 4, 1, 2, 3), rhs_spec=(4, 3, 0, 1, 2), out_spec=(0, 4, 1, 2, 3))
out shape:  (1, 30, 30, 30, 1)