Exercise $2\mathrm{.}2(5$ points$)$: implement the BlockDiagonalConv$2$d layer

Above, you were provided with the description of the block$-$diagonal Conv$2$d layer. Now, you are asked to implement it$.$ I.e$.$ you should

implement a reversible $11$ Conv$2$d layer which weight matrix is structure as above: as a block$-$diagonal matrix with $22$ blocks. This layer

should also compute the logarithm of the absolute value of the determinant of the jacobian $("$LADJ$")$ at a given point during the forward pass $-$

see model. py to figure out the interface.

Rules:

You are not allowed to store a full weight matrix: find a convenient way to store only $2n$ values instead of $nn$ matrix. For example, for

a diagonal matrix we could store only one vector of size $n$ instead of a $nn$ matrix.

You are not allowed to use functions like torch.slogdet $/$ torch.inverse $/$ torch.det $/$ torch. logdet for this implementation

You are not allowed to unpack your $2n$ values into $nn$ matrix $-$ otherwise it will be meaningless to have this compressed

representation. For example, a multiplication of a diagonal matrix with a vector $x$ can be computed as an element$-$

wise multiplication of two vectors: $x$ and $a=(a_1,$dots,$a_n).$

Note: a test for reversibility below is useful, but it does not check the correctness of your ladj computation. So be mindful about your ladj

implementation.

fron typing import Tuple

import torch

front torch import Tensor

fron torch inport nn

fron torch.nn import functional as F

fron math import log$,$ pi$,$ exp

import numpy as np

fron scipy inport linalg as la

fron model import Flow, Block, Glow

class BlockDiagonalConv$2$d$($nn$.$Module$)$:

$*"."$

A $1$x$1$ Conv$2$d layer with $2$x$2$ block$-$diagonal welght matrix as described above

$*"$n

def $($self$,$ in$\_$channel: int$)$:

super$().()$

$\backslash$ TODO: Define the parameters

$|$ Hint: think about how we should store the weight matrix in such a way$|$ Hint: we can store it just as $2$ vectors of length $\text{'}$in$\_$channel': for upper values

$|$l of each block and for lower values of each block. You can also store tham as $[$n$/2,2,2]$

$|$ Hint: what initialization should we use for then?