internal change

PiperOrigin-RevId: 758287280 Change-Id: I8b99a5d44910d5b2e7857b808b1696f244778a26
2025-05-13 10:51:49 -07:00
parent 627fffdef9
commit 6bc0bbf7df
1 changed files with 1 additions and 1 deletions
@@ -429,4 +429,4 @@ def collision(m: Model, d: Data) -> Data:
  contacts = sum([condim_groups[k] for k in sorted(condim_groups)], [])
  contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)

-  return d.replace(_impl=d._impl.replace(contact=contact))  # pytype: disable=attribute-error
+  return d.tree_replace({'_impl.contact': contact})