From 680fb3e5ddbf21fb1e08877f2d03589ad01968ea Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 18 Oct 2024 17:14:23 -0700
Subject: [PATCH] Add more public support functions for MJX.

PiperOrigin-RevId: 687469475
Change-Id: I4218cc2f6d51d3f2f63b1a9a96953d550a5503a5
---
 doc/changelog.rst              |  8 ++++++++
 mjx/mujoco/mjx/__init__.py     |  3 +++
 mjx/mujoco/mjx/_src/support.py | 10 +++++-----
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index a13e1e90..8d911190 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+MJX
+^^^
+
+1. Added ``apply_ft``, ``jac``, and  ``xfrc_accumulate`` as public functions.
+
 Version 3.2.4 (Oct 15, 2024)
 ----------------------------
 
diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py
index fef95c32..6a031b57 100644
--- a/mjx/mujoco/mjx/__init__.py
+++ b/mjx/mujoco/mjx/__init__.py
@@ -48,10 +48,13 @@ from mujoco.mjx._src.smooth import subtree_vel
 from mujoco.mjx._src.smooth import tendon
 from mujoco.mjx._src.smooth import transmission
 from mujoco.mjx._src.solver import solve
+from mujoco.mjx._src.support import apply_ft
 from mujoco.mjx._src.support import full_m
 from mujoco.mjx._src.support import id2name
 from mujoco.mjx._src.support import is_sparse
+from mujoco.mjx._src.support import jac
 from mujoco.mjx._src.support import mul_m
 from mujoco.mjx._src.support import name2id
+from mujoco.mjx._src.support import xfrc_accumulate
 from mujoco.mjx._src.test_util import benchmark
 from mujoco.mjx._src.types import *
diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py
index 1b8e6c23..622c6659 100644
--- a/mjx/mujoco/mjx/_src/support.py
+++ b/mjx/mujoco/mjx/_src/support.py
@@ -348,7 +348,7 @@ def contact_force_dim(
     raise ValueError(f'Unknown cone type: {m.opt.cone}.')
 
 
-def length_circle(
+def _length_circle(
     p0: jax.Array, p1: jax.Array, ind: jax.Array, rad: jax.Array
 ) -> jax.Array:
   """Compute length of circle."""
@@ -366,7 +366,7 @@ def length_circle(
   return rad * angle
 
 
-def is_intersect(
+def _is_intersect(
     p1: jax.Array, p2: jax.Array, p3: jax.Array, p4: jax.Array
 ) -> jax.Array:
   """Check for intersection between two lines defined by their endpoints."""
@@ -429,7 +429,7 @@ def wrap_circle(
     good = jp.where(sidesite, good0, good1)
 
     # penalize for intersection
-    intersect = is_intersect(d[:2], sol[0], d[2:], sol[1])
+    intersect = _is_intersect(d[:2], sol[0], d[2:], sol[1])
     good = jp.where(intersect, -10000, good)
 
     return sol, good
@@ -442,10 +442,10 @@ def wrap_circle(
   pnt = sol.reshape(-1)
 
   # check for intersection
-  intersect = is_intersect(d[:2], pnt[:2], d[2:], pnt[2:])
+  intersect = _is_intersect(d[:2], pnt[:2], d[2:], pnt[2:])
 
   # compute curve length
-  wlen = length_circle(sol[0], sol[1], i, rad)
+  wlen = _length_circle(sol[0], sol[1], i, rad)
 
   # check cases
   invalid = (