From ddf6f10cdc5b147bfdaf113adce1d013a46d56a9 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Tue, 20 Aug 2024 22:43:05 +0200 Subject: [PATCH 01/14] Type get_context correctly get_context returns an instance of a Model, not a ContextMeta object We don't need the typevar, since we don't use it for anything special --- pymc/model/core.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 48d2117eb26..b89c825029d 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -24,7 +24,6 @@ TYPE_CHECKING, Literal, Optional, - TypeVar, Union, cast, overload, @@ -92,9 +91,6 @@ ] -T = TypeVar("T", bound="ContextMeta") - - class ContextMeta(type): """Functionality for objects that put themselves in a context manager.""" @@ -125,13 +121,13 @@ def __init__(cls, name, bases, nmspc, context_class: type | None = None, **kwarg cls._context_class = context_class super().__init__(name, bases, nmspc) - def get_context(cls, error_if_none=True, allow_block_model_access=False) -> T | None: + def get_context(cls, error_if_none=True, allow_block_model_access=False) -> "Model | None": """Return the most recently pushed context object of type ``cls`` on the stack, or ``None``. If ``error_if_none`` is True (default), raise a ``TypeError`` instead of returning ``None``. """ try: - candidate: T | None = cls.get_contexts()[-1] + candidate: Model | None = cls.get_contexts()[-1] except IndexError: # Calling code expects to get a TypeError if the entity # is unfound, and there's too much to fix. @@ -142,7 +138,7 @@ def get_context(cls, error_if_none=True, allow_block_model_access=False) -> T | raise BlockModelAccessError(candidate.error_msg_on_access) return candidate - def get_contexts(cls) -> list[T]: + def get_contexts(cls) -> list["Model"]: """Return a stack of context instances for the ``context_class`` of ``cls``.""" # This lazily creates the context class's contexts # thread-local object, as needed. This seems inelegant to me, From 74f66161c43b24d5039a41777de4e77c82631d54 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Wed, 21 Aug 2024 08:41:14 +0200 Subject: [PATCH 02/14] Import from future to use delayed evaluation of annotations All of these are supported on python>=3.9. --- pymc/model/core.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index b89c825029d..5ef5a2df829 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations import functools import sys @@ -23,8 +24,6 @@ from typing import ( TYPE_CHECKING, Literal, - Optional, - Union, cast, overload, ) @@ -121,7 +120,7 @@ def __init__(cls, name, bases, nmspc, context_class: type | None = None, **kwarg cls._context_class = context_class super().__init__(name, bases, nmspc) - def get_context(cls, error_if_none=True, allow_block_model_access=False) -> "Model | None": + def get_context(cls, error_if_none=True, allow_block_model_access=False) -> Model | None: """Return the most recently pushed context object of type ``cls`` on the stack, or ``None``. If ``error_if_none`` is True (default), raise a ``TypeError`` instead of returning ``None``. @@ -138,7 +137,7 @@ def get_context(cls, error_if_none=True, allow_block_model_access=False) -> "Mod raise BlockModelAccessError(candidate.error_msg_on_access) return candidate - def get_contexts(cls) -> list["Model"]: + def get_contexts(cls) -> list[Model]: """Return a stack of context instances for the ``context_class`` of ``cls``.""" # This lazily creates the context class's contexts # thread-local object, as needed. This seems inelegant to me, @@ -199,7 +198,7 @@ def __call__(cls, *args, **kwargs): return instance -def modelcontext(model: Optional["Model"]) -> "Model": +def modelcontext(model: Model | None) -> Model: """Return the given model or, if None was supplied, try to find one in the context stack.""" if model is None: model = Model.get_context(error_if_none=False) @@ -499,7 +498,7 @@ def __enter__(self: Self) -> Self: def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None: """Exit the context manager.""" - def __new__(cls, *args, model: Union[Literal[UNSET], None, "Model"] = UNSET, **kwargs): + def __new__(cls, *args, model: Literal[UNSET] | None | Model = UNSET, **kwargs): # resolves the parent instance instance = super().__new__(cls) if model is UNSET: @@ -521,7 +520,7 @@ def __init__( check_bounds=True, *, coords_mutable=None, - model: Union[Literal[UNSET], None, "Model"] = UNSET, + model: Literal[UNSET] | None | Model = UNSET, ): del model # used in __new__ to define the parent of this model self.name = self._validate_name(name) From aa253ac6ee813c08043aab9031324f6506224865 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Wed, 21 Aug 2024 15:59:50 +0200 Subject: [PATCH 03/14] New ModelManager class for managing model contexts We create a global instance of it within this module, which is similar to how it worked before, where a `context_class` attribute was attached to the Model class. We inherit from threading.local to ensure thread safety when working with models on multiple threads. See #1552 for the reasoning. This is already tested in `test_thread_safety`. --- pymc/model/core.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/pymc/model/core.py b/pymc/model/core.py index 5ef5a2df829..d2bb378d3e5 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -90,6 +90,35 @@ ] +class ModelManager(threading.local): + """Keeps track of currently active model contexts. + + A global instance of this is created in this module on import. + Use that instance, `MODEL_MANAGER` to inspect current contexts. + + It inherits from threading.local so is thread-safe, if models + can be entered/exited within individual threads. + """ + + def __init__(self): + self.active_contexts: list[Model] = [] + + @property + def current_context(self) -> Model | None: + """Return the innermost context of any current contexts.""" + return self.active_contexts[-1] if self.active_contexts else None + + @property + def parent_context(self) -> Model | None: + """Return the parent context to the active context, if any.""" + return self.active_contexts[-2] if len(self.active_contexts) > 1 else None + + +# MODEL_MANAGER is instantiated at import, and serves as a truth for +# what any currently active model contexts are. +MODEL_MANAGER = ModelManager() + + class ContextMeta(type): """Functionality for objects that put themselves in a context manager.""" From d5beee6997b61e0c79a65cde458a44991a4aa784 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Wed, 21 Aug 2024 18:55:22 +0200 Subject: [PATCH 04/14] Model class is now the context manager directly --- pymc/model/core.py | 34 +++++++--------------------------- 1 file changed, 7 insertions(+), 27 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index d2bb378d3e5..e1f61fa67f5 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -22,7 +22,6 @@ from collections.abc import Iterable, Sequence from sys import modules from typing import ( - TYPE_CHECKING, Literal, cast, overload, @@ -40,7 +39,6 @@ from pytensor.tensor.random.op import RandomVariable from pytensor.tensor.random.type import RandomType from pytensor.tensor.variable import TensorConstant, TensorVariable -from typing_extensions import Self from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.data import is_valid_observed @@ -122,25 +120,6 @@ def parent_context(self) -> Model | None: class ContextMeta(type): """Functionality for objects that put themselves in a context manager.""" - def __new__(cls, name, bases, dct, **kwargs): - """Add __enter__ and __exit__ methods to the class.""" - - def __enter__(self): - self.__class__.context_class.get_contexts().append(self) - return self - - def __exit__(self, typ, value, traceback): - self.__class__.context_class.get_contexts().pop() - - dct[__enter__.__name__] = __enter__ - dct[__exit__.__name__] = __exit__ - - # We strip off keyword args, per the warning from - # StackExchange: - # DO NOT send "**kwargs" to "type.__new__". It won't catch them and - # you'll get a "TypeError: type() takes 1 or 3 arguments" exception. - return super().__new__(cls, name, bases, dct) - # FIXME: is there a more elegant way to automatically add methods to the class that # are instance methods instead of class methods? def __init__(cls, name, bases, nmspc, context_class: type | None = None, **kwargs): @@ -519,13 +498,14 @@ class Model(WithMemoization, metaclass=ContextMeta): """ - if TYPE_CHECKING: - - def __enter__(self: Self) -> Self: - """Enter the context manager.""" + def __enter__(self): + """Enter the context manager.""" + MODEL_MANAGER.active_contexts.append(self) + return self - def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None: - """Exit the context manager.""" + def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None: + """Exit the context manager.""" + _ = MODEL_MANAGER.active_contexts.pop() def __new__(cls, *args, model: Literal[UNSET] | None | Model = UNSET, **kwargs): # resolves the parent instance From e2cefc3e4aff2cff5cf0324f52640e2c29bb0daa Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Wed, 21 Aug 2024 19:34:59 +0200 Subject: [PATCH 05/14] Fix type of UNSET in type definition UNSET is the instance of the _UnsetType type. We should be typing the latter here. --- pymc/model/core.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index e1f61fa67f5..5a3ee338e4a 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -69,6 +69,7 @@ VarName, WithMemoization, _add_future_warning_tag, + _UnsetType, get_transformed_name, get_value_vars_from_user_vars, get_var_name, @@ -507,7 +508,7 @@ def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None: """Exit the context manager.""" _ = MODEL_MANAGER.active_contexts.pop() - def __new__(cls, *args, model: Literal[UNSET] | None | Model = UNSET, **kwargs): + def __new__(cls, *args, model: _UnsetType | None | Model = UNSET, **kwargs): # resolves the parent instance instance = super().__new__(cls) if model is UNSET: @@ -529,7 +530,7 @@ def __init__( check_bounds=True, *, coords_mutable=None, - model: Literal[UNSET] | None | Model = UNSET, + model: _UnsetType | None | Model = UNSET, ): del model # used in __new__ to define the parent of this model self.name = self._validate_name(name) From 5a1f0b8285f29fa0043124182b29f07bc139c924 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Wed, 21 Aug 2024 19:40:47 +0200 Subject: [PATCH 06/14] Set model parent in init rather than in __new__ We use the new ModelManager.parent_context property to reliably set any parent context, or else set it to None. --- pymc/model/core.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 5a3ee338e4a..4cc14e068c8 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -508,15 +508,6 @@ def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None: """Exit the context manager.""" _ = MODEL_MANAGER.active_contexts.pop() - def __new__(cls, *args, model: _UnsetType | None | Model = UNSET, **kwargs): - # resolves the parent instance - instance = super().__new__(cls) - if model is UNSET: - instance._parent = cls.get_context(error_if_none=False) - else: - instance._parent = model - return instance - @staticmethod def _validate_name(name): if name.endswith(":"): @@ -532,9 +523,9 @@ def __init__( coords_mutable=None, model: _UnsetType | None | Model = UNSET, ): - del model # used in __new__ to define the parent of this model self.name = self._validate_name(name) self.check_bounds = check_bounds + self._parent = model if not isinstance(model, _UnsetType) else MODEL_MANAGER.parent_context if coords_mutable is not None: warnings.warn( From 8846db93da3cc6a91976c9bcf62ac4ae77b5a234 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Thu, 22 Aug 2024 19:20:19 +0200 Subject: [PATCH 07/14] Replace get_context in metaclass with classmethod We set this directly on the class as a classmethod, which is clearer than going via the metaclass. --- pymc/model/core.py | 28 +++++++++++----------------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 4cc14e068c8..ee150554a67 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -129,23 +129,6 @@ def __init__(cls, name, bases, nmspc, context_class: type | None = None, **kwarg cls._context_class = context_class super().__init__(name, bases, nmspc) - def get_context(cls, error_if_none=True, allow_block_model_access=False) -> Model | None: - """Return the most recently pushed context object of type ``cls`` on the stack, or ``None``. - - If ``error_if_none`` is True (default), raise a ``TypeError`` instead of returning ``None``. - """ - try: - candidate: Model | None = cls.get_contexts()[-1] - except IndexError: - # Calling code expects to get a TypeError if the entity - # is unfound, and there's too much to fix. - if error_if_none: - raise TypeError(f"No {cls} on context stack") - return None - if isinstance(candidate, BlockModelAccess) and not allow_block_model_access: - raise BlockModelAccessError(candidate.error_msg_on_access) - return candidate - def get_contexts(cls) -> list[Model]: """Return a stack of context instances for the ``context_class`` of ``cls``.""" # This lazily creates the context class's contexts @@ -573,6 +556,17 @@ def __init__( functools.partial(str_for_model, formatting="latex"), self ) + @classmethod + def get_context( + cls, error_if_none: bool = True, allow_block_model_access: bool = False + ) -> Model | None: + model = MODEL_MANAGER.current_context + if isinstance(model, BlockModelAccess) and not allow_block_model_access: + raise BlockModelAccessError(model.error_msg_on_access) + if model is None and error_if_none: + raise TypeError("No model on context stack") + return model + @property def parent(self): return self._parent From 7318d7562d8e1dc302a7ae409f57fa81e1bb4707 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Thu, 22 Aug 2024 19:25:26 +0200 Subject: [PATCH 08/14] Remove get_contexts from metaclass The original function does not behave as I expected. In the following example I expected that it would return only the final model, not root. This method is not used anywhere in the pymc codebase, so I have dropped it from the codebase. I originally included the following code to replace it, but since it is not used anyway, it is better to remove it. ```python` @classmethod def get_contexts(cls) -> list[Model]: """Return a list of the currently active model contexts.""" return MODEL_MANAGER.active_contexts ``` Example for testing behaviour in current main branch: ```python import pymc as pm with pm.Model(name="root") as root: print([c.name for c in pm.Model.get_contexts()]) with pm.Model(name="first") as first: print([c.name for c in pm.Model.get_contexts()]) with pm.Model(name="m_with_model_None", model=None) as m_with_model_None: # This one doesn't make much sense: print([c.name for c in pm.Model.get_contexts()]) ``` --- pymc/model/core.py | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index ee150554a67..bd080031e83 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -129,29 +129,6 @@ def __init__(cls, name, bases, nmspc, context_class: type | None = None, **kwarg cls._context_class = context_class super().__init__(name, bases, nmspc) - def get_contexts(cls) -> list[Model]: - """Return a stack of context instances for the ``context_class`` of ``cls``.""" - # This lazily creates the context class's contexts - # thread-local object, as needed. This seems inelegant to me, - # but since the context class is not guaranteed to exist when - # the metaclass is being instantiated, I couldn't figure out a - # better way. [2019/10/11:rpg] - - # no race-condition here, contexts is a thread-local object - # be sure not to override contexts in a subclass however! - context_class = cls.context_class - assert isinstance( - context_class, type - ), f"Name of context class, {context_class} was not resolvable to a class" - if not hasattr(context_class, "contexts"): - context_class.contexts = threading.local() - - contexts = context_class.contexts - - if not hasattr(contexts, "stack"): - contexts.stack = [] - return contexts.stack - # the following complex property accessor is necessary because the # context_class may not have been created at the point it is # specified, so the context_class may be a class *name* rather From 560cd3bc5c1928f80164e27c4679cbe6d7dac5f9 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Thu, 22 Aug 2024 19:28:37 +0200 Subject: [PATCH 09/14] Simplify ContextMeta We only keep the __call__ method, which is necessary to keep the model context itself active during that model's __init__. --- pymc/model/core.py | 62 +++++++++------------------------------------- 1 file changed, 12 insertions(+), 50 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index bd080031e83..f3e0efa5306 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -20,7 +20,6 @@ import warnings from collections.abc import Iterable, Sequence -from sys import modules from typing import ( Literal, cast, @@ -118,55 +117,6 @@ def parent_context(self) -> Model | None: MODEL_MANAGER = ModelManager() -class ContextMeta(type): - """Functionality for objects that put themselves in a context manager.""" - - # FIXME: is there a more elegant way to automatically add methods to the class that - # are instance methods instead of class methods? - def __init__(cls, name, bases, nmspc, context_class: type | None = None, **kwargs): - """Add ``__enter__`` and ``__exit__`` methods to the new class automatically.""" - if context_class is not None: - cls._context_class = context_class - super().__init__(name, bases, nmspc) - - # the following complex property accessor is necessary because the - # context_class may not have been created at the point it is - # specified, so the context_class may be a class *name* rather - # than a class. - @property - def context_class(cls) -> type: - def resolve_type(c: type | str) -> type: - if isinstance(c, str): - c = getattr(modules[cls.__module__], c) - if isinstance(c, type): - return c - raise ValueError(f"Cannot resolve context class {c}") - - assert cls is not None - if isinstance(cls._context_class, str): - cls._context_class = resolve_type(cls._context_class) - if not isinstance(cls._context_class, str | type): - raise ValueError( - f"Context class for {cls.__name__}, {cls._context_class}, is not of the right type" - ) - return cls._context_class - - # Inherit context class from parent - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - cls.context_class = super().context_class - - # Initialize object in its own context... - # Merged from InitContextMeta in the original. - def __call__(cls, *args, **kwargs): - # We type hint Model here so type checkers understand that Model is a context manager. - # This metaclass is only used for Model, so this is safe to do. See #6809 for more info. - instance: Model = cls.__new__(cls, *args, **kwargs) - with instance: # appends context - instance.__init__(*args, **kwargs) - return instance - - def modelcontext(model: Model | None) -> Model: """Return the given model or, if None was supplied, try to find one in the context stack.""" if model is None: @@ -336,6 +286,18 @@ def profile(self): return self._pytensor_function.profile +class ContextMeta(type): + """A metaclass in order to apply a model's context during `Model.__init__``.""" + + # We want the Model's context to be active during __init__. In order for this + # to apply to subclasses of Model as well, we need to use a metaclass. + def __call__(cls: type[Model], *args, **kwargs): + instance = cls.__new__(cls, *args, **kwargs) + with instance: # applies context + instance.__init__(*args, **kwargs) + return instance + + class Model(WithMemoization, metaclass=ContextMeta): """Encapsulates the variables and likelihood factors of a model. From 6ff3d5d3c591e785882b74a1b2ca5767de97bfc5 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Wed, 28 Aug 2024 19:53:30 +0200 Subject: [PATCH 10/14] Type Model.register_rv for for downstream typing In pymc/distributions/distribution.py, this change allows the type checker to infer that `rv_out` can only be a TensorVariable. Thanks to @ricardoV94 for type hint on rv_var. --- pymc/model/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index f3e0efa5306..6ccf8a9b0dd 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1162,8 +1162,8 @@ def set_data( def register_rv( self, - rv_var, - name, + rv_var: RandomVariable, + name: str, *, observed=None, total_size=None, @@ -1171,7 +1171,7 @@ def register_rv( default_transform=UNSET, transform=UNSET, initval=None, - ): + ) -> TensorVariable: """Register an (un)observed random variable with the model. Parameters From 2ebc48f2209b8b9e8806de87770c618e42876870 Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Fri, 6 Sep 2024 15:14:10 +0200 Subject: [PATCH 11/14] Include np.ndarray as possible type for coord values I originally tried numpy's ArrayLike, replacing Sequence entirely, but then I realized that ArrayLike also allows non-sequences like integers and floats. I am not certain if `values="a string"` should be legal. With the type hint sequence, it is. Might be more accurate, but verbose to use `list | tuple | set | np.ndarray | None`. --- pymc/model/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 6ccf8a9b0dd..cf205aac6db 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -896,7 +896,7 @@ def shape_from_dims(self, dims): def add_coord( self, name: str, - values: Sequence | None = None, + values: Sequence | np.ndarray | None = None, mutable: bool | None = None, *, length: int | Variable | None = None, From 165a35f3e1d640b8139a81dfaa6b2863bfe4e91e Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Fri, 4 Oct 2024 11:12:45 +0200 Subject: [PATCH 12/14] Use function-scoped new_dims to handle type hint varying throughout function We don't want to allow the user to pass a `dims=[None, None]` to our function, but current behaviour set `dims=[None] * N` at the end of `determine_coords`. To handle this, I created a `new_dims` with a larger type scope which matches the return type of `dims` in `determine_coords`. Then I did the same within def Data to support this new type hint. --- pymc/data.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/pymc/data.py b/pymc/data.py index 247825981f1..4e34aaa9a7b 100644 --- a/pymc/data.py +++ b/pymc/data.py @@ -221,9 +221,9 @@ def Minibatch(variable: TensorVariable, *variables: TensorVariable, batch_size: def determine_coords( model, value: pd.DataFrame | pd.Series | xr.DataArray, - dims: Sequence[str | None] | None = None, + dims: Sequence[str] | None = None, coords: dict[str, Sequence | np.ndarray] | None = None, -) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None]]: +) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str] | Sequence[None]]: """Determine coordinate values from data or the model (via ``dims``).""" if coords is None: coords = {} @@ -268,9 +268,10 @@ def determine_coords( if dims is None: # TODO: Also determine dim names from the index - dims = [None] * np.ndim(value) - - return coords, dims + new_dims: Sequence[str] | Sequence[None] = [None] * np.ndim(value) + else: + new_dims = dims + return coords, new_dims def ConstantData( @@ -366,7 +367,7 @@ def Data( The name for this variable. value : array_like or pandas.Series, pandas.Dataframe A value to associate with this variable. - dims : str or tuple of str, optional + dims : str, tuple of str or tuple of None, optional Dimension names of the random variables (as opposed to the shapes of these random variables). Use this when ``value`` is a pandas Series or DataFrame. The ``dims`` will then be the name of the Series / DataFrame's columns. See ArviZ @@ -451,13 +452,16 @@ def Data( expected=x.ndim, ) + new_dims: Sequence[str] | Sequence[None] | None if infer_dims_and_coords: - coords, dims = determine_coords(model, value, dims) + coords, new_dims = determine_coords(model, value, dims) + else: + new_dims = dims - if dims: + if new_dims: xshape = x.shape # Register new dimension lengths - for d, dname in enumerate(dims): + for d, dname in enumerate(new_dims): if dname not in model.dim_lengths: model.add_coord( name=dname, @@ -467,6 +471,6 @@ def Data( length=xshape[d], ) - model.register_data_var(x, dims=dims) + model.register_data_var(x, dims=new_dims) return x From dbbb9a2404007a1e0acfb879bccfcd267e085dbc Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Fri, 4 Oct 2024 11:15:22 +0200 Subject: [PATCH 13/14] Fix case of dims = [None, None, ...] The only case where dims=[None, ...] is when the user has passed dims=None. Since the user passed dims=None, they shouldn't be expecting any coords to match that dimension. Thus we don't need to try to add any more coords to the model. --- pymc/data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/data.py b/pymc/data.py index 4e34aaa9a7b..22fc8717c3f 100644 --- a/pymc/data.py +++ b/pymc/data.py @@ -462,7 +462,7 @@ def Data( xshape = x.shape # Register new dimension lengths for d, dname in enumerate(new_dims): - if dname not in model.dim_lengths: + if dname not in model.dim_lengths and dname is not None: model.add_coord( name=dname, # Note: Coordinate values can't be taken from From 8565965da0cb27da66e482c4198391d1c037e29e Mon Sep 17 00:00:00 2001 From: Thomas Aarholt Date: Thu, 10 Oct 2024 11:12:28 +0200 Subject: [PATCH 14/14] Remove unused hack --- pymc/model/core.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index cf205aac6db..ad60a84dfb9 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -2003,11 +2003,6 @@ def to_graphviz( ) -# this is really disgusting, but it breaks a self-loop: I can't pass Model -# itself as context class init arg. -Model._context_class = Model - - class BlockModelAccess(Model): """Can be used to prevent user access to Model contexts."""