diff --git a/.github/workflows/ty.yml b/.github/workflows/ty.yml index cd44e9963..3d75ecbe3 100644 --- a/.github/workflows/ty.yml +++ b/.github/workflows/ty.yml @@ -3,10 +3,8 @@ name: ty on: workflow_dispatch: push: - branches: + branches: &ty_branches - "main" - - "220-roughpy-jax" - - "335-typecheck-ruff" paths: &ty_paths - ".github/workflows/ty.yml" - "ty.toml" @@ -14,7 +12,7 @@ on: - "roughpy_jax/**/*.py" - "benchmarks/**/*.py" pull_request: - branches: ["main", "220-roughpy-jax", "335-typecheck-ruff"] + branches: *ty_branches paths: *ty_paths permissions: @@ -41,6 +39,30 @@ jobs: shell: bash run: uv sync --extra typecheck --no-install-project + - name: Setup gha caching for vcpkg + uses: actions/github-script@v7 + with: + script: | + core.exportVariable('ACTIONS_CACHE_URL', process.env.ACTIONS_CACHE_URL || ''); + core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); + + - name: Checkout vcpkg + shell: bash + run: git clone --depth=1 https://github.com/Microsoft/vcpkg.git tools/vcpkg + + - name: Configure CMake + run: | + cmake -B "${{ github.workspace }}/build" \ + -DCMAKE_BUILD_TYPE="${{ env.BUILD_TYPE }}" \ + -DROUGHPY_BUILD_TESTS=ON \ + -DROUGHPY_JAX=ON + env: + CMAKE_TOOLCHAIN_FILE: ${{ github.workspace }}/tools/vcpkg/scripts/buildsystems/vcpkg.cmake + VCPKG_FORCE_SYSTEM_BINARIES: 1 + + - name: Build + run: cmake --build "${{ github.workspace }}/build" --config "${{ env.BUILD_TYPE }}" + - name: Ty type check shell: bash - run: uv run --no-sync ty check . --output-format=github \ No newline at end of file + run: uv run ty check . --output-format=github \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 7d1316e5f..5d7d182f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -156,6 +156,7 @@ lint = [ typecheck = [ "jax", "ty", + "pybind11<3.0.0", ] jax = [ "jax>=0.4.0", diff --git a/roughpy_jax/bases.py b/roughpy_jax/bases.py index 3f664507a..277830e49 100644 --- a/roughpy_jax/bases.py +++ b/roughpy_jax/bases.py @@ -1,27 +1,32 @@ import typing -from collections.abc import Hashable, Iterable -from typing import Literal, TypeVar +from collections.abc import Iterable +from typing import Literal, Protocol, TypeVar import jax import numpy as np +from numpy.typing import NDArray from roughpy import compute as rpc @typing.runtime_checkable -class Basis(typing.Protocol, Hashable): +class Basis(Protocol): """ Structural protocol shared by basis objects used in ``roughpy_jax``. Any object implementing this protocol provides the width, truncation depth, and degree offsets needed to construct compatible tensor or Lie bases. + + Note that this protocol is supposed to be hashable, hence the presence of + __hash__ and __eq__. """ - width: np.int32 - depth: np.int32 - degree_begin: np.ndarray[np.int64.dtype] + width: int + depth: int + degree_begin: NDArray[np.int64] def size(self) -> int: ... def __eq__(self, other: object) -> bool: ... + def __hash__(self) -> int: ... BasisT = TypeVar("BasisT", bound=Basis) diff --git a/roughpy_jax/intervals.py b/roughpy_jax/intervals.py index 9fb7ea01b..ab7c63e50 100644 --- a/roughpy_jax/intervals.py +++ b/roughpy_jax/intervals.py @@ -1,9 +1,11 @@ +from __future__ import annotations + import enum import math import numbers import typing from dataclasses import dataclass -from typing import Generic, Protocol, Self, TypeVar +from typing import Any, Generic, Protocol, Self, TypeVar, cast import jax import jax.numpy as jnp @@ -16,32 +18,6 @@ class IntervalType(enum.IntEnum): OpenCl = 1 -def _interval_dataclass(cls): - """ - Combined decorator for roughpy_jax interval objects - - Registers dataclass and JAX data class with dynamic inf and sup and static - interval_type - """ - cls = dataclass(cls, frozen=True) - return jax.tree_util.register_dataclass( - cls, data_fields=["_inf", "_sup"], meta_fields=["_interval_type"] - ) - - -def _partition_dataclass(cls): - """ - Combined decorator for roughpy_jax partition objects - - Registers dataclass and JAX data class with dynamic endpoints and - interval_type - """ - cls = dataclass(cls, frozen=True) - return jax.tree_util.register_dataclass( - cls, data_fields=["_endpoints"], meta_fields=["_interval_type"] - ) - - @typing.runtime_checkable class Interval(Protocol[RealT]): """ @@ -77,7 +53,7 @@ def sup(self) -> RealT: ... @property def length(self) -> RealT: ... - def intersection(self, other: Self) -> Self | None: + def intersection(self, other: Interval[RealT]) -> Interval[RealT] | None: """ Calculate the intersection of this interval with another interval. :param other: The other interval to intersect with. @@ -88,13 +64,13 @@ def intersection(self, other: Self) -> Self | None: ... -class BaseInterval(Interval): +class BaseInterval: # TODO: These don't need to be in a class, just have module-level functions for str, length, and intersection that # take Intervals as arguments. The only reason to have these in a class is if we want to use inheritance to share # code between different Interval implementations. @staticmethod - def __str__(interval) -> str: + def to_string(interval: Any) -> str: reprs = { IntervalType.ClOpen: "[{}, {})", IntervalType.OpenCl: "({}, {}]", @@ -102,19 +78,19 @@ def __str__(interval) -> str: return reprs[interval.interval_type].format(interval.inf, interval.sup) @staticmethod - def length(interval: Interval) -> RealT: + def length(interval: Any) -> Any: """ Calculate the length of the interval. :return: The length of the interval, calculated as sup - inf. :rtype: float """ - return jnp.maximum(0.0, interval.sup - interval.inf) + return jnp.maximum(0.0, jnp.asarray(interval.sup) - jnp.asarray(interval.inf)) @staticmethod def intersection( - left_interval: Interval, - right_interval: Interval, - ) -> Interval | None: + left_interval: Any, + right_interval: Any, + ) -> Interval[Any] | None: """ Calculate the intersection of this interval with another interval. :param other: The other interval to intersect with. @@ -140,14 +116,14 @@ def intersection( # if new_inf >= new_sup: # return None # No intersection - IntervalType = left_interval.__class__ + IntervalT: type = type(left_interval) interval_type = left_interval.interval_type # Intersections are defined by bounds, but not every Interval implementation # can be constructed from (inf, sup, interval_type) (e.g. DyadicInterval). # Use a canonical RealInterval result to avoid incorrect reconstruction. # NOTE: Do NOT call float() here — new_inf/new_sup may be JAX tracers during JIT. - return IntervalType(_inf=new_inf, _sup=new_sup, _interval_type=interval_type) + return IntervalT(_inf=new_inf, _sup=new_sup, _interval_type=interval_type) @dataclass(frozen=True) @@ -234,7 +210,7 @@ def sup(self) -> float: def length(self) -> float: return BaseInterval.length(self) - def intersection(self, other: Self) -> Interval | None: + def intersection(self, other: Self) -> RealInterval[float] | None: """ Calculate the intersection of this dyadic interval with another dyadic interval. :param other: The other dyadic interval to intersect with. @@ -254,7 +230,7 @@ def intersection(self, other: Self) -> Interval | None: ) -@_interval_dataclass +@dataclass(frozen=True) class RealInterval(Generic[RealT]): """ Represents a real interval with specified bounds and interval type. @@ -274,7 +250,7 @@ class RealInterval(Generic[RealT]): _interval_type: IntervalType def __str__(self) -> str: - return BaseInterval.__str__(self) + return BaseInterval.to_string(self) @property def interval_type(self) -> IntervalType: @@ -289,10 +265,10 @@ def sup(self) -> RealT: return self._sup @property - def length(self) -> RealT: + def length(self) -> Any: return BaseInterval.length(self) - def intersection(self, other: Self) -> Self | None: + def intersection(self, other: Interval[Any]) -> RealInterval[RealT] | None: """ Calculate the intersection of this real interval with another real interval. :param other: The other real interval to intersect with. @@ -300,10 +276,17 @@ def intersection(self, other: Self) -> Self | None: :return: A new RealInterval representing the intersection, or None if there is no intersection. :rtype: typing.Optional[RealInterval] """ - return BaseInterval.intersection(self, other) + return cast(RealInterval[RealT] | None, BaseInterval.intersection(self, other)) + +RealInterval = jax.tree_util.register_dataclass( + RealInterval, + data_fields=["_inf", "_sup"], + meta_fields=["_interval_type"], +) -@_partition_dataclass + +@dataclass(frozen=True) class Partition(Generic[RealT]): _endpoints: list[RealT] _interval_type: IntervalType @@ -312,22 +295,22 @@ def __len__(self) -> int: return len(self._endpoints) - 1 def __str__(self) -> str: - return BaseInterval.__str__(self) + return BaseInterval.to_string(self) @property def interval_type(self) -> IntervalType: return self._interval_type @property - def inf(self) -> RealT: + def inf(self) -> Any: return jnp.asarray(self._endpoints[0]) @property - def sup(self) -> RealT: + def sup(self) -> Any: return jnp.asarray(self._endpoints[-1]) @property - def length(self) -> RealT: + def length(self) -> Any: return BaseInterval.length(self) def to_real_interval(self) -> RealInterval[RealT]: @@ -342,7 +325,7 @@ def to_real_interval(self) -> RealInterval[RealT]: _interval_type=self._interval_type, ) - def intersection(self, other: Interval) -> Self | None: + def intersection(self, other: Interval[RealT]) -> Partition[RealT] | None: """ Calculate the intersection of this partition with another Interval. :param other: The other interval to intersect with. @@ -358,7 +341,12 @@ def intersection(self, other: Interval) -> Self | None: # Here we convert to RealInterval to perform the intersection logic intermediate_itvl = self.to_real_interval() - intersect_itvl = intermediate_itvl.intersection(other) + other_interval = RealInterval( + _inf=other.inf, + _sup=other.sup, + _interval_type=other.interval_type, + ) + intersect_itvl = intermediate_itvl.intersection(other_interval) if intersect_itvl is None: return None @@ -368,7 +356,9 @@ def intersection(self, other: Interval) -> Self | None: new_endpoints.append(intersect_itvl.inf) # 2) Include all inner points for ep in self._endpoints: - if intersect_itvl.inf <= ep <= intersect_itvl.sup: + if bool( + jnp.logical_and(intersect_itvl.inf <= ep, ep <= intersect_itvl.sup) + ): new_endpoints.append(ep) # 3) Add new sup if within bounds of old interval if intersect_itvl.sup < self.sup: @@ -379,7 +369,7 @@ def intersection(self, other: Interval) -> Self | None: _interval_type=self.interval_type, ) - def to_intervals(self) -> list[Interval[RealT]]: + def to_intervals(self) -> list[RealInterval[RealT]]: """ Convert the partition into a list of RealIntervals corresponding to the subintervals defined by the partition. @@ -397,3 +387,10 @@ def to_intervals(self) -> list[Interval[RealT]]: ) for i in range(len(self._endpoints) - 1) ] + + +Partition = jax.tree_util.register_dataclass( + Partition, + data_fields=["_endpoints"], + meta_fields=["_interval_type"], +) diff --git a/roughpy_jax/ops.py b/roughpy_jax/ops.py index 23ef4df41..668fbe388 100644 --- a/roughpy_jax/ops.py +++ b/roughpy_jax/ops.py @@ -1,7 +1,7 @@ import collections.abc as cabc -from collections.abc import Callable +from collections.abc import Callable, Mapping from functools import partial -from typing import Any, ClassVar, TypedDict, TypeVar +from typing import Any, ClassVar, TypedDict import jax import jax.numpy as jnp @@ -21,7 +21,6 @@ # Potentially useful to have cached versions to the l2t and t2l matrices as JAX # arrays, as these are used in many operations and converting from the C++ # buffers to JAX arrays can be expensive. -global _lie_sparse_matrix_cache _lie_sparse_matrix_cache: dict[ tuple, tuple[ @@ -52,9 +51,6 @@ def _get_lie_sparse_matrices(lie_basis, dtype): class EmptyStaticArgs(TypedDict): ... -OperationT = TypeVar("OperationT") - - def _batched_fallback_wrapper(single_tensor_fn): """ Generate a batched tensor function from a function that operates on a single tensor. @@ -188,7 +184,7 @@ class Operation: # when deriving from this class. # # Users should not interact with this directly - __all_operations: ClassVar[dict[tuple[str, str], type[OperationT]]] = {} + __all_operations: ClassVar[dict[tuple[str, str], type["Operation"]]] = {} # The supported layout for data for algebra objects. At the moment all # operations only support densely represented objects. In the future, @@ -223,7 +219,7 @@ class Operation: # of all the required and optional arguments. This will be passed to the # FFI calls by ** unpacking. Using a TypedDict gives some level of # argument checking - StaticArgs: ClassVar[type[TypedDict]] + StaticArgs: ClassVar[type[Mapping[str, object]]] ## The following instance attributes are used by the class upon call to ## select from available implementations and populate static arguments. @@ -237,7 +233,7 @@ class Operation: # the configuration of batching batch_dims: tuple[int, ...] # dictionary of static arguments - static_args: type[TypedDict] + static_args: Mapping[str, object] # For FFI calls, the shape of the output array(s) result_shape_dtypes: tuple[jax.ShapeDtypeStruct, ...] @@ -284,7 +280,7 @@ def register_all( @classmethod def get_operation( cls, fn_name: str, layout: str = "dense" - ) -> type[OperationT] | None: + ) -> type["Operation"] | None: """ Retrieves a registered operation class based on the function name and layout. @@ -353,7 +349,7 @@ def get_result_basis(cls, bases: tuple[Basis, ...], preferred_basis) -> Basis: if preferred_basis is not None: return result_basis(preferred_basis, *bases, strategy="first") - return result_basis(bases, strategy="max_depth") + return result_basis(*bases, strategy="max_depth") @classmethod def __init_subclass__(cls, **kwargs): @@ -386,7 +382,7 @@ def __init__( self.ffi_call_args = self.default_ffi_call_args | (ffi_call_args or {}) self.static_args = self.make_static_args(kwargs) - def make_static_args(self, kwargs) -> type[TypedDict]: + def make_static_args(self, kwargs) -> Mapping[str, object]: """ Construct static args from class kwargs. @@ -971,7 +967,7 @@ def get_result_basis(cls, bases: tuple[Basis, ...], preferred_basis) -> Basis: return basis - def make_static_args(self, kwargs) -> type[TypedDict]: + def make_static_args(self, kwargs) -> Mapping[str, object]: arg_basis = self.bases[0] tensor_basis = self.basis @@ -1026,7 +1022,7 @@ def get_result_basis(cls, bases: tuple[Basis, ...], preferred_basis) -> Basis: return basis - def make_static_args(self, kwargs) -> type[TypedDict]: + def make_static_args(self, kwargs) -> Mapping[str, object]: lie_basis = self.basis t2l_data, t2l_indices, t2l_indptr = _get_lie_sparse_matrices( diff --git a/roughpy_jax/streams/lie_increment_stream.py b/roughpy_jax/streams/lie_increment_stream.py index 4dfc96bfa..d5f3f190d 100644 --- a/roughpy_jax/streams/lie_increment_stream.py +++ b/roughpy_jax/streams/lie_increment_stream.py @@ -280,6 +280,8 @@ def __init__( ) cache_length = int(cache.shape[0]) + if resolution is None: + raise ValueError("resolution must be provided") expected_length = self._cache_length_from_resolution(resolution) if cache_length != expected_length: raise ValueError( @@ -291,7 +293,7 @@ def __init__( self._lie_basis = lie_basis self._group_basis = group_basis or to_tensor_basis(lie_basis) self._support = support or RealInterval(0.0, 1.0, IntervalType.ClOpen) - self._resolution = int(resolution) + self._resolution = resolution self._interval_type = interval_type self._zero_index = cache_length - 1 diff --git a/roughpy_jax/streams/piecewise_abelian_stream.py b/roughpy_jax/streams/piecewise_abelian_stream.py index 80b16513d..687371e36 100644 --- a/roughpy_jax/streams/piecewise_abelian_stream.py +++ b/roughpy_jax/streams/piecewise_abelian_stream.py @@ -18,17 +18,7 @@ from .concepts import GroupT, LieT, Stream -def _pas_dataclass(cls): - """Helper to apply the dataclass and register_dataclass decorators in the correct order.""" - cls = dataclass(cls, frozen=True) - return jax.tree_util.register_dataclass( - cls, - data_fields=["_data", "_partition"], - meta_fields=["_lie_basis", "_group_basis"], - ) - - -@_pas_dataclass +@dataclass(frozen=True) class PiecewiseAbelianStream(Stream[LieT, GroupT]): """A stream representing a piecewise abelian path.""" @@ -125,6 +115,13 @@ def signature(self, interval: Interval) -> GroupT: return ft_exp(tensor, self._group_basis) +PiecewiseAbelianStream = jax.tree_util.register_dataclass( + PiecewiseAbelianStream, + data_fields=["_data", "_partition"], + meta_fields=["_lie_basis", "_group_basis"], +) + + def to_piecewise_abelian( stream: Stream[LieT, GroupT], partition: Partition ) -> PiecewiseAbelianStream[LieT, GroupT]: diff --git a/ty.toml b/ty.toml index d9499e711..75105e30a 100644 --- a/ty.toml +++ b/ty.toml @@ -14,7 +14,7 @@ respect-ignore-files = true [analysis] respect-type-ignore-comments = true -replace-imports-with-any = ["roughpy._roughpy"] +replace-imports-with-any = ["roughpy._roughpy", "roughpy_jax._rpy_jax_internals"] [rules] all = "warn" @@ -28,4 +28,21 @@ unknown-argument = "ignore" include = ["tests/**", "benchmarks/**", "examples/**"] [overrides.rules] -possibly-unresolved-reference = "warn" \ No newline at end of file +possibly-unresolved-reference = "warn" + +[[overrides]] +include = ["roughpy_jax/intervals.py"] + +[overrides.rules] +unsupported-operator = "ignore" + +[[overrides]] +include = ["roughpy_jax/ops.py"] + +[overrides.rules] +invalid-assignment = "ignore" +invalid-argument-type = "ignore" +invalid-parameter-default = "ignore" +invalid-type-arguments = "ignore" +invalid-type-form = "ignore" +unresolved-import = "ignore" \ No newline at end of file