# Path: difflogic/nn/neural_logic/modules/_utils.py
# def exclude_mask(inputs, cnt=2, dim=1):
# def mask_value(inputs, mask, value):


import functools
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import jacinle.random as random
from difflogic.nn.neural_logic.modules._utils import meshgrid
from jactorch.quickstart.models import MLPModel
    id_shape = list(relations.size()[:-1])
    ids = [random.permutation(2 ** self.id_dim - 1)[:nr] + 1 \
        for i in range(batch_size)]
    ids = np.vstack(ids)
    binary_ids = self.binarize_code(ids)
    zeros = create_zeros(binary_ids.shape)
    binary_ids = torch.tensor(
        binary_ids, dtype=relations.dtype, device=relations.device)
    binary_ids2 = torch.cat(meshgrid(binary_ids, dim=1), dim=-1)
    padded_binary_ids = torch.cat([binary_ids, zeros], dim=-1)

    def embed(embed, x):
      input_size = x.size()[:-1]
      input_channel = x.size(-1)
      f = x.view(-1, input_channel)
      f = embed(f)
      return f.view(*input_size, -1)

    if attributes is None:
      rels = [binary_ids2, relations]
    else:
      padding = create_zeros(*binary_ids2.size()[:-1], attributes.size(-1))
      rels = [binary_ids2, padding, relations]
    rels = torch.cat(rels, dim=-1)
    memory = rels.view(batch_size, -1, rels.size(-1))
    if attributes is not None:
      assert nr == attributes.size(1)
      padding = create_zeros(*padded_binary_ids.size()[:-1], relations.size(-1))
      attributes = torch.cat([padded_binary_ids, attributes, padding], dim=-1)
      memory = torch.cat([memory, attributes], dim=1)

# Path: difflogic/nn/neural_logic/modules/_utils.py
# def exclude_mask(inputs, cnt=2, dim=1):
# def mask_value(inputs, mask, value):


# Path: difflogic/nn/baselines/memory_net.py
import functools
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import jacinle.random as random
from difflogic.nn.neural_logic.modules._utils import meshgrid
from jactorch.quickstart.models import MLPModel

    id_shape = list(relations.size()[:-1])
    ids = [random.permutation(2 ** self.id_dim - 1)[:nr] + 1 \
        for i in range(batch_size)]
    ids = np.vstack(ids)
    binary_ids = self.binarize_code(ids)
    zeros = create_zeros(binary_ids.shape)
    binary_ids = torch.tensor(
        binary_ids, dtype=relations.dtype, device=relations.device)
    binary_ids2 = torch.cat(meshgrid(binary_ids, dim=1), dim=-1)
    padded_binary_ids = torch.cat([binary_ids, zeros], dim=-1)

    def embed(embed, x):
      input_size = x.size()[:-1]
      input_channel = x.size(-1)
      f = x.view(-1, input_channel)
      f = embed(f)
      return f.view(*input_size, -1)

    if attributes is None:
      rels = [binary_ids2, relations]
    else:
      padding = create_zeros(*binary_ids2.size()[:-1], attributes.size(-1))
      rels = [binary_ids2, padding, relations]
    rels = torch.cat(rels, dim=-1)
    memory = rels.view(batch_size, -1, rels.size(-1))
    if attributes is not None:
      assert nr == attributes.size(1)
      padding = create_zeros(*padded_binary_ids.size()[:-1], relations.size(-1))
      attributes = torch.cat([padded_binary_ids, attributes, padding], dim=-1)
      memory = torch.cat([memory, attributes], dim=1)

    keys = embed(self.key_embed, memory).transpose(1, 2)