diff --git a/gnm/shape/gnm_jax.py b/gnm/shape/gnm_jax.py index c301b751..ab4881ef 100644 --- a/gnm/shape/gnm_jax.py +++ b/gnm/shape/gnm_jax.py @@ -38,11 +38,10 @@ from __future__ import annotations -from collections.abc import Mapping import dataclasses -from typing import Any, Self from absl import logging +from etils import enp from gnm.shape import gnm_landmarks from gnm.shape import gnm_xnp from gnm.shape.data.versions import gnm_specs @@ -117,12 +116,9 @@ class GNM(gnm_xnp.GNM): _shape_error_type = jt.TypeCheckError @classmethod - def _from_model_data( - cls, - model_data: Mapping[str, Any], - ) -> Self: - """Creates a JAX GNM instance from model data.""" - return cls._from_model_data_with_xnp(model_data, xnp=jnp) + def _get_np_module(cls) -> enp.NpModule: + """Returns the JAX NumPy backend module.""" + return jnp def compute_vertex_normals( self, diff --git a/gnm/shape/gnm_numpy.py b/gnm/shape/gnm_numpy.py index 96ebbe91..8064b12c 100644 --- a/gnm/shape/gnm_numpy.py +++ b/gnm/shape/gnm_numpy.py @@ -32,11 +32,10 @@ from __future__ import annotations -from collections.abc import Mapping import dataclasses -from typing import Any, Self from absl import logging +from etils import enp from gnm.shape import gnm_common from gnm.shape import gnm_landmarks from gnm.shape import gnm_xnp @@ -116,12 +115,9 @@ class GNM(gnm_xnp.GNM): """ @classmethod - def _from_model_data( - cls, - model_data: Mapping[str, Any], - ) -> Self: - """Creates a GNM instance from model data.""" - return cls._from_model_data_with_xnp(model_data, xnp=np) + def _get_np_module(cls) -> enp.NpModule: + """Returns the NumPy backend module.""" + return np def compute_vertex_normals( self, diff --git a/gnm/shape/gnm_pytorch.py b/gnm/shape/gnm_pytorch.py index 71b57657..05769bdd 100644 --- a/gnm/shape/gnm_pytorch.py +++ b/gnm/shape/gnm_pytorch.py @@ -34,11 +34,10 @@ from __future__ import annotations -from collections.abc import Mapping import dataclasses -from typing import Any, Self from absl import logging +from etils import enp from gnm.shape import gnm_landmarks from gnm.shape import gnm_xnp from gnm.shape.data.versions import gnm_specs @@ -125,12 +124,9 @@ def __post_init__(self): self.register_buffer(field.name, val, persistent=False) @classmethod - def _from_model_data( - cls, - model_data: Mapping[str, Any], - ) -> Self: - """Creates a PyTorch GNM instance from model data.""" - return cls._from_model_data_with_xnp(model_data, xnp=torch) + def _get_np_module(cls) -> enp.NpModule: + """Returns the PyTorch backend module.""" + return torch def compute_vertex_normals( self, diff --git a/gnm/shape/gnm_tensorflow.py b/gnm/shape/gnm_tensorflow.py index 3cd4200e..318f0c90 100644 --- a/gnm/shape/gnm_tensorflow.py +++ b/gnm/shape/gnm_tensorflow.py @@ -34,9 +34,7 @@ from __future__ import annotations -from collections.abc import Mapping import dataclasses -from typing import Any, Self from absl import logging from etils import enp @@ -119,12 +117,9 @@ class GNM(gnm_xnp.GNM): """ @classmethod - def _from_model_data( - cls, - model_data: Mapping[str, Any], - ) -> Self: - """Creates a TensorFlow GNM instance from model data.""" - return cls._from_model_data_with_xnp(model_data, xnp=enp.lazy.tnp) + def _get_np_module(cls) -> enp.NpModule: + """Returns the TensorFlow NumPy backend module.""" + return enp.lazy.tnp def compute_vertex_normals( self, diff --git a/gnm/shape/gnm_xnp.py b/gnm/shape/gnm_xnp.py index 1cf651d1..1c425b23 100644 --- a/gnm/shape/gnm_xnp.py +++ b/gnm/shape/gnm_xnp.py @@ -165,13 +165,23 @@ def __setstate__(self, state): self.__dict__.update(state) self._xnp = enp.get_np_module(self.template_vertex_positions) + @classmethod + @abc.abstractmethod + def _get_np_module(cls) -> enp.NpModule: + """Returns the array module (e.g., np, jnp, torch, tnp) for this backend.""" + raise NotImplementedError( + f'{cls.__name__} is an abstract backend-agnostic base class and' + ' does not define a backend array module. Use a concrete backend' + ' subclass (e.g., gnm_numpy.GNM, gnm_pytorch.GNM).' + ) + @classmethod def _prepare_init_kwargs( cls, model_data: Mapping[str, Any], - xnp: enp.NpModule, ) -> dict[str, Any]: """Prepares and casts GNM initialization arguments using xnp.""" + xnp = cls._get_np_module() init_kwargs = {} # Identify type convert functions. @@ -217,13 +227,12 @@ def as_original(val): return init_kwargs @classmethod - def _from_model_data_with_xnp( + def _from_model_data( cls, model_data: Mapping[str, Any], - xnp: enp.NpModule, ) -> Self: - """Creates a GNM instance from a model data dictionary and array module.""" - init_kwargs = cls._prepare_init_kwargs(model_data, xnp) + """Creates a GNM instance from a model data dictionary.""" + init_kwargs = cls._prepare_init_kwargs(model_data) # pylint: disable=no-value-for-parameter instance = super(GNM, cls).__new__(cls) # pylint: enable=no-value-for-parameter @@ -232,19 +241,6 @@ def _from_model_data_with_xnp( instance.__post_init__() return instance - @classmethod - @abc.abstractmethod - def _from_model_data( - cls, - model_data: Mapping[str, Any], - ) -> Self: - """Creates a GNM instance from a model data dictionary.""" - raise NotImplementedError( - f'{cls.__name__} is an abstract backend-agnostic base class and' - ' cannot be loaded directly. Use a concrete backend subclass (e.g.,' - ' gnm_numpy.GNM, gnm_pytorch.GNM).' - ) - def to_numpy_data_dict(self) -> dict[str, Any]: """Returns a dictionary of the GNM data represented as NumPy arrays.