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


import numpy as np
import torch
import torch.nn as nn
import jacinle.random as random
from difflogic.nn.neural_logic.modules._utils import meshgrid
from jactorch.functional.shape import broadcast
    self.feature_dim = feature_dim
    assert feature_dim == 1 or feature_dim == 2, ('only support attributes or '
                                                  'relations')
    self.num_layers = num_layers
    self.hidden_size = hidden_size
    self.code_length = code_length
    self.lstm = nn.LSTM(
        current_dim,
        hidden_size,
        num_layers,
        batch_first=True,
        bidirectional=True)

  def forward(self, relations, attributes=None):
    batch_size, nr = relations.size()[:2]
    assert nr == relations.size(2)

    id_shape = list(relations.size()[:-1])
    ids = [
        random.permutation(2**self.code_length - 1)[:nr] + 1
        for i in range(batch_size)
    ]
    ids = np.vstack(ids)
    binary_ids = self.binarize_code(ids)
    zeros = torch.tensor(
        np.zeros(binary_ids.shape),
        dtype=relations.dtype,
        device=relations.device)
    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/lstm.py
import numpy as np
import torch
import torch.nn as nn
import jacinle.random as random
from difflogic.nn.neural_logic.modules._utils import meshgrid
from jactorch.functional.shape import broadcast

    self.feature_dim = feature_dim
    assert feature_dim == 1 or feature_dim == 2, ('only support attributes or '
                                                  'relations')
    self.num_layers = num_layers
    self.hidden_size = hidden_size
    self.code_length = code_length
    self.lstm = nn.LSTM(
        current_dim,
        hidden_size,
        num_layers,
        batch_first=True,
        bidirectional=True)

  def forward(self, relations, attributes=None):
    batch_size, nr = relations.size()[:2]
    assert nr == relations.size(2)

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

    binary_ids2 = torch.cat(meshgrid(binary_ids, dim=1), dim=-1)