ltnjax.core ¶

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 LTNObject(nnx.Module):
    r"""Class representing a generic LTN object.

    An LTN object contains the results of an expression with variables
    `free_vars`.

    In LTNjax, LTN objects are constants, variables, and outputs of predicates,
    formulas, functions, connectives, and quantifiers.

    Internally this contains a tensor `value` and an ordered list `free_vars`.
    If the expression has $n$ `free_vars` $x_1,\dotsc,x_n$, then
    the first $n$ axes belong to these variables. The rest is reserved
    for `feature_dims` the result. This can $[,\infty)$ many axes.
    For each setting of the `free_vars`, the tensor will contain the
    corresponding result. This gets more clear by looking at the example below.

    The class extends `nnx.Module` since it may contain values.

    Attributes:
        value: The result of the expression. If there are $n$
            `free_vars`, then the first $n$ axes belong to these
            `free_vars` in the same order.
        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.

    Examples:
        We define the variables $x \in [1, 0]$ and $y \in [0, 1]$
        and the expression $x * y$.
        ```python
        >>> import ltn
        >>> x = ltn.Variable('x', [1., 0.])
        >>> y = ltn.Variable('y', [0., 1.])
        >>> prod = ltn.Connective(ltn.fuzzy_ops.AndProd(stable=False))
        >>> res = prod(x, y)
        ```

        If we choose x to be 1. (the value at index 0) and y to be 1. (the
        value at index 1), we get $1 * 1 = 1$.
        ```python
        >>> print(res.take('x', 0).take('y', 1))
        LTNObject(value=Array(1., dtype=float32), free_vars=[])
        ```

        If we only choose y to be 1 (the value at index 0), we get the
        LTN object $x * 1$. The result still depends on which x we chose.
        Either $x=0$ and we get $x * 1=0$ or $x=1$ and we
        get $x * 1=1$.
        ```python
        >>> print(res.take('y', 1))
        LTNObject(value=Array(shape=(2,), dtype=dtype('float32')),\
        free_vars=['x'])
        >>> print(res.take('y', 1).value)
        [1. 0.]
        ```

        If we choose no variable, then the LTN object will contain all
        possible results:
        ```python
        >>> print(res)
        LTNObject(value=Array(shape=(2, 2), dtype=dtype('float32')),\
        free_vars=['x', 'y'])
        >>> print(res.value)
        [[0. 1.]
        [0. 0.]]
        ```

        Note that LTN objects can also be tensors of any rank:
        ```python
        >>> x = ltn.Variable('x', [[[3.1415926535, 2.7182818284],\
        [1.414213562, .6931471805]], [[1., 0.], [0., 1.]]])
        >>> print(x)
        Variable(value=Array(shape=(2, 2, 2), dtype=dtype('float32')),\
        free_vars=['x'], var_label='x')
        >>> print(x.value)
        [[[3.1415927 2.7182817]
        [1.4142135 0.6931472]]

        [[1.        0.       ]
        [0.        1.       ]]]
        ```
        Here, the first axis belongs to free variable 'x' while the remaining
        axes are feature dimensions.
    """

    def __init__(
        self,
        value: ArrayLike,
        free_vars: list[VarLabel],
        trainable: bool = False,
    ) -> None:
        """Constructor.

        Args:
            value: The result of the expression. If there are $n$
                `free_vars`, then the first $n$ axes belong to these
                `free_vars` in the same order.
            free_vars: List of labels of 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.
        """
        if not trainable:
            self.value = jnp.asarray(value, dtype=jnp.float32)
        else:
            # We put a MyPy ignore here, because nnx.Param behaves as an
            # jnp.Array.
            self.value = nnx.Param(jnp.asarray(value, dtype=jnp.float32))  # type: ignore
        self.free_vars: list[VarLabel] = free_vars

    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})"
        )

    def _copy(self) -> LTNObject:
        """Copy function.

        Copy the LTN object but point to the same tensor, for
        gradient tracking.

        Returns:
            Copy of this LTN object instance.
        """
        return LTNObject(self.value, self.free_vars)

    def _get_axis_of_free_var(self, free_var: VarLabel) -> int:
        """Axis of `free_var`.

        Given a free variable `free_var`, returns the axis in attribute
        `value` that belongs to this free variable.

        Args:
            free_var: Label of the free variable, whose axis we want.

        Returns:
            Axis that belongs to the free variable `free_var`.

        Raises:
            ValueError: If `free_var` does not occur in `free_vars`.
        """
        if free_var not in self.free_vars:
            raise ValueError(
                f"{free_var} is not a free variable occurring in the LTN "
                "object."
            )
        return self.free_vars.index(free_var)

    def _get_dim_of_free_var(self, free_var: VarLabel) -> int:
        """Dimension of `free_var`.

        Given a label `free_var`, returns the corresp. dimension.

        Args:
            free_var: The label of the variable.

        Returns:
            The dimension that corresponds to variable `free_var`.

        Note:
            Do not confuse `axis` with `dimension`. The first is the position
            in the shape of `value` where the second is the value.
        """
        return jnp.shape(self.value)[self._get_axis_of_free_var(free_var)]

    def shape(self) -> tuple[int, ...]:
        """Returns the shape of the grounding of the LTN object.

        Returns:
            The shape of the grounding of the LTN object.
        """
        return self.value.shape

    def take(self, free_var: VarLabel, indices: int | list[int]) -> LTNObject:
        r"""Take elements along the axis that corresponds to `free_var`.

        Args:
            free_var: The variable we want to fix.
            indices: This is either an integer or a list of integers from
                $\{0, \dotsc, n-1\}$ where $n$ is the dimension
                of variable `free_var`. This parameter has the same effect as
                `jnp.array(indices)` in `jax.numpy.take`.

        Returns:
            The LTN object containing the elements along the axis that
            corresponds to `free_var`.

        Raises:
            ValueError: If `indices` is neither int nor list of ints.
        """
        indices_arr: Array = jnp.array(indices)
        remaining_free_vars = [v for v in self.free_vars if v != free_var]
        if jnp.ndim(indices_arr) not in [0, 1]:
            raise ValueError("Give a single indice or a list of indices.")

        result = self._copy()
        result.value = jnp.take(
            self.value, indices_arr, axis=self._get_axis_of_free_var(free_var)
        )
        result.free_vars = remaining_free_vars
        return result