Enable adding constraints to tensors

README.md

@@ -298,7 +298,30 @@ If you run into issues with a dltyped decorator and would like to see detailed s
 - In the current implementation, _every_ call will be checked, the performance overhead on most systems should be negligible (OTOO microseconds).
 - Pydantic default values are not checked.
 - Only symbolic, literal, and expressions are allowed for dimension specifiers, f-string syntax from `jaxtyping` is not supported.
-- Only torch tensors and numpy arrays are supported for now.
+- Only jax arrays, torch tensors and numpy arrays are supported for now.
 - Static shape checking is not supported, DLType only performs runtime checks, though some expression errors will be caught statically by construction if symbolic (i.e. non-string) shapes are used.
 - DLType does not support checkking inside unbounded container types (i.e. `list[TensorTypeBase]`) for performance reasons.
 - DLType does not support unions, but does support optionals.
+
+### A note about numpy.typing and pydantic BaseModels
+
+Starting with numpy 2.5.1 python 3.11 is deprecated which paves the way for use of the new PEP 695 type annotations.
+These annotations replace the old TypeAlias NDArray we had before with the following internal to numpy.typing:
+
+[Source](https://github.com/numpy/numpy/blame/5528ef840237704d55e06992e5c528f03a15a299/numpy/_typing/_array_like.py#L15)
+
+```
+# implementation of numpy.typing NDArray as of numpy 2.5.1
+type NDArray[ScalarT: np.generic] = np.ndarray[_AnyShape, np.dtype[ScalarT]]
+```
+
+Unfortunately for us, this means that the numpy.typing module is no longer natively compatible with pydantic model fields because `type`s defined through this syntax are not compatible with `isinstance`.
+See [PEP-695](https://peps.python.org/pep-0695/) for more information on why this is the case.
+Furthermore, we cannot add pydantic support dynamically to the type because numpy arrays are immutable classes so monkey patching `__get_pydantic_core_schema__` is not an option.
+
+Because of this limitation, we recommend using plain `np.ndarray` for numpy fields in classes.
+If you want to add static shapes to your tensors like you could with the `numpy.typing` module, you can do this directly on `np.ndarray` in newer versions of numpy.
+
+```
+array_argument: Annotated[np.ndarray[tuple[int, ...], np.dtype[np.float32 | np.float64]], dltype.FloatTensor["b c h w"]]
+```

dltype/__init__.py

@@ -17,6 +17,7 @@
 )
 from dltype._lib._dtypes import SUPPORTED_TENSOR_TYPES
 from dltype._lib._errors import (
+    DLTypeConstraintError,
     DLTypeDtypeError,
     DLTypeDuplicateError,
     DLTypeError,
@@ -121,6 +122,7 @@
     "BFloat16Tensor",
     "BoolTensor",
     "ConstantAxis",
+    "DLTypeConstraintError",
     "DLTypeDtypeError",
     "DLTypeDuplicateError",
     "DLTypeError",

dltype/_lib/_constants.py

@@ -1,5 +1,6 @@
 """Constants related to the dltype library."""
 
+import logging
 import typing
 import warnings
 
@@ -31,6 +32,7 @@ class _Env(BaseSettings):
 
 if DEBUG_MODE:
     warnings.warn("DLType debug mode enabled", UserWarning, stacklevel=1)
+    logging.getLogger("dltype").setLevel(logging.DEBUG)
 
 if GLOBAL_DISABLE:
     warnings.warn(

dltype/_lib/_errors.py

@@ -148,3 +148,20 @@ def __init__(
 
     def __str__(self) -> str:
         return f"Invalid scope provider {self._bad_scope_provider}, expected 'self' or a DLTypeScopeProvider"
+
+
+class DLTypeConstraintError(DLTypeError):
+    """An error raised when a constraint is violated."""
+
+    def __init__(
+        self,
+        tensor_name: str | None,
+        constraint: str | None,
+        error_ctx: str | None = None,
+    ) -> None:
+        self._tensor_name = tensor_name or "?"
+        self._constraint = constraint or "?"
+        super().__init__(error_ctx=error_ctx)
+
+    def __str__(self) -> str:
+        return f"Constraint violation, tensor={self._tensor_name} constraint={self._constraint}"

dltype/_lib/_parser.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import enum
 import math
 import re
@@ -55,6 +56,12 @@ class _DLTypeOperator(enum.Enum):
     MIN = "min"
     MAX = "max"
     ISQRT = "isqrt"
+    EQ = "=="
+    GT = ">"
+    GE = ">="
+    LT = "<"
+    LE = "<="
+    MOD = "%"
 
     def __repr__(self) -> str:
         return self.value
@@ -65,33 +72,55 @@ def evaluate_unary(self, a: int) -> int:
             return math.isqrt(a)
         raise NotImplementedError(self)
 
-    def evaluate(self, a: int, b: int) -> int:  # noqa: PLR0911
+    def evaluate(self, a: int, b: int) -> int:  # noqa: C901, PLR0911, PLR0912
         """Evaluate the operator."""
-        if self is _DLTypeOperator.ADD:
-            return a + b
-        if self is _DLTypeOperator.SUB:
-            return a - b
-        if self is _DLTypeOperator.MUL:
-            return a * b
-        if self is _DLTypeOperator.EXP:
-            return int(a**b)
-        if self is _DLTypeOperator.DIV:
-            return a // b
-        if self is _DLTypeOperator.MIN:
-            return min(a, b)
-        if self is _DLTypeOperator.MAX:
-            return max(a, b)
+        match self:
+            case _DLTypeOperator.GT:
+                return int(a > b)
+            case _DLTypeOperator.GE:
+                return int(a >= b)
+            case _DLTypeOperator.LT:
+                return int(a < b)
+            case _DLTypeOperator.LE:
+                return int(a <= b)
+            case _DLTypeOperator.EQ:
+                return int(a == b)
+            case _DLTypeOperator.ADD:
+                return a + b
+            case _DLTypeOperator.SUB:
+                return a - b
+            case _DLTypeOperator.MUL:
+                return a * b
+            case _DLTypeOperator.EXP:
+                return int(a**b)
+            case _DLTypeOperator.DIV:
+                return a // b
+            case _DLTypeOperator.MIN:
+                return min(a, b)
+            case _DLTypeOperator.MAX:
+                return max(a, b)
+            case _DLTypeOperator.MOD:
+                return a % b
+            case _DLTypeOperator.ISQRT:
+                msg = f"Invalid unary operator {self=}"
+                raise AssertionError(msg)
         raise NotImplementedError(self)
 
 
 _op_precedence: typing.Final = {
+    _DLTypeOperator.EQ: 0,
     _DLTypeOperator.ADD: 1,
     _DLTypeOperator.SUB: 1,
     _DLTypeOperator.MUL: 2,
     _DLTypeOperator.DIV: 2,
+    _DLTypeOperator.MOD: 2,
     _DLTypeOperator.EXP: 3,
     _DLTypeOperator.MIN: 4,
     _DLTypeOperator.MAX: 4,
+    _DLTypeOperator.GT: 5,
+    _DLTypeOperator.GE: 5,
+    _DLTypeOperator.LT: 5,
+    _DLTypeOperator.LE: 5,
     _DLTypeOperator.ISQRT: 5,
     _DLTypeGroupToken.LPAREN: 6,
 }
@@ -100,12 +129,20 @@ def evaluate(self, a: int, b: int) -> int:  # noqa: PLR0911
 _binary_functions: typing.Final = frozenset({_DLTypeOperator.MIN, _DLTypeOperator.MAX})
 _functional_operators: typing.Final = frozenset(_unary_functions.union(_binary_functions))
 _infix_operators: typing.Final = frozenset(
-    {_DLTypeOperator.ADD, _DLTypeOperator.SUB, _DLTypeOperator.MUL, _DLTypeOperator.DIV, _DLTypeOperator.EXP}
+    {
+        _DLTypeOperator.ADD,
+        _DLTypeOperator.SUB,
+        _DLTypeOperator.MUL,
+        _DLTypeOperator.DIV,
+        _DLTypeOperator.EXP,
+        _DLTypeOperator.EQ,
+        _DLTypeOperator.GT,
+        _DLTypeOperator.GE,
+        _DLTypeOperator.LT,
+        _DLTypeOperator.LE,
+        _DLTypeOperator.MOD,
+    }
 )
-_valid_operators: frozenset[str] = frozenset(
-    {op.value for op in _DLTypeOperator if op not in _functional_operators},
-)
-
 _VALID_IDENTIFIER_RX: typing.Final = re.compile(r"^[a-zA-Z][a-zA-Z0-9\_]*$")
 
 
@@ -130,6 +167,7 @@ def __init__(
         self.is_literal = not is_multiaxis_literal and all(
             isinstance(token, int) for token in postfix_expression
         )
+
         self.is_identifier = (
             is_multiaxis_literal or is_named_multiaxis or (postfix_expression == [identifier])
         )
@@ -142,6 +180,17 @@ def __init__(
         self.is_multiaxis_literal = is_multiaxis_literal
         self.is_anonymous = is_anonymous
         self.is_named_multiaxis = is_named_multiaxis
+
+        if (
+            not self.is_literal
+            and not self.is_anonymous
+            and not self.is_named_multiaxis
+            and not self.is_multiaxis_literal
+        ):
+            with contextlib.suppress(KeyError):
+                self.evaluate({})
+                self.is_literal = True
+
         _logger.debug(
             "Created new %s dimension expression %r", "multiaxis" if self.is_multiaxis_literal else "", self
         )
@@ -401,6 +450,24 @@ def _tokenize_string_expr(
             current_span += character
     if current_span:
         return_list.append(_span_to_tok(current_span) or _span_to_str_or_int(current_span))
+    # maybe combine repeated tokens like "==" or ">=" into a single token
+    for op in _DLTypeOperator:
+        if len(op.value) > 1:
+            combined_token = op.value
+            for idx in range(len(return_list) - len(combined_token) + 1):
+                if all(
+                    (
+                        (
+                            isinstance(return_list[idx + j], TokenT)
+                            and typing.cast("_DLTypeOperator | _DLTypeSpecifier", return_list[idx + j]).value
+                        )
+                        == combined_token[j]
+                    )
+                    or (isinstance(return_list[idx + j], str) and return_list[idx + j] == combined_token[j])
+                    for j in range(len(combined_token))
+                ):
+                    # replace the sequence with the combined token
+                    return_list[idx : idx + len(combined_token)] = [op]
     _assert_token_list_valid(return_list)
     return return_list
 
@@ -447,6 +514,9 @@ def expression_from_string(expression: str) -> DLTypeDimensionExpression:
     # split the expression into the identifier and the expression if it has a specifier
     identifier = expression
     if _DLTypeSpecifier.EQUALS.value in expression:
-        identifier, expression = expression.split(_DLTypeSpecifier.EQUALS.value, maxsplit=1)
+        identifier, _expression = expression.split(_DLTypeSpecifier.EQUALS.value, maxsplit=1)
+        if not _expression.startswith(_DLTypeSpecifier.EQUALS.value):
+            # we had an anonymous axis like a==3 so leave it alone
+            expression = _expression
     tokenized = _tokenize_string_expr(expression)
     return _postfix_from_infix(identifier, tokenized)

dltype/_lib/_tensor_type_base.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import typing
 
 from pydantic_core import core_schema
@@ -51,13 +52,41 @@ class TensorTypeBase:
     DTYPES: typing.ClassVar[tuple[_dtypes.DLtypeDtypeT, ...]] = ()
     """The torch dtypes that this tensor type asserts to contain. (empty for any dtype)."""
 
-    def __init__(self, shape: str | None, *, optional: bool = False) -> None:
+    def __init__(
+        self, shape: str | None, constraints: set[str] | None = None, *, optional: bool = False
+    ) -> None:
         """Create a new tensor type object."""
         self.multiaxis_index: int | None = None
         self.anonymous_multiaxis: bool = False
         self.multiaxis_name: str | None = None
         self.optional = optional
         self.expected_shape = self._parse_shape_string(shape)
+        self.constraints = (
+            [_parser.expression_from_string(cond) for cond in sorted(constraints)] if constraints else []
+        )
+        self.axis_identifiers = tuple(dim.identifier for dim in self.expected_shape)
+        _fake_scope_eval = {axis.identifier: -1 for axis in self.expected_shape}
+        for expression in self.constraints:
+            if expression.is_literal:
+                msg = (
+                    f"Invalid constraint syntax {expression=}: constraints must be expressions, not a literal"
+                )
+                raise SyntaxError(msg)
+            if expression.is_named_multiaxis or expression.is_anonymous or expression.is_multiaxis_literal:
+                msg = f"Invalid constraint syntax {expression=}: constraints cannot be multiaxis or anonymous"
+                raise SyntaxError(msg)
+            # Attempt to evaluate the expression with an empty scope, if we get anything other than a key error the expression is invalid.
+            with contextlib.suppress(Exception):
+                expression.evaluate({})
+                msg = f"Invalid constraint syntax {expression=}: constraints must reference at least one axis identifier"
+                raise SyntaxError(msg)
+
+            try:
+                assert expression.evaluate(_fake_scope_eval) in (0, 1)
+            except Exception as e:
+                msg = f"Invalid constraint syntax {expression=}: {e}"
+                raise SyntaxError(msg) from e
+
         # only include literal dimensions that aren't multiaxis
         self._literal_dims = tuple(
             (idx, dim.evaluate({}))
@@ -155,14 +184,15 @@ def validate_tensor(
             schema=core_schema.is_instance_schema(source_type),
         )
 
-    def check(
+    def check(  # noqa: C901, PLR0912
         self,
         tensor: _dtypes.DLtypeTensorT,
         tensor_name: str = "anonymous",
     ) -> None:
         """Check if the tensor matches this type."""
         # Basic validation for multi-axis dimensions
         __tracebackhide__ = not _constants.DEBUG_MODE
+        scope: dict[str, int] = {}
 
         if self.multiaxis_index is not None:
             # Min required dimensions = expected shape length + extra dimensions - 1 (the multi-axis placeholder)
@@ -203,3 +233,19 @@ def check(
                     expected_shape=dim,
                     actual=tensor.shape[adjusted_idx],
                 )
+
+            scope[self.axis_identifiers[idx]] = tensor.shape[adjusted_idx]
+
+        if self.constraints:
+            for idx, dim in enumerate(self.expected_shape):
+                # Adjust index if multiaxis exists and is before this dimension
+                adjusted_idx = idx
+                if self.multiaxis_index is not None and idx > self.multiaxis_index:
+                    adjusted_idx += len(tensor.shape) - len(self.expected_shape)
+
+                if dim.identifier not in scope:
+                    scope[dim.identifier] = tensor.shape[adjusted_idx]
+
+            for expression in self.constraints:
+                if not expression.evaluate(scope) == 1:
+                    raise _errors.DLTypeConstraintError(tensor_name=tensor_name, constraint=str(expression))