From 4eba85093a4b8061ae92275d2e11c7b0ef8802aa Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 23 Apr 2026 12:15:25 -0700
Subject: [PATCH] Add new flex_node fields to MJX.

PiperOrigin-RevId: 904576383
Change-Id: Ie0cb8e0f25bdeadbecc2ea229835b0f2a2957254
---
 mjx/mujoco/mjx/_src/types.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index ff08579c..faf47012 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -776,6 +776,10 @@ class Model(PyTreeNode):
   flex_vertadr: np.ndarray
   flex_vertnum: np.ndarray
   flex_vert0: np.ndarray
+  flex_nodeadr: np.ndarray
+  flex_nodenum: np.ndarray
+  flex_nodebodyid: np.ndarray
+  flex_node0: np.ndarray
   hfield_size: np.ndarray
   hfield_nrow: np.ndarray
   hfield_ncol: np.ndarray