# 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
    self.feature_dim = feature_dim
    assert feature_dim == 1 or feature_dim == 2, \
        'only support attributes or relations'
    self.queries = queries
    self.hidden_dim = hidden_dim
    self.key_dim = key_dim
    self.value_dim = value_dim
    self.id_dim = id_dim

    current_dim = id_dim * 2 + input_dim
    self.key_embed = MLPModel(current_dim, key_dim, [])
    self.value_embed = MLPModel(current_dim, value_dim, [])
    self.query_embed = MLPModel(id_dim * 2, key_dim, [])
    self.to_query = MLPModel(hidden_dim, key_dim, [])
    self.lstm_cell = nn.LSTMCell(value_dim, hidden_dim)

  def forward(self, relations, attributes=None):
    batch_size, nr = relations.size()[:2]
    assert nr == relations.size(2)
    create_zeros = functools.partial(
        torch.zeros, dtype=relations.dtype, device=relations.device)

    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)

# 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

    self.feature_dim = feature_dim
    assert feature_dim == 1 or feature_dim == 2, \
        'only support attributes or relations'
    self.queries = queries
    self.hidden_dim = hidden_dim
    self.key_dim = key_dim
    self.value_dim = value_dim
    self.id_dim = id_dim

    current_dim = id_dim * 2 + input_dim
    self.key_embed = MLPModel(current_dim, key_dim, [])
    self.value_embed = MLPModel(current_dim, value_dim, [])
    self.query_embed = MLPModel(id_dim * 2, key_dim, [])
    self.to_query = MLPModel(hidden_dim, key_dim, [])
    self.lstm_cell = nn.LSTMCell(value_dim, hidden_dim)

  def forward(self, relations, attributes=None):
    batch_size, nr = relations.size()[:2]
    assert nr == relations.size(2)
    create_zeros = functools.partial(
        torch.zeros, dtype=relations.dtype, device=relations.device)

    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)