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

Code before:

        Parameters
        ----------
        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:
            local_point = torch.zeros(3, device=gs.device, dtype=torch.float32)
        elif isinstance(local_point, np.ndarray):
            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 space.")

        p_local_ti = ti.Vector(local_point.tolist(), dt=ti.f32)
        self._kernel_get_jacobian(link.idx, p_local_ti)

        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, p_local: ti.types.vector(3, ti.f32)):
        ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL)
        for i_b in range(self._solver._B):
            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),
            )


Review comment (security):
Just call `local_point = torch.as_tensor(local_point, device=gs.device, dtype=torch.float32)` systematically if not None. Take inspiration from `RigidSolver._sanitize_1D_io_variables`.

Code after:

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