Improve docstring and README consistency in gnm.shape

gnm/shape/README.md

@@ -230,7 +230,7 @@ relevant for the GNM v3.x.
 
 ### Joint Parameters
 
-*   **Shape:** rotations: `[batch_size, 4x3 Rotation matrix]`, global translation: `[batch_size, 3]`
+*   **Shape:** rotations: `[batch_size, num_joints, 3]` (axis-angle), global translation: `[batch_size, 3]`
 *   **Description:** Controls the global head position and joint angles for head pose and eyeball orientation.
 
 ## Model Data

gnm/shape/gnm_jax.py

@@ -125,7 +125,14 @@ def compute_vertex_normals(
       self,
       vertices: jt.Float[jnp.ndarray, "... V 3"],
   ) -> jt.Float[jnp.ndarray, "... V 3"]:
-    """Compute vertex normals for GNM mesh."""
+    """Compute vertex normals for GNM mesh.
+
+    Args:
+      vertices: Vertex positions with shape `(..., V, 3)`.
+
+    Returns:
+      Vertex normals with shape `(..., V, 3)`.
+    """
     batch_dims = vertices.shape[:-2]
     num_vertices = vertices.shape[-2]
     vertices_flat = vertices.reshape(-1, num_vertices, 3)

gnm/shape/gnm_numpy.py

@@ -124,7 +124,14 @@ def compute_vertex_normals(
       self,
       vertices: npt.NDArray[np.floating],
   ) -> npt.NDArray[np.floating]:
-    """Compute vertex normals for GNM mesh."""
+    """Compute vertex normals for GNM mesh.
+
+    Args:
+      vertices: Vertex positions with shape `(..., V, 3)`.
+
+    Returns:
+      Vertex normals with shape `(..., V, 3)`.
+    """
     batch_dims = vertices.shape[:-2]
     num_vertices = vertices.shape[-2]
     vertices_flat = vertices.reshape(-1, num_vertices, 3)

gnm/shape/gnm_tensorflow.py

@@ -74,6 +74,7 @@ class GNM(gnm_xnp.GNM):
   * 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:
     joint_parent_indices: Parent's index for each joint in the skeleton, (J).
@@ -95,6 +96,21 @@ class GNM(gnm_xnp.GNM):
       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.
+    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.
+    variant: The variant of the loaded GNM model.
+    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.
+    edge_list: The quad topology represented as a list of directed edges (E, 2).
+    vertex_uvs: Per-vertex UV texture coordinates shaped (V, 2).
     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.
@@ -113,7 +129,14 @@ def compute_vertex_normals(
       self,
       vertices: tf.Tensor,
   ) -> tf.Tensor:
-    """Compute vertex normals for GNM mesh."""
+    """Compute vertex normals for GNM mesh.
+
+    Args:
+      vertices: Vertex positions with shape `(..., V, 3)`.
+
+    Returns:
+      Vertex normals with shape `(..., V, 3)`.
+    """
     batch_dims = tf.shape(vertices)[:-2]
     num_vertices = tf.shape(vertices)[-2]
     vertices_flat = tf.reshape(vertices, (-1, num_vertices, 3))