From 5e9a21b6466b7f22688853e542566b05ebba920e Mon Sep 17 00:00:00 2001 From: Addison Date: Mon, 21 Sep 2026 15:22:12 -0500 Subject: [PATCH 1/2] align_obj: keep param names updated on isl object Co-authored-by: Andreas Kloeckner --- namedisl/core.py | 12 ++++++++++++ namedisl/test/test_namedisl.py | 19 ++++++++++++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/namedisl/core.py b/namedisl/core.py index faafcd1..eaffcc2 100644 --- a/namedisl/core.py +++ b/namedisl/core.py @@ -360,6 +360,9 @@ def align_obj( *, allow_cross_dim_type: bool = False, obj_larger_than_space_ok: bool = False, ) -> NamedIslObjectT: + if named_obj.space.order_equals(space): + return named_obj + obj = named_obj._obj running_name_to_dim_id = dict(named_obj.space.name_to_dim) @@ -393,6 +396,11 @@ def align_obj( else: another_dim_type = DimType.in_ + if another_dim_type == DimType.param: + # Only promoted dimensions need consistent parameter IDs. + obj = _set_dim_name( + obj, old_dim_id.dim_type, old_dim_id.dim_index, name) + obj = obj.move_dims( another_dim_type.as_isl(), 0, old_dim_id.dim_type.as_isl(), old_dim_id.dim_index, @@ -406,6 +414,10 @@ def align_obj( if not allow_cross_dim_type: raise ValueError("moves across dim_types are not allowed") + if target_dt == DimType.param: + obj = _set_dim_name( + obj, old_dim_id.dim_type, old_dim_id.dim_index, name) + obj = obj.move_dims( target_dim_id.dim_type.as_isl(), target_dim_id.dim_index, old_dim_id.dim_type.as_isl(), old_dim_id.dim_index, diff --git a/namedisl/test/test_namedisl.py b/namedisl/test/test_namedisl.py index e769f4f..1387cea 100644 --- a/namedisl/test/test_namedisl.py +++ b/namedisl/test/test_namedisl.py @@ -33,7 +33,7 @@ import namedisl as nisl from .utils_for_tests import generate_random_named_set from namedisl import to_named -from namedisl.core import DimType +from namedisl.core import DimType, Space, align_obj @pytest.mark.parametrize("ndims", [2, 3, 4, 5]) @@ -228,3 +228,20 @@ def test_ticked_names_are_distinct_names() -> None: named_set = nisl.make_set(isl.Set.universe(space)) assert named_set.space.names == frozenset({"x", "x'"}) + + +def test_align_obj_returns_already_aligned_object() -> None: + named_set = nisl.make_set("{ [i] : i >= 0 }") + + assert align_obj(named_set, named_set.space) is named_set + + +def test_align_obj_promotes_renamed_dimension_with_its_named_id() -> None: + renamed_set = nisl.make_set("{ [i] : i = 4 }").rename_dims((("i", "n"),)) + parameter_space = Space.from_names(param=["n"], out=[]) + + aligned = align_obj( + renamed_set, parameter_space, allow_cross_dim_type=True) + + assert aligned._obj.get_dim_name(isl.dim_type.param, 0) == "n" + assert aligned.equals(nisl.make_set("[n] -> { [] : n = 4 }")) From 03d4b7a11a0d932b8deb01274fe3bfb12cdcd6dd Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 21 Sep 2026 15:36:02 -0500 Subject: [PATCH 2/2] Fix, test pickling Co-authored-by: Addison --- namedisl/core.py | 15 ++++++++++++++- namedisl/test/test_namedisl.py | 16 +++++++++++++++- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/namedisl/core.py b/namedisl/core.py index eaffcc2..c301578 100644 --- a/namedisl/core.py +++ b/namedisl/core.py @@ -40,7 +40,7 @@ import enum import re from collections.abc import Callable, Collection, Hashable, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, fields from functools import cached_property from importlib import metadata from typing import ( @@ -955,6 +955,19 @@ def as_isl(self) -> IslObjectT_co: object.__setattr__(self, "_isl_names_ok", True) # ruff: ignore[unnecessary-dunder-call] return res + @override + def __reduce__(self): + return (type(self), + ( + # islpy's pickle protocol requires unique raw dimension names. + self.as_isl(), + *( + getattr(self, f.name) + for f in fields(self)[1:] + ) + ) + ) + def involves_dims(self, names: Collection[str]) -> bool: """True if *self* involves any of the given dimensions.""" for dt, chunks in chunked_dims_by_type(names, self.space.name_to_dim).items(): diff --git a/namedisl/test/test_namedisl.py b/namedisl/test/test_namedisl.py index 1387cea..ca9dffc 100644 --- a/namedisl/test/test_namedisl.py +++ b/namedisl/test/test_namedisl.py @@ -26,6 +26,9 @@ """ +import pickle +from typing import cast + import pytest import islpy as isl @@ -243,5 +246,16 @@ def test_align_obj_promotes_renamed_dimension_with_its_named_id() -> None: aligned = align_obj( renamed_set, parameter_space, allow_cross_dim_type=True) - assert aligned._obj.get_dim_name(isl.dim_type.param, 0) == "n" assert aligned.equals(nisl.make_set("[n] -> { [] : n = 4 }")) + + +def test_pickling_restores_unique_raw_dimension_names() -> None: + raw_space = isl.Space.set_alloc(isl.DEFAULT_CONTEXT, 0, 2) + named_set = nisl.Set( + isl.Set.universe(raw_space), Space.from_names(param=[], out=["i", "j"])) + + restored = cast("nisl.Set", pickle.loads(pickle.dumps(named_set))) + + assert restored.equals(named_set) + assert restored._obj.get_dim_name(isl.dim_type.set, 0) == "i" + assert restored._obj.get_dim_name(isl.dim_type.set, 1) == "j"