Deductive Systems, Rule Schemas, and Span Components¶

@abstractmethod
def forward(
    self,
    input: torch.Tensor,
    semiring: ChartSemiring,
) -> torch.Tensor:
    """Create and populate the initial chart.

    Parameters
    ----------
    input : torch.Tensor
        Raw input (e.g. token indices).
    semiring : ChartSemiring
        The scoring semiring.

    Returns
    -------
    torch.Tensor
        The chart tensor with initial items filled in.
    """
    ...

@abstractmethod
def forward(
    self,
    chart: torch.Tensor,
    semiring: ChartSemiring,
    **context,
) -> torch.Tensor:
    """Apply this deduction step.

    Parameters
    ----------
    chart : torch.Tensor
        Current chart state.
    semiring : ChartSemiring
        The scoring semiring.
    **context
        Schedule-provided context (e.g. ``span_length``,
        ``span_start``, ``split`` for CKY).

    Returns
    -------
    torch.Tensor
        Updated chart.
    """
    ...

@abstractmethod
def forward(self, chart: torch.Tensor) -> torch.Tensor:
    """Extract goal items from the chart.

    Parameters
    ----------
    chart : torch.Tensor
        Completed chart.

    Returns
    -------
    torch.Tensor
        Goal item weights.
    """
    ...

@abstractmethod
def run(
    self,
    chart: torch.Tensor,
    deductions: nn.ModuleList,
    semiring: ChartSemiring,
) -> torch.Tensor:
    """Execute deductions on chart to fixpoint.

    Parameters
    ----------
    chart : torch.Tensor
        Chart with axioms filled in.
    deductions : nn.ModuleList
        The deduction steps.
    semiring : ChartSemiring
        The scoring semiring.

    Returns
    -------
    torch.Tensor
        Completed chart.
    """
    ...

def forward(self, input: torch.Tensor) -> torch.Tensor:
    """Run the deductive system on input.

    Parameters
    ----------
    input : torch.Tensor
        Raw input (e.g. token indices).

    Returns
    -------
    torch.Tensor
        Goal item weights.
    """
    chart = self.axiom(input, self._semiring)
    chart = self._schedule.run(chart, self.deductions, self._semiring)
    return self.goal(chart)

@abstractmethod
def __call__(self, system: CategorySystem) -> RuleSystem:
    """Instantiate this schema over a category system.

    Parameters
    ----------
    system : CategorySystem
        The finite category inventory.

    Returns
    -------
    RuleSystem
        Concrete rules for this schema.
    """
    ...

def __or__(self, other: RuleSchema) -> RuleSchema:
    """Union of two schemas."""
    return UnionSchema(self, other)

def __add__(self, other: RuleSchema) -> RuleSchema:
    """Alias for union."""
    return self | other

def weighted(self, weight: float) -> RuleSchema:
    """Attach a default weight to all rules from this schema."""
    return WeightedSchema(self, weight)

@abstractmethod
def match(
    self,
    left: Category,
    right: Category,
) -> Category | None:
    """Test whether (left, right) triggers this rule.

    Parameters
    ----------
    left : Category
        Left antecedent.
    right : Category
        Right antecedent.

    Returns
    -------
    Category or None
        Result category if the pair matches, else None.
    """
    ...

@abstractmethod
def match(
    self,
    cat: Category,
    system: CategorySystem,
) -> list[Category]:
    """Find all valid results for a given input category.

    Parameters
    ----------
    cat : Category
        Input category.
    system : CategorySystem
        The full category inventory (needed for some schemas
        that must check existence of intermediate categories).

    Returns
    -------
    list[Category]
        Valid result categories.
    """
    ...

def forward(self, chart, semiring, **context):
    """Apply binary rules for one span cell (i, j).

    Expected context keys: ``i``, ``j``, ``chart_cells``.
    Returns the (batch, C) tensor for cell (i, j).
    """
    i = context["i"]
    j = context["j"]
    chart_cells = context["chart_cells"]
    n_categories = context["n_categories"]

    _results = cast(torch.Tensor, self._results)
    _lefts = cast(torch.Tensor, self._lefts)
    _rights = cast(torch.Tensor, self._rights)

    if self._n_rules == 0:
        return torch.full(
            (chart_cells[(i, i + 1)].shape[0], n_categories),
            semiring.zero,
            device=_results.device,
        )

    parts = []

    for k in range(i + 1, j):
        left_cell = chart_cells[(i, k)]
        right_cell = chart_cells[(k, j)]

        left_scores = left_cell[:, _lefts]
        right_scores = right_cell[:, _rights]
        combined = semiring.times(left_scores, right_scores)

        if self._n_rules > 0:
            combined = semiring.times(
                combined,
                self.weights.unsqueeze(0),
            )

        parts.append(combined)

    stacked = torch.stack(parts, dim=0)
    split_scores = semiring.plus(stacked, dim=0)

    return _scatter_semiring(
        split_scores,
        _results,
        n_categories,
        split_scores.shape[0],
        device=_results.device,
        semiring=semiring,
    )

def forward(self, chart, semiring, **context):
    """Apply unary rules to a cell tensor.

    Expects ``chart`` to be a (batch, C) cell tensor.
    Returns updated (batch, C) tensor.
    """
    cell = chart

    if not self.has_rules:
        return cell

    _results = cast(torch.Tensor, self._results)
    _inputs = cast(torch.Tensor, self._inputs)
    n_categories = cell.shape[1]

    for _ in range(self._iterations):
        input_scores = cell[:, _inputs]

        if self._n_rules > 0:
            input_scores = semiring.times(
                input_scores,
                self.weights.unsqueeze(0),
            )

        additions = _scatter_semiring(
            input_scores,
            _results,
            n_categories,
            cell.shape[0],
            device=cell.device,
            semiring=semiring,
        )

        cell = semiring.plus_pair(cell, additions)

    return cell