class_docstr: Module that implements a linear layer with block matrices with positive diagonal
  blocks. Moreover, it uses Weight Normalization
  (https://arxiv.org/abs/1602.07868) for stability.
class_docstr_tok: ["Module","that","implements","a","linear","layer","with","block","matrices","with","positive","diagonal","blocks",".","Moreover",",","it","uses","Weight","Normalization","(","https",":","\/\/arxiv.org\/abs\/1602.07868",")","for","stability","."]
code: def __init__(self, in_features, out_features, dim, bias=True):
      super().__init__()
      self.in_features, self.out_features, self.dim = (
          in_features,
          out_features,
          dim,
      )
  
      weight = torch.zeros(out_features, in_features)
  
      # Fill in non-zero entries of block weight matrix, going from top
      # to bottom.
      for i in range(dim):
          weight[
              i * out_features // dim : (i + 1) * out_features // dim,
              0 : (i + 1) * in_features // dim,
          ] = torch.nn.init.xavier_uniform_(
              torch.Tensor(
                  out_features // dim, (i + 1) * in_features // dim
              )
          )
  
      self._weight = torch.nn.Parameter(weight)
      self._diag_weight = torch.nn.Parameter(
          torch.nn.init.uniform_(torch.Tensor(out_features, 1)).log()
      )
  
      self.bias = (
          torch.nn.Parameter(
              torch.nn.init.uniform_(
                  torch.Tensor(out_features),
                  -1 / math.sqrt(out_features),
                  1 / math.sqrt(out_features),
              )
          )
          if bias
          else 0
      )
  
      # Diagonal block mask
      mask_d = (
          torch.eye(dim)
          .unsqueeze(-1)
          .repeat(1, out_features // dim, in_features // dim)
          .view(out_features, in_features)
      )
      self.register_buffer("mask_d", mask_d)
  
      # Off-diagonal block mask for lower triangular weight matrix
      mask_o = torch.tril(torch.ones(dim, dim), diagonal=-1).unsqueeze(
          -1
      )
      mask_o = mask_o.repeat(
          1, out_features // dim, in_features // dim
      ).view(out_features, in_features)
      self.register_buffer("mask_o", mask_o)
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
  from ..None import constraints
  from ..torch_transform import TransformModule
  from ..util import copy_docs_from
  from .neural_autoregressive import ELUTransform, LeakyReLUTransform, TanhTransform
label_desc: Use image node_id 1 to create a new MaskedBlockLinear object from inherited base classes: torch.nn with example: obj = MaskedBlockLinear(in_features, out_features, dim, bias)