ltnjax ¶

class Connective:
    """Class representing an LTN connective.

    Wrapper for connectives that aggregates given
    [LTNObject][ltnjax.core.LTNObject] objects according a given aggregator
    operation `connective_op` and also broadcasts variables, see
    [_broadcast_exprs][ltnjax.core._broadcast_exprs].

    Attributes:
        connective_op: Aggregation function.
    """

    def __init__(
        self, connective_op: ltn.fuzzy_ops.ConnectiveOperator
    ) -> None:
        """Constructor.

        Args:
            connective_op: Aggregation function.
        """
        self.connective_op = connective_op

    def __call__(self, *wffs: LTNObject, **kwargs: Any) -> LTNObject:
        """Applies the connective using the given `connective_op`.

        Args:
            wffs: Tuple of LTN objects.
            kwargs: Further arguments to pass to `connective_op`.

        Returns:
            The resulting [LTNObject][ltnjax.core.LTNObject] object that
            combines the given `wffs` into one joint LTN objects.

        Raises:
            TypeError: If `wffs` are not of type
                [LTNObject][ltnjax.core.LTNObject].
            ValueError: If number of `wffs` does not fit to `connective_op`.
        """
        wffs = list(wffs)  # type: ignore
        for x in wffs:
            if not isinstance(x, LTNObject):
                raise TypeError(
                    "The operands of a LTN connective should be "
                    f"instances of {LTNObject}. Got an instance "
                    f"of {type(x)} instead."
                )

        wffs = _broadcast_exprs(wffs)  # type: ignore
        if isinstance(
            self.connective_op, ltn.fuzzy_ops.UnaryConnectiveOperator
        ):
            if len(wffs) != 1:
                raise ValueError(
                    "wffs must have length 1 since connective_op "
                    "is an UnaryConnectiveOperator."
                )
            t_result = self.connective_op(*_as_arrays(wffs), **kwargs)
        elif isinstance(
            self.connective_op, ltn.fuzzy_ops.BinaryConnectiveOperator
        ):
            if len(wffs) != 2:
                raise ValueError(
                    "wffs must have length 2 since connective_op "
                    "is an BinaryConnectiveOperator."
                )
            t_result = self.connective_op(*_as_arrays(wffs), **kwargs)
        elif isinstance(self.connective_op, ltn.fuzzy_ops.AggregationOperator):
            if len(wffs) < 1:
                raise ValueError("wffs must have length at least 1.")
            t_result = self.connective_op(
                jnp.stack(_as_arrays(wffs)), axis=0, **kwargs
            )
        result = LTNObject(t_result, wffs[0].free_vars)
        return result

class Constant(LTNObject):
    """The class representing constants.

    A constant can be a tensor of any rank.

    Attributes:
        value: Value of the constant that is a array of an arbitrary rank.
        free_vars: The free variables that are contained in the expression.
            `free_vars` is ordered in a way that if we have $n$
            `free_vars`, the first $n$ axes of `value` belong to these
            variables.
        trainable: Flag indicating whether the LTN constant is trainable
            (embedding) or not.
    """

    def __init__(self, value: Any, trainable: bool = False) -> None:
        """Constructor.

        Args:
            value: An object that is convertible to an array. This includes JAX
                arrays, NumPy arrays, Python scalars, Python collections like
                lists and tuples, objects with an ``__array__`` method, and
                objects supporting the Python buffer protocol.
            trainable: Flag indicating whether the LTN constant is trainable
                (embedding) or not.
        """
        # This is necessary as the input value could be scalars, lists,
        # numpy arrays, etc.
        value = jnp.asarray(value, dtype=jnp.float32)

        free_vars: list[VarLabel] = []
        super().__init__(value, free_vars=free_vars, trainable=trainable)

    def __repr__(self) -> str:
        """Representation function.

        Called by the repr() built-in function to compute the "official"
        string representation of an object.
        """
        return (
            f"ltn.{self.__class__.__name__}(value={self.value}, "
            f"free_vars={self.free_vars})"
        )

class Function(nnx.Module):
    """Class representing LTN functions.

    A function that maps $n$ tensors of any rank to
    one single tensor of any rank.

    Attributes:
        model: A `nnx.Module` that evaluates this function.

    Note:
        - `model` will be called with a tensor that has a batch dimension and
        optionally feature dimensions. The batch dimension is the axis that
        results from flattening the axes of the free variables. See
        [_flatten_free_dims][ltnjax.core._flatten_free_dims].
    """

    def __init__(
        self,
        model: nnx.Module | None = None,
        func: types.LambdaType | None = None,
    ) -> None:
        """Constructor.

        Initializes the LTN predicate in two different ways:
        1. if `model` is not None, it initializes the predicate with the given
        nnx.Module;
        2. if `model` is None, it uses the `func` as a function to define
        the LTN predicate. Note that, in this case, the LTN predicate is not
        learnable. So, the lambda function has to be used only for simple
        predicates.

        Args:
            model: (default=None) A `nnx.Module` that evaluates this
                function.
            func: (default=None) A lambda_expression.

        Raises:
            ValueError: If either both `model` and `func` is given or not
                `model` nor `func` is given.
            TypeError: If `model` is given and not an `nnx.Module` or
                if `func` is given and not an `types.LambdaType` object.
        """
        if model is not None and func is not None:
            raise ValueError(
                "Both model and func parameters have been "
                "specified. Expected only one of the two "
                "parameters to be specified."
            )

        if model is None and func is None:
            raise ValueError(
                "Both model and func parameters have not been "
                "specified. Expected one of the two parameters "
                "to be specified."
            )

        if model is not None:
            if not isinstance(model, nnx.Module):
                raise TypeError(
                    "Function() : argument 'model' (position 1) "
                    "must be a nnx.Module, not " + str(type(model))
                )
            self.model = model
        else:  # func is not None
            if not isinstance(func, types.LambdaType):
                raise TypeError(
                    "Function() : argument 'func' (position 2) "
                    "must be a function, not " + str(type(model))
                )
            self.model = LambdaModel(func)

    def __call__(self, *inputs: LTNObject, **kwargs: Any) -> LTNObject:
        """Evaluates the `model` of the given `inputs` and `kwargs`.

        Args:
            inputs: tuple of [LTNObject][ltnjax.core.LTNObject] to apply on
                `model`.
            kwargs: Further arguments to pass to `model`.

        Returns:
            `model(inputs)`.

        Raises:
            TypeError: If `inputs` are not of type
                [LTNObject][ltnjax.core.LTNObject].
        """
        # check input
        inputs_as_list = list(inputs)
        for x in inputs_as_list:
            if not isinstance(x, LTNObject):
                raise TypeError(
                    "The input to a LTN Function should be "
                    f"instances of {LTNObject}. Got an instance "
                    f"of {type(x)}  instead."
                )

        # forward input
        inputs_as_list = _broadcast_exprs(inputs_as_list)
        # invariant: flat_inputs has shape: (flat) + feature_dims
        flat_inputs = _flatten_free_dims(inputs_as_list)
        # invariant: t_outputs has shape: (batch) + model_output
        t_outputs = self.model(*_as_arrays(flat_inputs), **kwargs)

        # recover shape
        free_vars = (
            inputs_as_list[0].free_vars if len(inputs_as_list) > 0 else []
        )
        free_dims = (
            jnp.shape(inputs_as_list[0].value)[: len(free_vars)]
            if len(inputs_as_list) > 0
            else ()
        )

        # Case: Function
        # This line differs in function vs predicate.
        t_outputs = jnp.reshape(
            t_outputs, tuple(list(free_dims) + list(jnp.shape(t_outputs)[1::]))
        )

        # ensure that values are float
        t_outputs = jnp.astype(t_outputs, jnp.float32)
        wff = LTNObject(t_outputs, free_vars)
        return wff

class Predicate(nnx.Module):
    """Class representing an LTN predicate.

    An LTN predicate is grounded as a mathematical
    function (either pre-defined or learnable) that maps from some n-ary
    domain of individuals to a real number in [0,1] (fuzzy), which can be
    interpreted as a truth value.

    In LTNtorch, the inputs of a predicate are automatically broadcasted
    before the computation of the predicate, if necessary. Moreover, the
    output is organized in a tensor where each dimension is related to one
    variable given in input.

    Attributes:
        model: A `nnx.Module` that evaluates this function.

    Note:
        - `model` will be called with a tensor that has a batch dimension and
        optionally feature dimensions. The batch dimension is the axis that
        results from flattening the axes of the free variables. See
        [_flatten_free_dims][ltnjax.core._flatten_free_dims].
    """

    def __init__(
        self, model: nnx.Module | None = None, func: Callable | None = None
    ) -> None:
        """Constructor.

        Initializes the LTN predicate in two different ways:
        1. if `model` is not None, it initializes the predicate with the given
        nnx.Module;
        2. if `model` is None, it uses the `func` as a function to define
        the LTN predicate. Note that, in this case, the LTN predicate is not
        learnable. So, the lambda function has to be used only for simple
        predicates.

        Args:
            model: (default=None) A `nnx.Module` that evaluates this
                function.
            func: (default=None) A lambda_expression.

        Raises:
            ValueError: If either both `model` and `func` is given or not
                `model` nor `func` is given.
            TypeError: If `model` is given and not an `nnx.Module` or
                if `func` is given and not an `types.LambdaType` object.
        """
        if model is not None and func is not None:
            raise ValueError(
                "Both model and func parameters have been "
                "specified. Expected only one of the two "
                "parameters to be specified."
            )

        if model is None and func is None:
            raise ValueError(
                "Both model and func parameters have not been "
                "specified. Expected one of the two parameters "
                "to be specified."
            )

        if model is not None:
            if not isinstance(model, nnx.Module):
                raise TypeError(
                    "Predicate() : argument 'model' (position 1) "
                    "must be a nnx.Module, not " + str(type(model))
                )
            self.model = model
        else:  # func is not None
            if not isinstance(func, types.LambdaType):
                raise TypeError(
                    "Predicate() : argument 'func' (position 2) "
                    "must be a function, not " + str(type(model))
                )
            self.model = LambdaModel(func)

    def __call__(self, *inputs: LTNObject, **kwargs: Any) -> LTNObject:
        """Evaluates the `model` of the given `inputs` and `kwargs`.

        Args:
            inputs: tuple of [LTNObject][ltnjax.core.LTNObject] to apply on
                `model`.
            kwargs: Further arguments to pass to `model`.

        Returns:
            `model(inputs)`.

        Raises:
            TypeError: If `inputs` are not of type
                [LTNObject][ltnjax.core.LTNObject].
        """
        # check input
        inputs_as_list = list(inputs)
        for x in inputs_as_list:
            if not isinstance(x, LTNObject):
                raise TypeError(
                    "The input to a LTN Predicate should be "
                    f"instances of {LTNObject}. Got an instance "
                    f"of {type(x)}  instead."
                )

        # forward input
        inputs_as_list = _broadcast_exprs(inputs_as_list)
        # invariant: flat_inputs has shape: (flat) + feature_dims
        flat_inputs = _flatten_free_dims(inputs_as_list)
        # invariant: t_outputs has shape: (batch) + model_output
        t_outputs = self.model(*_as_arrays(flat_inputs), **kwargs)

        # recover shape
        free_vars = (
            inputs_as_list[0].free_vars if len(inputs_as_list) > 0 else []
        )
        free_dims = (
            jnp.shape(inputs_as_list[0].value)[: len(free_vars)]
            if len(inputs_as_list) > 0
            else ()
        )

        # Case: Predicate
        # This line differs in function vs predicate.
        t_outputs = jnp.reshape(t_outputs, free_dims)

        # ensure that values are float
        t_outputs = jnp.astype(t_outputs, jnp.float32)
        wff = LTNObject(t_outputs, free_vars)
        return wff

class Quantifier:
    r"""Class representing an LTN quantifier.

    Wrapper for Quantifiers. This evaluates a given LTN object `wff` for
    all variable combinations for that the condition `mask` is true. Then, the
    results will be aggregated with the aggregation operator `aggreg_op`.

    Attributes:
        aggreg_op: Aggregation operator.
        quantifier: (str = "f" | "e"`) Decides whether this is a
            "forall"- or "exists"-quantifier. This has no effect on the
            aggregator but is important for cases, where we aggregate over
            $\emptyset$. This may happen if the variables are empty
            or the mask masks each variable-combination. In these cases,
            "forall" expressions are true while "exists"-quantifiers are
            false.
            If the Quantifier is used with non-truth values, the quantifier
            can be used as the <b>neural elements</b> like for `sum` or
            `prod`.

    Raises:
        TypeError: If `aggreg_op` is not of type
            [ConnectiveOperator][ltnjax.fuzzy_ops.ConnectiveOperator].
        ValueError: If `quantifier` is not one of the strings `forall` or
            `exists`.

    Note:
        It is possible that the variable-combinations are empty or that the
        condition `mask` will mask every variable-combination. In both cases,
        a "forall"-statement will always be true while an "exists"-statement
        will always be false.
    """

    def __init__(
        self, aggreg_op: ltn.fuzzy_ops.ConnectiveOperator, quantifier: str
    ) -> None:
        r"""Constructor.

        Args:
            aggreg_op: Aggregation operator.
            quantifier: (str = "f" | "e"`) Decides whether this is a
                "forall"- or "exists"-quantifier. This has no effect on the
                aggregator but is important for cases, where we aggregate over
                $\emptyset$. This may happen if the variables are empty
                or the mask masks each variable-combination. In these cases,
                "forall" expressions are true while "exists"-quantifiers are
                false.
                If the Quantifier is used with non-truth values, the quantifier
                can be used as the <b>neural elements</b> like for `sum` or
                `prod`.
        """
        self.aggreg_op = aggreg_op
        if not isinstance(aggreg_op, ltn.fuzzy_ops.ConnectiveOperator):
            raise TypeError(
                "The aggregation operator for the quantifier "
                "should be an instance of "
                f"{ltn.fuzzy_ops.ConnectiveOperator}. Got an "
                f"instance of {type(aggreg_op)} instead."
            )
        if quantifier not in ["f", "e"]:
            raise ValueError(
                '`quantifier` for the quantifier should be "f" or "e".'
            )
        self.quantifier = quantifier

    def __call__(
        self,
        variables: Variable | list[Variable],
        wff: LTNObject,
        mask: LTNObject | None = None,
        **kwargs: Any,
    ) -> LTNObject:
        """Applies the quantification and outputs the resulting LTN object.

        As a side-effect, this removes the `diagonal quantification` from the
        given `variables`.
        Refer to [undiag][ltnjax.core.undiag].

        Args:
            variables: Variable or list of variables.
            wff: LTN object.
            mask: (default=None) Condition operation.
            kwargs: Further arguments to pass to `connective_op`.

        Returns:
            The resulting LTN object.

        Raises:
            TypeError: If the values of `variables` are not instances of
                [Variable][ltnjax.core.Variable] or if `wff` is not of type
                [LTNObject][ltnjax.core.LTNObject].
        """
        # check inputs
        variables = (
            [variables] if not isinstance(variables, list) else variables
        )
        for x in variables:
            if not isinstance(x, Variable):
                raise TypeError(
                    "The quantified variables should be "
                    f"instances of {Variable}. Got an instance of "
                    "{type(x)} instead."
                )
        if not isinstance(wff, LTNObject):
            raise TypeError(
                "The quantified LTN object should be an instance "
                f"of {LTNObject}. Got an instance of {type(x)} "
                "instead ."
            )
        if mask is not None and not isinstance(mask, LTNObject):
            raise TypeError(
                "The mask argument should be an instance of "
                f"{LTNObject}. Got an instance of {type(mask)} "
                "instead."
            )

        # Note: For the edge-case that variables are empty.
        # Since the empty variables are already passed into predicates that
        # return size-0 arrays, we have to check wff.
        # This will be done before self.aggreg_op is applied.

        aggreg_vars = {var.free_vars[0] for var in variables}

        if mask is not None:
            # This block is for broadcasting variables in wff and the
            # variables in mask.
            # Important to put aggreg dims last, to keep other dims in the
            # ragged result.
            mask = Quantifier._transpose_free_vars(
                mask,
                new_var_order=[
                    var for var in mask.free_vars if var not in aggreg_vars
                ]
                + [var for var in mask.free_vars if var in aggreg_vars],
            )
            wff = Quantifier._broadcast_wff_and_mask(wff, mask)

            mask.value = jnp.astype(mask.value, jnp.bool)

            # Ignore vars in variables that do not occur in wff.free_vars.
            aggreg_axes = [
                wff.free_vars.index(var)
                for var in aggreg_vars
                if var in wff.free_vars
            ]
            if wff.value.size == 0:  # Check edge case, when wff is empty:
                t_result = (
                    jnp.prod(wff.value, axis=aggreg_axes, **kwargs)
                    if self.quantifier == "f"
                    else jnp.sum(wff.value, axis=aggreg_axes, **kwargs)
                )
            else:
                t_result = self.aggreg_op(
                    wff.value, axis=aggreg_axes, mask=mask.value, **kwargs
                )

            aggreg_axes_in_mask = [
                mask.free_vars.index(var)
                for var in aggreg_vars
                if var in mask.free_vars
            ]
            # empty_vars are the variable-combinations that are completely
            # masked and we need to apply replacement value `rep_value`.
            non_empty_vars = (
                jnp.sum(
                    jnp.astype(mask.value, jnp.int32), axis=aggreg_axes_in_mask
                )
                != 0
            )
            rep_value = 1.0 if self.quantifier == "f" else 0

            t_result = jnp.where(non_empty_vars, t_result, rep_value)
        else:
            # ignore vars in variables that do not occur in wff.free_vars
            aggreg_axes = [
                wff.free_vars.index(var)
                for var in aggreg_vars
                if var in wff.free_vars
            ]
            if wff.value.size == 0:  # Check edge case, when wff is empty:
                t_result = (
                    jnp.prod(wff.value, axis=aggreg_axes, **kwargs)
                    if self.quantifier == "f"
                    else jnp.sum(wff.value, axis=aggreg_axes, **kwargs)
                )
            else:
                t_result = self.aggreg_op(
                    wff.value, axis=aggreg_axes, **kwargs
                )
        free_vars_remaining = [
            var for var in wff.free_vars if var not in aggreg_vars
        ]
        result = LTNObject(t_result, free_vars_remaining)
        undiag(*variables)
        return result

    @staticmethod
    def _broadcast_wff_and_mask(
        wff: LTNObject, mask: LTNObject, in_place: bool = False
    ) -> LTNObject:
        """Broadcasts the free variables from `mask` to `wff`.

        The variables of `mask` are put in the first axes.

        Args:
            wff: LTN object.
            mask: LTN object.
            in_place: (default=False) Boolean that decides whether we perform
                the operation on `wff` or a new copy.

        Returns:
            The LTN object `wff` with the vars from `mask` added.
        """
        if not in_place:
            wff = wff._copy()
        # 1. Broadcast wff with vars that are in the mask but not yet in the
        # LTN object.
        mask_vars_not_in_wff = [
            var for var in mask.free_vars if var not in wff.free_vars
        ]
        for var in mask_vars_not_in_wff:
            new_idx = len(wff.free_vars)
            wff.value = jnp.expand_dims(wff.value, axis=new_idx)
            wff.value = jnp.repeat(
                wff.value, mask._get_dim_of_free_var(var), axis=new_idx
            )
            wff.free_vars.append(var)
        # 2. Transpose wff so that the masked vars on the first axes.
        vars_not_in_mask = [
            var for var in wff.free_vars if var not in mask.free_vars
        ]
        wff = Quantifier._transpose_free_vars(
            wff, new_var_order=mask.free_vars + vars_not_in_mask
        )
        return wff

    @staticmethod
    def _transpose_free_vars(
        expr: LTNObject, new_var_order: list[VarLabel], in_place: bool = False
    ) -> LTNObject:
        """Transposes free variables.

        This changes the order of variables in `expr.free_vars` and the
        axes of `expr.value` will be transposed accordingly.

        Args:
            expr: The LTN object whose variables will be transposed.
            new_var_order: List of variables that defines the new order.
            in_place: (default=False) Boolean that decides whether we perform
                the operation on a new copy or on the same LTN objects.

        Returns:
            The transposed LTN object.
        """
        permutation = [expr.free_vars.index(var) for var in new_var_order]
        if not in_place:
            expr = expr._copy()
        expr.value = jnp.transpose(expr.value, permutation)
        expr.free_vars = new_var_order
        return expr