Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/entities/rigid_entity/rigid_entity.py (Python)

Code before:

        Returns
        -------
        jacobian : torch.Tensor
            The Jacobian matrix of shape (n_envs, 6, entity.n_dofs) or (6, entity.n_dofs) if n_envs == 0.
        """
        if not self._requires_jac_and_IK:
            gs.raise_exception(
                "Inverse kinematics and jacobian are disabled for this entity. Set `morph.requires_jac_and_IK` to True if you need them."
            )

        if self.n_dofs == 0:
            gs.raise_exception("Entity has zero dofs.")

        if local_point is None:
            self._kernel_get_jacobian(link.idx, 0.0, 0.0, 0.0)
        else:
            local_point = torch.as_tensor(
                local_point,
                device=gs.device,
                dtype=torch.float32,
            )
            if local_point.shape != (3,):
                gs.raise_exception("`local_point` must be a 3-vector in link frame.")

            x, y, z = local_point.tolist()
            self._kernel_get_jacobian(link.idx, x, y, z)

        jacobian = self._jacobian.to_torch(gs.device).permute(2, 0, 1)
        if self._solver.n_envs == 0:
            jacobian = jacobian.squeeze(0)

        return jacobian

    @ti.kernel
    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, px: ti.f32, py: ti.f32, pz: ti.f32):
        for i_b in range(self._solver._B):
            p_local = ti.Vector([px, py, pz])  # build once inside Taichi
            self._func_get_jacobian(
                tgt_link_idx,
                i_b,
                p_local,
                ti.Vector.one(gs.ti_int, 3),
                ti.Vector.one(gs.ti_int, 3),
            )

    @ti.func
    def _func_get_jacobian(self, tgt_link_idx, i_b, p_local, pos_mask, rot_mask):
        for i_row, i_d in ti.ndrange(6, self.n_dofs):
            self._jacobian[i_row, i_d, i_b] = 0.0


Review comment (suggestion):
Converting as tensor to finally convert back to list is useless. Anyway, just pass vector. 

By the way, you should prefer this signature:
```python
def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.template()):
```

To support both `ti.types.ndarray()` and `ti.Vector` as input argument (hopefully?). If it does not work, let's reconsider the best approach.


Code after:

        Returns
        -------
        jacobian : torch.Tensor
            The Jacobian matrix of shape (n_envs, 6, entity.n_dofs) or (6, entity.n_dofs) if n_envs == 0.
        """
        if not self._requires_jac_and_IK:
            gs.raise_exception(
                "Inverse kinematics and jacobian are disabled for this entity. Set `morph.requires_jac_and_IK` to True if you need them."
            )

        if self.n_dofs == 0:
            gs.raise_exception("Entity has zero dofs.")

        if local_point is None:
            self._kernel_get_jacobian_zero(link.idx)
        else:
            p_local = torch.as_tensor(local_point, dtype=gs.tc_float, device=gs.device)
            if p_local.shape != (3,):
                gs.raise_exception("Must be a vector of length 3")
            self._kernel_get_jacobian(link.idx, p_local)

        jacobian = self._jacobian.to_torch(gs.device).permute(2, 0, 1)
        if self._solver.n_envs == 0:
            jacobian = jacobian.squeeze(0)

        return jacobian

    @ti.func
    def _impl_get_jacobian(self, tgt_link_idx, i_b, p_vec):
        self._func_get_jacobian(
            tgt_link_idx,
            i_b,
            p_vec,
            ti.Vector.one(gs.ti_int, 3),
            ti.Vector.one(gs.ti_int, 3),
        )

    @ti.kernel
    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.types.ndarray()):
        p_vec = ti.Vector([p_local[0], p_local[1], p_local[2]], dt=gs.ti_float)
        for i_b in range(self._solver._B):
            self._impl_get_jacobian(tgt_link_idx, i_b, p_vec)

    @ti.kernel
    def _kernel_get_jacobian_zero(self, tgt_link_idx: ti.i32):
        for i_b in range(self._solver._B):
            self._impl_get_jacobian(tgt_link_idx, i_b, ti.Vector.zero(gs.ti_float, 3))

    @ti.func
    def _func_get_jacobian(self, tgt_link_idx, i_b, p_local, pos_mask, rot_mask):