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

Code before:

        if to_torch:
            buf = torch.full((out_size, 3), -1, dtype=gs.tc_int, device=gs.device)
        else:
            buf = np.full((out_size, 3), -1, dtype=np.int32)

        if n_max > 0:
            kernel_collect_welds(
                as_tensor,
                buf,
                self.constraint_solver.constraint_state,
                self.equalities_info,
                self._static_rigid_sim_config,
            )

        if as_tensor:
            if n_envs > 0:
                buf = buf.view(n_envs, n_max, 3) if to_torch else buf.reshape(n_envs, n_max, 3)
            env_idx, obj_a, obj_b = buf[..., 0], buf[..., 1], buf[..., 2]
            return {"env": env_idx, "obj_a": obj_a, "obj_b": obj_b}

        if n_envs == 0:
            return {"obj_a": [], "obj_b": []}

        parts = torch.split(buf, n_eqs) if to_torch else np.split(buf, splits)
        a = [p[:, 1] for p in parts]
        b = [p[:, 2] for p in parts]
        if n_envs == 1:
            return {"obj_a": a[0], "obj_b": b[0]}
        return {"obj_a": a, "obj_b": b}

    # ------------------------------------------------------------------------------------
    # ----------------------------------- properties -------------------------------------
    # ------------------------------------------------------------------------------------

    @property
    def links(self):
        if self.is_built:
            return self._links
        return gs.List(link for entity in self._entities for link in entity.links)

    @property
    def joints(self):
        if self.is_built:
            return self._joints
        return gs.List(joint for entity in self._entities for joint in entity.joints)

    @property
    def geoms(self):
        if self.is_built:
            return self._geoms

Review comment (suggestion):
`a` and `b` are not suitable variable names. Just use `obj_a, obj_b`.

Create tuple to keep the data immutable. And generate both output at once:
```python
obj_a, obj_b = tuple(zip(*((data[:, 1], data[:, 2]) for data in data_chunks)))
```

Code after:

        if n_max > 0:
            kernel_collect_welds(
                as_tensor,
                buf,
                self.constraint_solver.constraint_state,
                self.equalities_info,
                self._static_rigid_sim_config,
            )

        if n_envs > 0:
            if as_tensor:
                buf = buf.reshape((n_envs, n_max, 3))
                obj_a = buf[..., 1]
                obj_b = buf[..., 2]
            else:
                if to_torch:
                    data_chunks = torch.split(buf, n_eqs)
                else:
                    splits = list(np.cumsum(n_eqs, dtype=np.int32)[:-1])
                    data_chunks = np.split(buf, splits)
                obj_a, obj_b = tuple(zip(*((data[:, 1], data[:, 2]) for data in data_chunks)))
        else:
            if to_torch:
                obj_a = torch.empty((0,), dtype=gs.tc_int, device=gs.device)
                obj_b = torch.empty((0,), dtype=gs.tc_int, device=gs.device)
            else:
                obj_a = []
                obj_b = []

        return {"obj_a": obj_a, "obj_b": obj_b}

    # ------------------------------------------------------------------------------------
    # ----------------------------------- properties -------------------------------------
    # ------------------------------------------------------------------------------------

    @property
    def links(self):
        if self.is_built:
            return self._links
        return gs.List(link for entity in self._entities for link in entity.links)

    @property
    def joints(self):
        if self.is_built:
            return self._joints
        return gs.List(joint for entity in self._entities for joint in entity.joints)

    @property
    def geoms(self):
        if self.is_built: