Enable adding constraints to tensors

README.md

@@ -259,10 +259,22 @@ class MyProvider:
         }
 
 @dltyped(provider=MyProvider())
-def free_function(tensor: FloatTensor["batch dim1"]) -> None:
+def free_function(tensor: Annotated[torch.Tensor, FloatTensor["batch dim1"]]) -> None:
     # ... implementation details, dim1 provided by the external scope
 ```
 
+## Providing Tensor Constraints
+
+You may need to apply additional constraints to a tensor that cannot be expressed through the regular axis scope syntax.
+
+To provide additional constraints on a tensor you may provide any number of expressions which will be evaluated during a check.
+
+```python
+# NOTE: to use constraints you must use IntTensor() not IntTensor[] as python class_getitem does not support keyword arguments.
+def func(tensor: Annotated[torch.Tensor, IntTensor("batch chan feat", constraints={"chan%6==0", "batch>=1"})]) -> None:
+    ...
+```
+
 ## Supported Types
 
 - `FloatTensor`: For any precision floating point tensor. Is a superset of the following:
@@ -298,7 +310,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
@@ -45,16 +46,27 @@ def __repr__(self) -> str:
 
 
 class _DLTypeOperator(enum.Enum):
-    """An enum representing a mathematical operator for a dimension expression."""
+    """
+    An enum representing a mathematical operator for a dimension expression.
+
+    NOTE: the ordering here has to be maintained such that any operator that contains another operator appears first.
+    """
 
+    MIN = "min"
+    MAX = "max"
+    ISQRT = "isqrt"
+    EQ = "=="
+    NE = "!="
+    GE = ">="
+    LE = "<="
     ADD = "+"
     SUB = "-"
     MUL = "*"
     EXP = "^"
     DIV = "/"
-    MIN = "min"
-    MAX = "max"
-    ISQRT = "isqrt"
+    GT = ">"
+    LT = "<"
+    MOD = "%"
 
     def __repr__(self) -> str:
         return self.value
@@ -65,33 +77,58 @@ 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.NE:
+                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.NE: 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 +137,21 @@ 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.NE,
+        _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 +176,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 +189,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
         )
@@ -331,15 +389,6 @@ def _postfix_from_infix(identifier: str, expression: list[TokenT | str | int]) -
 TokenT: typing.TypeAlias = _DLTypeSpecifier | _DLTypeGroupToken | _DLTypeOperator
 
 
-def _span_to_tok(character: str) -> TokenT | None:
-    maybe_operator = typing.cast("_DLTypeOperator | None", _DLTypeOperator._value2member_map_.get(character))
-    maybe_specifier = typing.cast(
-        "_DLTypeSpecifier | None", _DLTypeSpecifier._value2member_map_.get(character)
-    )
-    maybe_group = typing.cast("_DLTypeGroupToken | None", _DLTypeGroupToken._value2member_map_.get(character))
-    return maybe_operator or maybe_specifier or maybe_group
-
-
 def _span_to_str_or_int(span: str) -> str | int:
     if span.isnumeric():
         return int(span)
@@ -382,27 +431,74 @@ def _assert_token_list_valid(tokenized: list[str | int | TokenT]) -> None:
         raise SyntaxError("Invalid expression syntax: " + "".join(map(repr, tokenized)))
 
 
-def _tokenize_string_expr(
+def _tokenize_string_expr(  # noqa: C901, PLR0912
     expression: str,
 ) -> list[str | int | TokenT]:
-    return_list: list[str | int | TokenT] = []
+    return_list: list[str] = []
     current_span = ""
+
+    # first pass just split by character, then we can combine repeated tokens like "==" or ">=" into a single token
     for character in expression:
         if character == " ":
             msg = "Spaces not permitted in dimension expressions"
             raise SyntaxError(msg)
 
-        if token := _span_to_tok(character):
+        return_list.append(character)
+
+    # second pass combine repeated tokens like "==" or ">=" into a single token
+    second_pass_return_list: list[str | TokenT] = [""] * len(return_list)
+    current_span = ""
+    offset = 0
+    found_operator = False
+
+    for _idx in range(len(return_list)):
+        idx = _idx + offset
+        if idx >= len(return_list):
+            second_pass_return_list = second_pass_return_list[: len(return_list) - offset]
+            break
+        character = return_list[idx]
+        found_operator = False
+        assert isinstance(character, str)
+        for op in _DLTypeOperator:
+            if tuple(expression[idx : idx + len(op.value)]) == tuple(op.value):
+                second_pass_return_list[idx - offset] = op
+                offset += len(op.value) - 1
+                found_operator = True
+                break
+        if not found_operator:
+            for spec in _DLTypeSpecifier:
+                if tuple(expression[idx : idx + len(spec.value)]) == tuple(spec.value):
+                    second_pass_return_list[idx - offset] = spec
+                    offset += len(spec.value) - 1
+                    found_operator = True
+                    break
+        if not found_operator:
+            for group in _DLTypeGroupToken:
+                if tuple(expression[idx : idx + len(group.value)]) == tuple(group.value):
+                    second_pass_return_list[idx - offset] = group
+                    offset += len(group.value) - 1
+                    found_operator = True
+                    break
+        if not found_operator:
+            second_pass_return_list[idx - offset] = character
+    second_pass_return_list = [tok for tok in second_pass_return_list if tok]
+
+    # third pass convert to tokens or strings/ints
+    final_return_list: list[str | int | TokenT] = []
+    current_span = ""
+    for token in second_pass_return_list:
+        if isinstance(token, _DLTypeOperator | _DLTypeSpecifier | _DLTypeGroupToken):
             if current_span:
-                return_list.append(_span_to_tok(current_span) or _span_to_str_or_int(current_span))
-            current_span = ""
-            return_list.append(token)
+                final_return_list.append(_span_to_str_or_int(current_span))
+                current_span = ""
+            final_return_list.append(token)
         else:
-            current_span += character
+            current_span += token
     if current_span:
-        return_list.append(_span_to_tok(current_span) or _span_to_str_or_int(current_span))
-    _assert_token_list_valid(return_list)
-    return return_list
+        final_return_list.append(_span_to_str_or_int(current_span))
+
+    _assert_token_list_valid(final_return_list)
+    return final_return_list
 
 
 def _get_group_indices(expr: list[TokenT | int | str], offset: int) -> tuple[int, list[int], int]:
@@ -447,6 +543,10 @@ 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 _VALID_IDENTIFIER_RX.match(_identifier) and _expression[0].isalnum():
+            # we had an assignment expression, so we can use the identifier and expression as-is
+            identifier = _identifier
+            expression = _expression
     tokenized = _tokenize_string_expr(expression)
     return _postfix_from_infix(identifier, tokenized)