From eb77b59441173e5034fc40ef1df02a21efcb5032 Mon Sep 17 00:00:00 2001
From: Mirco Mazzoni <mircom@google.com>
Date: Mon, 27 Jul 2026 05:10:08 -0700
Subject: [PATCH] Internal changes.

PiperOrigin-RevId: 954565492
---
 gnm/shape/data/versions/v3_0/gnm_head.npz | Bin 53305389 -> 53305389 bytes
 gnm/shape/gnm_common.py                   |   8 ++--
 gnm/shape/gnm_xnp.py                      |  46 +++++++++++++++++-----
 gnm/shape/visualization/render_gnm.py     |   4 +-
 4 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/gnm/shape/data/versions/v3_0/gnm_head.npz b/gnm/shape/data/versions/v3_0/gnm_head.npz
index 1ddcb04c27d8eb7dc2d3f82bdca1dd335c9fd29b..f76d97e659c33afdc4c084c06246134bbf72d680 100644
GIT binary patch
delta 5644
zcmajjcXW;S1IO{)+=Q4xh6Eu*jKq!DF+zgcBWe?SZ)&gFLW3YG6(eS?rbZAuwnmNE
zyJqdZwdn8jc2Dl<&!af!{W|&M^L+1da`PnjiQl!+GhSs3t8I30wR3WEvg=r8M#d^`
zmX2Mzr?>OA<DbdS!>(`OWRH}%20PtNF41-JxX+l{)Lz3|rM@&wsA5{{rCIdL+(iw$
zuznN&b>KE;>+-uq=kj(oOhan@d+HJ0!#j-m4sX9NO%t5+E;L(@&1cvJ`uzWqF1~e{
zw)ZXkrD;Nf@5xT=CPfx9?7ZJiNc|5boQn)9&o0|Ri($8L?|9p;%fSG9b}`8X|6_N2
zx4X%8-0kFu+sP(nuguCpIVRptPLtxV>tnXOOq1VmfOq7$)EA0ye3eF}O>}(aRBkC(
zgvO@xF<f!=(O6ya@g*as#->y0Rffc=u^GEB=Bl?9GI<-Wn(oh4AKw}gp~0CFLxVGq
zTgUpqqs;z>b>GWlQrC%|k2CLDqGjB*1@0z~nCCN{vm5qdQMNsAyqFr%&V2V{Vxyf4
zd}Zu0dzSM+j?Xg3-Q?xdzFq%r{R4XU9oo2&qn+u~H+V=_+v`Len&qsrsH`fRa#60z
zO}VS=%0qdo94e>ErMy&bl}CB2yegmaQTdgxDxeCgLaMOxQ~s)mvZw%6R0XPHs<<kl
zN~%(-v?`;5R9RI{l~)y1u&SsksSp*aDyu50s`^G%Q`J=sRa4bcwN)KeSJhMXRhVj^
zzEurXBh^?nQB74d)m*huEmgQ`rCO^ts;z3L+N%z#qw1tOt1ham>ZZD@9;&D6rFyGA
zs;}y&zEl0x05wnzQiIhHHB=2#!_^2iQhl#RsnKeT8mq>s@oIvas3O!PHCg?jBGnW%
zRZUY-YPy=CW~yj4OZ}*3s~Gi@nxkUXTs2S4R|`~}TBsJO#cGLKs+Osr)pE5$tyHVj
zYW0hXSHG$?YOPwQ)~f`yL2Xo<RHE9fwy3RYo7%2+sGTZF?NYnd9<^8PQ~T8cbx{4L
zlGP#gyE?3nsH5tbN>Rtv33XDPQm54!byl5I=hX#uQC(7()gS7Lx~i_J>*|KOss2>A
z)NOS~-BtJ0ef2;+RFBkS^+Y{Y&(w4ELcLV4)NA!dy;XmyztubSUVTviL>`)Dic(+?
zW^e#UNCRoX3DQA&$N(836J!Qw$O2g*8@PZgxS5i4Q!rI^NWIngyA}E5w`Yh?z0{jt
zQU~3PMmVP4^@@IR8x-emI-i@Zi+=MK{`@wHTW?u@)x}NGZ;L>~D)3!hkCZ5G>moaN
zfG6aDoRAB=AUEUzZ^#Szzz6b!FBE`+PzVZxANWHNus{G5g+M3<#i0b0gi=r%%0Lj5
zg>q0HDnKw)gh~(sp->sBKvnn#szG(A0X3l()P_1x7wSQM2!jUjEi{Bi&={IPQ)mXw
zp#`*raA*asp$)W!cF-O=Ku72VouLbKg>KLtdO%O;1-+pU^o4%#9rT9*Fc1d8U>E{J
zVHgaD5ik<Ihfy#Z#=uw@2jgJ^OoRxS1e4(hh=eIH6{bNHOotgT6QW@j{0Orl27ZD$
z5DRl*9?XXY5C;ok5iEu!uoRZT&#)X;z)DyJtKk=jhhJe0tc7*39ui;!Y=lja2%BLG
zY=v#G9d^J@NP=Ck8}`6n*a!RJ033whAQ=w9?{FB7z)?5`DR3N4z)3g-r{N5og>!Ho
zF2F^&1ef6txB^$<8eE4Pa1;K7TW}lhz+Jcp_u&CNgh%igp1@Oh2G8LIyo6Wq8s5NL
z_zV7qckmuQz(1zQ<Gbzkg$eAz3=ZH3X&@~)K{`ke86YEMg3RCySs*K90~c@wH*kmS
z-~pbH19Czx@Pgcs2fQIK<O3hb557<U3PK?$41VAbMZf|9P!s~87!-#RP!dW(X($6h
zP!`HTd8h!vP!TFY2!ujqr~*~t8>j}=p$621T2LG6KwYQ@^&t!zz_-v48bM=d0!^VA
zG=~<@62hSsw1zg&7TQ63=l~s|6Lf|y&=tBtcjy5<p%?UqKF}BX!FSLf2EafV1cPA+
z425AZ97e!M_#Q^VXcz-yVH}Ky2`~{NU=mD*A0QH@z*Lw9Q7|25z)XmSS@0vwh8Xw>
z=0GgWg?TU^7C;;<ghj9zmcUY220z1cSOF_x6|9C|ARc~&HLw=e!Fou54X_b5K_YC1
zEwB}~!FJdIJ0S^n!EV?Cdto2!hXZgBeuHE<1i!;!I08rE7^J{)H~}Z&6r6@La2C$N
zdAI-<;SyYiKi~>ng==sfZop0W6K=t6xC3|L9^8iq@DLusV|W5j;Tb%K7w{5Z!E1N}
zZ{aWa8{WZt_yGUd>-J~X7bdU=GdO@Fq=B^H1nD3>WPps22{MB-WPz-Z4P3w#+`t{O
zg9ms*4#)|)zzcFi9`J^|kPmzyKlnlcC<ukjNxJ=+J?HQz=z|u1<`RZKK|==Er~U!V
zr}(&5E&a?RnP1WuB@Oeyy0&>#P`IC2_?wf0!i$*izU3;PCKhu8!>51iGpgOKEap6n
zR%!HK53+0{pQaXbu;C=@!ztTnwnl-5k@c~ZZS+o~ps$U(H~ZX4)~8XnlTtJ)V>s!P
zQE+pMxwK(qeF$YcX{AQx4I}FlDBCDU3yazK9Pg9S1dWOtPO?5?vYqr$qhf}U^{J9=
z)TpJ!9AFq(A0pXCn>8wH7+Iek*+vD!Kl{D=_`t|Eiq)v3;Uw!*CELhB?-}FESsw}6
zM*TmIKL63qHoC4+u;Kk|{rqPeHEzu}_$w#v(Ws(f^zpM?ob9C2Z7k+c!^ry6W*fz8
TRM{}H{){nli}P%2F}wZ`Tv@)_

delta 5644
zcmajjcT|mk0LSrew+InR3K7XD*&`!l@4fflJA3bt9U>!p?~t8MX7=7Ad++)Cyxrp-
z$K&`p=lweU@p+zm=z4B_*X-HszD7A#t7x`!v2k#4uqmE$iet$XD}38_NND3}!yl84
zo6W6-_1mwFsJ`3P<Q!Tlo$Hhdb!|0#kk2rT9bg?smNc#R&@A*_+8l=6`ki%tyWtLI
z8!|eFrtx(8pQ&4DN6%`^cX|34rt_y;?;7cpe!kgq+YE-?<Sy2o{yJ`<&fb-ow(!pS
zuW4kY_u1C$#)M=s?1q-L-ZJu1h<{;r$q)M(b~aA6fBnTH+aAtt%PuS?vtjo#u9kME
z_qv)~M#RMg$HkbGtuiY+Wgit66ED_R*N1I;7cZmX2175@{PjY?_V44V_)+%n9SSYx
ziooy$UWO}*7vhSaFBv=`Jb_B65=Bi2ckHl`t3DM;>}j~_W$zl+*9i<r92FRlWW)y6
zhn^?#HLUlTtm`QE7fJRk(lX-qJXe!j*qf<NDGmGeSE~QIJ8!(Oa-mLq_rt<NoihE;
z*lor%r(WDX=`2^1hgXZ{LG6RGcj?xzb}f4w)1Pnf`H|Mw2|hZ_NhMXuRC494TvQ6>
zs!}R9<*rhx)GCegP-#^<<*CxE49ZJoRNg9+%B-@etjb6Es%*+nWmh>=PL)gLR(Vuj
zl~3hY1(d%ks0yjVs)!0uMO86XTm`BUs-!BVN~<!etSYC<s|u>3s-!BbDypigrmCwN
zs-~)?YO6Y`uBxZ%s|Ko}YNQ&gCaS4wrkbl3s-<eBTB|myt!k&*s}8E8>ZCfWE~=~Q
zrn;*hDoFKIy;N`2NA*?xRDU%<4OD~FU^PSyRm0S9HA0P4qts{>tj4IZYMcsD<JAN;
zQB6{l)f6>Vg{o<4x|*TF)J!!?g{#?Wj+(3HsR%V+El>;9BDGj8QA^b_wOp-GE7dBs
zTCGuQ)jG9aZBQFkq}rr5t1T)@ZB^UUcC|z8RJ+t}6|MHDy=tG@uMVh#>X16Dj;I)Q
zR2@^t)d_V{ol>#tv^t~Cs&neRx}YwqOX{+^qOPiI>bkn2ZmL`Awz{M4s(b3bdZ6Oe
zL-j~KR!`JZ^-MiiFVsu*O1)NZ)LZpVy;mR9NA*d4R$tUt^-X<OKSGX9Gfh%p3udqb
zdx!_|!2uFLLP!LTkQkDH6C{OXkQ|)B1yY!z^-wUCvitR@F+TK)kM#KU5Q}c9htWX$
zUr#;b2jlNs{|Hyp<+N;_^~r1eY|7bK_l@RLxgZ~P5##k~k<;*0{JEOMV_l>KH*kkk
zkQ&l}2c(5`;0fs=19(A3@P<r~8L~iD@Bv@Q27ZtoazIYV1-T&)<b`~Y9}0j!6of)h
z7>Ymu6oq0?90H*Pl!Q`H8p=RfC<o=C0#t-bP#LN~Rj3Blp$621T2LG6KwYQ@^`QYY
zghtR9nm|)%2F;-bw1igB8rncxXb0_~19XH=&>6ZwSLg=ap$7y(Pv`}`p%3(he$XEV
zz(5!TgJB2^g<&upM!-lI1*0Jt#=uw@2O%&XCcs3P1e0M3OodRG2Gd~%guzUh1>rCo
z=D=K-2N5tI7QjMS1dCw_EQMvT99F<eSOu$L4XlNAupTzRMu>z>uo<>M6l{fUupM^5
zPS^#zAsY6;Uf2iw;Q$<jLvR?5Knxs(V{jZ!z)3g-v2YsBz*#s4=ivfegiCN4uE15e
z2G`*R+=N?j8}7hexCi&)0mQ*Wcm$8(2|R^o@El&iOLzsZ;SIcnckmuQz(@E5pWzF9
zg>Ud3ewac|@3qwzCa?uF*nvI7gZSV82_PXP0!K&;Nx%t`LNZ7W&fo$mz!g%08@NL%
zNDXPg1JXh|@Pzb`0lXk1cta-03|Sy6_<%2D13$<PIUpzGg4~b?@<Kky4+X#<3PK?$
z3`HOSib63c4uMbtN<t|p4P~G#l!Nk60V+Zzs0>x0DpZ5&Py=d0EvOB3pf1#d`p^Ix
zLL+DlO`s_>gXYizT0$#m4Q-$;w1f800Xjk_=nP$;D|CbI&;x>?C-j2e&<FZLKj;qw
zU?2>F!7v1d!Y~*PBVZ(qg3%BRV_+<dgAf=G6JR1tg2^xira~x8gXu5>!eA!Mf^e7(
zb6_sag9w-p3t%BEg2k`|mclYv4l7_Ktb*0B2G+tlSPvUuBSgX`*bG}B3bw*F*bX~j
zC+vdV5Dj}^FYJT;Z~zX%Avg?2AO?=YF*pt<;3S-aSU3%5;4GYj^Kbz!!X>y2SKumK
zgX?euZo)0N4R_!!+=Kh@0OH^wJc7sY1fIe(cn&Y%CA@;y@CM$(J9rNt;3Is3&+rAl
z!Z-L1KWz2*GwTZz*n%1Cz#ig3d~kpSkPs4qBP50--~>q_86*d1aDf!y3Ms)2+#waD
zhBV*-X(1hWLVCymUXT&IAroYVEaqrE{><*P_!0DJ10Qo9!;hdtZDQN~c6=1kxS@}E
z5c6yLBJaQEll&X`n1!!7+P_gY^P^8(<yFVeT-|V$<&0`?V?T2`Mk_Tk9%Nb0r>rBd
zx_;&W!<{ULQ`XT8jf_85|1tWik-y<imIEp4ojTO}+np?@QPxqcMg<Ia`ePJO-_M-i
zFtQv%S?{z$qr!%f<pjz)O4Y#6Y&^%a96wn{qcqBGxRd3G$vS$bQ7*&Ca;jt<)oSQx
z&Tbf44w0;*tr`_Hj4WqI)=}n0fBoM5JTS73!Zk9!!7Qgr){&jwGsc&*90^%RL4O@t
zeqFPU?r3CuKm9SP-Q;iI;Ge&yMOg2&PotuSe}X?o`J4Kg0}Ug~|2FHL)@W4XAEOiz
K?#=wnF8=}ih|O{U

diff --git a/gnm/shape/gnm_common.py b/gnm/shape/gnm_common.py
index dd5a3109..2cd3263e 100644
--- a/gnm/shape/gnm_common.py
+++ b/gnm/shape/gnm_common.py
@@ -191,10 +191,10 @@ def axis_angle_to_rotation_matrix(
   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),
-  )
+  eyem = eye(3, dtype=angle.dtype, reference_array=axis_angle, xnp=xnp)
+  for _ in range(axis_angle.ndim - 1):
+    eyem = eyem[None, ...]
+  matrix = eyem + xnp.zeros_like(axis_angle[..., None])
 
   skew_01 = -axis[..., 2]
   skew_02 = axis[..., 1]
diff --git a/gnm/shape/gnm_xnp.py b/gnm/shape/gnm_xnp.py
index 396384e9..af89f677 100644
--- a/gnm/shape/gnm_xnp.py
+++ b/gnm/shape/gnm_xnp.py
@@ -179,6 +179,11 @@ def as_int_array(val):
     def as_original(val):
       return val
 
+    def as_string_tuple(val):
+      if val is None:
+        return ()
+      return tuple(str(item) for item in val)
+
     field_converters = {
         'version': as_original,
         'variant': as_original,
@@ -188,16 +193,16 @@ def as_original(val):
         '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,
+        'identity_names': as_string_tuple,
+        'joint_names': as_string_tuple,
+        'expression_names': as_string_tuple,
+        'joint_parent_indices': as_int_array,
         '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,
+        'mesh_component_names': as_string_tuple,
         'mirror_indices': as_int_array,
         'joint_regressor': as_float_array,
         'pose_correctives_regressor': as_float_array,
@@ -208,6 +213,18 @@ def as_original(val):
 
     for field_name, converter in field_converters.items():
       val = model_data.get(field_name, None)
+      if (
+          val is None
+          and field_name == 'bone_aligned_template_joint_orientations'
+      ):
+        val = model_data.get('bone_aligned_orientations', None)
+        if val is None and 'joint_names' in model_data:
+          n_joints = len(model_data['joint_names'])
+          val = xnp.tile(xnp.eye(3, dtype=xnp.float32), (n_joints, 1, 1))
+      if val is None and field_name == 'mesh_component_names':
+        val = model_data.get('part_names', None)
+        if val is None:
+          val = ('skin',)
       init_kwargs[field_name] = converter(val)
 
     return init_kwargs
@@ -218,7 +235,7 @@ def _from_model_data_with_xnp(
       model_data: Mapping[str, Any],
       xnp: enp.NpModule,
   ) -> Self:
-    """Creates a GNM instance from a model data dictionary and array module."""    
+    """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)
@@ -257,6 +274,11 @@ def to_numpy_data_dict(self) -> dict[str, Any]:
       data_dict[field.name] = val
     return data_dict
 
+  @property
+  def bone_aligned_orientations(self) -> enpt.FloatArray:
+    """Alias for bone_aligned_template_joint_orientations."""
+    return self.bone_aligned_template_joint_orientations
+
   def _check_parameter_shapes(
       self,
       identity: enpt.FloatArray['...'] | None = None,
@@ -390,6 +412,7 @@ def vertex_positions_world(
       joints: enpt.FloatArray['A1 ... An J 3'],
       rotations: enpt.FloatArray['A1 ... An J 3'],
       translation: enpt.FloatArray['A1 ... An 3'],
+      **_kwargs,
   ) -> enpt.FloatArray['A1 ... An V 3']:
     """Applies linear blend skinning to GNM vertices."""
     return gnm_common.linear_blend_skinning(
@@ -500,6 +523,7 @@ def vertex_positions_bind_pose(
       self,
       identity: enpt.FloatArray['A1 ... An I'],
       expression: enpt.FloatArray['A1 ... An E'],
+      **_kwargs,
   ) -> enpt.FloatArray['A1 ... An V 3']:
     return gnm_common.vertex_positions_bind_pose(
         identity,
@@ -511,7 +535,7 @@ def vertex_positions_bind_pose(
 
   @enp.check_and_normalize_arrays(strict=False)
   def joint_positions_bind_pose(
-      self, identity: enpt.FloatArray['A1 ... An I']
+      self, identity: enpt.FloatArray['A1 ... An I'], **_kwargs
   ) -> enpt.FloatArray['A1 ... An J 3']:
     return gnm_common.joint_positions_bind_pose(
         identity,
@@ -521,7 +545,7 @@ def joint_positions_bind_pose(
 
   @enp.check_and_normalize_arrays(strict=False)
   def compute_pose_correctives(
-      self, rotations: enpt.FloatArray['A1 ... An J 3']
+      self, rotations: enpt.FloatArray['A1 ... An J 3'], **_kwargs
   ) -> enpt.FloatArray['A1 ... An V 3']:
     return gnm_common.compute_pose_correctives(
         rotations,
@@ -537,6 +561,7 @@ def joint_transforms_world(
       joints: enpt.FloatArray['A1 ... An J 3'],
       rotations: enpt.FloatArray['A1 ... An J 3'],
       translation: enpt.FloatArray['A1 ... An 3'],
+      **_kwargs,
   ) -> enpt.FloatArray['A1 ... An J 4 4']:
     return gnm_common.joint_transforms_world(
         joints, rotations, translation, self.joint_parent_indices
@@ -579,7 +604,10 @@ def vertex_group_mask(
         operator, name = name[0], name[1:]
       if name[0] == '~':
         inverse, name = True, name[1:]
-      group_mask = self.vertex_group(name) > threshold
+      try:
+        group_mask = self.vertex_group(name) > threshold
+      except KeyError:
+        continue
       if inverse:
         group_mask = ~group_mask
       match operator:
diff --git a/gnm/shape/visualization/render_gnm.py b/gnm/shape/visualization/render_gnm.py
index d63805c2..78926e69 100644
--- a/gnm/shape/visualization/render_gnm.py
+++ b/gnm/shape/visualization/render_gnm.py
@@ -421,8 +421,8 @@ def get_look_at_world_to_camera(
     share_camera: np.ndarray | bool = True,
     y_up: np.ndarray | bool = False,
     look_at_vertex_groups: Sequence[str] = ('hockey_mask',),
-    left_vertex_groups: Sequence[str] = ('ears', '&left'),
-    right_vertex_groups: Sequence[str] = ('ears', '&right'),
+    left_vertex_groups: Sequence[str] = ('left_ear', 'ears', '&left'),
+    right_vertex_groups: Sequence[str] = ('right_ear', 'ears', '&right'),
     forward_vertex_groups: Sequence[str] = ('nose_region',),
 ) -> np.ndarray:
   """Compute world-to-camera matrices for a 'look-at' transform.