Make _from_model_data concrete in gnm_xnp by introducing abstract _get_np_module method.

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,

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,

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,

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,

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.