chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adding type suppressions for pyrefly

Browse Source 
PiperOrigin-RevId: 950703735
Change-Id: I6b6d5826a0dd1f4c0a1cbb2a46aee944cd60d884
This commit is contained in:
Hana Joo
2026-07-20 02:39:22 -07:00
committed by Copybara-Service
parent 00abc5465f
commit 78946ca94e
30 changed files with 193 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_convex.py
+8 -8
View File
@@ -56,27 +56,27 @@ def collider(ncon: int):
fn = collision_fn
for i in [0, 1]:
if key.types[i] == GeomType.BOX:
infos[i] = mesh.box(infos[i])
infos[i] = mesh.box(infos[i]) # pyrefly: ignore[unsupported-operation]
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0, size=0, face=0, vert=0
)
elif key.types[i] == GeomType.MESH:
c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]]
infos[i] = ConvexInfo(**vars(c), **vars(cm))
infos[i] = ConvexInfo(**vars(c), **vars(cm)) # pyrefly: ignore[unsupported-operation]
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0, size=0
)
elif key.types[i] == GeomType.HFIELD:
hfield_info = mesh.hfield(m, key.data_ids[i])
infos[i] = hfield_info.replace(pos=infos[i].pos, mat=infos[i].mat)
in_axes[i] = hfield_info.replace(pos=0, mat=0, data=None)
infos[i] = hfield_info.replace(pos=infos[i].pos, mat=infos[i].mat) # pyrefly: ignore[unsupported-operation]
in_axes[i] = hfield_info.replace(pos=0, mat=0, data=None) # pyrefly: ignore[unsupported-operation]
fn = functools.partial(fn, subgrid_size=key.subgrid_size)
dist, pos, frame = jax.vmap(fn, in_axes=in_axes)(*infos) # pytype: disable=wrong-keyword-args
if ncon > 1:
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame
collide.ncon = ncon
collide.ncon = ncon # pyrefly: ignore[missing-attribute]
return collide
return wrapper
@@ -1051,7 +1051,7 @@ def hfield_sphere(
) -> Collision:
"""Calculates contacts between a hfield and a sphere."""
rbound = jp.max(s.size)
dist, pos, n = _hfield_collision(_sphere_convex, h, s, rbound, subgrid_size)
dist, pos, n = _hfield_collision(_sphere_convex, h, s, rbound, subgrid_size) # pyrefly: ignore[bad-argument-type]
n_mean = jp.mean(n, axis=0)
mask = dist < jp.minimum(0, dist.min() + 1e-3)
@@ -1075,7 +1075,7 @@ def hfield_capsule(
) -> Collision:
"""Calculates contacts between a hfield and a capsule."""
rbound = c.size[0] + c.size[1]
dist, pos, n = _hfield_collision(_capsule_convex, h, c, rbound, subgrid_size)
dist, pos, n = _hfield_collision(_capsule_convex, h, c, rbound, subgrid_size) # pyrefly: ignore[bad-argument-type]
n_mean = jp.mean(n, axis=0)
mask = dist < jp.minimum(0, dist.min() + 1e-3)
@@ -1099,7 +1099,7 @@ def hfield_convex(
) -> Collision:
"""Calculates contacts between a hfield and a capsule."""
rbound = jp.max(c.size)
dist, pos, n = _hfield_collision(_convex_convex, h, c, rbound, subgrid_size)
dist, pos, n = _hfield_collision(_convex_convex, h, c, rbound, subgrid_size) # pyrefly: ignore[bad-argument-type]
n_mean = jp.mean(n, axis=0)
mask = dist < jp.minimum(0, dist.min() + 1e-3)
mjx/mujoco/mjx/_src/collision_driver.py
+3 -3
View File
@@ -384,7 +384,7 @@ def make_condim(
if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF:
ncon = sdf_initpoints
else:
func = _COLLISION_FUNC.get(k.types, None)
func = _COLLISION_FUNC.get(k.types, None) # pyrefly: ignore[no-matching-overload]
if func is not None:
ncon = func.ncon # pytype: disable=attribute-error
else:
@@ -430,7 +430,7 @@ def collision(m: Model, d: Data) -> Data:
contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact)
# run the collision function specified by the grouping key
func = _COLLISION_FUNC[key.types]
func = _COLLISION_FUNC[key.types] # pyrefly: ignore[bad-index]
ncon = func.ncon # pytype: disable=attribute-error
dist, pos, frame = func(m, d, key, contact.geom)
@@ -457,4 +457,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.tree_replace({'_impl.contact': contact})
return d.tree_replace({'_impl.contact': contact}) # pyrefly: ignore[bad-return]
mjx/mujoco/mjx/_src/collision_primitive.py
+1 -1
View File
@@ -40,7 +40,7 @@ def collider(ncon: int):
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame
collide.ncon = ncon
collide.ncon = ncon # pyrefly: ignore[missing-attribute]
return collide
return wrapper
mjx/mujoco/mjx/_src/collision_sdf.py
+1 -1
View File
@@ -54,7 +54,7 @@ def collider(ncon: int):
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame
collide.ncon = ncon
collide.ncon = ncon # pyrefly: ignore[missing-attribute]
return collide
return wrapper
mjx/mujoco/mjx/_src/constraint.py
+4 -4
View File
@@ -716,8 +716,8 @@ def make_constraint(m: Model, d: Data) -> Data:
if not efcs:
z = jp.empty(0)
d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))})
d = d.tree_replace({
d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))}) # pyrefly: ignore[bad-assignment]
d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'_impl.efc_D': z,
'_impl.efc_aref': z,
'_impl.efc_frictionloss': z,
@@ -736,13 +736,13 @@ def make_constraint(m: Model, d: Data) -> Data:
return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss
aref, r, pos, margin, frictionloss = fn(efc)
d = d.tree_replace({
d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'_impl.efc_J': efc.J,
'_impl.efc_D': 1 / r,
'_impl.efc_aref': aref,
'_impl.efc_pos': pos,
'_impl.efc_margin': margin,
})
d = d.tree_replace({'_impl.efc_frictionloss': frictionloss})
d = d.tree_replace({'_impl.efc_frictionloss': frictionloss}) # pyrefly: ignore[bad-assignment]
return d
mjx/mujoco/mjx/_src/dataclasses.py
+4 -4
View File
@@ -75,7 +75,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
Returns:
the resulting dataclass, registered with Jax
"""
data_clz = dataclasses.dataclass(frozen=True)(clz)
data_clz = dataclasses.dataclass(frozen=True)(clz) # pyrefly: ignore[bad-argument-type]
data_clz.replace = dataclasses.replace
if register_as_pytree:
@@ -137,7 +137,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
data_clz, iterate_clz_with_keys, clz_from_iterable
)
return data_clz
return data_clz # pyrefly: ignore[bad-return]
TNode = TypeVar('TNode', bound='PyTreeNode')
@@ -163,7 +163,7 @@ class PyTreeNode:
@classmethod
def fields(cls) -> Tuple[dataclasses.Field, ...]: # pylint: disable=g-bare-generic
return dataclasses.fields(cls)
return dataclasses.fields(cls) # pyrefly: ignore[bad-argument-type]
def tree_replace(
self, params: Dict[str, Optional[jax.typing.ArrayLike]]
@@ -190,7 +190,7 @@ def _tree_replace(
for i, g in enumerate(lst):
if not hasattr(g, attr[1]):
continue
v = val if not hasattr(val, '__iter__') else val[i]
v = val if not hasattr(val, '__iter__') else val[i] # pyrefly: ignore[bad-index, unsupported-operation]
lst[i] = _tree_replace(g, attr[1:], v)
return base.replace(**{attr[0]: lst})
mjx/mujoco/mjx/_src/forward.py
+2 -2
View File
@@ -89,7 +89,7 @@ def fwd_velocity(m: Model, d: Data) -> Data:
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('fwd_velocity requires JAX backend implementation.')
d = d.tree_replace({
d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel,
'_impl.ten_velocity': d._impl.ten_J @ d.qvel,
})
@@ -357,7 +357,7 @@ def euler(m: Model, d: Data) -> Data:
else:
M = d._impl.M + jp.diag(m.opt.timestep * m.dof_damping)
dh = d.tree_replace({'_impl.M': M})
dh = smooth.factor_m(m, dh)
dh = smooth.factor_m(m, dh) # pyrefly: ignore[bad-argument-type]
qfrc = d.qfrc_smooth + d.qfrc_constraint
qacc = smooth.solve_m(m, dh, qfrc)
return _advance(m, d, d.act_dot, qacc)
mjx/mujoco/mjx/_src/inverse.py
+1 -1
View File
@@ -76,7 +76,7 @@ def inv_constraint(m: Model, d: Data) -> Data:
# update
ctx = solver.Context.create(m, d, grad=False)
return d.tree_replace({
return d.tree_replace({ # pyrefly: ignore[bad-return]
'qfrc_constraint': ctx.qfrc_constraint,
'_impl.efc_force': ctx.efc_force,
})
mjx/mujoco/mjx/_src/io.py
+25 -25
View File
@@ -151,15 +151,15 @@ def _resolve_impl_and_device(
if (has_impl, has_device) == (True, True):
pass
elif (has_impl, has_device) == (True, False):
device = _resolve_device(impl)
device = _resolve_device(impl) # pyrefly: ignore[bad-argument-type]
elif (has_impl, has_device) == (False, True):
impl = _resolve_impl(device)
impl = _resolve_impl(device) # pyrefly: ignore[bad-argument-type]
else:
device = jax.devices()[0]
logging.info('Using JAX default device: %s.', device)
impl = _resolve_impl(device)
_check_impl_device_compatibility(impl, device)
_check_impl_device_compatibility(impl, device) # pyrefly: ignore[bad-argument-type]
return impl, device # pytype: disable=bad-return-type
@@ -724,7 +724,7 @@ def _make_data_jax(
qpos=jp.array(m.qpos0, dtype=float_),
eq_active=m.eq_active0,
_impl=impl,
**_make_data_public_fields(m),
**_make_data_public_fields(m), # pyrefly: ignore[bad-argument-type]
)
if m.nmocap:
@@ -773,7 +773,7 @@ def _make_data_warp(
nvmax=nvmax,
) # pylint: disable=undefined-variable
fields = _make_data_public_fields(m)
fields = _make_data_public_fields(m) # pyrefly: ignore[bad-argument-type]
for k in fields:
if k in {'userdata', 'plugin_state', 'history'}:
continue
@@ -1287,14 +1287,14 @@ def _get_data_into_warp(
if batched
else d
)
result_i = result[i] if batched else result
result_i = result[i] if batched else result # pyrefly: ignore[bad-index]
ncon = d_i._impl.nacon[0]
nefc = int(d_i._impl.nefc)
# nj = int(d_i._impl.nj[0])
nj = 0 # TODO(btaba): add nj back
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ:
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: # pyrefly: ignore[missing-attribute]
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
all_fields = types.Data.fields() + mjxw.types.DataWarp.fields()
for field in all_fields:
@@ -1368,14 +1368,14 @@ def _get_data_into(
for i in range(batch_size):
d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
result_i = result[i] if batched else result
result_i = result[i] if batched else result # pyrefly: ignore[bad-index]
ncon = (d_i._impl.contact.dist <= 0).sum()
efc_active = (d_i._impl.efc_J != 0).any(axis=1)
nefc = int(efc_active.sum())
nj = (d_i._impl.efc_J != 0).sum() if support.is_sparse(m) else nefc * m.nv
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ:
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: # pyrefly: ignore[missing-attribute]
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
if d.impl == types.Impl.JAX:
all_fields = types.Data.fields() + types.DataJAX.fields()
@@ -1389,10 +1389,10 @@ def _get_data_into(
continue
if field.name == 'contact':
_get_contact(result_i.contact, d_i._impl.contact)
_get_contact(result_i.contact, d_i._impl.contact) # pyrefly: ignore[missing-attribute]
# efc_address must be updated because rows were deleted above:
efc_map = np.cumsum(efc_active) - 1
result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address]
result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] # pyrefly: ignore[missing-attribute]
continue
# MuJoCo actuator_moment is sparse, MJX uses a dense representation.
@@ -1412,10 +1412,10 @@ def _get_data_into(
)
else:
actuator_moment = d_i._impl.actuator_moment
result_i.moment_rownnz[:] = moment_rownnz
result_i.moment_rowadr[:] = moment_rowadr
result_i.moment_colind[:] = moment_colind
result_i.actuator_moment[:] = actuator_moment
result_i.moment_rownnz[:] = moment_rownnz # pyrefly: ignore[missing-attribute]
result_i.moment_rowadr[:] = moment_rowadr # pyrefly: ignore[missing-attribute]
result_i.moment_colind[:] = moment_colind # pyrefly: ignore[missing-attribute]
result_i.actuator_moment[:] = actuator_moment # pyrefly: ignore[missing-attribute]
continue
# MuJoCo ten_J is sparse, MJX uses a dense representation.
@@ -1435,7 +1435,7 @@ def _get_data_into(
)
else:
ten_j = d_i._impl.ten_J
result_i.ten_J[:] = ten_j
result_i.ten_J[:] = ten_j # pyrefly: ignore[missing-attribute]
continue
if hasattr(d_i._impl, field.name):
@@ -1461,9 +1461,9 @@ def _get_data_into(
efc_J_rowadr,
efc_J_colind,
)
result_i.efc_J_rownnz[:] = efc_J_rownnz
result_i.efc_J_rowadr[:] = efc_J_rowadr
result_i.efc_J_colind[:] = efc_J_colind
result_i.efc_J_rownnz[:] = efc_J_rownnz # pyrefly: ignore[missing-attribute]
result_i.efc_J_rowadr[:] = efc_J_rowadr # pyrefly: ignore[missing-attribute]
result_i.efc_J_colind[:] = efc_J_colind # pyrefly: ignore[missing-attribute]
value = efc_J
else:
value = value.reshape(-1)
@@ -1496,7 +1496,7 @@ def _get_data_into(
# recalculate qLD and qLDiagInv as MJX and MuJoCo have different
# representations of the Cholesky decomposition.
mujoco.mj_factorM(m, result_i)
mujoco.mj_factorM(m, result_i) # pyrefly: ignore[bad-argument-type]
# TODO(josechenf): Iterate on the keepalive implementation to make it easier to
@@ -1533,7 +1533,7 @@ def _get_data_into_cpp(
d_i: types.Data = (
jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
)
result_i = result[i] if batched else result
result_i = result[i] if batched else result # pyrefly: ignore[bad-index]
if batched:
addr_i = int(d_impl.pointer_lo[i]) | (int(d_impl.pointer_hi[i]) << 32)
@@ -1562,7 +1562,7 @@ def _get_data_into_cpp(
src_data.mocap_quat[:] = d_i.mocap_quat
mujoco.mj_kinematics(m, src_data)
mujoco.mj_copyData(result_i, m, src_data)
mujoco.mj_copyData(result_i, m, src_data) # pyrefly: ignore[bad-argument-type]
def get_data_into(
@@ -1810,5 +1810,5 @@ def create_render_context(
_check_warp_installed()
from mujoco.mjx.warp import io as mjxw_io # pylint: disable=g-import-not-at-top # pytype: disable=import-error
return mjxw_io.create_render_context(
mjm, nworld=nworld, devices=devices, **kwargs
mjm, nworld=nworld, devices=devices, **kwargs # pyrefly: ignore[bad-argument-type]
)
mjx/mujoco/mjx/_src/ray.py
+4 -4
View File
@@ -263,8 +263,8 @@ def ray(
for bodyid in bodyexclude:
geom_filter &= (m.geom_bodyid != bodyid)
if geomgroup:
geomgroup = np.array(geomgroup, dtype=bool)
geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)]
geomgroup = np.array(geomgroup, dtype=bool) # pyrefly: ignore[bad-assignment]
geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)] # pyrefly: ignore[bad-index]
# map ray to local geom frames
geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos)
@@ -281,7 +281,7 @@ def ray(
args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
if geom_type == GeomType.MESH:
dist, id_ = fn(m, id_, *args)
dist, id_ = fn(m, id_, *args) # pyrefly: ignore[bad-argument-count, bad-argument-type]
else:
dist = jax.vmap(fn)(*args)
@@ -314,4 +314,4 @@ def ray_geom(
Returns:
dist: distance from ray origin to geom surface
"""
return _RAY_FUNC[geomtype](size, pnt, vec)
return _RAY_FUNC[geomtype](size, pnt, vec) # pyrefly: ignore[bad-argument-type, bad-return, missing-argument]
mjx/mujoco/mjx/_src/scan.py
+3 -3
View File
@@ -331,7 +331,7 @@ def flat(
reordered_ys.append(_take(y, _index(ids, input_order)))
y = reordered_ys if f_ret_is_seq else reordered_ys[0]
return y
return y # pyrefly: ignore[bad-return]
def body_tree(
@@ -443,7 +443,7 @@ def body_tree(
if i < len(in_types):
key_in_take.setdefault(key, []).append(ids)
else:
key_y_take.setdefault(key, []).append(np.hstack(ids))
key_y_take.setdefault(key, []).append(np.hstack(ids)) # pyrefly: ignore[no-matching-overload]
# use this grouping to take the right data subsets and call vmap(f)
keys = sorted(key_body_ids, reverse=reverse)
@@ -494,4 +494,4 @@ def body_tree(
y = y[0] if len(out_types) == 1 else y
return y
return y # pyrefly: ignore[bad-return]
mjx/mujoco/mjx/_src/sensor.py
+12 -12
View File
@@ -339,9 +339,9 @@ def sensor_vel(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.JOINTVEL:
sensor = d.qvel[m.jnt_dofadr[objid]]
elif sensor_type == SensorType.TENDONVEL:
sensor = d._impl.ten_velocity[objid]
sensor = d._impl.ten_velocity[objid] # pyrefly: ignore[missing-attribute]
elif sensor_type == SensorType.ACTUATORVEL:
sensor = d._impl.actuator_velocity[objid]
sensor = d._impl.actuator_velocity[objid] # pyrefly: ignore[missing-attribute]
elif sensor_type == SensorType.BALLANGVEL:
jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
sensor = d.qvel[jnt_dotadr]
@@ -403,10 +403,10 @@ def sensor_vel(m: Model, d: Data) -> Data:
adrs.append(adrt.reshape(-1))
continue # avoid adding to sensors/adrs list a second time
elif sensor_type == SensorType.SUBTREELINVEL:
sensor = d._impl.subtree_linvel[objid]
sensor = d._impl.subtree_linvel[objid] # pyrefly: ignore[missing-attribute]
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.SUBTREEANGMOM:
sensor = d._impl.subtree_angmom[objid]
sensor = d._impl.subtree_angmom[objid] # pyrefly: ignore[missing-attribute]
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
else:
# TODO(taylorhowell): raise error after adding sensor check to io.py
@@ -502,7 +502,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# compute conray, flip if second body
conray = jax.vmap(
lambda frame, force: math.normalize(frame[0] * force[0])
)(d._impl.contact.frame, contact_force)
)(d._impl.contact.frame, contact_force) # pyrefly: ignore[unbound-name]
conray = jp.where(conbody1[..., None], -conray, conray)
# compute distance, mapping over sites and contacts
@@ -547,7 +547,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# compute force magnitude for each contact
force_mag = jax.vmap(
lambda forcetorque: jp.dot(forcetorque[:3], forcetorque[:3])
)(contact_force)
)(contact_force) # pyrefly: ignore[unbound-name]
def _reduce(reduction, mask):
if reduction == 1: # mindist
@@ -662,7 +662,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
slot.append(jp.repeat(nfound, num)[:, None])
if dataspec & (1 << 1): # force
slot.append(flip * contact_force[cid, :3])
slot.append(flip * contact_force[cid, :3]) # pyrefly: ignore[unbound-name]
if dataspec & (1 << 2): # torque
slot.append(flip * contact_force[cid, 3:])
@@ -699,14 +699,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
bodyid = m.site_bodyid[objid]
rot = d.site_xmat[objid]
cvel = d.cvel[bodyid]
cacc = d._impl.cacc[bodyid]
cacc = d._impl.cacc[bodyid] # pyrefly: ignore[missing-attribute]
dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]]
sensor = _accelerometer(cvel, cacc, dif, rot)
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.FORCE:
bodyid = m.site_bodyid[objid]
cfrc_int = d._impl.cfrc_int[bodyid]
cfrc_int = d._impl.cfrc_int[bodyid] # pyrefly: ignore[missing-attribute]
site_xmat = d.site_xmat[objid]
sensor = jax.vmap(lambda mat, vec: mat.T @ vec)(
site_xmat, cfrc_int[:, 3:]
@@ -715,7 +715,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.TORQUE:
bodyid = m.site_bodyid[objid]
rootid = m.body_rootid[bodyid]
cfrc_int = d._impl.cfrc_int[bodyid]
cfrc_int = d._impl.cfrc_int[bodyid] # pyrefly: ignore[missing-attribute]
site_xmat = d.site_xmat[objid]
dif = d.site_xpos[objid] - d.subtree_com[rootid]
sensor = jax.vmap(
@@ -744,7 +744,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
pos, bodyid = objtype_data[ot]
pos = pos[objidt]
bodyid = bodyid[objidt]
cacc = d._impl.cacc[bodyid]
cacc = d._impl.cacc[bodyid] # pyrefly: ignore[missing-attribute]
if sensor_type == SensorType.FRAMELINACC:
@@ -774,7 +774,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# TODO(taylorhowell): raise error after adding sensor check to io.py
continue # unsupported sensor type
sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) # pyrefly: ignore[bad-argument-type]
adrs.append(adr)
if not adrs:
mjx/mujoco/mjx/_src/smooth.py
+14 -14
View File
@@ -173,7 +173,7 @@ def com_pos(m: Model, d: Data) -> Data:
root_com = subtree_com[m.body_rootid]
offset = d.xipos - root_com
cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
d = d.tree_replace({'_impl.cinert': cinert})
d = d.tree_replace({'_impl.cinert': cinert}) # pyrefly: ignore[bad-assignment]
# map motion dofs to global frame centered at subtree_com
def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis):
@@ -212,7 +212,7 @@ def com_pos(m: Model, d: Data) -> Data:
d.xanchor,
d.xaxis,
)
d = d.tree_replace({'cdof': cdof})
d = d.tree_replace({'cdof': cdof}) # pyrefly: ignore[bad-assignment]
return d
@@ -302,12 +302,12 @@ def crb(m: Model, d: Data) -> Data:
crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True)
crb_body = crb_body.at[0].set(0.0)
d = d.tree_replace({'_impl.crb': crb_body})
d = d.tree_replace({'_impl.crb': crb_body}) # pyrefly: ignore[bad-assignment]
crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.tree_replace({'_impl.M': qm})
d = d.tree_replace({'_impl.M': qm}) # pyrefly: ignore[bad-assignment]
return d
@@ -318,7 +318,7 @@ def factor_m(m: Model, d: Data) -> Data:
if not support.is_sparse(m):
qh, _ = jax.scipy.linalg.cho_factor(d._impl.M)
d = d.tree_replace({'_impl.qLD': qh})
d = d.tree_replace({'_impl.qLD': qh}) # pyrefly: ignore[bad-assignment]
return d
depth = []
@@ -367,7 +367,7 @@ def factor_m(m: Model, d: Data) -> Data:
qld_diag = qld[diag_adr]
qld = (qld / qld[np.array(diag_ds)]).at[diag_adr].set(qld_diag)
d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag})
d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag}) # pyrefly: ignore[bad-assignment]
return d
@@ -448,7 +448,7 @@ def com_vel(m: Model, d: Data) -> Data:
d.qvel,
)
d = d.tree_replace({'cvel': cvel, 'cdof_dot': cdof_dot})
d = d.tree_replace({'cvel': cvel, 'cdof_dot': cdof_dot}) # pyrefly: ignore[bad-assignment]
return d
@@ -544,7 +544,7 @@ def subtree_vel(m: Model, d: Data) -> Data:
reverse=True,
)
return d.tree_replace({
return d.tree_replace({ # pyrefly: ignore[bad-return]
'_impl.subtree_linvel': subtree_linvel,
'_impl.subtree_angmom': subtree_angmom,
})
@@ -841,7 +841,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
)
# update data
return d.tree_replace({
return d.tree_replace({ # pyrefly: ignore[bad-return]
'_impl.cacc': cacc,
'_impl.cfrc_int': cfrc_int,
'_impl.cfrc_ext': cfrc_ext,
@@ -1172,7 +1172,7 @@ def tendon(m: Model, d: Data) -> Data:
[wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)]
).reshape((m.nwrap, 2))
return d.tree_replace({
return d.tree_replace({ # pyrefly: ignore[bad-return]
'ten_length': ten_length,
'_impl.ten_J': ten_moment,
'_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
@@ -1276,7 +1276,7 @@ def transmission(m: Model, d: Data) -> Data:
moment = jac @ wrench
elif trntype == TrnType.TENDON:
length = d.ten_length[trnid[0]] * gear[:1]
moment = d._impl.ten_J[trnid[0]] * gear[0]
moment = d._impl.ten_J[trnid[0]] * gear[0] # pyrefly: ignore[missing-attribute]
else:
raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
@@ -1308,7 +1308,7 @@ def transmission(m: Model, d: Data) -> Data:
length = length.reshape((m.nu,))
moment = moment.reshape((m.nu, m.nv))
d = d.tree_replace(
d = d.tree_replace( # pyrefly: ignore[bad-assignment]
{'actuator_length': length, '_impl.actuator_moment': moment}
)
return d
@@ -1332,7 +1332,7 @@ def tendon_armature(m: Model, d: Data) -> Data:
j = m.M_colind
JTAJ = JTAJ[(i, j)]
return d.tree_replace({'_impl.M': d._impl.M + JTAJ})
return d.tree_replace({'_impl.M': d._impl.M + JTAJ}) # pyrefly: ignore[bad-return]
def tendon_dot(m: Model, d: Data) -> jax.Array:
@@ -1457,7 +1457,7 @@ def tendon_bias(m: Model, d: Data) -> Data:
# add bias term: qfrc += ten_J * armature * ten_Jdot @ qvel
coef = m.tendon_armature * jp.dot(ten_Jdot, d.qvel)
return d.tree_replace({
return d.tree_replace({ # pyrefly: ignore[bad-return]
'qfrc_bias': (
d.qfrc_bias
+ jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0)
mjx/mujoco/mjx/_src/solver.py
+2 -2
View File
@@ -482,7 +482,7 @@ def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
vv = jp.sum(v[:, 1:] * v[:, 1:], axis=1)
point_fn = lambda a: _LSPoint.create(
m, d, ctx, a, jv, quad, quad_gauss, uu, v0, uv, vv
m, d, ctx, a, jv, quad, quad_gauss, uu, v0, uv, vv # pyrefly: ignore[bad-argument-type]
)
def cond(ctx: _LSContext) -> jax.Array:
@@ -601,7 +601,7 @@ def solve(m: Model, d: Data) -> Data:
else:
ctx = jax.lax.while_loop(cond, body, ctx)
d = d.tree_replace({
d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'qfrc_constraint': ctx.qfrc_constraint,
'qacc': ctx.qacc,
'_impl.efc_force': ctx.efc_force,
mjx/mujoco/mjx/_src/support.py
+6 -6
View File
@@ -501,15 +501,15 @@ class BindData(object):
num = sum((typ == jt) * jt.dof_width() for jt in JointType)
if isinstance(self.id, list):
idx = []
for a, n in zip(adr, num):
for a, n in zip(adr, num): # pyrefly: ignore[bad-argument-type]
idx.extend(a + j for j in range(n))
return self._slice(self.__getname(name), idx)
elif num > 1:
return self._slice(self.__getname(name), slice(adr, adr + num))
else:
return self._slice(self.__getname(name), adr)
return self._slice(self.__getname(name), adr) # pyrefly: ignore[bad-argument-type]
elif name in ('mocap_pos', 'mocap_quat'):
return self._slice(self.__getname(name), self.model.body_mocapid[self.id])
return self._slice(self.__getname(name), self.model.body_mocapid[self.id]) # pyrefly: ignore[bad-argument-type]
return self._slice(self.__getname(name), self.id)
def set(self, name: str, value: jax.Array) -> Data:
@@ -521,7 +521,7 @@ class BindData(object):
try:
iter(value)
except TypeError:
value = [value]
value = [value] # pyrefly: ignore[bad-assignment]
if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'):
adr = num = 0
if name == 'qpos':
@@ -549,7 +549,7 @@ class BindData(object):
num = [dim]
i = 0
value = jax.numpy.array(value).flatten()
for a, n in zip(adr, num):
for a, n in zip(adr, num): # pyrefly: ignore[bad-argument-type]
shape = array.shape
array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)
i += n
@@ -623,7 +623,7 @@ def contact_force_dim(
if m.opt.cone == ConeType.PYRAMIDAL:
efc_address = (
d._impl.contact.efc_address[idx_dim, None] # pytype: disable=attribute-error
+ np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None]
+ np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None] # pyrefly: ignore[no-matching-overload]
)
efc_force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error
force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
mjx/mujoco/mjx/_src/test_util.py
+4 -4
View File
@@ -63,13 +63,13 @@ def benchmark(
xla_flags += ' --xla_gpu_triton_gemm_any=True'
os.environ['XLA_FLAGS'] = xla_flags
m.opt.solver = {
m.opt.solver = { # pyrefly: ignore[bad-assignment]
'cg': mujoco.mjtSolver.mjSOL_CG,
'newton': mujoco.mjtSolver.mjSOL_NEWTON,
}[solver.lower()]
m.opt.iterations = iterations
m.opt.ls_iterations = ls_iterations
m = io.put_model(m)
m = io.put_model(m) # pyrefly: ignore[bad-assignment]
@jax.pmap
def init(key):
@@ -449,7 +449,7 @@ def create_mjcf(
if p(80):
continue
geom1, geom2 = np.random.choice(geom_names, replace=False, size=2)
geom1, geom2 = np.random.choice(geom_names, replace=False, size=2) # pyrefly: ignore[no-matching-overload]
if geom1 > geom2:
geom1, geom2 = geom2, geom1
@@ -467,7 +467,7 @@ def create_mjcf(
if p(50):
continue
body1, body2 = np.random.choice(body_names, replace=False, size=2)
body1, body2 = np.random.choice(body_names, replace=False, size=2) # pyrefly: ignore[no-matching-overload]
ET.SubElement(contact, 'exclude', {'body1': body1, 'body2': body2})
# ElementTree.indent is not available before Python 3.9
mjx/mujoco/mjx/_src/types.py
+59 -59
View File
@@ -105,10 +105,10 @@ class JointType(enum.IntEnum):
HINGE: rotation angle (rad) around body-fixed axis (1,)
"""
FREE = mujoco.mjtJoint.mjJNT_FREE
BALL = mujoco.mjtJoint.mjJNT_BALL
SLIDE = mujoco.mjtJoint.mjJNT_SLIDE
HINGE = mujoco.mjtJoint.mjJNT_HINGE
FREE = mujoco.mjtJoint.mjJNT_FREE # pyrefly: ignore[bad-assignment]
BALL = mujoco.mjtJoint.mjJNT_BALL # pyrefly: ignore[bad-assignment]
SLIDE = mujoco.mjtJoint.mjJNT_SLIDE # pyrefly: ignore[bad-assignment]
HINGE = mujoco.mjtJoint.mjJNT_HINGE # pyrefly: ignore[bad-assignment]
def dof_width(self) -> int:
return {0: 6, 1: 3, 2: 1, 3: 1}[self.value]
@@ -126,9 +126,9 @@ class IntegratorType(enum.IntEnum):
IMPLICITFAST: implicit in velocity, no rne derivative
"""
EULER = mujoco.mjtIntegrator.mjINT_EULER
RK4 = mujoco.mjtIntegrator.mjINT_RK4
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
EULER = mujoco.mjtIntegrator.mjINT_EULER # pyrefly: ignore[bad-assignment]
RK4 = mujoco.mjtIntegrator.mjINT_RK4 # pyrefly: ignore[bad-assignment]
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST # pyrefly: ignore[bad-assignment]
# unsupported: IMPLICIT
@@ -147,14 +147,14 @@ class GeomType(enum.IntEnum):
SDF: signed distance field
"""
PLANE = mujoco.mjtGeom.mjGEOM_PLANE
HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD
SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE
CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE
ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID
CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER
BOX = mujoco.mjtGeom.mjGEOM_BOX
MESH = mujoco.mjtGeom.mjGEOM_MESH
PLANE = mujoco.mjtGeom.mjGEOM_PLANE # pyrefly: ignore[bad-assignment]
HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD # pyrefly: ignore[bad-assignment]
SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE # pyrefly: ignore[bad-assignment]
CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE # pyrefly: ignore[bad-assignment]
ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID # pyrefly: ignore[bad-assignment]
CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER # pyrefly: ignore[bad-assignment]
BOX = mujoco.mjtGeom.mjGEOM_BOX # pyrefly: ignore[bad-assignment]
MESH = mujoco.mjtGeom.mjGEOM_MESH # pyrefly: ignore[bad-assignment]
# unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
@@ -184,8 +184,8 @@ class ConeType(enum.IntEnum):
ELLIPTIC: elliptic
"""
PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL
ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC
PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL # pyrefly: ignore[bad-assignment]
ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC # pyrefly: ignore[bad-assignment]
class JacobianType(enum.IntEnum):
@@ -197,9 +197,9 @@ class JacobianType(enum.IntEnum):
AUTO: sparse if nv>60 and device is TPU, dense otherwise
"""
DENSE = mujoco.mjtJacobian.mjJAC_DENSE
SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
AUTO = mujoco.mjtJacobian.mjJAC_AUTO
DENSE = mujoco.mjtJacobian.mjJAC_DENSE # pyrefly: ignore[bad-assignment]
SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE # pyrefly: ignore[bad-assignment]
AUTO = mujoco.mjtJacobian.mjJAC_AUTO # pyrefly: ignore[bad-assignment]
class SolverType(enum.IntEnum):
@@ -211,8 +211,8 @@ class SolverType(enum.IntEnum):
"""
# unsupported: PGS
CG = mujoco.mjtSolver.mjSOL_CG
NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
CG = mujoco.mjtSolver.mjSOL_CG # pyrefly: ignore[bad-assignment]
NEWTON = mujoco.mjtSolver.mjSOL_NEWTON # pyrefly: ignore[bad-assignment]
class EqType(enum.IntEnum):
@@ -225,10 +225,10 @@ class EqType(enum.IntEnum):
TENDON: couple the lengths of two tendons with cubic
"""
CONNECT = mujoco.mjtEq.mjEQ_CONNECT
WELD = mujoco.mjtEq.mjEQ_WELD
JOINT = mujoco.mjtEq.mjEQ_JOINT
TENDON = mujoco.mjtEq.mjEQ_TENDON
CONNECT = mujoco.mjtEq.mjEQ_CONNECT # pyrefly: ignore[bad-assignment]
WELD = mujoco.mjtEq.mjEQ_WELD # pyrefly: ignore[bad-assignment]
JOINT = mujoco.mjtEq.mjEQ_JOINT # pyrefly: ignore[bad-assignment]
TENDON = mujoco.mjtEq.mjEQ_TENDON # pyrefly: ignore[bad-assignment]
# unsupported: DISTANCE
@@ -243,11 +243,11 @@ class WrapType(enum.IntEnum):
CYLINDER: wrap around (infinite) cylinder
"""
JOINT = mujoco.mjtWrap.mjWRAP_JOINT
PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY
SITE = mujoco.mjtWrap.mjWRAP_SITE
SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE
CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER
JOINT = mujoco.mjtWrap.mjWRAP_JOINT # pyrefly: ignore[bad-assignment]
PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY # pyrefly: ignore[bad-assignment]
SITE = mujoco.mjtWrap.mjWRAP_SITE # pyrefly: ignore[bad-assignment]
SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE # pyrefly: ignore[bad-assignment]
CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER # pyrefly: ignore[bad-assignment]
class TrnType(enum.IntEnum):
@@ -260,10 +260,10 @@ class TrnType(enum.IntEnum):
SITE: force on site
"""
JOINT = mujoco.mjtTrn.mjTRN_JOINT
JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT
SITE = mujoco.mjtTrn.mjTRN_SITE
TENDON = mujoco.mjtTrn.mjTRN_TENDON
JOINT = mujoco.mjtTrn.mjTRN_JOINT # pyrefly: ignore[bad-assignment]
JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT # pyrefly: ignore[bad-assignment]
SITE = mujoco.mjtTrn.mjTRN_SITE # pyrefly: ignore[bad-assignment]
TENDON = mujoco.mjtTrn.mjTRN_TENDON # pyrefly: ignore[bad-assignment]
# unsupported: SLIDERCRANK, BODY
@@ -278,11 +278,11 @@ class DynType(enum.IntEnum):
MUSCLE: piece-wise linear filter with two time constants
"""
NONE = mujoco.mjtDyn.mjDYN_NONE
INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR
FILTER = mujoco.mjtDyn.mjDYN_FILTER
FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT
MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE
NONE = mujoco.mjtDyn.mjDYN_NONE # pyrefly: ignore[bad-assignment]
INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR # pyrefly: ignore[bad-assignment]
FILTER = mujoco.mjtDyn.mjDYN_FILTER # pyrefly: ignore[bad-assignment]
FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT # pyrefly: ignore[bad-assignment]
MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE # pyrefly: ignore[bad-assignment]
# unsupported: USER
@@ -295,9 +295,9 @@ class GainType(enum.IntEnum):
MUSCLE: muscle FLV curve computed by muscle_gain
"""
FIXED = mujoco.mjtGain.mjGAIN_FIXED
AFFINE = mujoco.mjtGain.mjGAIN_AFFINE
MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE
FIXED = mujoco.mjtGain.mjGAIN_FIXED # pyrefly: ignore[bad-assignment]
AFFINE = mujoco.mjtGain.mjGAIN_AFFINE # pyrefly: ignore[bad-assignment]
MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE # pyrefly: ignore[bad-assignment]
# unsupported: USER
@@ -310,9 +310,9 @@ class BiasType(enum.IntEnum):
MUSCLE: muscle passive force computed by muscle_bias
"""
NONE = mujoco.mjtBias.mjBIAS_NONE
AFFINE = mujoco.mjtBias.mjBIAS_AFFINE
MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE
NONE = mujoco.mjtBias.mjBIAS_NONE # pyrefly: ignore[bad-assignment]
AFFINE = mujoco.mjtBias.mjBIAS_AFFINE # pyrefly: ignore[bad-assignment]
MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE # pyrefly: ignore[bad-assignment]
# unsupported: USER
@@ -327,14 +327,14 @@ class ConstraintType(enum.IntEnum):
CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
"""
EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF
FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON
LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT
LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON
CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS
CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC
EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY # pyrefly: ignore[bad-assignment]
FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF # pyrefly: ignore[bad-assignment]
FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON # pyrefly: ignore[bad-assignment]
LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT # pyrefly: ignore[bad-assignment]
LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON # pyrefly: ignore[bad-assignment]
CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS # pyrefly: ignore[bad-assignment]
CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL # pyrefly: ignore[bad-assignment]
CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC # pyrefly: ignore[bad-assignment]
class CamLightType(enum.IntEnum):
@@ -348,11 +348,11 @@ class CamLightType(enum.IntEnum):
TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
"""
FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED
TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK
TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM
TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY
TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM
FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED # pyrefly: ignore[bad-assignment]
TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK # pyrefly: ignore[bad-assignment]
TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM # pyrefly: ignore[bad-assignment]
TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY # pyrefly: ignore[bad-assignment]
TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM # pyrefly: ignore[bad-assignment]
class SensorType(enum.IntEnum):
mjx/mujoco/mjx/codegen/generate_warp_shim.py
+2 -2
View File
@@ -264,8 +264,8 @@ def _warp_function(
# create a dummy output if there are no output fields
needs_dummy_output = not field_usage.data_out_fields
if needs_dummy_output and fn_name != 'render':
fn_args_raw.append('# Dummy output')
fn_args_raw.append('dummy: wp.array[int],')
fn_args_raw.append('# Dummy output') # pyrefly: ignore[bad-argument-type]
fn_args_raw.append('dummy: wp.array[int],') # pyrefly: ignore[bad-argument-type]
return fn_args_raw, fn_assignments, fn_call
mjx/mujoco/mjx/codegen/generate_warp_types.py
+1 -1
View File
@@ -440,7 +440,7 @@ def write_core_cls(
shape_property = _DATA_SHAPE_PROPERTY_FIELD
new_class_body = _build_new_class_body_ast(
keys,
keys, # pyrefly: ignore[bad-argument-type]
cls_name,
annotations,
shape_property=shape_property,
mjx/mujoco/mjx/codegen/trace.py
+7 -7
View File
@@ -162,7 +162,7 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
called_fn_name = node.func.id
key = (hash(self._current_fpath), called_fn_name)
if key not in self._visited_fns:
self._visited_fns.add(key)
self._visited_fns.add(key) # pyrefly: ignore[bad-argument-type]
next_fpath = self._module_fpaths.get(
called_fn_name, self._current_fpath
)
@@ -181,8 +181,8 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
if len(node.args) != 2:
raise ValueError(f'wp.copy() must have 2 arguments, got {node.args}.')
out_node, in_node = node.args
self.add_field_usage(out_node, True)
self.add_field_usage(in_node, False)
self.add_field_usage(out_node, True) # pyrefly: ignore[bad-argument-type]
self.add_field_usage(in_node, False) # pyrefly: ignore[bad-argument-type]
return
# do not trace array creation
@@ -201,7 +201,7 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
key = (hash(self._current_fpath), called_fn_name)
next_fpath = self._module_fpaths.get(parts[0])
if next_fpath and key not in self._visited_fns:
self._visited_fns.add(key)
self._visited_fns.add(key) # pyrefly: ignore[bad-argument-type]
self.recurse_trace(next_fpath, called_fn_name)
self.generic_visit(node)
@@ -223,10 +223,10 @@ def trace_function(
) -> FieldUsage:
"""Traces the function statically to find usages of model and data fields."""
base_path = file.get_base_path()
fpath = base_path / fpath
fpath = base_path / fpath # pyrefly: ignore[bad-assignment]
logging.info('Tracing function: "%s" in "%s"', fn, fpath)
src = fpath.read_text()
src = fpath.read_text() # pyrefly: ignore[missing-attribute]
parsed_ast = ast.parse(src, filename=str(fpath))
target_fn_nodes = (
@@ -262,7 +262,7 @@ def trace_function(
if visited_fns is None:
visited_fns = set()
visitor = _FunctionFieldUsageVisitor(fpath, visited_fns, mjwarp_field_info)
visitor = _FunctionFieldUsageVisitor(fpath, visited_fns, mjwarp_field_info) # pyrefly: ignore[bad-argument-type]
for body in target_fn_node.body:
visitor.visit(body)
mjx/mujoco/mjx/warp/ffi.py
+6 -6
View File
@@ -119,7 +119,7 @@ def jax_callable_variadic_tuple(
# Provide a flattened signature for the Warp callable machinery.
new_signature = flatten_signature(inspect.signature(func), args)
func_wrapper.__signature__ = new_signature
func_wrapper.__signature__ = new_signature # pyrefly: ignore[missing-attribute]
func_wrapper.__annotations__ = {
p.name: p.annotation
for p in new_signature.parameters.values()
@@ -205,7 +205,7 @@ def _expand_dim_from_path(
if ndim is None or ndim < 0:
return leaf
if ndim > leaf.ndim:
leaf = jp.expand_dims(leaf, axis=np.arange(ndim - leaf.ndim))
leaf = jp.expand_dims(leaf, axis=np.arange(ndim - leaf.ndim)) # pyrefly: ignore[bad-argument-type]
if ndim != leaf.ndim:
raise AssertionError(
f'Leaf node ndim ({leaf.ndim}) and expected ndim ({ndim}) do not match'
@@ -221,7 +221,7 @@ def _squeeze_dim(leaf_expanded: Any, leaf: Any) -> Any:
f' ndim {leaf.ndim}'
)
if leaf_expanded.ndim > leaf.ndim:
return jp.squeeze(leaf_expanded, np.arange(leaf_expanded.ndim - leaf.ndim))
return jp.squeeze(leaf_expanded, np.arange(leaf_expanded.ndim - leaf.ndim)) # pyrefly: ignore[bad-argument-type]
return leaf_expanded
@@ -288,7 +288,7 @@ def _maybe_broadcast_to(
"""Broadcasts fields that are used in MuJoCo Warp."""
ndim = _get_mapping_from_tree_path(path, mjx_warp_types._NDIM[cls_str])
needs_batch_dim = _get_mapping_from_tree_path(
path, mjx_warp_types._BATCH_DIM[cls_str] # pylint: disable=protected-access
path, mjx_warp_types._BATCH_DIM[cls_str] # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
)
needs_batch_dim = bool(needs_batch_dim) and (ndim is not None and ndim > 0)
if needs_batch_dim and not is_batched:
@@ -305,7 +305,7 @@ def _check_leading_dim(
):
"""Asserts that the batch dimension of a leaf node matches the expected batch dimension."""
has_batch_dim = _get_mapping_from_tree_path(
path, mjx_warp_types._BATCH_DIM['Data']
path, mjx_warp_types._BATCH_DIM['Data'] # pyrefly: ignore[bad-argument-type]
)
attr = tree_path_to_attr_str(path)
if has_batch_dim and leaf.shape[0] != expected_batch_dim:
@@ -396,7 +396,7 @@ def marshal_custom_vmap(
out_batched = jax.tree.map_with_path(
# NB: if a field is not in MuJoCo Warp, we let JAX do its magic.
lambda path, x: _get_mapping_from_tree_path(
path, mjx_warp_types._BATCH_DIM['Data'] # pylint: disable=protected-access
path, mjx_warp_types._BATCH_DIM['Data'] # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
)
or x,
out_batched,