# Path: difflogic/nn/neural_logic/modules/_utils.py
# def exclude_mask(inputs, cnt=2, dim=1):
#   """Produce an exclusive mask.
# 
#   Specifically, for cnt=2, given an array a[i, j] of n * n, it produces
#   a mask with size n * n where only a[i, j] = 1 if and only if (i != j).
# 
#   Args:
#     inputs: The tensor to be masked.
#     cnt: The operation is performed over [dim, dim + cnt) axes.
#     dim: The starting dimension for the exclusive mask.
# 
#   Returns:
#     A mask that make sure the coordinates are mutually exclusive.
#   """
#   assert cnt > 0
#   if dim < 0:
#     dim += inputs.dim()
#   n = inputs.size(dim)
#   for i in range(1, cnt):
#     assert n == inputs.size(dim + i)
# 
#   rng = torch.arange(0, n, dtype=torch.long, device=inputs.device)
#   q = []
#   for i in range(cnt):
#     p = rng
#     for j in range(cnt):
#       if i != j:
#         p = p.unsqueeze(j)
#     p = p.expand((n,) * cnt)
#     q.append(p)
#   mask = q[0] == q[0]
#   # Mutually Exclusive
#   for i in range(cnt):
#     for j in range(cnt):
#       if i != j:
#         mask *= q[i] != q[j]
#   for i in range(dim):
#     mask.unsqueeze_(0)
#   for j in range(inputs.dim() - dim - cnt):
#     mask.unsqueeze_(-1)
# 
#   return mask.expand(inputs.size()).float()
# 
# def mask_value(inputs, mask, value):
#   assert inputs.size() == mask.size()
#   return inputs * mask + value * (1 - mask)


import itertools
import torch
import torch.nn as nn
from jactorch.functional import broadcast
from ._utils import exclude_mask, mask_value
    super().__init__()
    self.dim = dim

  def forward(self, inputs, n=None):
    if self.dim == 0:
      assert n is not None
    elif n is None:
      n = inputs.size(self.dim)
    dim = self.dim + 1
    return broadcast(inputs.unsqueeze(dim), dim, n)

  def get_output_dim(self, input_dim):
    return input_dim


class Reducer(nn.Module):
  """Reduce out a variable via quantifiers (exists/forall), implemented by max/min-pooling."""

  def __init__(self, dim, exclude_self=True, exists=True):
    super().__init__()
    self.dim = dim
    self.exclude_self = exclude_self
    self.exists = exists

  def forward(self, inputs):
    shape = inputs.size()
    inp0, inp1 = inputs, inputs
    if self.exclude_self:
      mask = exclude_mask(inputs, cnt=self.dim, dim=-1 - self.dim)
      inp0 = mask_value(inputs, mask, 0.0)

# Path: difflogic/nn/neural_logic/modules/_utils.py
# def exclude_mask(inputs, cnt=2, dim=1):
#   """Produce an exclusive mask.
# 
#   Specifically, for cnt=2, given an array a[i, j] of n * n, it produces
#   a mask with size n * n where only a[i, j] = 1 if and only if (i != j).
# 
#   Args:
#     inputs: The tensor to be masked.
#     cnt: The operation is performed over [dim, dim + cnt) axes.
#     dim: The starting dimension for the exclusive mask.
# 
#   Returns:
#     A mask that make sure the coordinates are mutually exclusive.
#   """
#   assert cnt > 0
#   if dim < 0:
#     dim += inputs.dim()
#   n = inputs.size(dim)
#   for i in range(1, cnt):
#     assert n == inputs.size(dim + i)
# 
#   rng = torch.arange(0, n, dtype=torch.long, device=inputs.device)
#   q = []
#   for i in range(cnt):
#     p = rng
#     for j in range(cnt):
#       if i != j:
#         p = p.unsqueeze(j)
#     p = p.expand((n,) * cnt)
#     q.append(p)
#   mask = q[0] == q[0]
#   # Mutually Exclusive
#   for i in range(cnt):
#     for j in range(cnt):
#       if i != j:
#         mask *= q[i] != q[j]
#   for i in range(dim):
#     mask.unsqueeze_(0)
#   for j in range(inputs.dim() - dim - cnt):
#     mask.unsqueeze_(-1)
# 
#   return mask.expand(inputs.size()).float()
# 
# def mask_value(inputs, mask, value):
#   assert inputs.size() == mask.size()
#   return inputs * mask + value * (1 - mask)


# Path: difflogic/nn/neural_logic/modules/dimension.py
import itertools
import torch
import torch.nn as nn
from jactorch.functional import broadcast
from ._utils import exclude_mask, mask_value

    super().__init__()
    self.dim = dim

  def forward(self, inputs, n=None):
    if self.dim == 0:
      assert n is not None
    elif n is None:
      n = inputs.size(self.dim)
    dim = self.dim + 1
    return broadcast(inputs.unsqueeze(dim), dim, n)

  def get_output_dim(self, input_dim):
    return input_dim


class Reducer(nn.Module):
  """Reduce out a variable via quantifiers (exists/forall), implemented by max/min-pooling."""

  def __init__(self, dim, exclude_self=True, exists=True):
    super().__init__()
    self.dim = dim
    self.exclude_self = exclude_self
    self.exists = exists

  def forward(self, inputs):
    shape = inputs.size()
    inp0, inp1 = inputs, inputs
    if self.exclude_self:
      mask = exclude_mask(inputs, cnt=self.dim, dim=-1 - self.dim)
      inp0 = mask_value(inputs, mask, 0.0)

      inp1 = mask_value(inputs, mask, 1.0)