Upload folder using huggingface_hub

README.md

@@ -1,13 +1,39 @@
 ---
-title: Gnm
-emoji: 🏃
-colorFrom: purple
+title: GNM Head — Interactive 3D Face Model
+emoji: 🧬
+colorFrom: green
 colorTo: blue
 sdk: gradio
 sdk_version: 6.20.0
-python_version: '3.13'
 app_file: app.py
-pinned: false
+short_description: Reactive 3D demo of Google's GNM parametric head model
+python_version: "3.12"
+startup_duration_timeout: 30m
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# GNM Head — Interactive 3D Face Model
+
+A reactive Gradio demo of **[GNM](https://github.com/google/GNM)** (Generative
+aNthropometric Model), Google's state-of-the-art parametric 3D statistical model
+of the human head. This Space replicates the official interactive slider demo
+from [`gnm/shape/demos/gnm_head_demo.ipynb`](https://github.com/google/GNM/blob/main/gnm/shape/README.md#demo):
+move any slider and the 3D head mesh regenerates live.
+
+## Controls
+
+- **Identity** — first 10 components of the linear identity basis (head shape).
+- **Expression** — selected blendshape components per region (left/right eye,
+  lower face, tongue, pupil).
+- **Pose** — joint rotations for neck and head (degrees), plus a gaze direction
+  control with vergence.
+- **Translation** — global position.
+
+The model is pure NumPy (a few matrix multiplies, ~40 ms per evaluation), so the
+Space runs on `cpu-basic` with no GPU required.
+
+## Credits
+
+Model and reference demo: [google/GNM](https://github.com/google/GNM),
+Apache 2.0. This Space vendors the upstream `gnm.shape` NumPy backend and its
+model data (`gnm_head.npz`).

app.py

@@ -0,0 +1,339 @@
+"""Reactive Gradio demo for the GNM (Generative aNthropometric Model) head model.
+
+Replicates the interactive slider demo from
+`gnm/shape/demos/gnm_head_demo.ipynb` in the upstream repo
+(https://github.com/google/GNM). Moving any slider re-evaluates the GNM
+mesh-generating function and updates a live 3D viewer.
+
+GNM Head is a pure-NumPy parametric statistical model of the human head.
+Evaluating it is a few matrix multiplies (~40 ms on CPU), so this Space runs
+happily on `cpu-basic` with no GPU.
+"""
+
+import re
+import tempfile
+
+import gradio as gr
+import numpy as np
+import trimesh
+from scipy.spatial.transform import Rotation
+
+from gnm.shape import gnm_numpy
+from gnm.shape.visualization import vertex_colors as vertex_colors_module
+
+# ---------------------------------------------------------------------------
+# Load the GNM head model once at module scope (fast, CPU-only).
+# ---------------------------------------------------------------------------
+GNM = gnm_numpy.GNM.from_local(
+    version=gnm_numpy.GNMMajorVersion.V3,
+    variant=gnm_numpy.GNMVariant.HEAD,
+)
+
+# Mesh topology used by the upstream viewer: all triangles except the outer
+# eye shells, so the irises/scleras are visible.
+TRIANGLES = np.asarray(GNM.triangles_group("~eye_exteriors"))
+
+# Per-vertex colors highlighting eyes, teeth, tongue (matches the notebook).
+VERTEX_COLORS = vertex_colors_module.get_vertex_colors(gnm_np=GNM)
+if "pupils" in GNM.vertex_group_names:
+    VERTEX_COLORS[GNM.vertex_group_indices("pupils")] = 0.0
+VERTEX_RGBA = (
+    np.clip(
+        np.concatenate(
+            [VERTEX_COLORS, np.ones((VERTEX_COLORS.shape[0], 1))], axis=1
+        ),
+        0.0,
+        1.0,
+    )
+    * 255
+).astype(np.uint8)
+
+
+def _find_component_indices(names, region, suffix=r"_[0-9][0-9][0-9]"):
+    """Indices of basis components whose name matches `<region><suffix>`."""
+    regex = re.compile(f"{region}{suffix}")
+    return [i for i, n in enumerate(names) if regex.match(n)]
+
+
+# ---------------------------------------------------------------------------
+# Which basis components each slider group controls (mirrors the notebook).
+# ---------------------------------------------------------------------------
+_id_names = list(GNM.identity_names)
+_ex_names = list(GNM.expression_names)
+
+IDENTITY_INDICES = _find_component_indices(_id_names, "head")[:10]
+
+_tongue_mean = _find_component_indices(_ex_names, "tongue_mean", suffix="")
+_tongue = _find_component_indices(_ex_names, "tongue")
+_all_tongue = (_tongue_mean + _tongue)
+
+EXPRESSION_GROUPS = {
+    "Left eye": _find_component_indices(_ex_names, "left_eye_region")[:3],
+    "Right eye": _find_component_indices(_ex_names, "right_eye_region")[:3],
+    "Lower face": _find_component_indices(_ex_names, "lower_face_region")[:7],
+    "Tongue": _all_tongue[:4],
+    "Pupil": _find_component_indices(_ex_names, "pupils")[:1],
+}
+# Flat ordered list of expression indices exposed as sliders.
+EXPRESSION_INDICES = [i for g in EXPRESSION_GROUPS.values() for i in g]
+
+# Pose slider limits (degrees), as in the notebook.
+POSE_LIMITS = {
+    "neck": {"X": 90, "Y": 90, "Z": 90},
+    "head": {"X": 45, "Y": 45, "Z": 15},
+    "gaze": {"X": 25, "Y": 30, "Vergence": 10},
+}
+
+NUM_IDENTITY = len(IDENTITY_INDICES)          # 10
+NUM_EXPRESSION = len(EXPRESSION_INDICES)      # 16
+NUM_POSE = 9                                  # neck(3) + head(3) + gaze(3)
+NUM_TRANSLATION = 3
+
+TOTAL_SLIDERS = NUM_IDENTITY + NUM_EXPRESSION + NUM_POSE + NUM_TRANSLATION
+
+
+# ---------------------------------------------------------------------------
+# Parameter assembly + mesh generation.
+# ---------------------------------------------------------------------------
+def _build_rotations(pose_vals):
+    """Port of the notebook's pose->axis-angle conversion.
+
+    pose_vals is a flat list of 9 degrees:
+    [neck X,Y,Z, head X,Y,Z, gaze X, gaze Y, gaze Vergence].
+    """
+    neck = pose_vals[0:3]
+    head = pose_vals[3:6]
+    gaze_x, gaze_y, verg = pose_vals[6:9]
+
+    rotations_euler = np.zeros((GNM.num_joints, 3))
+    # Joint 0 = neck, joint 1 = head.
+    rotations_euler[0] = np.radians(neck)
+    rotations_euler[1] = np.radians(head)
+
+    # Both eyes (joints 2, 3) share gaze X/Y; vergence splits them apart.
+    rotations_euler[2, 0] = np.radians(gaze_x)
+    rotations_euler[3, 0] = np.radians(gaze_x)
+    rotations_euler[2, 1] = np.radians(gaze_y) + np.radians(verg)
+    rotations_euler[3, 1] = np.radians(gaze_y) - np.radians(verg)
+
+    return np.array(
+        [Rotation.from_euler("XYZ", r).as_rotvec() for r in rotations_euler]
+    )
+
+
+def generate_head(*slider_values) -> str:
+    """Generate a GNM head mesh (.glb) from the slider parameters.
+
+    Args:
+        slider_values: flat sequence of slider values in this order:
+            10 identity, 16 expression, 9 pose (degrees), 3 translation.
+
+    Returns:
+        Path to a temporary .glb file containing the colored head mesh.
+    """
+    vals = list(slider_values)
+    idx = 0
+    id_vals = vals[idx:idx + NUM_IDENTITY]; idx += NUM_IDENTITY
+    ex_vals = vals[idx:idx + NUM_EXPRESSION]; idx += NUM_EXPRESSION
+    pose_vals = vals[idx:idx + NUM_POSE]; idx += NUM_POSE
+    trans_vals = vals[idx:idx + NUM_TRANSLATION]
+
+    identity = np.zeros(GNM.identity_dim)
+    for i, v in zip(IDENTITY_INDICES, id_vals):
+        identity[i] = float(v)
+
+    expression = np.zeros(GNM.expression_dim)
+    for i, v in zip(EXPRESSION_INDICES, ex_vals):
+        expression[i] = float(v)
+
+    rotations = _build_rotations([float(v) for v in pose_vals])
+    translation = np.array([float(v) for v in trans_vals])
+
+    vertices = np.asarray(GNM(identity, expression, rotations, translation))
+
+    mesh = trimesh.Trimesh(
+        vertices=vertices,
+        faces=TRIANGLES,
+        vertex_colors=VERTEX_RGBA,
+        process=False,
+    )
+    # GNM is Y-up / +Z-forward; rotate so the face looks toward the camera in
+    # the default Model3D view.
+    out = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
+    mesh.export(out.name)
+    return out.name
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+INTRO = """
+# 🧬 GNM — Generative aNthropometric Model (Head)
+
+Interactive replica of the [official reactive demo](https://github.com/google/GNM/blob/main/gnm/shape/README.md#demo)
+(`gnm_head_demo.ipynb`). **GNM Head** is a state-of-the-art parametric 3D
+statistical model of the human head from [google/GNM](https://github.com/google/GNM),
+with disentangled control over **identity**, **expression**, **pose** and
+**translation**. Move any slider to see it re-generate the 3D mesh live.
+
+The exposed parameters mirror the notebook: the first 10 identity (head)
+components, a few expression components per region (eyes, lower face, tongue,
+pupil), joint rotations (neck / head / gaze), and global translation.
+"""
+
+
+def _slider(label, minimum, maximum, value=0.0, step=None):
+    kw = dict(label=label, minimum=minimum, maximum=maximum, value=value)
+    if step is not None:
+        kw["step"] = step
+    return gr.Slider(**kw)
+
+
+with gr.Blocks(title="GNM Head") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(INTRO)
+
+        with gr.Row():
+            # --- Controls -------------------------------------------------
+            with gr.Column(scale=3):
+                identity_sliders = []
+                expression_sliders = []
+                pose_sliders = []
+                translation_sliders = []
+
+                with gr.Tabs():
+                    with gr.Tab("Identity"):
+                        gr.Markdown(
+                            "First 10 components of the linear **identity** "
+                            "basis (head shape). Typical range −3 … +3."
+                        )
+                        with gr.Row():
+                            for k in range(NUM_IDENTITY):
+                                identity_sliders.append(
+                                    _slider(f"i{k}", -3.0, 3.0, 0.0, 0.05)
+                                )
+
+                    with gr.Tab("Expression"):
+                        gr.Markdown(
+                            "Selected **expression** blendshape components "
+                            "per region. Typical range −3 … +3."
+                        )
+                        for group_name, group_idx in EXPRESSION_GROUPS.items():
+                            with gr.Row():
+                                gr.Markdown(f"**{group_name}**")
+                            with gr.Row():
+                                for j in range(len(group_idx)):
+                                    expression_sliders.append(
+                                        _slider(
+                                            f"{group_name[:2].lower()}{j}",
+                                            -3.0, 3.0, 0.0, 0.05,
+                                        )
+                                    )
+
+                    with gr.Tab("Pose"):
+                        gr.Markdown(
+                            "Joint **rotations** in degrees (neck / head) and "
+                            "**gaze** direction (both eyes locked, with a "
+                            "vergence control)."
+                        )
+                        for joint, limits in POSE_LIMITS.items():
+                            with gr.Row():
+                                gr.Markdown(f"**{joint.capitalize()}**")
+                            with gr.Row():
+                                for name, limit in limits.items():
+                                    pose_sliders.append(
+                                        _slider(
+                                            name, -limit, limit, 0, 1
+                                        )
+                                    )
+
+                    with gr.Tab("Translation"):
+                        gr.Markdown("Global **translation** (meters).")
+                        with gr.Row():
+                            for d in ["X", "Y", "Z"]:
+                                translation_sliders.append(
+                                    _slider(d, -0.2, 0.2, 0.0, 0.01)
+                                )
+
+                reset_btn = gr.Button("↺ Reset all sliders", variant="secondary")
+
+            # --- Viewer ---------------------------------------------------
+            with gr.Column(scale=2):
+                model_out = gr.Model3D(
+                    label="GNM head mesh",
+                    clear_color=[0.94, 0.94, 0.94, 1.0],
+                    height=560,
+                    camera_position=(0, 90, None),
+                )
+
+        all_sliders = (
+            identity_sliders
+            + expression_sliders
+            + pose_sliders
+            + translation_sliders
+        )
+
+        # Reactive: any slider change regenerates the mesh.
+        for s in all_sliders:
+            s.change(
+                fn=generate_head,
+                inputs=all_sliders,
+                outputs=model_out,
+                show_progress="hidden",
+            )
+
+        def _reset():
+            return [0.0] * len(all_sliders)
+
+        reset_btn.click(_reset, inputs=None, outputs=all_sliders).then(
+            fn=generate_head, inputs=all_sliders, outputs=model_out
+        )
+
+        # Presets: full rows (every slider gets a concrete value, no None).
+        # Offsets into the flat slider vector:
+        _EX0 = NUM_IDENTITY                                   # expression start
+        _PO0 = NUM_IDENTITY + NUM_EXPRESSION                  # pose start
+        _TR0 = NUM_IDENTITY + NUM_EXPRESSION + NUM_POSE       # translation start
+
+        def _preset(identity=None, expression=None, pose=None, translation=None):
+            row = [0.0] * TOTAL_SLIDERS
+            for i, v in (identity or {}).items():
+                row[i] = v
+            for i, v in (expression or {}).items():
+                row[_EX0 + i] = v
+            for i, v in (pose or {}).items():
+                row[_PO0 + i] = v
+            for i, v in (translation or {}).items():
+                row[_TR0 + i] = v
+            return row
+
+        presets = [
+            _preset(),  # neutral template face
+            _preset(identity={0: 2.5, 1: -2.0, 2: 1.5, 4: -1.5}),  # identity A
+            _preset(identity={0: -2.5, 3: 2.0, 5: -1.8, 6: 1.5}),  # identity B
+            # Pose: neck Y (index 1), head X (index 3), gaze X (index 6).
+            _preset(pose={1: 35, 3: -15, 6: 15}),                  # looking aside
+            _preset(pose={0: -25, 4: 20}),                         # looking down/up
+        ]
+
+        gr.Examples(
+            label="Presets (click to load)",
+            examples=presets,
+            inputs=all_sliders,
+            outputs=model_out,
+            fn=generate_head,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        # Render the neutral template on load.
+        demo.load(fn=generate_head, inputs=all_sliders, outputs=model_out)
+
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)

gnm/shape/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+

gnm/shape/data/landmarks/head_sparse_68.txt

@@ -0,0 +1,68 @@
+8777 0.068 8841 0.079 11165 0.853
+9838 0.151 9835 0.045 9836 0.804
+9986 0.037 9984 0.216 9982 0.746
+10087 0.099 10090 0.828 9996 0.074
+6846 0.866 6849 0.102 8028 0.031
+8738 0.084 8736 0.158 8737 0.758
+10085 0.084 10084 0.057 9845 0.859
+9747 0.146 9754 0.407 9740 0.447
+9975 0.156 12257 0.138 12258 0.706
+3612 0.447 3626 0.407 3619 0.146
+3854 0.746 3856 0.216 3858 0.037
+3868 0.074 3962 0.828 3959 0.099
+1900 0.031 721 0.102 718 0.866
+2609 0.758 2608 0.158 2610 0.084
+3717 0.859 3956 0.057 3957 0.084
+3708 0.804 3707 0.045 3710 0.151
+5037 0.853 2713 0.079 2649 0.068
+7635 0.443 7636 0.143 7458 0.414
+7578 0.141 7575 0.065 7574 0.794
+7093 0.226 7090 0.351 7671 0.423
+7572 0.053 7566 0.041 7565 0.907
+7111 0.004 7108 0.996 7640 0.000
+1512 0.000 980 0.996 983 0.004
+1437 0.907 1438 0.041 1444 0.053
+1543 0.423 962 0.351 965 0.226
+1446 0.794 1447 0.065 1450 0.141
+1330 0.414 1508 0.143 1507 0.443
+12319 0.995 9970 0.005 9971 0.000
+12329 0.995 6766 0.004 6764 0.000
+12310 0.996 10070 0.003 10052 0.000
+12287 0.000 12296 0.997 10055 0.003
+10105 0.995 10106 0.003 9931 0.002
+9478 0.126 10221 0.355 10225 0.519
+12297 0.000 12298 0.998 10219 0.002
+3350 0.170 3354 0.588 4097 0.242
+3803 0.002 3978 0.003 3977 0.995
+7426 0.992 7425 0.002 6708 0.006
+7288 0.996 7289 0.002 7290 0.003
+7164 0.004 7165 0.001 7404 0.995
+11028 0.008 11027 0.974 6758 0.017
+6736 0.005 6739 0.988 7181 0.008
+6722 0.002 7307 0.994 7308 0.004
+630 0.017 4899 0.974 4900 0.008
+1276 0.995 1037 0.001 1036 0.004
+1162 0.003 1161 0.002 1160 0.996
+580 0.006 1297 0.002 1298 0.992
+1180 0.004 1179 0.994 594 0.002
+1053 0.008 611 0.988 608 0.005
+10683 0.409 10653 0.584 10652 0.006
+10720 0.000 10719 0.996 10673 0.004
+10732 0.000 10731 0.997 10672 0.003
+10711 0.002 12276 0.998 12279 0.000
+4544 0.003 4603 0.997 4604 0.000
+4545 0.004 4591 0.996 4592 0.000
+4524 0.006 4525 0.584 4555 0.409
+4685 0.003 4598 0.001 4599 0.996
+4691 0.002 4593 0.995 4594 0.002
+10702 0.000 12270 1.000 12284 0.000
+10722 0.002 10721 0.995 10819 0.002
+10727 0.996 10726 0.001 10813 0.003
+11285 0.558 11286 0.312 10694 0.129
+10805 0.997 10806 0.002 10801 0.000
+10709 0.003 12285 0.997 12272 0.000
+4673 0.000 4678 0.002 4677 0.997
+4566 0.129 5158 0.312 5157 0.558
+4689 0.998 4690 0.001 4687 0.001
+12271 0.996 10704 0.004 10703 0.000
+10815 0.001 10818 0.001 10817 0.998

gnm/shape/data/versions/gnm_catalog.py

@@ -0,0 +1,21 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""GNM catalog constants."""
+
+MAINTAINED_MAJOR_VERSIONS = ["v3"]
+MAJOR_VERSION_TO_VARIANTS_MAP = {"v3": ["head"]}
+ALL_VARIANTS = ["head"]
+VARIANT_TO_MODEL_FILE_NAME_MAP = {"head": "gnm_head"}
+VERSION_TO_FULL_VERSION_MAP = {"v3": "v3_0"}

gnm/shape/data/versions/gnm_specs.py

@@ -0,0 +1,34 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""GNM versions and variants specifications."""
+
+import enum
+
+
+class GNMVersion(enum.StrEnum):
+  V3_0 = '3.0'
+
+class GNMMajorVersion(enum.StrEnum):
+  V3 = '3'
+
+class GNMVariant(enum.StrEnum):
+  HEAD = 'head'
+
+class GNMBodyPart(enum.StrEnum):
+  HEAD = 'head'
+
+GNM_VARIANT_TO_BODY_PART_MAP = {
+  GNMVariant.HEAD: GNMBodyPart.HEAD,
+}

gnm/shape/data/versions/v3_0/gnm_head.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e50a702789af51347531e13721df1e5a26dcfc30ced9e243c7cede9fcac5db43
+size 53305389

gnm/shape/fitting_utils/project_on_pca.py

@@ -0,0 +1,261 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Projection on a linear vertex basis."""
+
+import dataclasses
+
+from gnm.shape.fitting_utils import regularized_least_squares
+import numpy as np
+
+_EPSILON = 1e-16
+
+
+@dataclasses.dataclass(frozen=True)
+class PCAProjectionResult:
+  """The result of projecting a mesh on a linear vertex basis.
+
+  Attributes:
+    coefficients: The coefficients of the vertex basis. Note that if
+      `num_components` is given, then the coefficients will be truncated to that
+      number of components.
+    reconstruction: The reconstructed vertices.
+  """
+
+  coefficients: np.ndarray
+  reconstruction: np.ndarray | None = None
+
+
+def project_on_linear_vertex_basis(
+    vertices: np.ndarray,
+    mean_vertex_positions: np.ndarray,
+    vertex_basis: np.ndarray,
+    indices: np.ndarray | None = None,
+    regularization: float = 0.0,
+    num_components: int | None = None,
+    compute_reconstruction: bool = True,
+) -> PCAProjectionResult:
+  """Finds the coefficients of the vertex basis that best fit the vertices.
+
+  Solves a least-squares problem to find the coefficients of the vertex basis
+  that best fit the vertices. The formulation is:
+
+  || mu + B @ c - v ||_2^2 + lambda * ||c||_2^2
+  || B @ c - (v - mu) ||_2^2 + lambda * ||c||_2^2
+
+  where mu contains the mean vertex positions, (V, 3), B is the vertex basis we
+  are projecting on, (I, V, 3), c contains the desired coefficients, (I,), and v
+  are the target vertices, (V, 3).
+
+  Args:
+    vertices: The vertices to project, ([A1, ..., An], V, 3).
+    mean_vertex_positions: The mean vertex positions, (V, 3).
+    vertex_basis: The vertex basis, (I, V, 3).
+    indices: The indices of the vertices to project, (V,). If None, all vertices
+      are used to compute the least-squares error.
+    regularization: The regularization weight for the least-squares problem.
+    num_components: The number of components to use for the projection. If None,
+      all components are used.
+    compute_reconstruction: If True, compute the reconstructed vertices using
+      the estimated coefficients and given vertex basis.
+
+  Returns:
+    An instance of `PCAProjectionResult` containing the coefficients and the
+    reconstructed vertices. Note that the coefficients and reconstructions in
+    `PCAProjectionResult` are *always* batched.
+  """
+  if indices is None:
+    indices = np.arange(len(mean_vertex_positions))
+
+  # If the number of components is not given, use all components.
+  if num_components is None:
+    num_components = vertex_basis.shape[0]
+
+  # Add a batch dimension if the vertices are not batched.
+  if len(vertices.shape) == 2:
+    vertices = np.expand_dims(vertices, axis=0)
+
+  # Compute the offset from the mean vertex positions, ([A1, ..., An], V, 3]).
+  offsets_from_mean = vertices - mean_vertex_positions
+  offsets_from_mean = np.take(offsets_from_mean, indices, axis=-2)
+
+  # Flatten the shape offsets to (prod(A1, ..., An), V * 3).
+  num_batch_elements = np.prod(offsets_from_mean.shape[:-2])
+  offsets_from_mean = np.reshape(
+      offsets_from_mean, [num_batch_elements, len(indices) * 3]
+  )
+
+  # Transpose the shape offsets to (V * 3, prod(A1, ..., An)).
+  offsets_from_mean = np.swapaxes(offsets_from_mean, 0, 1)
+
+  # Index the basis using the selected vertex indices and the desired number of
+  # components, flatten the vertex and coordinate axis, and transpose to
+  # (V*3, I).
+  vertex_basis_subset = vertex_basis[:num_components]
+  basis_for_fitting = np.reshape(
+      vertex_basis_subset[:, indices], [num_components, len(indices) * 3]
+  )
+  basis_for_fitting = np.swapaxes(basis_for_fitting, 0, 1)
+
+  coefficients, *_ = regularized_least_squares.regularized_least_squares(
+      basis_for_fitting, offsets_from_mean, regularization
+  )
+
+  # Reshape the coefficients to ([A1, ..., An], I).
+  coefficients = np.swapaxes(coefficients, 0, 1)
+  coefficients = np.reshape(
+      coefficients, [*vertices.shape[:-2], num_components]
+  )
+
+  reconstruction = None
+  if compute_reconstruction:
+    reconstruction = mean_vertex_positions + np.einsum(
+        '...i,ivm->...vm', coefficients, vertex_basis_subset
+    )
+
+  return PCAProjectionResult(
+      coefficients=coefficients, reconstruction=reconstruction
+  )
+
+
+class PCABasisProjection:
+  """Projects vertices on a linear vertex basis.
+
+  This class pre-computes the matrix used to project vertices on a linear vertex
+  basis for repeated solves.
+  """
+
+  def __init__(
+      self,
+      mean_vertex_positions: np.ndarray,
+      vertex_basis: np.ndarray,
+      vertex_indices: np.ndarray | None = None,
+      regularization: float = 0.0,
+      num_components: int | None = None,
+      compute_reconstruction: bool = True,
+  ):
+    """Builds the instance and pre-computes the projection matrix.
+
+    Args:
+      mean_vertex_positions: The mean vertex positions for the basis, (V, 3).
+      vertex_basis: The vertex basis, (I, V, 3).
+      vertex_indices: The indices of the vertices to project, (V,). If None, all
+        vertices are used to compute the least-squares error.
+      regularization: The regularization weight for the least-squares problem.
+      num_components: The number of components to use for the projection. If
+        None, all components are used.
+      compute_reconstruction: If True, compute the reconstructed vertices using
+        the estimated coefficients and given vertex basis.
+    """
+
+    self._compute_reconstruction = compute_reconstruction
+
+    if vertex_indices is None:
+      vertex_indices = np.arange(len(mean_vertex_positions))
+
+    # If the number of components is not given, use all components.
+    if num_components is None:
+      num_components = vertex_basis.shape[0]
+
+    self._num_components = num_components
+    self._vertex_indices = vertex_indices
+    self._mean_vertex_positions = mean_vertex_positions.astype(np.float32)
+
+    regularization_sqrt = np.sqrt(regularization + _EPSILON)
+    self._regularization_sqrt = regularization_sqrt
+
+    # Index the basis using the selected vertex indices and the desired number
+    # of components, flatten the vertex and coordinate axis, and transpose to
+    # (V*3, I).
+    vertex_basis_subset = vertex_basis[:num_components]
+    self._vertex_basis_subset = vertex_basis_subset
+
+    basis_for_fitting = np.reshape(
+        vertex_basis_subset[:, vertex_indices],
+        [num_components, len(vertex_indices) * 3],
+    )
+    basis_for_fitting = np.swapaxes(basis_for_fitting, 0, 1)
+
+    # Append a diagonal array multiplied with the regularization weight.
+    left_hand_side_regularization_term = regularization_sqrt * np.eye(
+        basis_for_fitting.shape[0],
+        basis_for_fitting.shape[1],
+        dtype=basis_for_fitting.dtype,
+    )
+    left_hand_side = np.concatenate(
+        [basis_for_fitting, left_hand_side_regularization_term]
+    ).astype(np.float32)
+
+    # Contains the array used to project the vertices on the basis.
+    self._basis_projection_matrix = np.linalg.pinv(left_hand_side).astype(
+        np.float32
+    )
+
+  def __call__(
+      self,
+      vertices: np.ndarray,
+  ) -> PCAProjectionResult:
+    """Finds the coefficients of the vertex basis that best fit the vertices.
+
+    Args:
+      vertices: The vertices to project, ([A1, ..., An], V, 3).
+
+    Returns:
+      An instance of `PCAProjectionResult` containing the coefficients and the
+      reconstructed vertices.
+    """
+    # Add a batch dimension if the vertices are not batched.
+    if len(vertices.shape) == 2:
+      vertices = np.expand_dims(vertices, axis=0)
+
+    # Compute the offset from the mean vertex positions, ([A1, ..., An], V, 3]).
+    shape_offsets = (vertices - self._mean_vertex_positions).astype(np.float32)
+    shape_offsets = np.take(shape_offsets, self._vertex_indices, axis=-2)
+
+    # Flatten the shape offsets to (prod(A1, ..., An), V * 3).
+    num_batch_elements = np.prod(shape_offsets.shape[:-2])
+    shape_offsets = np.reshape(
+        shape_offsets, [num_batch_elements, len(self._vertex_indices) * 3]
+    )
+
+    # Transpose the shape offsets to (V * 3, prod(A1, ..., An)).
+    shape_offsets = np.swapaxes(shape_offsets, 0, 1)
+
+    # Append an array full of zeros to the right-hand side for the
+    # regularization term.
+    right_hand_side_regularization_term = np.zeros_like(shape_offsets)
+    right_hand_side = np.concatenate(
+        [shape_offsets, right_hand_side_regularization_term]
+    )
+
+    # Compute the coefficients, (I, [A1, ..., An]]).
+    coefficients = np.einsum(
+        'iv,v...->i...',
+        self._basis_projection_matrix,
+        right_hand_side,
+    )
+    coefficients = np.swapaxes(coefficients, 0, 1)
+    coefficients = np.reshape(
+        coefficients, [*vertices.shape[:-2], self._num_components]
+    )
+
+    reconstruction = None
+    if self._compute_reconstruction:
+      reconstruction = self._mean_vertex_positions + np.einsum(
+          '...i,ivm->...vm', coefficients, self._vertex_basis_subset
+      )
+
+    return PCAProjectionResult(
+        coefficients=coefficients, reconstruction=reconstruction
+    )

gnm/shape/fitting_utils/regularized_least_squares.py

@@ -0,0 +1,131 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""A wrapper of `np.linalg.lstsq` that supports regularization."""
+
+import numpy as np
+import numpy.typing as npt
+
+_EPSILON = 1e-16
+
+
+def regularized_least_squares(
+    left_hand_side_data_term: npt.NDArray[np.floating],
+    right_hand_side_data_term: npt.NDArray[np.floating],
+    regularization: float = 0.0,
+    rcond: float | None = None,
+) -> tuple[
+    npt.NDArray[np.floating],
+    npt.NDArray[np.floating],
+    int,
+    npt.NDArray[np.floating],
+]:
+  """Solves a regularized least-squares problem.
+
+  Returns the solution of a problem of the form:
+
+  min_x || A x - b ||_2^2 + gamma * ||x||_2^2.
+
+  To solve this problem using a standard least-squares function, we convert it
+  to:
+
+  min_x || [A; sqrt(gamma) * I] x - [b; 0] ||_2^2
+
+  where [;] means that we concatenate the two arrays along the rows.
+
+  Args:
+    left_hand_side_data_term: The left hand side array,  A, (N, M).
+    right_hand_side_data_term: The right hand side term,  b, with shape (N,) or
+      (N, K).
+    regularization: The scalar regularization that will be used.
+    rcond: Cut-off ratio for small singular values of A.
+
+  Returns:
+    A tuple that contains the least-squares solution, the residuals for each
+    data term, the rank of the left hand side matrix and its singular values.
+  """
+
+  regularization_sqrt = np.sqrt(regularization + _EPSILON)
+
+  # Append a diagonal array multiplied with the regularization weight.
+  left_hand_side_regularization_term = regularization_sqrt * np.eye(
+      left_hand_side_data_term.shape[0],
+      left_hand_side_data_term.shape[1],
+      dtype=left_hand_side_data_term.dtype,
+  )
+  left_hand_side = np.concatenate(
+      [left_hand_side_data_term, left_hand_side_regularization_term]
+  )
+
+  # Append an array full of zeros to the right-hand side for the regularization
+  # term.
+  right_hand_side_regularization_term = np.zeros_like(right_hand_side_data_term)
+  right_hand_side = np.concatenate(
+      [right_hand_side_data_term, right_hand_side_regularization_term]
+  )
+
+  return np.linalg.lstsq(left_hand_side, right_hand_side, rcond)
+
+
+class RegularizedLeastSquares:
+  """A class for repeated solves using the same left-hand side."""
+
+  def __init__(
+      self,
+      left_hand_side_data_term: npt.NDArray[np.floating],
+      regularization: float = 0.0,
+  ):
+    """Pre-computes the left-hand side of the regularized least-squares problem.
+
+    Args:
+      left_hand_side_data_term: The left hand side array,  A, (N, M).
+      regularization: The scalar regularization that will be used.
+    """
+    regularization_sqrt = np.sqrt(regularization + _EPSILON)
+
+    # Append a diagonal array multiplied with the regularization weight.
+    left_hand_side_regularization_term = regularization_sqrt * np.eye(
+        left_hand_side_data_term.shape[0],
+        left_hand_side_data_term.shape[1],
+        dtype=left_hand_side_data_term.dtype,
+    )
+    left_hand_side = np.concatenate(
+        [left_hand_side_data_term, left_hand_side_regularization_term]
+    )
+
+    # Compute the pseudo-inverse of the left-hand side.
+    self._left_hand_side = (
+        np.linalg.pinv(left_hand_side.T @ left_hand_side) @ left_hand_side.T
+    )
+
+  def solve(
+      self, right_hand_side_data_term: npt.NDArray[np.floating]
+  ) -> npt.NDArray[np.floating]:
+    """Solves the least-squares problem for a given target.
+
+    Args:
+      right_hand_side_data_term: The right hand side term,  b, with shape (N,)
+        or (N, K).
+
+    Returns:
+      The least-squares solution, x.
+    """
+    right_hand_side_regularization_term = np.zeros_like(
+        right_hand_side_data_term
+    )
+    right_hand_side = np.concatenate(
+        [right_hand_side_data_term, right_hand_side_regularization_term],
+        axis=0,
+    )
+    return self._left_hand_side @ right_hand_side

gnm/shape/gnm_base.py

@@ -0,0 +1,73 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Base GNM class."""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Mapping
+import dataclasses
+from typing import Any, Self
+
+from gnm.shape import gnm_data_loader
+from gnm.shape.data.versions import gnm_specs
+
+
+@dataclasses.dataclass(init=False)
+class GNMBase(abc.ABC):
+  """Base GNM class."""
+
+  version: gnm_specs.GNMVersion
+  variant: gnm_specs.GNMVariant
+
+  @classmethod
+  def from_local(
+      cls,
+      version: gnm_specs.GNMMajorVersion,
+      variant: gnm_specs.GNMVariant,
+  ) -> Self:
+    """Creates a GNM instance from a local model file."""
+    data_dict = gnm_data_loader.load_model_from_runfile(version, variant)
+    return cls._from_model_data(data_dict)  # pyrefly: ignore[bad-return]
+
+  @classmethod
+  def from_gnm(cls, gnm: GNMBase) -> Self:
+    """Creates a GNM instance from another GNM instance."""
+    data_dict = gnm.to_numpy_data_dict()
+    return cls._from_model_data(data_dict)  # pyrefly: ignore[bad-return]
+
+  @abc.abstractmethod
+  def to_numpy_data_dict(self) -> dict[str, Any]:
+    """Returns a dictionary of the GNM data represented as NumPy arrays."""
+    raise NotImplementedError()
+
+  @classmethod
+  @abc.abstractmethod
+  def _from_model_data(
+      cls,
+      data_dict: Mapping[str, Any],
+  ) -> Self:
+    """Creates a GNM instance from a model data."""
+    pass
+
+  @property
+  def major_version(self) -> str:
+    """Returns the major version of the model."""
+    return gnm_data_loader.full_version_to_major(self.version)
+
+  @property
+  def body_part(self) -> gnm_specs.GNMBodyPart:
+    """Returns the body part of the GNM model."""
+    return gnm_specs.GNM_VARIANT_TO_BODY_PART_MAP[self.variant]

gnm/shape/gnm_common.py

@@ -0,0 +1,505 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Backend-agnostic core math functions of the GNM model using etils.enp."""
+
+from typing import Any, Sequence
+from etils import enp
+
+enpt = enp.typing
+
+_EPSILON = 1e-8
+
+
+def take(
+    array: enpt.FloatArray['...'],
+    indices: enpt.IntArray['...'],
+    axis: int = 0,
+    xnp: Any = None,
+) -> enpt.FloatArray['...']:
+  """Extracts elements from an array along a specified axis.
+
+  This is a backend-agnostic wrapper around backend specific operations.
+  Specifically, on PyTorch it uses index_select, while on other backends it uses
+  the backend's take.
+
+  Args:
+    array: The source array to select values from.
+    indices: 1D or multi-dimensional array of indices to extract.
+    axis: The axis along which to select values.
+    xnp: Optional array module (e.g. numpy, jax.numpy, torch) to use. If not
+      provided, it is auto-detected from the array argument.
+
+  Returns:
+    The array with selected indices along the specified axis.
+  """
+  if xnp is None:
+    xnp = enp.lazy.get_xnp(array)
+  if enp.lazy.is_torch_xnp(xnp):
+    axis = axis % array.ndim
+    indices_shape = indices.shape
+    if len(indices_shape) > 1:
+      indices_flat = indices.reshape(-1)
+      taken = xnp.index_select(array, axis, indices_flat)
+      out_shape = array.shape[:axis] + indices_shape + array.shape[axis + 1 :]
+      return taken.reshape(out_shape)
+    else:
+      return xnp.index_select(array, axis, indices)
+  else:
+    return xnp.take(array, indices, axis=axis)
+
+
+def eye(
+    size: int,
+    dtype: Any,
+    reference_array: enpt.FloatArray['...'],
+    xnp: Any = None,
+) -> enpt.FloatArray['...']:
+  """Creates a 2D identity matrix.
+
+  Args:
+    size: The number of rows (and columns) in the output matrix.
+    dtype: The data type of the output matrix.
+    reference_array: An array used to determine the device (e.g. for PyTorch).
+    xnp: Optional array module (e.g. numpy, jax.numpy, torch) to use. If not
+      provided, it is auto-detected from reference_array.
+
+  Returns:
+    A 2D identity matrix of shape (size, size).
+  """
+  if xnp is None:
+    xnp = enp.lazy.get_xnp(reference_array)
+  if enp.lazy.is_torch(reference_array):
+    return xnp.eye(size, dtype=dtype, device=reference_array.device)
+  else:
+    return xnp.eye(size, dtype=dtype)
+
+
+def reshape_with_batch_dims(
+    array: enpt.FloatArray['...'],
+    target_suffix: tuple[int, ...],
+    reference_array: enpt.FloatArray['...'],
+    num_reference_non_batch_dims: int,
+) -> enpt.FloatArray['...']:
+  """Reshapes an array by prepending the batch shape of a reference array.
+
+  This function extracts the batch dimensions of `reference_array` (all but the
+  last `num_reference_non_batch_dims` dimensions) and prepends them to the
+  `target_suffix` shape. The input `array` is then reshaped to this combined
+  shape.
+
+  Args:
+    array: The array to reshape.
+    target_suffix: The shape suffix to append after the batch dimensions.
+    reference_array: The array from which to extract the batch dimensions.
+    num_reference_non_batch_dims: The number of trailing dimensions in
+      `reference_array` that represent non-batch dimensions.
+
+  Returns:
+    The reshaped array with shape `(*batch_shape, *target_suffix)`.
+  """
+  xnp = enp.lazy.get_xnp(reference_array)
+  shape = _graph_shape(reference_array)
+  batch_shape = shape[:-num_reference_non_batch_dims]
+
+  if enp.lazy.is_tf(reference_array):
+    output_shape = enp.lazy.tf.concat(
+        [batch_shape, list(target_suffix)], axis=0
+    )
+  else:
+    output_shape = batch_shape + target_suffix
+  return xnp.reshape(array, output_shape)
+
+
+def zeros_with_batch_dims(
+    reference_array: enpt.FloatArray['...'],
+    num_reference_non_batch_dims: int,
+    suffix_shape: tuple[int, ...],
+    dtype: Any,
+) -> enpt.FloatArray['...']:
+  """Creates a zeros array with the batch shape of a reference array.
+
+  This function extracts the batch dimensions of `reference_array` (all but the
+  last `num_reference_non_batch_dims` dimensions) and creates an array of zeros
+  with the combined shape `(*batch_shape, *suffix_shape)`.
+
+  Args:
+    reference_array: The array from which to extract the batch dimensions.
+    num_reference_non_batch_dims: The number of trailing dimensions in
+      `reference_array` that represent non-batch dimensions.
+    suffix_shape: The shape suffix of the output array.
+    dtype: The data type of the output array.
+
+  Returns:
+    An array of zeros of shape `(*batch_shape, *suffix_shape)` on the same
+    device as `reference_array`.
+  """
+  xnp = enp.lazy.get_xnp(reference_array)
+
+  shape = _graph_shape(reference_array)
+  batch_shape = shape[:-num_reference_non_batch_dims]
+  array_kwargs = dict(dtype=dtype)
+
+  if enp.lazy.is_tf_xnp(xnp):
+    full_shape = enp.lazy.tf.concat([batch_shape, list(suffix_shape)], axis=0)
+  else:
+    full_shape = batch_shape + suffix_shape
+  if enp.lazy.is_torch(reference_array):
+    array_kwargs['device'] = reference_array.device
+  return xnp.zeros(full_shape, **array_kwargs)
+
+
+def axis_angle_to_rotation_matrix(
+    axis_angle: enpt.FloatArray['... 3'],
+    epsilon: float = _EPSILON,
+) -> enpt.FloatArray['... 3 3']:
+  """Builds a 3x3 rotation matrix from an axis-angle vector.
+
+  Uses Rodrigues' rotation formula:
+    R = I + sin(theta) * K + (1 - cos(theta)) * K^2
+  where K is the skew-symmetric cross-product matrix of the unit rotation axis.
+
+  Args:
+    axis_angle: The rotation vector where the direction represents the rotation
+      axis and the magnitude represents the rotation angle in radians. Shape:
+      `(..., 3)`.
+    epsilon: A small value to avoid division by zero when the axis-angle vector
+      has zero magnitude.
+
+  Returns:
+    A batch of 3x3 rotation matrices. Shape: `(..., 3, 3)`.
+  """
+  xnp = enp.get_np_module(axis_angle)
+
+  norm_squared = xnp.sum(xnp.square(axis_angle), axis=-1, keepdims=True)
+  angle = xnp.sqrt(
+      xnp.maximum(norm_squared, xnp.asarray(epsilon, dtype=norm_squared.dtype))
+  )
+  axis = axis_angle / angle
+
+  sin_angle, cos_angle = xnp.sin(angle), xnp.cos(angle)
+  sin_angle, cos_angle = sin_angle[..., None], cos_angle[..., None]
+
+  matrix = xnp.broadcast_to(
+      eye(3, dtype=angle.dtype, reference_array=axis_angle, xnp=xnp),
+      (*axis_angle.shape[:-1], 3, 3),
+  )
+
+  skew_01 = -axis[..., 2]
+  skew_02 = axis[..., 1]
+  skew_10 = axis[..., 2]
+  skew_12 = -axis[..., 0]
+  skew_20 = -axis[..., 1]
+  skew_21 = axis[..., 0]
+  row0 = xnp.stack([xnp.zeros_like(skew_01), skew_01, skew_02], axis=-1)
+  row1 = xnp.stack([skew_10, xnp.zeros_like(skew_12), skew_12], axis=-1)
+  row2 = xnp.stack([skew_20, skew_21, xnp.zeros_like(skew_21)], axis=-1)
+  skew_symmetric_axis_matrix = xnp.stack([row0, row1, row2], axis=-2)
+
+  matrix = (
+      matrix
+      + sin_angle * skew_symmetric_axis_matrix
+      + (1.0 - cos_angle)
+      * xnp.matmul(skew_symmetric_axis_matrix, skew_symmetric_axis_matrix)
+  )
+
+  return matrix
+
+
+def joint_transforms_world(
+    joints: enpt.FloatArray['... J 3'],
+    rotations: enpt.FloatArray['... J 3'],
+    translation: enpt.FloatArray['... 3'],
+    joint_parent_indices: Sequence[int],
+) -> enpt.FloatArray['... J 4 4']:
+  """Computes the world-space transformation matrices for each joint.
+
+  This function traverses the joint hierarchy and applies forward kinematics
+  to compute the global/world space transforms of each joint, combining
+  joint rotations and joint translations.
+
+  Args:
+    joints: The bind-pose joints positions. Shape: `(..., J, 3)`.
+    rotations: The joint rotations as axis-angle vectors. Shape: `(..., J, 3)`.
+    translation: The translation vector for the root joint. Shape: `(..., 3)`.
+    joint_parent_indices: A sequence of parent indices for each joint. The root
+      joint must have itself or a sentinel as parent (usually index 0 has parent
+      0).
+
+  Returns:
+    A batch of 4x4 transformation matrices in world space for each joint.
+    Shape: `(..., J, 4, 4)`.
+  """
+  xnp = enp.get_np_module(joints)
+  num_joints = len(joint_parent_indices)
+
+  # Gather the parent's joints position for each joint.
+  # Under JAX/TF/NumPy, we can index with parents.
+  # But we must exclude the root (index 0 has parent 0, but is not used).
+  parent_indices_except_root = xnp.asarray(
+      joint_parent_indices[1:], dtype=xnp.int32
+  )
+  joint_parents = take(joints, parent_indices_except_root, axis=-2, xnp=xnp)
+
+  # Since we cannot do in-place assignments, we construct the local
+  # transformation matrices by combining rotation matrices and translation
+  # vectors.
+  local_rotations = axis_angle_to_rotation_matrix(rotations)
+
+  # Construct local translations.
+  root_translation = joints[..., 0, :] + translation
+  non_root_translations = joints[..., 1:, :] - joint_parents
+  local_translations = xnp.concatenate(
+      [root_translation[..., None, :], non_root_translations], axis=-2
+  )
+
+  # Combine rotation and translation.
+  # local_transforms: (..., J, 3, 4)
+  local_transforms_3x4 = xnp.concatenate(
+      [local_rotations, local_translations[..., None]], axis=-1
+  )
+
+  # Append [0, 0, 0, 1] row.
+  bottom_row = zeros_with_batch_dims(
+      joints, 2, (num_joints, 1, 4), dtype=joints.dtype
+  )
+  # Set the last element to 1.0.
+  bottom_row = bottom_row + xnp.asarray(
+      [0.0, 0.0, 0.0, 1.0], dtype=joints.dtype
+  )
+
+  local_transforms = xnp.concatenate(
+      [local_transforms_3x4, bottom_row], axis=-2
+  )
+
+  # Compute forward kinematics.
+  # We cannot loop with direct array assignment in TF/JAX graph mode.
+  # We must stack them sequentially.
+  # FK: T_world[i] = T_world[parent[i]] @ T_local[i]
+  # Since skeleton is a tree, we can build the list of transforms sequentially.
+  world_transforms_list = [local_transforms[..., 0, :, :]]
+  for joint_index in range(1, num_joints):
+    parent_transform = world_transforms_list[joint_parent_indices[joint_index]]
+    current_transform = (
+        parent_transform @ local_transforms[..., joint_index, :, :]
+    )
+    world_transforms_list.append(current_transform)
+
+  return xnp.stack(world_transforms_list, axis=-3)
+
+
+def linear_blend_skinning(
+    local_vertices: enpt.FloatArray['... V 3'],
+    joint_positions: enpt.FloatArray['... J 3'],
+    rotations: enpt.FloatArray['... J 3'],
+    translation: enpt.FloatArray['... 3'],
+    skinning_weights: enpt.FloatArray['J V'],
+    joint_parent_indices: Sequence[int],
+) -> enpt.FloatArray['... V 3']:
+  """Poses the GNM vertices using Linear Blend Skinning (LBS).
+
+  Linear Blend Skinning computes posed vertex positions by taking a weighted
+  sum of the transformations applied to each joint, where the weights are
+  defined per vertex and per joint.
+
+  Args:
+    local_vertices: The vertex positions in the local character space, with
+      shape `(..., V, 3)`.
+    joint_positions: The joint positions in bind pose, with shape `(..., J, 3)`.
+    rotations: Joint rotations as axis-angle vectors, with shape `(..., J, 3)`.
+    translation: Translation vector for the root joint, with shape: `(..., 3)`.
+    skinning_weights: Skinning weights defining joint influence per vertex with
+      shape `(J, V)`.
+    joint_parent_indices: Sequence of parent indices for each joint.
+
+  Returns:
+    The posed skinned vertices with shape `(..., V, 3)`.
+  """
+  xnp = enp.get_np_module(local_vertices)
+  num_joints = skinning_weights.shape[0]
+
+  # The local-to-world transforms of each joint, after posing, (N, J, 4, 4).
+  joint_transforms_world_local = joint_transforms_world(
+      joint_positions, rotations, translation, joint_parent_indices
+  )
+
+  # deltas = T_world[..., :3, :3] @ joints
+  deltas = xnp.einsum(
+      '...jik,...jk->...ji',
+      joint_transforms_world_local[..., :3, :3],
+      joint_positions,
+  )
+  deltas = deltas[..., None]  # (..., J, 3, 1)
+
+  offset_3x3 = zeros_with_batch_dims(
+      local_vertices, 2, (num_joints, 3, 3), dtype=local_vertices.dtype
+  )
+  bottom_row = zeros_with_batch_dims(
+      local_vertices, 2, (num_joints, 1, 4), dtype=local_vertices.dtype
+  )
+
+  offset = xnp.concatenate([offset_3x3, deltas], axis=-1)
+  offset = xnp.concatenate([offset, bottom_row], axis=-2)
+  joint_transforms = joint_transforms_world_local - offset
+
+  # Convert the vertices to homogeneous coordinates.
+  ones = (
+      zeros_with_batch_dims(
+          local_vertices,
+          2,
+          (local_vertices.shape[-2], 1),
+          dtype=local_vertices.dtype,
+      )
+      + 1.0
+  )
+  vertices_h = xnp.concatenate([local_vertices, ones], axis=-1)
+
+  # Perform Linear Blend Skinning.
+  vertices_skinned = xnp.einsum(
+      'jv,...jmn,...vn->...vm',
+      skinning_weights,
+      joint_transforms,
+      vertices_h,
+  )[..., :3]
+
+  return vertices_skinned
+
+
+def vertex_positions_bind_pose(
+    identity: enpt.FloatArray['... I'] | None,
+    expression: enpt.FloatArray['... E'] | None,
+    template_vertex_positions: enpt.FloatArray['V 3'],
+    vertex_identity_basis: enpt.FloatArray['I V 3'],
+    expression_basis: enpt.FloatArray['E V 3'],
+) -> enpt.FloatArray['... V 3']:
+  """Computes vertices in the bind pose, with identity and expression applied.
+
+  This function computes the neutral shape of the model by applying the identity
+  (shape) blendshapes and expression (expression) blendshapes to the template
+  vertices.
+
+  Args:
+    identity: The identity parameter vector. Shape: `(..., I)`.
+    expression: The expression parameter vector. Shape: `(..., E)`.
+    template_vertex_positions: The base template vertices. Shape: `(V, 3)`.
+    vertex_identity_basis: The identity displacement basis. Shape: `(I, V, 3)`.
+    expression_basis: The expression displacement basis. Shape: `(E, V, 3)`.
+
+  Returns:
+    The template vertices with shape and expression offsets applied. Shape:
+    `(..., V, 3)`.
+  """
+  xnp = enp.get_np_module(template_vertex_positions)
+
+  # Apply linear identity and expression bases.
+  identity_deltas = 0.0
+  if identity is not None:
+    identity_deltas = xnp.einsum(
+        '...i,ijk->...jk', identity, vertex_identity_basis
+    )
+
+  expression_deltas = 0.0
+  if expression is not None:
+    expression_deltas = xnp.einsum(
+        '...i,ijk->...jk', expression, expression_basis
+    )
+
+  return template_vertex_positions + identity_deltas + expression_deltas
+
+
+def joint_positions_bind_pose(
+    identity: enpt.FloatArray['... I'] | None,
+    template_joint_positions: enpt.FloatArray['J 3'],
+    joint_identity_basis: enpt.FloatArray['I J 3'],
+) -> enpt.FloatArray['... J 3']:
+  """Joint positions in the bind pose, with identity basis applied.
+
+  Computes joint locations in the bind pose by applying the identity shape
+  regressor basis to the base template joint positions.
+
+  Args:
+    identity: The identity parameter vector. Shape: `(..., I)`.
+    template_joint_positions: Base template joint positions. Shape: `(J, 3)`.
+    joint_identity_basis: Identity regressor basis for joints. Shape: `(I, J,
+      3)`.
+
+  Returns:
+    The template joints with shape offsets applied. Shape: `(..., J, 3)`.
+  """
+  xnp = enp.get_np_module(template_joint_positions)
+  deltas = 0.0
+  if identity is not None:
+    deltas = xnp.einsum('...i,ijk->...jk', identity, joint_identity_basis)
+
+  return template_joint_positions + deltas
+
+
+def compute_pose_correctives(
+    rotations: enpt.FloatArray['... J 3'] | None,
+    pose_correctives_regressor: enpt.FloatArray['F F2'] | None,
+    template_vertex_positions: enpt.FloatArray['V 3'],
+    num_joints: int,
+    num_vertices: int,
+) -> enpt.FloatArray['... V 3']:
+  """Applies pose-dependent corrective shape offsets to vertices.
+
+  Pose correctives modify the mesh vertices based on joint rotation angles to
+  fix skinning artifacts (e.g. at elbows, knees).
+
+  Args:
+    rotations: Joint rotations as axis-angle vectors, with shape `(..., J, 3)`.
+    pose_correctives_regressor: The regressor mapping pose features to vertex
+      offsets, with shape `(F, F2)`.
+    template_vertex_positions: The base template vertices, used as reference for
+      device and shape, with shape: `(V, 3)`.
+    num_joints: The total number of joints.
+    num_vertices: The total number of vertices.
+
+  Returns:
+    The corrective vertex offsets. Shape: `(..., V, 3)`.
+  """
+  xnp = enp.get_np_module(template_vertex_positions)
+
+  if pose_correctives_regressor is None or rotations is None:
+    return zeros_with_batch_dims(
+        template_vertex_positions,
+        2,
+        (num_vertices, 3),
+        dtype=template_vertex_positions.dtype,
+    )
+
+  rotation_matrices = axis_angle_to_rotation_matrix(rotations)
+
+  # Construct eye matrices.
+  eye_matrix = eye(3, dtype=rotations.dtype, reference_array=rotations, xnp=xnp)
+  # Broadcast to matching shape.
+  eye_reshaped = xnp.reshape(eye_matrix, (1,) * (rotations.ndim - 1) + (3, 3))
+  pose_features = rotation_matrices - eye_reshaped
+
+  pose_features = reshape_with_batch_dims(
+      pose_features, (num_joints * 9,), rotations, 2
+  )
+
+  pose_deltas = xnp.einsum(
+      '...f,fv->...v', pose_features, pose_correctives_regressor
+  )
+  return reshape_with_batch_dims(pose_deltas, (num_vertices, 3), rotations, 2)
+
+
+def _graph_shape(array: enpt.FloatArray['...']) -> tuple[int, ...]:
+  """Returns the shape of an array, supporting graph-mode arrays."""
+  if enp.lazy.is_tf(array):
+    return enp.lazy.tf.shape(array)
+  else:
+    return array.shape

gnm/shape/gnm_data_loader.py

@@ -0,0 +1,247 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""GNM data loader."""
+
+# from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
+import functools
+from typing import Any
+
+from absl import logging
+from etils import epath
+from gnm.shape import gnm_data_schema
+from gnm.shape.data.versions import gnm_catalog
+from gnm.shape.data.versions import gnm_specs
+import numpy as np
+
+_pkg = __package__ or 'gnm.shape'
+_MODELS_VERSIONS_DIR = epath.resource_path(f'{_pkg}.data.versions')
+_VARIANT_TO_MODEL_FILE_NAME_MAP = gnm_catalog.VARIANT_TO_MODEL_FILE_NAME_MAP
+
+
+class GNMModelDataNotLinkedError(Exception):
+  """Raised when a GNM model data is not linked into the binary."""
+
+  pass
+
+
+def _get_model_path_from_version_and_variant(
+    version: gnm_specs.GNMMajorVersion,
+    variant: gnm_specs.GNMVariant,
+) -> epath.Path:
+  """Returns the GNM model runfiles path for given variant and version."""
+  version_value = major_to_newest_full_version(version).value.replace('.', '_')
+  version_dir_name = f'v{version_value}'
+  model_file_name = f'{_VARIANT_TO_MODEL_FILE_NAME_MAP[variant]}.npz'
+  return _MODELS_VERSIONS_DIR / version_dir_name / model_file_name
+
+
+def major_to_newest_full_version(
+    major: gnm_specs.GNMMajorVersion,
+) -> gnm_specs.GNMVersion:
+  """Returns the newest GNMVersion for a given GNMMajorVersion."""
+  minors = [
+      e for e in gnm_specs.GNMVersion if e.value.split('.')[0] == major.value
+  ]
+  return sorted(minors, key=lambda e: int(e.value.split('.')[1]))[-1]
+
+
+def full_version_to_major(
+    version: gnm_specs.GNMVersion,
+) -> gnm_specs.GNMMajorVersion:
+  """Returns the major version of a GNMVersion."""
+  return gnm_specs.GNMMajorVersion(version.value.split('.')[0])
+
+
+@functools.lru_cache
+def load_model_from_runfile(
+    version: gnm_specs.GNMMajorVersion, variant: gnm_specs.GNMVariant
+) -> dict[str, Any]:
+  """Loads GNM model data from a runfile for the given version/variant."""
+  model_file = _get_model_path_from_version_and_variant(version, variant)
+
+  logging.info(
+      'Loading GNM model version %s, variant %s from runfiles: %s',
+      version,
+      variant,
+      model_file,
+  )
+  with model_file.open('rb') as f:
+    data_dict = dict(np.load(f))
+
+  # Validate the data.
+  valid, missing, extra = _validate_gnm_data(data_dict)
+  if not valid:
+    raise ValueError(
+        f'Validation failed for version {version}, variant {variant}.'
+        f' Missing: {missing}, Extra: {extra}'
+    )
+
+  return _standardize_gnm_data_types(data_dict)
+
+
+def _validate_gnm_data(
+    data: dict[str, Any],
+) -> tuple[bool, Sequence[str], Sequence[str]]:
+  """Validates the GNM data dict.
+
+  It returns any extra or missing fields and a boolean indicating if the data
+  dict has exactly the expected fields.
+
+  Args:
+    data: The GNM data dict to validate.
+
+  Returns:
+    A tuple of (bool, Sequence[str], Sequence[str]) indicating if the data dict
+    has exactly the expected fields, the missing fields and the extra fields.
+  """
+  expected_fields = gnm_data_schema.GNM_DATA_ATTRIBUTES
+  missing_fields = list(set(expected_fields) - set(data.keys()))
+  extra_fields = list(set(data.keys()) - set(expected_fields))
+  return not missing_fields and not extra_fields, missing_fields, extra_fields
+
+
+def _standardize_gnm_data_types(data: dict[str, Any]) -> dict[str, Any]:
+  """Standardizes the GNM data data types in-place.
+
+  The data loaded from the .npz model files are defined as Numpy arrays. This
+  function converts the items to their expected Python types.
+
+  Args:
+    data: The GNM data dict to standardize.
+
+  Returns:
+    The GNM data dict with standardized data types.
+  """
+  keys_to_standardize = (
+      'version',
+      'variant',
+      'identity_names',
+      'joint_names',
+      'expression_names',
+      'mesh_component_names',
+      'vertex_group_names',
+  )
+  for k in keys_to_standardize:
+    if k not in data:
+      raise ValueError(f'Required attribute {k} not found in GNM data.')
+
+  try:
+    data['version'] = gnm_specs.GNMVersion(str(data['version']))
+  except ValueError as e:
+    version_val = data['version']
+    raise ValueError(f'Unknown GNM version: {version_val}') from e
+  try:
+    data['variant'] = gnm_specs.GNMVariant(str(data['variant']))
+  except ValueError as e:
+    variant_val = data['variant']
+    raise ValueError(f'Unknown GNM variant: {variant_val}') from e
+  for key in (
+      'identity_names',
+      'joint_names',
+      'expression_names',
+      'mesh_component_names',
+      'vertex_group_names',
+  ):
+    data[key] = [str(v) for v in data[key]]
+
+  return data
+
+
+def _rename_legacy_basis_names(data: dict[str, Any]) -> None:
+  """Renames legacy identity and expression basis names to their new names."""
+  if 'identity_names' in data:
+    identity_renames = [
+        ('eyeballs_', 'eyes_'),
+    ]
+    new_identity_names = []
+    for name in data['identity_names']:
+      name_str = str(name)
+      for old_prefix, new_prefix in identity_renames:
+        if name_str.startswith(old_prefix) and not name_str.startswith(
+            new_prefix
+        ):
+          name_str = new_prefix + name_str[len(old_prefix) :]
+          break
+      new_identity_names.append(name_str)
+    data['identity_names'] = new_identity_names
+
+  if 'expression_names' in data:
+    expression_renames = [
+        ('left_eye_', 'left_eye_region_'),
+        ('right_eye_', 'right_eye_region_'),
+        ('mouth_', 'lower_face_region_'),
+        ('eyeballs_', 'pupils_'),
+    ]
+    new_expression_names = []
+    for name in data['expression_names']:
+      name_str = str(name)
+      for old_prefix, new_prefix in expression_renames:
+        if name_str.startswith(old_prefix) and not name_str.startswith(
+            new_prefix
+        ):
+          name_str = new_prefix + name_str[len(old_prefix) :]
+          break
+      new_expression_names.append(name_str)
+    data['expression_names'] = new_expression_names
+
+
+def _populate_legacy_vertex_group_aliases(data: dict[str, Any]) -> None:
+  """Populates standardized aliases for legacy vertex groups."""
+  if 'vertex_group_names' not in data or 'vertex_groups' not in data:
+    return
+
+  vertex_group_names_list = [str(v) for v in data['vertex_group_names']]
+  vertex_group_weights_array = np.array(data['vertex_groups'], dtype=np.float32)
+  if (
+      vertex_group_weights_array.ndim != 2
+      or len(vertex_group_names_list) != vertex_group_weights_array.shape[0]
+  ):
+    return
+
+  index_lookup = {name: i for i, name in enumerate(vertex_group_names_list)}
+  extra_group_names = []
+  extra_group_weights = []
+
+  # Mapping to retrieve new vertex group name by combinating legacy ones.
+  # Note that this list is not exhaustive, and mostly here to support the
+  # standard API calls in third_party/py/gnm.
+  vertex_group_mappings = {
+      'upper_teeth_and_gums': ('upper_teeth',),
+      'lower_teeth_and_gums': ('lower_teeth',),
+      'eyes': ('left_eye', 'right_eye'),
+      'eye_interiors': ('eyeball_interior',),
+      'eye_exteriors': ('eyeball_exterior',),
+      'scleras': ('sclera',),
+      'irises': ('iris',),
+      'pupils': ('pupil',),
+      'ears': ('left_ear', 'right_ear'),
+  }
+  for target_group_name, source_group_names in vertex_group_mappings.items():
+    if target_group_name not in index_lookup and all(
+        name in index_lookup for name in source_group_names
+    ):
+      source_weights_list = [
+          vertex_group_weights_array[index_lookup[name]]
+          for name in source_group_names
+      ]
+      extra_group_names.append(target_group_name)
+      extra_group_weights.append(np.maximum.reduce(source_weights_list))
+
+  if extra_group_names:
+    data['vertex_group_names'] = vertex_group_names_list + extra_group_names
+    data['vertex_groups'] = np.concatenate(
+        (vertex_group_weights_array, np.stack(extra_group_weights)), axis=0
+    )

gnm/shape/gnm_data_schema.py

@@ -0,0 +1,41 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""GNM data schema."""
+
+GNM_DATA_ATTRIBUTES = [
+    'version',
+    'variant',
+    'template_vertex_positions',
+    'template_joint_positions',
+    'vertex_identity_basis',
+    'joint_identity_basis',
+    'expression_basis',
+    'identity_names',
+    'joint_names',
+    'expression_names',
+    'joint_parent_indices',
+    'skinning_weights',
+    'quads',
+    'triangles',
+    'quad_uvs',
+    'triangle_uvs',
+    'mesh_component_names',
+    'mirror_indices',
+    'joint_regressor',
+    'pose_correctives_regressor',
+    'bone_aligned_template_joint_orientations',
+    'vertex_groups',
+    'vertex_group_names',
+]

gnm/shape/gnm_landmarks.py

@@ -0,0 +1,102 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""GNM landmarks definitions and loaders."""
+
+import dataclasses
+import enum
+from etils import epath
+from gnm.shape.data.versions import gnm_specs
+import numpy as np
+
+_pkg = __package__ or 'gnm.shape'
+_LANDMARKS_DIR = epath.resource_path(_pkg) / 'data' / 'landmarks'
+
+
+class GNMLandmarksType(enum.StrEnum):
+  """Available GNM landmarks types."""
+
+  HEAD_SPARSE_68 = 'head_sparse_68'
+
+
+_LANDMARKS_TYPE_TO_BODY_PART_MAP = {
+    GNMLandmarksType.HEAD_SPARSE_68: gnm_specs.GNMBodyPart.HEAD,
+}
+
+
+@dataclasses.dataclass(frozen=True)
+class LandmarksConfiguration:
+  """Configuration holding landmark definition indices and barycentric weights.
+  """
+
+  indices: np.ndarray
+  weights: np.ndarray
+
+
+class GNMLandmarksDataNotLinkedError(Exception):
+  """Raised when GNM landmark definition data is not linked into the binary."""
+
+  pass
+
+
+def _check_landmarks_data_linked(landmarks_type: GNMLandmarksType) -> None:
+  """Validates that the GNM landmarks file exists in the binary package."""
+  landmark_file = _get_landmark_path(landmarks_type)
+  if not landmark_file.exists():
+    file_name = f'{landmarks_type}.txt'
+    build_target = f'//third_party/py/gnm/shape/data/landmarks:{file_name}'
+    raise GNMLandmarksDataNotLinkedError(
+        f'The GNM landmark data for {landmarks_type} is not linked into the'
+        ' binary package. Please check the BUILD file and make sure to include'
+        f' \'data = ["{build_target}"]\'.'
+    )
+
+
+def _get_landmark_path(landmarks_type: GNMLandmarksType) -> epath.Path:
+  """Returns the path to the landmark definition file."""
+  file_name = f'{landmarks_type}.txt'
+  return _LANDMARKS_DIR / file_name
+
+
+def check_body_part_compatibility(
+    landmarks_type: GNMLandmarksType,
+    body_part: gnm_specs.GNMBodyPart,
+) -> None:
+  """Checks if the landmark type is compatible with the model body part."""
+  expected_body_part = _LANDMARKS_TYPE_TO_BODY_PART_MAP[landmarks_type]
+  if body_part != expected_body_part:
+    raise ValueError(
+        f'Landmark type {landmarks_type} is only compatible with GNM models'
+        f' of body part {expected_body_part}, but got'
+        f' body part {body_part}.'
+    )
+
+
+def load_landmarks(
+    landmarks_type: GNMLandmarksType,
+) -> LandmarksConfiguration:
+  """Loads landmark indices and weights for the given landmark type.
+
+  Args:
+    landmarks_type: The type of landmarks to load.
+
+  Returns:
+    A LandmarksConfiguration instance containing indices and weights arrays.
+  """
+  landmark_file = _get_landmark_path(landmarks_type)
+  with landmark_file.open('r') as f:
+    data = np.loadtxt(f)
+    indices = data[:, ::2].astype(np.int32)
+    weights = data[:, 1::2].astype(np.float32)
+  return LandmarksConfiguration(indices=indices, weights=weights)

gnm/shape/gnm_numpy.py

@@ -0,0 +1,152 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""NumPy implementation of the GNM model .
+
+Example usage:
+  ```
+  gnm = gnm_numpy.from_local(version=GNMVersion.V3, variant=GNMVariant.HEAD)
+
+  # Generate random identity, expression, rotations, translation parameters
+  identity = np.random.normal(size=gnm.identity_dim)
+  expression = np.zeros(gnm.expression_dim)
+  rotations = np.random.uniform(-1, 1, size=(gnm.num_joints, 3)) * 0.15
+  translation = np.random.uniform(-1, 1, size=(3,)) * 0.15
+
+  vertices = gnm(identity, expression, rotations, translation)
+  ```
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+import dataclasses
+from typing import Any
+
+from absl import logging
+from gnm.shape import gnm_common
+from gnm.shape import gnm_landmarks
+from gnm.shape import gnm_xnp
+from gnm.shape.data.versions import gnm_specs
+import numpy as np
+import numpy.typing as npt
+
+GNMVersion = gnm_specs.GNMVersion
+GNMMajorVersion = gnm_specs.GNMMajorVersion
+GNMVariant = gnm_specs.GNMVariant
+GNMBodyPart = gnm_specs.GNMBodyPart
+GNMLandmarksType = gnm_landmarks.GNMLandmarksType
+
+_rotation_matrix = gnm_common.axis_angle_to_rotation_matrix
+
+
+@dataclasses.dataclass(frozen=False, kw_only=True, init=False)
+class GNM(gnm_xnp.GNM):
+  """NumPy implementation of the GNM parametric model.
+
+  GNM is a mesh-generating function. Given identity, expression, joint
+  rotation, and translation parameters, it produces vertices of a mesh.
+
+  The GNM class also surfaces useful data for down-stream users, e.g. the
+  names of each expression dimension, and the topology of the mesh.
+
+  Shape dimensions are denoted:
+  * V: Number of vertices.
+  * J: Number of joints.
+  * I: Identity basis dimensionality.
+  * E: Expression basis dimensionality.
+  * Q: The number of quads in the mesh topology.
+  * T: The number of triangles, in a triangulated version of the mesh topology.
+  * G: Number of vertex groups.
+  * P: Number of separate parts of the mesh.
+
+  Attributes:
+    version: The version of the loaded GNM model.
+    variant: The variant of the loaded GNM model.
+    template_vertex_positions: Vertex positions in the template mesh, (V, 3).
+    template_joint_positions: Joint positions in the template GNM, (J, 3).
+    vertex_identity_basis: The vertex identity basis of the model, (I, V, 3).
+    joint_identity_basis: The joint identity basis of the model, (I, J, 3).
+    expression_basis: The vertex expression basis, aka blend-shapes, (E, V, 3).
+    identity_names: The name of each identity in the identity basis.
+    joint_names: The name of each joint in the skeleton.
+    expression_names: The name of each expression in the expression basis.
+    joint_parent_indices: Parent's index for each joint in the skeleton, (J).
+    skinning_weights: The model's skinning weights, (J, V).
+    quads: The mesh topology as quads, (Q, 4).
+    triangles: The mesh topology as triangles, (T, 3).
+    quad_uvs: Texture coordinates per quad, (Q, 4, 2).
+    triangle_uvs: Texture coordinates per triangle, (T, 3, 2).
+    mesh_component_names: The vertex group name corresponding to each separate
+      mesh part.
+    mirror_indices: The index of each vertex on the other side of the mesh.
+    joint_regressor: Mapping from vertices to joints, (J, V).
+    pose_correctives_regressor: Matrix for pose correctives, (9*J, 3*V).
+    bone_aligned_template_joint_orientations: The bone-aligned rotations for
+      each joint, (J, 3, 3). If they do not exist in the GNM npz, they are set
+      to the identity matrix. Note that these are not used to compute the GNM
+      joint and vertex positions.
+    vertex_groups: The weights in each vertex group, (G, V).
+    vertex_group_names: The name of each vertex group, (G,).
+    joint_connections: The joint connections of the GNM rig.
+    joint_children_indices: A dictionary that contains the joint indices of the
+      children of each joint.
+    skinning_segmentation: A decomposition of the vertices into separate parts
+      based on skinning weights.
+    num_vertices: The number of vertices in the mesh V.
+    num_joints: The number of joints in the skeleton J.
+    identity_dim: The dimensionality of the linear identity basis I.
+    expression_dim: The dimensionality of the linear expression basis E.
+    edge_list: The quad topology represented as a list of directed edges (E, 2).
+    vertex_uvs: Per-vertex UV texture coordinates shaped (V, 2).
+  """
+
+  @classmethod
+  def _from_model_data(
+      cls,
+      model_data: Mapping[str, Any],
+  ) -> GNM:
+    """Creates a GNM instance from model data."""
+    return cls._from_model_data_with_xnp(model_data, xnp=np)
+
+  def compute_vertex_normals(
+      self,
+      vertices: npt.NDArray[np.floating],
+  ) -> npt.NDArray[np.floating]:
+    """Compute vertex normals for GNM mesh."""
+    batch_dims = vertices.shape[:-2]
+    num_vertices = vertices.shape[-2]
+    vertices_flat = vertices.reshape(-1, num_vertices, 3)
+
+    face_vertices = vertices_flat[:, self.triangles, :]
+    v0 = face_vertices[..., 0, :]
+    v1 = face_vertices[..., 1, :]
+    v2 = face_vertices[..., 2, :]
+    face_normals_area = np.cross(v1 - v0, v2 - v0, axis=-1)
+
+    vertex_normals = np.zeros_like(vertices_flat)
+    np.add.at(
+        vertex_normals,
+        (slice(None), self.triangles, slice(None)),
+        face_normals_area[:, :, None, :],
+    )
+
+    normal_magnitudes = np.linalg.norm(vertex_normals, axis=-1, keepdims=True)
+    if np.any(np.isclose(normal_magnitudes, 0.0)):
+      logging.warning(
+          'Some vertex normals have zero magnitude. This is unexpected and'
+          ' likely indicates triangle collapse.'
+      )
+    vertex_normals = vertex_normals / np.maximum(normal_magnitudes, 1e-8)
+    return vertex_normals.reshape(batch_dims + (num_vertices, 3))

gnm/shape/gnm_test_utils.py

@@ -0,0 +1,101 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Test utilities for GNM, primarily for mocking TFHub access."""
+
+from gnm.shape import gnm_numpy
+import numpy as np
+import numpy.typing as npt
+
+DTypeLike = npt.DTypeLike
+
+
+def default_gnm_parameters(
+    gnm: gnm_numpy.GNM,
+    batch_shape: tuple[int, ...] | None = None,
+    dtype: DTypeLike = np.float32,
+) -> dict[str, npt.NDArray[np.floating]]:
+  """Returns default GNM parameters.
+
+  Args:
+    gnm: The GNM model to use.
+    batch_shape: The batch shape to use for the parameters.
+    dtype: The dtype to use for the parameters.
+
+  Returns:
+    A dictionary of GNM parameters.
+  """
+  if batch_shape is None:
+    batch_shape = tuple()
+
+  parameters = {
+      'identity': np.zeros(batch_shape + (gnm.identity_dim,)),
+      'expression': np.zeros(batch_shape + (gnm.expression_dim,)),
+      'rotations': np.zeros(batch_shape + (gnm.num_joints, 3)),
+      'translation': np.zeros(batch_shape + (3,)),
+  }
+
+  return {k: v.astype(dtype) for k, v in parameters.items()}
+
+
+def random_gnm_parameters(
+    gnm: gnm_numpy.GNM,
+    batch_shape: tuple[int, ...] | None = None,
+    seed: int | np.random.Generator | None = None,
+    identity_range: tuple[float, float] = (-1.0, 1.0),
+    expression_range: tuple[float, float] = (-1.0, 1.0),
+    rotation_range_degrees: tuple[float, float] = (-30.0, 30.0),
+    translation_range_m: tuple[float, float] = (-0.01, 0.01),
+    dtype: DTypeLike = np.float32,
+) -> dict[str, npt.NDArray[np.floating]]:
+  """Returns random GNM parameters sampled from a uniform distribution.
+
+  Args:
+    gnm: The GNM model to use.
+    batch_shape: The batch shape to use for the parameters.
+    seed: The random seed to use. If None, a default seed is used.
+    identity_range: The range of values to sample identity from.
+    expression_range: The range of values to sample expression from.
+    rotation_range_degrees: The range of values to sample rotation from (in
+      degrees).
+    translation_range_m: The range of values to sample translation from (in
+      meters).
+    dtype: The dtype to use for the parameters.
+
+  Returns:
+    A dictionary of GNM parameters.
+  """
+  if seed is None:
+    seed = np.random.default_rng(0)
+  if batch_shape is None:
+    batch_shape = tuple()
+
+  rng = np.random.default_rng(seed)
+
+  rotation_range_radians = np.deg2rad(rotation_range_degrees)
+
+  parameters = {
+      'identity': rng.uniform(
+          *identity_range, size=batch_shape + (gnm.identity_dim,)
+      ),
+      'expression': rng.uniform(
+          *expression_range, size=batch_shape + (gnm.expression_dim,)
+      ),
+      'rotations': rng.uniform(
+          *rotation_range_radians, size=batch_shape + (gnm.num_joints, 3)
+      ),
+      'translation': rng.uniform(*translation_range_m, size=batch_shape + (3,)),
+  }
+
+  return {k: v.astype(dtype) for k, v in parameters.items()}

gnm/shape/gnm_utils.py

@@ -0,0 +1,550 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Utility functions for manipulating GNM identity and expression basis."""
+
+from collections.abc import Mapping, Sequence
+import enum
+from typing import Any
+
+from gnm.shape import gnm_numpy
+import numpy as np
+import numpy.typing as npt
+from scipy import linalg
+
+
+class BasisType(enum.StrEnum):
+  IDENTITY = 'identity'
+  EXPRESSION = 'expression'
+
+
+BASIS_DIM_ATTRIBUTE_MAP = {
+    BasisType.IDENTITY: 'identity_dim',
+    BasisType.EXPRESSION: 'expression_dim',
+}
+
+_BASIS_NAMES_ATTRIBUTE_MAP = {
+    BasisType.IDENTITY: 'identity_names',
+    BasisType.EXPRESSION: 'expression_names',
+}
+
+_BASIS_BASIS_ATTRIBUTE_MAP = {
+    BasisType.IDENTITY: 'vertex_identity_basis',
+    BasisType.EXPRESSION: 'expression_basis',
+}
+
+
+# Used to compute the expression components stddevs.
+_EXPRESSION_REGION_VERTEX_GROUP_MAP = {
+    'left_eye_region': ['expression_basis_left_eye'],
+    'right_eye_region': ['expression_basis_right_eye'],
+    'lower_face_region': [
+        'expression_basis_mouth_nose_ears',
+        'lower_teeth_and_gums',
+        'tongue',
+    ],
+    'tongue': ['tongue'],
+    'pupils': ['eye_interiors'],
+}
+
+# Used to compute the identity components stddevs.
+_IDENTITY_REGION_VERTEX_GROUP_MAP = {}
+
+_BASIS_REGION_VERTEX_GROUP_MAP = {
+    BasisType.IDENTITY: _IDENTITY_REGION_VERTEX_GROUP_MAP,
+    BasisType.EXPRESSION: _EXPRESSION_REGION_VERTEX_GROUP_MAP,
+}
+
+_TONGUE_MEAN_NAME = 'tongue_mean'
+
+
+# gnm_utils supports only the following variants.
+_SUPPORTED_HEAD_VARIANTS = frozenset([
+    gnm_numpy.GNMVariant.HEAD,
+])
+
+_SUPPORTED_GNM_VARIANTS = frozenset(_SUPPORTED_HEAD_VARIANTS)
+
+
+def _get_major_version(version_str: str) -> int:
+  """Extracts the major version from a version string."""
+  try:
+    v = version_str.lower().lstrip('v')
+    return int(v.split('.')[0])
+  except ValueError:
+    return 1
+
+
+def validate_gnm(gnm: gnm_numpy.GNM) -> None:
+  """Validates that the GNM model is supported by gnm_utils."""
+  variant = getattr(gnm, 'variant', None)
+  if variant not in _SUPPORTED_GNM_VARIANTS:
+    raise ValueError(f"GNM variant '{variant}' is not supported by gnm_utils.")
+
+
+def _validate_conversion(
+    from_gnm: gnm_numpy.GNM,
+    to_gnm: gnm_numpy.GNM,
+    basis_type: BasisType,
+) -> None:
+  """Validates that the GNM models can be converted between each other."""
+  del basis_type
+
+  validate_gnm(from_gnm)
+  validate_gnm(to_gnm)
+
+
+def get_gnm_attribute(gnm: gnm_numpy.GNM, attribute: str) -> Any:
+  """Returns the attribute from the GNM model."""
+  if (data := getattr(gnm, attribute, None)) is None:
+    raise ValueError(
+        f'No attribute {attribute} in GNM model version {gnm.version} and'
+        f' variant {gnm.variant}.'
+    )
+  return data
+
+
+def _get_corresponding_indices(
+    basis_type: BasisType,
+    from_gnm: gnm_numpy.GNM,
+    to_gnm: gnm_numpy.GNM,
+) -> tuple[np.ndarray, np.ndarray]:
+  """Get indices mapping common names subset from `from_basis` to `to_basis`."""
+  attribute_name = _BASIS_NAMES_ATTRIBUTE_MAP[basis_type]
+  from_names = np.array(get_gnm_attribute(from_gnm, attribute_name))
+  to_names = np.array(get_gnm_attribute(to_gnm, attribute_name))
+
+  # Check uniqueness of names.
+  from_names_unique = np.unique(from_names).size == from_names.size
+  to_names_unique = np.unique(to_names).size == to_names.size
+  if not (from_names_unique and to_names_unique):
+    raise ValueError('Names must be unique.')
+
+  from_indices, to_indices = np.where(from_names[:, None] == to_names[None])
+  return from_indices.astype(np.int32), to_indices.astype(np.int32)
+
+
+def components_names(
+    basis_type: BasisType, gnm: gnm_numpy.GNM
+) -> Sequence[str]:
+  """Returns the names of the basis components."""
+  attribute_name = _BASIS_NAMES_ATTRIBUTE_MAP[basis_type]
+  return get_gnm_attribute(gnm, attribute_name)
+
+
+def region_names_from_components_names(names: Sequence[str]) -> Sequence[str]:
+  """Returns the names of the basis regions keeping the components order."""
+  regions_names = ['_'.join(n.split('_')[:-1]) for n in names]
+  _, indices = np.unique(regions_names, return_index=True)
+  return [regions_names[i] for i in np.sort(indices)]
+
+
+def region_to_component_indices_map(
+    basis_type: BasisType, gnm: gnm_numpy.GNM
+) -> Mapping[str, np.ndarray]:
+  """Return a map from basis region name to component indices."""
+  names = components_names(basis_type, gnm)
+  region_names = region_names_from_components_names(names)
+
+  region_indices_map = {}
+  for r in region_names:
+    indices = [index for index, name in enumerate(names) if name.startswith(r)]
+    region_indices_map[r] = np.array(indices, dtype=np.int32)
+  return region_indices_map
+
+
+def convert_coefficients(
+    coefficients: npt.NDArray[np.floating],
+    *,
+    basis_type: BasisType,
+    from_gnm: gnm_numpy.GNM,
+    to_gnm: gnm_numpy.GNM,
+) -> npt.NDArray[np.floating]:
+  """Convert coefficients from one GNM model to another.
+
+  Args:
+    coefficients: Coefficients corresponding to the basis in `from_gnm`, shape
+      ([A1, ..., An], C_from).
+    basis_type: Type of the basis, either identity or expression.
+    from_gnm: GNM model instance to convert from.
+    to_gnm: GNM model instance to convert to.
+
+  Returns:
+    Coefficients corresponding to the basis in `to_gnm`, shape
+    ([A1, ..., An], C_to).
+  """
+  if from_gnm is None or to_gnm is None:
+    raise ValueError('Both from_gnm and to_gnm must be provided.')
+
+  _validate_conversion(from_gnm, to_gnm, basis_type)
+
+  batch_shape = coefficients.shape[:-1]
+
+  attribute_name = BASIS_DIM_ATTRIBUTE_MAP[basis_type]
+  from_dim = get_gnm_attribute(from_gnm, attribute_name)
+  to_dim = get_gnm_attribute(to_gnm, attribute_name)
+
+  if coefficients.shape[-1] != from_dim:
+    raise ValueError(
+        f'Dimension mismatch: {coefficients.shape[-1]} vs {from_dim}.'
+    )
+
+  to_indices, from_indices = _get_corresponding_indices(
+      basis_type, from_gnm=to_gnm, to_gnm=from_gnm
+  )
+
+  # Convert.
+  converted = np.zeros(batch_shape + (to_dim,), dtype=coefficients.dtype)
+  converted[..., to_indices] = coefficients[..., from_indices]
+  return converted
+
+
+def _coefficients_to_regions(
+    coefficients: np.ndarray, basis_type: BasisType, gnm: gnm_numpy.GNM
+) -> dict[str, np.ndarray]:
+  """Splits the coefficients ([A1, ..., An], C) to face regions."""
+  validate_gnm(gnm)
+  attribute_name = BASIS_DIM_ATTRIBUTE_MAP[basis_type]
+  coeffs_dim = get_gnm_attribute(gnm, attribute_name)
+  if coefficients.shape[-1] != coeffs_dim:
+    raise ValueError(
+        f'Unexpected coefficients dimension of {coefficients.shape[-1]}.'
+        f' Expected {coeffs_dim} for the basis type {basis_type} and GNM model'
+        f' version {gnm.version} and variant {gnm.variant}.'
+    )
+
+  region_indices = region_to_component_indices_map(basis_type, gnm)
+  regions = {
+      r: coefficients[..., indices] for r, indices in region_indices.items()
+  }
+  for region_name, region_coefficients in regions.items():
+    assert region_coefficients.base is None or (
+        region_coefficients.base is not coefficients
+        and region_coefficients.base is not coefficients.base
+    ), (
+        f'Region {region_name} is a view of the original array. This is most'
+        ' likely a bug and may cause unexpected behavior.'
+    )
+  return regions
+
+
+def _regions_to_coefficients(
+    regions: Mapping[str, np.ndarray],
+    basis_type: BasisType,
+    gnm: gnm_numpy.GNM,
+) -> np.ndarray:
+  """Combines face regions to a single basis array ([A1, ..., An], C])."""
+  validate_gnm(gnm)
+  attribute_name = BASIS_DIM_ATTRIBUTE_MAP[basis_type]
+  coeffs_dim = get_gnm_attribute(gnm, attribute_name)
+
+  if not regions:
+    return np.zeros((coeffs_dim,), dtype=np.float32)
+
+  region_indices = region_to_component_indices_map(basis_type, gnm)
+
+  tmp_region_name = list(regions.keys())[0]
+  batch_shape = regions[tmp_region_name].shape[:-1]
+  dtype = regions[tmp_region_name].dtype
+
+  coeffs = np.zeros(batch_shape + (coeffs_dim,), dtype=dtype)
+  for region_name, values in regions.items():
+    if region_name not in region_indices:
+      gnm_type_name = getattr(gnm, 'model_type', 'unknown')
+      raise ValueError(f'No region {region_name} in model {gnm_type_name}.')
+    else:
+      coeffs[..., region_indices[region_name]] = values
+
+  return coeffs
+
+
+def _region_components(
+    basis_type: BasisType, gnm: gnm_numpy.GNM
+) -> dict[str, np.ndarray]:
+  """Return per-region basis components, each region shape (Ri, V, 3)."""
+  validate_gnm(gnm)
+  attribute_name = _BASIS_BASIS_ATTRIBUTE_MAP[basis_type]
+  basis = get_gnm_attribute(gnm, attribute_name)
+  region_indices = region_to_component_indices_map(basis_type, gnm)
+  return {r: basis[indices] for r, indices in region_indices.items()}
+
+
+def compute_scaling_for_basis(
+    vertex_basis: npt.NDArray[np.floating],
+) -> npt.NDArray[np.floating]:
+  """Computes the scaling matrix to make basis equivalent to vertex losses.
+
+  Args:
+    vertex_basis: The basis to compute the scaling for, shape (D, V, 3).
+
+  Returns:
+    A scaling matrix to apply to differences between basis components,
+    shape (D, D).
+  """
+  basis_dim = vertex_basis.shape[0]
+  vertex_basis = vertex_basis.reshape(basis_dim, -1).astype(np.float64)
+
+  basis_dot_transpose = vertex_basis @ vertex_basis.T
+  return linalg.sqrtm(basis_dot_transpose).astype(np.float32)
+
+
+def scaling_for_expression_loss(
+    gnm: gnm_numpy.GNM,
+) -> npt.NDArray[np.floating]:
+  """Scaling matrix to make expression losses equivalent to vertex losses.
+
+  The loss between expression components is: || e1 - e2 ||^2.
+  The loss between vertices is: || v1 - v2 ||^2.
+  || v1 - v2 ||^2 = || (template + B_v @ e1 - (template + B_v @ c2) ||^2 =
+  || B_v @ (c1 - c2) ||^2 =
+  (c1 - c2)^T @ B_v^T @ B_v @ (c1 - c2) =
+  (c1 - c2)^T @ M @ (c1 - c2), M = B_v^T @ B_v (1)
+
+  This if we want to make the expression loss equivalent to the vertex loss,
+  we need to scale the expression loss by: M = matrix_sqrt(B_v^T @ B_v).
+
+
+  Args:
+    gnm: The GNM model instance to compute the scaling for.
+
+  Returns:
+    A scaling matrix to apply to differences between expression components,
+    shape (E, E).
+  """
+  validate_gnm(gnm)
+  return compute_scaling_for_basis(
+      get_gnm_attribute(gnm, _BASIS_BASIS_ATTRIBUTE_MAP[BasisType.EXPRESSION])
+  )
+
+
+def scaling_for_identity_loss(
+    gnm: gnm_numpy.GNM,
+) -> npt.NDArray[np.floating]:
+  """Scaling matrix to make identity losses equivalent to vertex losses."""
+  validate_gnm(gnm)
+  return compute_scaling_for_basis(
+      get_gnm_attribute(gnm, _BASIS_BASIS_ATTRIBUTE_MAP[BasisType.IDENTITY])
+  )
+
+
+def identity_to_regions(
+    identity: np.ndarray, gnm: gnm_numpy.GNM
+) -> dict[str, np.ndarray]:
+  """Splits the identity ([A1, ..., An], I]) to face regions."""
+  return _coefficients_to_regions(
+      identity, basis_type=BasisType.IDENTITY, gnm=gnm
+  )
+
+
+def expression_to_regions(
+    expression: np.ndarray, gnm: gnm_numpy.GNM
+) -> dict[str, np.ndarray]:
+  """Splits the expression ([A1, ..., An], E]) to face regions."""
+  return _coefficients_to_regions(
+      expression, basis_type=BasisType.EXPRESSION, gnm=gnm
+  )
+
+
+def expression_regions_indices(
+    gnm: gnm_numpy.GNM,
+) -> dict[str, np.ndarray]:
+  """Returns mapping from region name to expression component indices."""
+  validate_gnm(gnm)
+  return dict(region_to_component_indices_map(BasisType.EXPRESSION, gnm))
+
+
+def identity_regions_indices(
+    gnm: gnm_numpy.GNM,
+) -> dict[str, np.ndarray]:
+  """Returns mapping from region name to identity component indices."""
+  validate_gnm(gnm)
+  return dict(region_to_component_indices_map(BasisType.IDENTITY, gnm))
+
+
+def identity_dim(gnm: gnm_numpy.GNM) -> int:
+  """Returns the identity dimension of the GNM model."""
+  validate_gnm(gnm)
+  return get_gnm_attribute(gnm, 'identity_dim')
+
+
+def expression_dim(gnm: gnm_numpy.GNM) -> int:
+  """Returns the expression dimension of the GNM model."""
+  validate_gnm(gnm)
+  return get_gnm_attribute(gnm, 'expression_dim')
+
+
+def expression_regions_dims(gnm: gnm_numpy.GNM) -> dict[str, int]:
+  """Returns mapping from region name to expression component dimension."""
+  validate_gnm(gnm)
+  region_indices = region_to_component_indices_map(BasisType.EXPRESSION, gnm)
+  return {k: v.size for k, v in region_indices.items()}
+
+
+def identity_regions_dims(gnm: gnm_numpy.GNM) -> dict[str, int]:
+  """Returns mapping from region name to identity component dimension."""
+  validate_gnm(gnm)
+  region_indices = region_to_component_indices_map(BasisType.IDENTITY, gnm)
+  return {k: v.size for k, v in region_indices.items()}
+
+
+def joint_rotations_to_regions(
+    joint_rotations: np.ndarray, gnm: gnm_numpy.GNM
+) -> dict[str, np.ndarray]:
+  """Splits the joint rotations ([A1, ..., An], J, 3]) to invidual regions."""
+  validate_gnm(gnm)
+  joint_names: list[str] = get_gnm_attribute(gnm, 'joint_names')
+  if len(joint_names) != joint_rotations.shape[-2]:
+    raise ValueError(
+        f'Number of joints in the GNM model ({len(joint_names)}) does not'
+        ' match the number of joints in the input rotations'
+        f' ({joint_rotations.shape[-2]}).'
+    )
+  return {
+      joint_name: np.copy(joint_rotations[..., joint_index, :])
+      for joint_index, joint_name in enumerate(joint_names)
+  }
+
+
+def regions_to_identity(
+    regions: Mapping[str, np.ndarray], gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Combines face regions to single identity array ([A1, ..., An], I])."""
+  return _regions_to_coefficients(regions, BasisType.IDENTITY, gnm)
+
+
+def regions_to_expression(
+    regions: Mapping[str, np.ndarray], gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Combines face regions to single expression array ([A1, ..., An], E])."""
+  return _regions_to_coefficients(regions, BasisType.EXPRESSION, gnm)
+
+
+def regions_to_joint_rotations(
+    regions: Mapping[str, np.ndarray], gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Combines joint rotations to single rotations array.
+
+  Returns array of shape ([A1, ..., An], J, 3]).
+  """
+  validate_gnm(gnm)
+  joint_names: list[str] = get_gnm_attribute(gnm, 'joint_names')
+  if not regions:
+    raise ValueError('No regions provided.')
+
+  tmp_region_name = list(regions.keys())[0]
+  batch_shape = regions[tmp_region_name].shape[:-1]
+  dtype = regions[tmp_region_name].dtype
+
+  missing_joint_value = np.zeros((*batch_shape, 3), dtype=dtype)
+  return np.stack(
+      [
+          regions.get(
+              joint_name,
+              missing_joint_value,
+          )
+          for joint_name in joint_names
+      ],
+      axis=-2,
+  )
+
+
+def region_identity_components(
+    gnm: gnm_numpy.GNM,
+) -> dict[str, np.ndarray]:
+  """Return per-region identity components, each region shape (Ri, V, 3)."""
+  return _region_components(BasisType.IDENTITY, gnm)
+
+
+def region_expression_components(
+    gnm: gnm_numpy.GNM,
+) -> dict[str, np.ndarray]:
+  """Return per-region expression components, each region shape (Ri, V, 3)."""
+  return _region_components(BasisType.EXPRESSION, gnm)
+
+
+def _sigmas(basis_type: BasisType, gnm: gnm_numpy.GNM) -> np.ndarray:
+  """Return the standard deviations of basis components, shape (C, )."""
+  validate_gnm(gnm)
+
+  if basis_type == BasisType.IDENTITY:
+    raise NotImplementedError('Identity sigmas are not implemented yet.')
+
+  region_group = _BASIS_REGION_VERTEX_GROUP_MAP[basis_type]
+  regions = {}
+  for name, components in _region_components(basis_type, gnm).items():
+    vertex_indices = gnm.vertex_group_indices(*region_group[name])
+    regions[name] = np.linalg.norm(components[:, vertex_indices], axis=(1, 2))
+  region_sigmas = _regions_to_coefficients(regions, basis_type, gnm)
+
+  # Treat tongue mean in expressions individually, it is not a real component.
+  if basis_type == BasisType.EXPRESSION:
+    expression_names = gnm.expression_names
+    if _TONGUE_MEAN_NAME in expression_names:
+      # Convert to list to use .index()
+      expression_names_list = list(expression_names)
+      region_sigmas[expression_names_list.index(_TONGUE_MEAN_NAME)] = 1.0
+
+  if np.any(np.isclose(region_sigmas, 0.0, atol=1e-8)):
+    raise RuntimeError('Expression components have zero standard deviation.')
+
+  return region_sigmas
+
+
+def expression_sigmas(gnm: gnm_numpy.GNM) -> np.ndarray:
+  """Return the standard deviations of expression components, shape (E, )."""
+  return _sigmas(BasisType.EXPRESSION, gnm)
+
+
+def _normalize_coefficients(
+    coefficients: np.ndarray, basis_type: BasisType, gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Normalize coefficients ([A1, ..., An], C) by corresponding sigmas."""
+  s = _sigmas(basis_type, gnm)
+  return coefficients / s
+
+
+def _denormalize_coefficients(
+    coefficients: np.ndarray, basis_type: BasisType, gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Denormalize coefficients ([A1, ..., An], C) by corresponding sigmas."""
+  s = _sigmas(basis_type, gnm)
+  return coefficients * s
+
+
+def normalize_expression(
+    expression: np.ndarray, gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Normalize expression coefficients ([A1, ..., An], E)."""
+  return _normalize_coefficients(expression, BasisType.EXPRESSION, gnm)
+
+
+def denormalize_expression(
+    expression: np.ndarray, gnm: gnm_numpy.GNM
+) -> np.ndarray:
+  """Denormalize expression coefficients ([A1, ..., An], E)."""
+  return _denormalize_coefficients(expression, BasisType.EXPRESSION, gnm)
+
+
+def identity_region_names(gnm: gnm_numpy.GNM) -> Sequence[str]:
+  """Gets identity basis region names ordered as they appear in the basis."""
+  validate_gnm(gnm)
+  names = components_names(BasisType.IDENTITY, gnm)
+  return region_names_from_components_names(names)
+
+
+def expression_region_names(gnm: gnm_numpy.GNM) -> Sequence[str]:
+  """Gets expression basis region names ordered as they appear in the basis."""
+  validate_gnm(gnm)
+  names = components_names(BasisType.EXPRESSION, gnm)
+  return region_names_from_components_names(names)

gnm/shape/gnm_xnp.py

@@ -0,0 +1,738 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Unified backend-agnostic implementation of the GNM model using etils.enp."""
+
+from __future__ import annotations
+
+import abc
+import collections
+from collections.abc import Mapping, Sequence
+import dataclasses
+import functools
+from typing import Any, Self
+
+from absl import logging
+from etils import enp
+from gnm.shape import gnm_base
+from gnm.shape import gnm_common
+from gnm.shape import gnm_landmarks
+from gnm.shape.data.versions import gnm_specs
+import numpy as np
+import numpy.typing as npt
+
+enpt = enp.typing
+
+_NONZERO_THRESHOLD = 1e-4
+_EPSILON = 1e-8
+
+
+@dataclasses.dataclass(frozen=False, kw_only=True, init=False)
+class GNM(gnm_base.GNMBase):
+  """Backend-agnostic implementation of the GNM parametric model.
+
+  This class is abstract. Use a concrete backend subclass (gnm_numpy.GNM,
+  gnm_pytorch.GNM, or gnm_jax.GNM).
+
+  GNM is a mesh-generating function. Given identity, expression, joint
+  rotation, and translation parameters, it produces vertices of a mesh.
+
+  The GNM class also surfaces useful data for down-stream users, e.g. the
+  names of each expression dimension, and the topology of the mesh.
+
+  Shape dimensions are denoted:
+  * V: Number of vertices.
+  * J: Number of joints.
+  * I: Identity basis dimensionality.
+  * E: Expression basis dimensionality.
+  * Q: The number of quads in the mesh topology.
+  * T: The number of triangles, in a triangulated version of the mesh topology.
+  * G: Number of vertex groups.
+  * P: Number of separate parts of the mesh.
+
+  Attributes:
+    version: The version of the loaded GNM model.
+    variant: The variant of the loaded GNM model.
+    cl_number: The CL used to create this model.
+    template_vertex_positions: Vertex positions in the template mesh, (V, 3).
+    template_joint_positions: Joint positions in the template GNM, (J, 3).
+    vertex_identity_basis: The vertex identity basis of the model, (I, V, 3).
+    joint_identity_basis: The joint identity basis of the model, (I, J, 3).
+    expression_basis: The vertex expression basis, aka blend-shapes, (E, V, 3).
+    identity_names: The name of each identity in the identity basis.
+    joint_names: The name of each joint in the skeleton.
+    expression_names: The name of each expression in the expression basis.
+    joint_parent_indices: Parent's index for each joint in the skeleton, (J).
+    skinning_weights: The model's skinning weights, (J, V).
+    quads: The mesh topology as quads, (Q, 4).
+    triangles: The mesh topology as triangles, (T, 3).
+    quad_uvs: Texture coordinates per quad, (Q, 4, 2).
+    triangle_uvs: Texture coordinates per triangle, (T, 3, 2).
+    mesh_component_names: The vertex group name corresponding to each separate
+      mesh part.
+    mirror_indices: The index of each vertex on the other side of the mesh.
+    joint_regressor: Mapping from vertices to joints, (J, V).
+    pose_correctives_regressor: Matrix for pose correctives, (9*J, 3*V).
+    bone_aligned_template_joint_orientations: The bone-aligned rotations for
+      each joint, (J, 3, 3). If they do not exist in the GNM npz, they are set
+      to the identity matrix. Note that these are not used to compute the GNM
+      joint and vertex positions.
+    vertex_groups: The weights in each vertex group, (G, V).
+    vertex_group_names: The name of each vertex group, (G,).
+    joint_connections: The joint connections of the GNM rig.
+    joint_children_indices: A dictionary that contains the joint indices of the
+      children of each joint.
+    skinning_segmentation: A decomposition of the vertices into separate parts
+      based on skinning weights.
+    num_vertices: The number of vertices in the mesh V.
+    num_joints: The number of joints in the skeleton J.
+    identity_dim: The dimensionality of the linear identity basis I.
+    expression_dim: The dimensionality of the linear expression basis E.
+    edge_list: The quad topology represented as a list of directed edges (E, 2).
+    vertex_uvs: Per-vertex UV texture coordinates shaped (V, 2).
+  """
+
+  _shape_error_type = ValueError
+
+  version: gnm_specs.GNMVersion
+  variant: gnm_specs.GNMVariant
+  template_vertex_positions: enpt.FloatArray
+  template_joint_positions: enpt.FloatArray
+  vertex_identity_basis: enpt.FloatArray
+  joint_identity_basis: enpt.FloatArray
+  expression_basis: enpt.FloatArray
+  identity_names: Sequence[str]
+  joint_names: Sequence[str]
+  expression_names: Sequence[str]
+  joint_parent_indices: Sequence[int]
+  skinning_weights: enpt.FloatArray
+  quads: enpt.IntArray
+  triangles: enpt.IntArray
+  quad_uvs: enpt.FloatArray
+  triangle_uvs: enpt.FloatArray
+  mesh_component_names: Sequence[str]
+  mirror_indices: enpt.IntArray
+  joint_regressor: enpt.FloatArray
+  pose_correctives_regressor: enpt.FloatArray
+  bone_aligned_template_joint_orientations: enpt.FloatArray
+  vertex_groups: enpt.FloatArray
+  vertex_group_names: Sequence[str]
+
+  def __init__(self, *args: Any, **kwargs: Any) -> None:
+    del args, kwargs
+    raise TypeError(
+        f'{self.__class__.__name__} cannot be instantiated directly via'
+        ' constructor. Please use a class factory method such as from_local()'
+        ' or from_model_data().'
+    )
+
+  def __post_init__(self):
+    self._vertex_group_names_lookup: dict[str, int] = {
+        name: i for i, name in enumerate(self.vertex_group_names)
+    }
+    self._xnp = enp.get_np_module(self.template_vertex_positions)
+    self._landmarks: dict[
+        gnm_landmarks.GNMLandmarksType, tuple[enpt.IntArray, enpt.FloatArray]
+    ] = {}
+
+  @property
+  def xnp(self) -> enp.NpModule:
+    """Returns the backend module dynamically from G-Nome array state."""
+    return self._xnp
+
+  def __getstate__(self):
+    state = self.__dict__.copy()
+    state.pop('_xnp', None)
+    return state
+
+  def __setstate__(self, state):
+    self.__dict__.update(state)
+    self._xnp = enp.get_np_module(self.template_vertex_positions)
+
+  @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."""
+    init_kwargs = {}
+
+    # Identify type convert functions.
+    def as_float_array(val):
+      return xnp.asarray(val, dtype=xnp.float32) if val is not None else None
+
+    def as_int_array(val):
+      return xnp.asarray(val, dtype=xnp.int32) if val is not None else None
+
+    def as_original(val):
+      return val
+
+    field_converters = {
+        'version': as_original,
+        'variant': as_original,
+        'cl_number': as_original,
+        'template_vertex_positions': as_float_array,
+        'template_joint_positions': as_float_array,
+        'vertex_identity_basis': as_float_array,
+        'joint_identity_basis': as_float_array,
+        'expression_basis': as_float_array,
+        'identity_names': as_original,
+        'joint_names': as_original,
+        'expression_names': as_original,
+        'joint_parent_indices': as_original,
+        'skinning_weights': as_float_array,
+        'quads': as_int_array,
+        'triangles': as_int_array,
+        'quad_uvs': as_float_array,
+        'triangle_uvs': as_float_array,
+        'mesh_component_names': as_original,
+        'mirror_indices': as_int_array,
+        'joint_regressor': as_float_array,
+        'pose_correctives_regressor': as_float_array,
+        'bone_aligned_template_joint_orientations': as_float_array,
+        'vertex_groups': as_float_array,
+        'vertex_group_names': as_original,
+    }
+
+    for field_name, converter in field_converters.items():
+      val = model_data.get(field_name, None)
+      init_kwargs[field_name] = converter(val)
+
+    return init_kwargs
+
+  @classmethod
+  def _from_model_data_with_xnp(
+      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)
+    # pylint: disable=no-value-for-parameter
+    instance = super(GNM, cls).__new__(cls)
+    # pylint: enable=no-value-for-parameter
+    for k, v in init_kwargs.items():
+      object.__setattr__(instance, k, v)
+    instance.__post_init__()
+    return instance
+
+  @classmethod
+  @abc.abstractmethod
+  def _from_model_data(
+      cls,
+      data_dict: 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."""
+    data_dict = {}
+    for field in dataclasses.fields(self):
+      val = getattr(self, field.name)
+      if enp.lazy.is_array(val):
+        # Convert JAX, PyTorch or TF tensors to numpy.
+        if enp.lazy.is_tf(val):
+          val = val.numpy()
+        elif enp.lazy.is_jax(val):
+          val = np.array(val)
+        elif enp.lazy.is_torch(val):
+          val = val.detach().cpu().numpy()
+      data_dict[field.name] = val
+    return data_dict
+
+  def _check_parameter_shapes(
+      self,
+      identity: enpt.FloatArray['...'] | None = None,
+      expression: enpt.FloatArray['...'] | None = None,
+      rotations: enpt.FloatArray['...'] | None = None,
+      translation: enpt.FloatArray['...'] | None = None,
+  ) -> None:
+    error_type = self._shape_error_type
+    if identity is not None:
+      if identity.ndim < 1 or identity.shape[-1] != self.identity_dim:
+        raise error_type(
+            f'identity shape mismatch: expected last dim {self.identity_dim},'
+            f' got {identity.shape}'
+        )
+    if expression is not None:
+      if expression.ndim < 1 or expression.shape[-1] != self.expression_dim:
+        raise error_type(
+            'expression shape mismatch: expected last dim'
+            f' {self.expression_dim}, got {expression.shape}'
+        )
+    if rotations is not None:
+      if (
+          rotations.ndim < 2
+          or rotations.shape[-2] != self.num_joints
+          or rotations.shape[-1] != 3
+      ):
+        raise error_type(
+            'rotations shape mismatch: expected last dims'
+            f' {(self.num_joints, 3)}, got {rotations.shape}'
+        )
+    if translation is not None:
+      if translation.ndim < 1 or translation.shape[-1] != 3:
+        raise error_type(
+            'translation shape mismatch: expected last dim 3, got'
+            f' {translation.shape}'
+        )
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def __call__(
+      self,
+      identity: enpt.FloatArray['A1 ... An I'] | None = None,
+      expression: enpt.FloatArray['A1 ... An E'] | None = None,
+      rotations: enpt.FloatArray['A1 ... An J 3'] | None = None,
+      translation: enpt.FloatArray['A1 ... An 3'] | None = None,
+  ) -> enpt.FloatArray['A1 ... An V 3']:
+    """Evaluates the GNM mesh-generating function."""
+    self._check_parameter_shapes(identity, expression, rotations, translation)
+    batch_shape = _check_batch_dims(
+        identity, expression, rotations, translation
+    )
+
+    # Fill in missing parameters.
+    if identity is None:
+      identity = gnm_common.zeros_with_batch_dims(
+          self.template_vertex_positions,
+          2,
+          (*batch_shape, self.identity_dim),
+          dtype=self.template_vertex_positions.dtype,
+      )
+    if expression is None:
+      expression = gnm_common.zeros_with_batch_dims(
+          self.template_vertex_positions,
+          2,
+          (*batch_shape, self.expression_dim),
+          dtype=self.template_vertex_positions.dtype,
+      )
+    if rotations is None:
+      rotations = gnm_common.zeros_with_batch_dims(
+          self.template_vertex_positions,
+          2,
+          (*batch_shape, self.num_joints, 3),
+          dtype=self.template_vertex_positions.dtype,
+      )
+    if translation is None:
+      translation = gnm_common.zeros_with_batch_dims(
+          self.template_vertex_positions,
+          2,
+          (*batch_shape, 3),
+          dtype=self.template_vertex_positions.dtype,
+      )
+
+    # Bind pose vertex positions.
+    vertices = self.vertex_positions_bind_pose(identity, expression)
+
+    # Bind pose joint positions.
+    joints = self.joint_positions_bind_pose(identity)
+
+    # Pose correctives.
+    pose_correctives = self.compute_pose_correctives(rotations)
+    vertices = vertices + pose_correctives
+
+    return self.vertex_positions_world(vertices, joints, rotations, translation)
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def vertices_and_landmarks(
+      self,
+      landmarks_type: gnm_landmarks.GNMLandmarksType,
+      identity: enpt.FloatArray['A1 ... An I'] | None = None,
+      expression: enpt.FloatArray['A1 ... An E'] | None = None,
+      rotations: enpt.FloatArray['A1 ... An J 3'] | None = None,
+      translation: enpt.FloatArray['A1 ... An 3'] | None = None,
+  ) -> tuple[
+      enpt.FloatArray['A1 ... An V 3'],
+      enpt.FloatArray['A1 ... An L 3'],
+  ]:
+    """Evaluates the GNM mesh function and extracts 3D landmarks."""
+    gnm_landmarks.check_body_part_compatibility(landmarks_type, self.body_part)
+    if landmarks_type not in self._landmarks:
+      config = gnm_landmarks.load_landmarks(landmarks_type)
+      self._landmarks[landmarks_type] = (
+          self.xnp.asarray(config.indices, dtype=self.xnp.int32),
+          self.xnp.asarray(config.weights, dtype=self.xnp.float32),
+      )
+    indices, weights = self._landmarks[landmarks_type]
+
+    vertices = self(
+        identity=identity,
+        expression=expression,
+        rotations=rotations,
+        translation=translation,
+    )
+    weights = enp.compat.astype(weights, vertices.dtype)
+
+    face_vertices = gnm_common.take(vertices, indices, axis=-2, xnp=self.xnp)
+    landmarks = self.xnp.sum(face_vertices * weights[..., None], axis=-2)
+    return vertices, landmarks
+
+  def vertex_positions_world(
+      self,
+      vertices: enpt.FloatArray['A1 ... An V 3'],
+      joints: enpt.FloatArray['A1 ... An J 3'],
+      rotations: enpt.FloatArray['A1 ... An J 3'],
+      translation: enpt.FloatArray['A1 ... An 3'],
+  ) -> enpt.FloatArray['A1 ... An V 3']:
+    """Applies linear blend skinning to GNM vertices."""
+    return gnm_common.linear_blend_skinning(
+        local_vertices=vertices,
+        joint_positions=joints,
+        rotations=rotations,
+        translation=translation,
+        skinning_weights=self.skinning_weights,
+        joint_parent_indices=self.joint_parent_indices,
+    )
+
+  apply_linear_blend_skinning = vertex_positions_world
+
+  @property
+  def num_vertices(self) -> int:
+    return self.template_vertex_positions.shape[0]
+
+  @property
+  def num_joints(self) -> int:
+    return self.skinning_weights.shape[0]
+
+  @property
+  def identity_dim(self) -> int:
+    return self.vertex_identity_basis.shape[0]
+
+  @property
+  def expression_dim(self) -> int:
+    return self.expression_basis.shape[0]
+
+  @functools.cached_property
+  def edge_list(self) -> npt.NDArray[np.integer]:
+    """The quad topology represented as a list of directed edges."""
+    quads = np.array(self.quads)
+    e1 = quads.ravel()
+    e2 = np.roll(quads, -1, axis=1).ravel()
+    edges = np.stack([np.minimum(e1, e2), np.maximum(e1, e2)], axis=1)
+    edge_keys = e1 + self.num_vertices * e2
+    _, unique_indices = np.unique(edge_keys, return_index=True)
+    unique_undirected_edges = edges[unique_indices]
+    return np.vstack(
+        [unique_undirected_edges, np.fliplr(unique_undirected_edges)]
+    )
+
+  @functools.cached_property
+  def joint_connections(self) -> npt.NDArray[np.int32]:
+    """The bones of the GNM rig."""
+    return np.stack(
+        [
+            np.array(self.joint_parent_indices),
+            np.arange(self.num_joints, dtype=np.int32),
+        ],
+        axis=1,
+    )[1:]
+
+  @functools.cached_property
+  def joint_children_indices(self) -> Mapping[int, Sequence[int]]:
+    children = collections.defaultdict(list)
+    for joint, parent in enumerate(self.joint_parent_indices):
+      if joint == 0:
+        continue
+      children[parent].append(joint)
+    return dict(children)
+
+  @functools.cached_property
+  def vertex_uvs(self) -> npt.NDArray[np.floating]:
+    """Per-vertex UV texture coordinates."""
+    logging.warning(
+        'These do not allow correct texturing across UV-seams. Please use '
+        'per-quad or per-triangle UV coordinates when possible.'
+    )
+    quads = np.array(self.quads)
+    quad_uvs = np.array(self.quad_uvs)
+    indices = quads.ravel()
+    uvs = quad_uvs.reshape(-1, 2)
+    _, last_indices_reversed = np.unique(indices[::-1], return_index=True)
+    last_indices = len(indices) - 1 - last_indices_reversed
+    vertex_uvs = np.zeros((self.num_vertices, 2), dtype=uvs.dtype)
+    vertex_uvs[indices[last_indices]] = uvs[last_indices]
+    return vertex_uvs
+
+  def add_vertex_group(self, name: str, value: enpt.FloatArray['V']) -> None:
+    """Adds a new vertex group.
+
+    Args:
+      name: The name of the vertex group.
+      value: The vertex group values, shape (V,).
+    """
+    xnp = self.xnp
+    if name in self.vertex_group_names:
+      raise ValueError(f'Vertex group {name} already exists.')
+    if value.ndim != 1 or value.shape[0] != self.num_vertices:
+      raise ValueError(
+          f'Vertex group must be 1D array of length {self.num_vertices}.'
+      )
+
+    self.vertex_groups = xnp.concatenate(
+        [self.vertex_groups, value[None]], axis=0
+    )
+    self.vertex_group_names = list(self.vertex_group_names) + [name]
+    self._vertex_group_names_lookup[name] = len(self.vertex_group_names) - 1
+
+  @property
+  def skinning_segmentation(self) -> npt.NDArray[np.int32]:
+    return np.array(self.skinning_weights).argmax(axis=0)
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def vertex_positions_bind_pose(
+      self,
+      identity: enpt.FloatArray['A1 ... An I'],
+      expression: enpt.FloatArray['A1 ... An E'],
+  ) -> enpt.FloatArray['A1 ... An V 3']:
+    return gnm_common.vertex_positions_bind_pose(
+        identity,
+        expression,
+        self.template_vertex_positions,
+        self.vertex_identity_basis,
+        self.expression_basis,
+    )
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def joint_positions_bind_pose(
+      self, identity: enpt.FloatArray['A1 ... An I']
+  ) -> enpt.FloatArray['A1 ... An J 3']:
+    return gnm_common.joint_positions_bind_pose(
+        identity,
+        self.template_joint_positions,
+        self.joint_identity_basis,
+    )
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def compute_pose_correctives(
+      self, rotations: enpt.FloatArray['A1 ... An J 3']
+  ) -> enpt.FloatArray['A1 ... An V 3']:
+    return gnm_common.compute_pose_correctives(
+        rotations,
+        self.pose_correctives_regressor,
+        self.template_vertex_positions,
+        self.num_joints,
+        self.num_vertices,
+    )
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def joint_transforms_world(
+      self,
+      joints: enpt.FloatArray['A1 ... An J 3'],
+      rotations: enpt.FloatArray['A1 ... An J 3'],
+      translation: enpt.FloatArray['A1 ... An 3'],
+  ) -> enpt.FloatArray['A1 ... An J 4 4']:
+    return gnm_common.joint_transforms_world(
+        joints, rotations, translation, self.joint_parent_indices
+    )
+
+  @enp.check_and_normalize_arrays(strict=False)
+  def get_posed_joint_transforms(
+      self,
+      identity: enpt.FloatArray['A1 ... An I'],
+      rotations: enpt.FloatArray['A1 ... An J 3'],
+      translation: enpt.FloatArray['A1 ... An 3'],
+  ) -> enpt.FloatArray['A1 ... An J 4 4']:
+    self._check_parameter_shapes(
+        identity=identity, rotations=rotations, translation=translation
+    )
+    _check_batch_dims(
+        identity=identity, rotations=rotations, translation=translation
+    )
+    return self.joint_transforms_world(
+        joints=self.joint_positions_bind_pose(identity),
+        rotations=rotations,
+        translation=translation,
+    )
+
+  def vertex_group(self, name: str) -> npt.NDArray[np.floating]:
+    try:
+      return np.array(self.vertex_groups[self._vertex_group_names_lookup[name]])
+    except KeyError as exc:
+      raise KeyError(
+          f'Vertex group {name} not found in {self.vertex_group_names}.'
+      ) from exc
+
+  def vertex_group_mask(
+      self, *names: str, threshold: float = _NONZERO_THRESHOLD
+  ) -> npt.NDArray[bool]:
+    result_mask = np.zeros(self.num_vertices, dtype=bool)
+    for name in names:
+      operator, inverse = '|', False
+      if name[0] in '|&-':
+        operator, name = name[0], name[1:]
+      if name[0] == '~':
+        inverse, name = True, name[1:]
+      group_mask = self.vertex_group(name) > threshold
+      if inverse:
+        group_mask = ~group_mask
+      match operator:
+        case '|':
+          result_mask |= group_mask
+        case '&':
+          result_mask &= group_mask
+        case '-':
+          result_mask &= ~group_mask
+    return result_mask
+
+  def vertex_group_indices(
+      self, *names: str, threshold: float = _NONZERO_THRESHOLD
+  ) -> npt.NDArray[np.integer]:
+    return np.where(self.vertex_group_mask(*names, threshold=threshold))[0]
+
+  def quad_indices_for_group(self, *names: str) -> npt.NDArray[np.integer]:
+    vertex_indices = self.vertex_group_indices(*names)
+    quads = np.array(self.quads)
+    return np.where(np.all(np.isin(quads, vertex_indices), axis=-1))[0]
+
+  def triangle_indices_for_group(self, *names: str) -> npt.NDArray[np.integer]:
+    vertex_indices = self.vertex_group_indices(*names)
+    triangles = np.array(self.triangles)
+    return np.where(np.all(np.isin(triangles, vertex_indices), axis=-1))[0]
+
+  def quads_group(self, *names: str) -> npt.NDArray[np.integer]:
+    return np.array(self.quads)[self.quad_indices_for_group(*names)]
+
+  def triangles_group(self, *names: str) -> npt.NDArray[np.integer]:
+    return np.array(self.triangles)[self.triangle_indices_for_group(*names)]
+
+  def quad_uvs_group(self, *names: str) -> npt.NDArray[np.floating]:
+    return np.array(self.quad_uvs)[self.quad_indices_for_group(*names)]
+
+  def triangle_uvs_group(self, *names: str) -> npt.NDArray[np.floating]:
+    return np.array(self.triangle_uvs)[self.triangle_indices_for_group(*names)]
+
+  def vertex_uvs_group(self, *names: str) -> npt.NDArray[np.floating]:
+    return self.vertex_uvs[self.vertex_group_indices(*names)]
+
+  def compute_vertex_normals(
+      self, vertices: enpt.FloatArray['A1 ... An V 3']
+  ) -> enpt.FloatArray['A1 ... An V 3']:
+    """Computes vertex normals. Must be overridden by subclasses."""
+    raise NotImplementedError(
+        'Subclasses must implement compute_vertex_normals.'
+    )
+
+  def prune_vertices(self, keep_vertices: enpt.IntArray['V_pruned']) -> None:
+    """Prunes model vertices in-place."""
+    xnp = self.xnp
+    num_vertices = self.num_vertices
+    keep_vertices = xnp.asarray(keep_vertices, dtype=xnp.int32)
+
+    self.template_vertex_positions = gnm_common.take(
+        self.template_vertex_positions, keep_vertices, axis=0, xnp=xnp
+    )
+    self.vertex_identity_basis = gnm_common.take(
+        self.vertex_identity_basis, keep_vertices, axis=1, xnp=xnp
+    )
+    self.expression_basis = gnm_common.take(
+        self.expression_basis, keep_vertices, axis=1, xnp=xnp
+    )
+    self.skinning_weights = gnm_common.take(
+        self.skinning_weights, keep_vertices, axis=1, xnp=xnp
+    )
+
+    mapper = _scatter_indices(keep_vertices, num_vertices, xnp)
+
+    quads = gnm_common.take(mapper, self.quads, xnp=xnp)
+    triangles = gnm_common.take(mapper, self.triangles, xnp=xnp)
+
+    # Remove quads/triangles containing negative indices.
+    quad_mask = xnp.all(quads >= 0, axis=-1)
+    triangle_mask = xnp.all(triangles >= 0, axis=-1)
+
+    quad_indices = xnp.where(quad_mask)[0]
+    triangle_indices = xnp.where(triangle_mask)[0]
+
+    self.quads = gnm_common.take(quads, quad_indices, axis=0, xnp=xnp)
+    self.triangles = gnm_common.take(
+        triangles, triangle_indices, axis=0, xnp=xnp
+    )
+
+    if self.pose_correctives_regressor is not None:
+      pose_correctives = xnp.reshape(
+          self.pose_correctives_regressor,
+          (self.num_joints * 9, num_vertices, 3),
+      )
+      pose_correctives = gnm_common.take(
+          pose_correctives, keep_vertices, axis=1, xnp=xnp
+      )
+      self.pose_correctives_regressor = xnp.reshape(
+          pose_correctives, (-1, keep_vertices.shape[0] * 3)
+      )
+
+
+def _check_batch_dims(
+    identity: enpt.FloatArray['... I'] | None = None,
+    expression: enpt.FloatArray['... E'] | None = None,
+    rotations: enpt.FloatArray['... J 3'] | None = None,
+    translation: enpt.FloatArray['... 3'] | None = None,
+) -> tuple[int, ...]:
+  """Ensures that the leading batch dimensions of all inputs are the same."""
+  shapes = {}
+  if identity is not None:
+    shapes['identity'] = identity.shape[:-1]
+  if expression is not None:
+    shapes['expression'] = expression.shape[:-1]
+  if rotations is not None:
+    shapes['rotations'] = rotations.shape[:-2]
+  if translation is not None:
+    shapes['translation'] = translation.shape[:-1]
+
+  if not shapes:
+    return ()
+
+  first_name, first_shape = next(iter(shapes.items()))
+  for name, shape in shapes.items():
+    if shape != first_shape:
+      raise ValueError(
+          f'Mismatched batch dimensions: {first_name} has {first_shape}, '
+          f'but {name} has {shape}.'
+      )
+  return first_shape
+
+
+def _scatter_indices(keep_vertices, num_vertices, xnp) -> Any:
+  """Computes a mapping array to translate pruned vertex indices."""
+  if enp.lazy.is_tf_xnp(xnp):
+    tf = enp.lazy.tf
+    return (
+        tf.scatter_nd(
+            indices=tf.expand_dims(keep_vertices, axis=-1),
+            updates=tf.range(tf.shape(keep_vertices)[0], dtype=tf.int32) + 1,
+            shape=(num_vertices,),
+        )
+        - 1
+    )
+  elif enp.lazy.is_jax_xnp(xnp):
+    # pytype: disable=import-error
+    import jax.numpy as jnp  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+    # pytype: enable=import-error
+    mapper = jnp.full((num_vertices,), -1, dtype=jnp.int32)
+    return mapper.at[keep_vertices].set(
+        jnp.arange(len(keep_vertices), dtype=jnp.int32)
+    )
+  elif enp.lazy.is_torch_xnp(xnp):
+    mapper = xnp.full((num_vertices,), -1, dtype=xnp.int32)
+    mapper[keep_vertices] = xnp.arange(len(keep_vertices), dtype=xnp.int32)
+    return mapper
+  else:
+    mapper = np.full((num_vertices,), -1, dtype=np.int32)
+    mapper[keep_vertices] = np.arange(len(keep_vertices), dtype=np.int32)
+    return mapper

gnm/shape/pyproject.toml

@@ -0,0 +1,79 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "gnm-shape"
+version = "3.0.0"
+description = "GNM Shape Library"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "Apache-2.0"}
+dependencies = [    
+    "absl-py",
+    "etils",
+    "imageio",
+    "immutabledict",
+    "importlib_resources",
+    "ipywidgets",
+    "notebook",
+    "numpy",    
+    "opencv-python",
+    "opt-einsum",
+    "rtree",
+    "scipy",
+    "tensorflow",
+    "tqdm",
+    "trimesh",
+    "typeguard",    
+]
+
+[project.optional-dependencies]
+jax = [
+    "jax",
+    "jaxtyping",
+]
+pytorch = [
+    "torch",
+]
+all = [
+    "gnm-shape[jax,pytorch]",
+]
+dev = [
+    "pytest",
+    "pylint",
+    "mediapy",
+    "tensorflow-graphics",
+]
+
+[tool.setuptools]
+packages = [
+    "gnm.shape",
+    "gnm.shape.fitting_utils",
+    "gnm.shape.visualization",
+]
+package-dir = {"gnm.shape" = "."}
+
+[tool.setuptools.package-data]
+"gnm.shape" = [
+    "data/**/*.npz",
+    "data/**/*.png",
+    "data/**/*.jpg",
+    "data/**/*.md",
+    "data/**/*.h5",
+    "assets/**/*",
+]

gnm/shape/visualization/vertex_colors.py

@@ -0,0 +1,100 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Per-vertex colors (V, 3) for GNM, highlighting eyes and teeth."""
+
+from collections.abc import Sequence
+
+from gnm.shape import gnm_numpy
+import numpy as np
+
+# UINT8 color values kept for easier pasting into Color Picker tools.
+DEFAULT_COLOR = tuple([c / 255 for c in [50, 156, 237]])
+ORANGE = tuple([c / 255 for c in [232, 138, 14]])
+GREEN = tuple([c / 255 for c in [83, 227, 39]])
+CYAN = tuple([c / 255 for c in [39, 227, 227]])
+
+# Defines a mapping for shading particular regions. It is from a vertex group to
+# a tuple of modifiers that scales and adds an offset to the given color value.
+_VERTEX_GROUP_COLOR_MODIFIERS = {
+    'skin': (1.0, 0.0),
+    'scleras': (0.6, 0.4),
+    'irises': (0.6, 0.0),
+    'gums': (0.7, 0.0),
+    'teeth': (0.6, 0.4),
+    'tongue': (0.7, 0.0),
+    'mouth_sock': (0.7, 0.0),
+}
+
+
+def get_vertex_colors(
+    gnm_np: gnm_numpy.GNM,
+    color: Sequence[float] = DEFAULT_COLOR,
+) -> np.ndarray:
+  """Per-vertex colors (V, 3) for GNM, highlighting eyes and teeth.
+
+  Args:
+    gnm_np: The GNM instance for which the colors will be generated. If None,
+      then use the GNM face model.
+    color: The RGB color of the skin. Color values should be in [0.0, 1.0]. The
+      irises will be darker, and the scleras and teeth will be brighter.
+
+  Returns:
+    Per-vertex colors (V, 3) for GNM, highlighting eyes and teeth.
+  """
+  color = np.array(color)  # pyrefly: ignore[bad-assignment]
+  colors = np.zeros((gnm_np.num_vertices, 3))
+
+  for region, (scale, offset) in _VERTEX_GROUP_COLOR_MODIFIERS.items():
+    if region in gnm_np.vertex_group_names:
+      colors[gnm_np.vertex_group_indices(region)] = (
+          color * scale + offset  # pyrefly: ignore[unsupported-operation]
+      )
+  return colors
+
+
+def get_vertex_colors_for_inner_head(
+    gnm_np: gnm_numpy.GNM,
+    base_color: Sequence[float] = DEFAULT_COLOR,
+) -> np.ndarray:
+  """Per-vertex GNM colors, highlighting mouth sock and teeth/tongue.
+
+  Useful for visualizing the inner part of the head where we need to distinguish
+  between the mouth sock, upper teeth, lower teeth, and tongue.
+
+  Args:
+    gnm_np: The GNM instance for which the colors will be generated.
+    base_color: The RGB color of the skin (see `get_vertex_colors`).
+
+  Returns:
+    Per-vertex colors (V, 3) for GNM.
+
+  Raises:
+    ValueError: The GNM model does not include a head.
+  """
+  # The function only makes sense for GNM models includling a head.
+  colored_groups = [
+      'mouth_sock',
+      'upper_teeth_and_gums',
+      'lower_teeth_and_gums',
+      'tongue',
+  ]
+  if not np.all([g in gnm_np.vertex_group_names for g in colored_groups]):
+    raise ValueError('The GNM model does not include a head.')
+
+  colors = get_vertex_colors(color=base_color, gnm_np=gnm_np)
+  colors[gnm_np.vertex_group_indices('mouth_sock')] = ORANGE
+  colors[gnm_np.vertex_group_indices('upper_teeth_and_gums')] = GREEN
+  colors[gnm_np.vertex_group_indices('lower_teeth_and_gums', 'tongue')] = CYAN
+  return colors

requirements.txt

@@ -0,0 +1,9 @@
+numpy
+scipy
+trimesh
+absl-py
+etils
+immutabledict
+opt-einsum
+typeguard
+importlib_resources