Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/solvers/rigid/rigid_solver_decomp.py (Python)

Code before:

    def set_geom_friction(self, friction, geoms_idx):
        self._kernel_set_geom_friction(geoms_idx, friction)

    @ti.kernel
    def _kernel_set_geom_friction(self, geoms_idx: ti.i32, friction: ti.f32):
        self.geoms_info[geoms_idx].friction = friction

    def set_geoms_friction(self, friction, geoms_idx=None, *, unsafe=False):
        friction, geoms_idx, _ = self._sanitize_1D_io_variables(
            friction, geoms_idx, self.n_geoms, None, batched=False, idx_name="geoms_idx", unsafe=unsafe
        )
        self._kernel_set_geoms_friction(friction, geoms_idx)

    @ti.kernel
    def _kernel_set_geoms_friction(
        self,
        friction: ti.types.ndarray(),
        geoms_idx: ti.types.ndarray(),
    ):
        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
        for i_g_ in ti.ndrange(geoms_idx.shape[0]):
            self.geoms_info[geoms_idx[i_g_]].friction = friction[i_g_]

    def add_weld_constraint(self, link1_idx, link2_idx, envs_idx=None, *, unsafe=False):
        link1_idx, link2_idx, envs_idx = self._sanitize_1D_io_variables(
            link1_idx, link2_idx, self.n_links, envs_idx, idx_name="links_idx", unsafe=unsafe
        )
        link1_idx = link1_idx.to(gs.tc_int)

        self._kernel_add_weld_constraint(link1_idx, link2_idx, envs_idx)

    @ti.kernel
    def _kernel_add_weld_constraint(
        self,
        link1_idx: ti.types.ndarray(),
        link2_idx: ti.types.ndarray(),
        envs_idx: ti.types.ndarray(),
    ):
        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
        for i_b_ in ti.ndrange(envs_idx.shape[0]):
            i_b = envs_idx[i_b_]
            i_e = self.constraint_solver.ti_n_equalities[i_b]
            l1 = link1_idx[i_b]
            l2 = link2_idx[i_b]

            shared_pos = self.links_state[l1, i_b].pos
            pos1 = gu.ti_inv_transform_by_trans_quat(
                shared_pos, self.links_state[l1, i_b].pos, self.links_state[l1, i_b].quat
            )
            pos2 = gu.ti_inv_transform_by_trans_quat(

Review comment (security):
I think it would be better to refactor `_sanitize_1D_io_variables` to support your use case. Here you are abusing the output tensor, which is not great.

```python
_, link1_idx, envs_idx = self._sanitize_1D_io_variables(
    None, link1_idx, self.n_links, envs_idx, idx_name="links_idx", skip_allocation=True, unsafe=unsafe
)
_, link2_idx, _ = self._sanitize_1D_io_variables(
    None, link2_idx, self.n_links, None, idx_name="links_idx", skip_allocation=True, unsafe=unsafe
)
```

Code after:
    def set_geom_friction(self, friction, geoms_idx):
        self._kernel_set_geom_friction(geoms_idx, friction)

    @ti.kernel
    def _kernel_set_geom_friction(self, geoms_idx: ti.i32, friction: ti.f32):
        self.geoms_info[geoms_idx].friction = friction

    def set_geoms_friction(self, friction, geoms_idx=None, *, unsafe=False):
        friction, geoms_idx, _ = self._sanitize_1D_io_variables(
            friction, geoms_idx, self.n_geoms, None, batched=False, idx_name="geoms_idx", unsafe=unsafe
        )
        self._kernel_set_geoms_friction(friction, geoms_idx)

    @ti.kernel
    def _kernel_set_geoms_friction(
        self,
        friction: ti.types.ndarray(),
        geoms_idx: ti.types.ndarray(),
    ):
        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
        for i_g_ in ti.ndrange(geoms_idx.shape[0]):
            self.geoms_info[geoms_idx[i_g_]].friction = friction[i_g_]

    def add_weld_constraint(self, link1_idx, link2_idx, envs_idx=None, *, unsafe=False):
        _, link1_idx, _ = self._sanitize_1D_io_variables(
            None, link1_idx, self.n_links, envs_idx, idx_name="links_idx", skip_allocation=True, unsafe=unsafe
        )
        _, link2_idx, envs_idx = self._sanitize_1D_io_variables(
            None, link2_idx, self.n_links, envs_idx, idx_name="links_idx", skip_allocation=True, unsafe=unsafe
        )
        self._kernel_add_weld_constraint(link1_idx, link2_idx, envs_idx)

    @ti.kernel
    def _kernel_add_weld_constraint(
        self,
        link1_idx: ti.types.ndarray(),
        link2_idx: ti.types.ndarray(),
        envs_idx: ti.types.ndarray(),
    ):
        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
        for i_b_ in ti.ndrange(envs_idx.shape[0]):
            i_b = envs_idx[i_b_]
            if self.constraint_solver.ti_n_equalities[i_b] >= self.n_equalities_candidate:
                self.constraint_solver.ti_n_equalities[i_b] = self.n_equalities_candidate - 1
                print("Warning: too many constraints, delete the last one.")
            i_e = self.constraint_solver.ti_n_equalities[i_b]

            l1 = link1_idx[i_b]
            l2 = link2_idx[i_b]

            shared_pos = self.links_state[l1, i_b].pos