d, k = 2,4
m = torch.randint(0, d-1, (k,))
i = torch.arange(d)
j = torch.arange(k)
input_mask = torch.where(m[j][..., None] >= i[None], 1., 0.)
input_mask
output_mask = torch.where(i[..., None] > m[j][None], 1., 0.)
output_mask
output_mask @ input_mask

tensor([[0., 0.],
        [4., 0.]])

make_masks[source]

make_masks(d:int, ks:Sequence[T_co])

masks = make_masks(2, [4]*3)
torch.linalg.multi_dot(masks[::-1]).bool().float()
for mask in masks:
    print(mask.shape)
fig, axs = plt.subplots(1, len(masks))
for ax, mask in zip(axs, masks):
    ax.matshow(mask)
plt.show()

torch.Size([4, 2])
torch.Size([4, 4])
torch.Size([4, 4])
torch.Size([2, 4])

class MaskedLinear[source]

MaskedLinear(d_in:int, d_out:int, mask=None) :: Module

Masked linear layer

m = MaskedLinear(d,k)
m.set_mask(masks[0])
m.mask
x = torch.rand(5, d)
out1 = m(x)

x[:, 1] = torch.rand(5)
out2 = m(x)
assert torch.all(out1 == out2)
out2

tensor([[-0.8757, -0.0918,  0.1242,  0.4327],
        [-0.8374, -0.0877,  0.1188,  0.4137],
        [-0.2710, -0.0284,  0.0385,  0.1339],
        [-0.0614, -0.0064,  0.0087,  0.0303],
        [-0.5964, -0.0625,  0.0846,  0.2947]], grad_fn=<AddmmBackward>)

class SimpleMADE[source]

SimpleMADE(d_in:int, d_h:int, n_layers:int, add_direct=True) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

model = SimpleMADE(2, 4, 2, True)
model
x = torch.rand(5, 2)
model(x)

tensor([[ 0.0000, -0.0571],
        [ 0.0000, -0.0635],
        [ 0.0000, -0.4485],
        [ 0.0000, -0.1129],
        [ 0.0000, -0.0245]], grad_fn=<AddBackward0>)

class MADE[source]

MADE(d_in:int, d_h:int, n_layers:int, add_direct=True, outs_per_input=1) :: Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call :meth:to, etc.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

x=torch.arange(10).reshape(2,5)
x.repeat(1,2)
x = torch.rand(5, 2)
model = MADE(2, 4, 3, outs_per_input=2)
model(x)
d_in = 4
outs_per_input = 2
direct_mask = torch.repeat_interleave(
    torch.tril(torch.ones(d_in, d_in), diagonal=-1), outs_per_input, 0
)
direct = MaskedLinear(d_in, d_in*outs_per_input, direct_mask)
direct_mask

tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 0., 0., 0.],
        [1., 1., 0., 0.],
        [1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [1., 1., 1., 0.]])