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

Code before:
        link : RigidLink
            The target link.
        local_point : torch.Tensor or None, shape (3,)
            Coordinates of the point in the link’s *local* frame.
            If None, the link origin is used (back-compat).

        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, device=gs.device, dtype=gs.tc_float)
            if p_local.numel() != 3:
                gs.raise_exception("`local_point` must have exactly 3 elements.")
            p_local = ti.Vector([float(p_local[0]), float(p_local[1]), float(p_local[2])])
            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_local):
        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.kernel
    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.template()):
        for i_b in range(self._solver._B):
            self._impl_get_jacobian(tgt_link_idx, i_b, p_local)

    @ti.kernel
    def _kernel_get_jacobian_zero(self, tgt_link_idx: ti.i32):

Review comment (performance):
Define a kernel that accepts `ti.types.ndarray()` directly so that you don't need to allocate a `ti.Vector` in Python-scope, which is slow.

Code after:
        link : RigidLink
            The target link.
        local_point : torch.Tensor or None, shape (3,)
            Coordinates of the point in the link’s *local* frame.
            If None, the link origin is used (back-compat).

        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):