jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/rbda/rnea.py

@@ -0,0 +1,235 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Cross, StandardGravity
+
+from . import utils
+
+
+def rnea(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.Vector,
+    base_quaternion: jtp.Vector,
+    joint_positions: jtp.Vector,
+    base_linear_velocity: jtp.Vector,
+    base_angular_velocity: jtp.Vector,
+    joint_velocities: jtp.Vector,
+    base_linear_acceleration: jtp.Vector | None = None,
+    base_angular_acceleration: jtp.Vector | None = None,
+    joint_accelerations: jtp.Vector | None = None,
+    link_forces: jtp.Matrix | None = None,
+    standard_gravity: jtp.FloatLike = StandardGravity,
+) -> tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute inverse dynamics using the Recursive Newton-Euler Algorithm (RNEA).
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity:
+            The linear velocity of the base link in inertial-fixed representation.
+        base_angular_velocity:
+            The angular velocity of the base link in inertial-fixed representation.
+        joint_velocities: The velocities of the joints.
+        base_linear_acceleration:
+            The linear acceleration of the base link in inertial-fixed representation.
+        base_angular_acceleration:
+            The angular acceleration of the base link in inertial-fixed representation.
+        joint_accelerations: The accelerations of the joints.
+        link_forces:
+            The forces applied to the links expressed in the world frame.
+        standard_gravity: The standard gravity constant.
+
+    Returns:
+        A tuple containing the 6D force applied to the base link expressed in the
+        world frame and the joint forces that, when applied respectively to the base
+        link and joints, produce the given base and joint accelerations.
+    """
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, W_v̇_WB, s̈, _, W_f, W_g = utils.process_inputs(
+        model=model,
+        base_position=base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=joint_positions,
+        base_linear_velocity=base_linear_velocity,
+        base_angular_velocity=base_angular_velocity,
+        joint_velocities=joint_velocities,
+        base_linear_acceleration=base_linear_acceleration,
+        base_angular_acceleration=base_angular_acceleration,
+        joint_accelerations=joint_accelerations,
+        link_forces=link_forces,
+        standard_gravity=standard_gravity,
+    )
+
+    W_g = jnp.atleast_2d(W_g).T
+    W_v_WB = jnp.atleast_2d(W_v_WB).T
+    W_v̇_WB = jnp.atleast_2d(W_v̇_WB).T
+
+    # Get the 6D spatial inertia matrices of all links.
+    M = js.model.link_spatial_inertia_matrices(model=model)
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute the base transform.
+    W_H_B = jaxlie.SE3.from_rotation_and_translation(
+        rotation=jaxlie.SO3(wxyz=W_Q_B),
+        translation=W_p_B,
+    )
+
+    # Compute 6D transforms of the base velocity.
+    W_X_B = W_H_B.adjoint()
+    B_X_W = W_H_B.inverse().adjoint()
+
+    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # These transforms define the relative kinematics of the entire model, including
+    # the base transform for both floating-base and fixed-base models.
+    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+        joint_positions=s, base_transform=W_H_B.as_matrix()
+    )
+
+    # Allocate buffers.
+    v = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    a = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    f = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+
+    # Allocate the buffer of transforms link -> base.
+    i_X_0 = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
+
+    # Initialize the acceleration of the base link.
+    a_0 = -B_X_W @ W_g
+    a = a.at[0].set(a_0)
+
+    if model.floating_base():
+
+        # Base velocity v₀ in body-fixed representation.
+        v_0 = B_X_W @ W_v_WB
+        v = v.at[0].set(v_0)
+
+        # Base acceleration a₀ in body-fixed representation w/o gravity.
+        a_0 = B_X_W @ (W_v̇_WB - W_g)
+        a = a.at[0].set(a_0)
+
+        # Force applied to the base link that produce the base acceleration w/o gravity.
+        f_0 = (
+            M[0] @ a[0]
+            + Cross.vx_star(v[0]) @ M[0] @ v[0]
+            - W_X_B.T @ jnp.vstack(W_f[0])
+        )
+        f = f.at[0].set(f_0)
+
+    # ======
+    # Pass 1
+    # ======
+
+    ForwardPassCarry = tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix]
+    forward_pass_carry: ForwardPassCarry = (v, a, i_X_0, f)
+
+    def forward_pass(
+        carry: ForwardPassCarry, i: jtp.Int
+    ) -> tuple[ForwardPassCarry, None]:
+
+        ii = i - 1
+        v, a, i_X_0, f = carry
+
+        # Project the joint velocity into its motion subspace.
+        vJ = S[i] * ṡ[ii]
+
+        # Propagate the link velocity.
+        v_i = i_X_λi[i] @ v[λ[i]] + vJ
+        v = v.at[i].set(v_i)
+
+        # Propagate the link acceleration.
+        a_i = i_X_λi[i] @ a[λ[i]] + S[i] * s̈[ii] + Cross.vx(v[i]) @ vJ
+        a = a.at[i].set(a_i)
+
+        # Compute the link-to-base transform.
+        i_X_0_i = i_X_λi[i] @ i_X_0[λ[i]]
+        i_X_0 = i_X_0.at[i].set(i_X_0_i)
+
+        # Compute link-to-world transform for the 6D force.
+        i_Xf_W = Adjoint.inverse(i_X_0[i] @ B_X_W).T
+
+        # Compute the force acting on the link.
+        f_i = (
+            M[i] @ a[i]
+            + Cross.vx_star(v[i]) @ M[i] @ v[i]
+            - i_Xf_W @ jnp.vstack(W_f[i])
+        )
+        f = f.at[i].set(f_i)
+
+        return (v, a, i_X_0, f), None
+
+    (v, a, i_X_0, f), _ = (
+        jax.lax.scan(
+            f=forward_pass,
+            init=forward_pass_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(v, a, i_X_0, f), None]
+    )
+
+    # ======
+    # Pass 2
+    # ======
+
+    τ = jnp.zeros_like(s)
+
+    BackwardPassCarry = tuple[jtp.Vector, jtp.Matrix]
+    backward_pass_carry: BackwardPassCarry = (τ, f)
+
+    def backward_pass(
+        carry: BackwardPassCarry, i: jtp.Int
+    ) -> tuple[BackwardPassCarry, None]:
+
+        ii = i - 1
+        τ, f = carry
+
+        # Project the 6D force to the DoF of the joint.
+        τ_i = S[i].T @ f[i]
+        τ = τ.at[ii].set(τ_i.squeeze())
+
+        # Propagate the force to the parent link.
+        def update_f(f: jtp.Matrix) -> jtp.Matrix:
+
+            f_λi = f[λ[i]] + i_X_λi[i].T @ f[i]
+            f = f.at[λ[i]].set(f_λi)
+
+            return f
+
+        f = jax.lax.cond(
+            pred=jnp.logical_or(λ[i] != 0, model.floating_base()),
+            true_fun=update_f,
+            false_fun=lambda f: f,
+            operand=f,
+        )
+
+        return (τ, f), None
+
+    (τ, f), _ = (
+        jax.lax.scan(
+            f=backward_pass,
+            init=backward_pass_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(τ, f), None]
+    )
+
+    # ==============
+    # Adjust outputs
+    # ==============
+
+    # Express the base 6D force in the world frame.
+    W_f0 = B_X_W.T @ f[0]
+
+    return W_f0.squeeze(), jnp.atleast_1d(τ.squeeze())

jaxsim/rbda/utils.py

@@ -0,0 +1,160 @@
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim import exceptions
+from jaxsim.math import StandardGravity
+
+
+def process_inputs(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.VectorLike | None = None,
+    base_quaternion: jtp.VectorLike | None = None,
+    joint_positions: jtp.VectorLike | None = None,
+    base_linear_velocity: jtp.VectorLike | None = None,
+    base_angular_velocity: jtp.VectorLike | None = None,
+    joint_velocities: jtp.VectorLike | None = None,
+    base_linear_acceleration: jtp.VectorLike | None = None,
+    base_angular_acceleration: jtp.VectorLike | None = None,
+    joint_accelerations: jtp.VectorLike | None = None,
+    joint_forces: jtp.VectorLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+    standard_gravity: jtp.ScalarLike | None = None,
+) -> tuple[
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Matrix,
+    jtp.Vector,
+]:
+    """
+    Adjust the inputs to rigid-body dynamics algorithms.
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity: The linear velocity of the base link.
+        base_angular_velocity: The angular velocity of the base link.
+        joint_velocities: The velocities of the joints.
+        base_linear_acceleration: The linear acceleration of the base link.
+        base_angular_acceleration: The angular acceleration of the base link.
+        joint_accelerations: The accelerations of the joints.
+        joint_forces: The forces applied to the joints.
+        link_forces: The forces applied to the links.
+        standard_gravity: The standard gravity constant.
+
+    Returns:
+        The adjusted inputs.
+    """
+
+    dofs = model.dofs()
+    nl = model.number_of_links()
+
+    # Floating-base position.
+    W_p_B = base_position
+    W_Q_B = base_quaternion
+    s = joint_positions
+
+    # Floating-base velocity in inertial-fixed representation.
+    W_vl_WB = base_linear_velocity
+    W_ω_WB = base_angular_velocity
+    ṡ = joint_velocities
+
+    # Floating-base acceleration in inertial-fixed representation.
+    W_v̇l_WB = base_linear_acceleration
+    W_ω̇_WB = base_angular_acceleration
+    s̈ = joint_accelerations
+
+    # System dynamics inputs.
+    f = link_forces
+    τ = joint_forces
+
+    # Fill missing data and adjust dimensions.
+    s = jnp.atleast_1d(s.squeeze()) if s is not None else jnp.zeros(dofs)
+    ṡ = jnp.atleast_1d(ṡ.squeeze()) if ṡ is not None else jnp.zeros(dofs)
+    s̈ = jnp.atleast_1d(s̈.squeeze()) if s̈ is not None else jnp.zeros(dofs)
+    τ = jnp.atleast_1d(τ.squeeze()) if τ is not None else jnp.zeros(dofs)
+    W_vl_WB = jnp.atleast_1d(W_vl_WB.squeeze()) if W_vl_WB is not None else jnp.zeros(3)
+    W_v̇l_WB = jnp.atleast_1d(W_v̇l_WB.squeeze()) if W_v̇l_WB is not None else jnp.zeros(3)
+    W_p_B = jnp.atleast_1d(W_p_B.squeeze()) if W_p_B is not None else jnp.zeros(3)
+    W_ω_WB = jnp.atleast_1d(W_ω_WB.squeeze()) if W_ω_WB is not None else jnp.zeros(3)
+    W_ω̇_WB = jnp.atleast_1d(W_ω̇_WB.squeeze()) if W_ω̇_WB is not None else jnp.zeros(3)
+    f = jnp.atleast_2d(f.squeeze()) if f is not None else jnp.zeros(shape=(nl, 6))
+    W_Q_B = (
+        jnp.atleast_1d(W_Q_B.squeeze())
+        if W_Q_B is not None
+        else jnp.array([1.0, 0, 0, 0])
+    )
+    standard_gravity = (
+        jnp.array(standard_gravity).squeeze()
+        if standard_gravity is not None
+        else StandardGravity
+    )
+
+    if s.shape != (dofs,):
+        raise ValueError(s.shape, dofs)
+
+    if ṡ.shape != (dofs,):
+        raise ValueError(ṡ.shape, dofs)
+
+    if s̈.shape != (dofs,):
+        raise ValueError(s̈.shape, dofs)
+
+    if τ.shape != (dofs,):
+        raise ValueError(τ.shape, dofs)
+
+    if W_p_B.shape != (3,):
+        raise ValueError(W_p_B.shape, (3,))
+
+    if W_vl_WB.shape != (3,):
+        raise ValueError(W_vl_WB.shape, (3,))
+
+    if W_ω_WB.shape != (3,):
+        raise ValueError(W_ω_WB.shape, (3,))
+
+    if W_v̇l_WB.shape != (3,):
+        raise ValueError(W_v̇l_WB.shape, (3,))
+
+    if W_ω̇_WB.shape != (3,):
+        raise ValueError(W_ω̇_WB.shape, (3,))
+
+    if f.shape != (nl, 6):
+        raise ValueError(f.shape, (nl, 6))
+
+    if W_Q_B.shape != (4,):
+        raise ValueError(W_Q_B.shape, (4,))
+
+    # Check that the quaternion is unary since our RBDAs make this assumption in order
+    # to prevent introducing additional normalizations that would affect AD.
+    exceptions.raise_value_error_if(
+        condition=~jnp.allclose(W_Q_B.dot(W_Q_B), 1.0),
+        msg="A RBDA received a quaternion that is not normalized.",
+    )
+
+    # Pack the 6D base velocity and acceleration.
+    W_v_WB = jnp.hstack([W_vl_WB, W_ω_WB])
+    W_v̇_WB = jnp.hstack([W_v̇l_WB, W_ω̇_WB])
+
+    # Create the 6D gravity acceleration.
+    W_g = jnp.zeros(6).at[2].set(-standard_gravity)
+
+    return (
+        W_p_B.astype(float),
+        W_Q_B.astype(float),
+        s.astype(float),
+        W_v_WB.astype(float),
+        ṡ.astype(float),
+        W_v̇_WB.astype(float),
+        s̈.astype(float),
+        τ.astype(float),
+        f.astype(float),
+        W_g.astype(float),
+    )

jaxsim/terrain/__init__.py

@@ -0,0 +1,2 @@
+from . import terrain
+from .terrain import FlatTerrain, PlaneTerrain, Terrain

jaxsim/terrain/terrain.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+import abc
+import dataclasses
+
+import jax.numpy as jnp
+import jax_dataclasses
+import numpy as np
+
+import jaxsim.math
+import jaxsim.typing as jtp
+from jaxsim import exceptions
+
+
+class Terrain(abc.ABC):
+    """
+    Base class for terrain models.
+
+    Attributes:
+        delta: The delta value used for numerical differentiation.
+    """
+
+    delta = 0.010
+
+    @abc.abstractmethod
+    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
+        """
+        Compute the height of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The height of the terrain at the specified location.
+        """
+
+        pass
+
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+        """
+        Compute the normal vector of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The normal vector of the terrain surface at the specified location.
+        """
+
+        # https://stackoverflow.com/a/5282364
+        h_xp = self.height(x=x + self.delta, y=y)
+        h_xm = self.height(x=x - self.delta, y=y)
+        h_yp = self.height(x=x, y=y + self.delta)
+        h_ym = self.height(x=x, y=y - self.delta)
+
+        n = jnp.array(
+            [(h_xm - h_xp) / (2 * self.delta), (h_ym - h_yp) / (2 * self.delta), 1.0]
+        )
+
+        return n / jaxsim.math.safe_norm(n)
+
+
+@jax_dataclasses.pytree_dataclass
+class FlatTerrain(Terrain):
+    """
+    Represents a terrain model with a flat surface and a constant height.
+    """
+
+    _height: float = dataclasses.field(default=0.0, kw_only=True)
+
+    @staticmethod
+    def build(height: jtp.FloatLike = 0.0) -> FlatTerrain:
+        """
+        Create a FlatTerrain instance with a specified height.
+
+        Args:
+            height: The height of the flat terrain.
+
+        Returns:
+            FlatTerrain: A FlatTerrain instance.
+        """
+
+        return FlatTerrain(_height=float(height))
+
+    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
+        """
+        Compute the height of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The height of the terrain at the specified location.
+        """
+
+        return jnp.array(self._height, dtype=float)
+
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+        """
+        Compute the normal vector of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The normal vector of the terrain surface at the specified location.
+        """
+
+        return jnp.array([0.0, 0.0, 1.0], dtype=float)
+
+    def __hash__(self) -> int:
+
+        return hash(self._height)
+
+    def __eq__(self, other: FlatTerrain) -> bool:
+
+        if not isinstance(other, FlatTerrain):
+            return False
+
+        return self._height == other._height
+
+
+@jax_dataclasses.pytree_dataclass
+class PlaneTerrain(FlatTerrain):
+    """
+    Represents a terrain model with a flat surface defined by a normal vector.
+    """
+
+    _normal: tuple[float, float, float] = jax_dataclasses.field(
+        default=(0.0, 0.0, 1.0), kw_only=True
+    )
+
+    @staticmethod
+    def build(height: jtp.FloatLike = 0.0, *, normal: jtp.VectorLike) -> PlaneTerrain:
+        """
+        Create a PlaneTerrain instance with a specified plane normal vector.
+
+        Args:
+            normal: The normal vector of the terrain plane.
+            height: The height of the plane over the origin.
+
+        Returns:
+            PlaneTerrain: A PlaneTerrain instance.
+        """
+
+        normal = jnp.array(normal, dtype=float)
+        height = jnp.array(height, dtype=float)
+
+        if normal.shape != (3,):
+            msg = "Expected a 3D vector for the plane normal, got '{}'."
+            raise ValueError(msg.format(normal.shape))
+
+        # Make sure that the plane normal is a unit vector.
+        normal = normal / jnp.linalg.norm(normal)
+
+        return PlaneTerrain(
+            _height=height.item(),
+            _normal=tuple(normal.tolist()),
+        )
+
+    def normal(
+        self, x: jtp.FloatLike | None = None, y: jtp.FloatLike | None = None
+    ) -> jtp.Vector:
+        """
+        Compute the normal vector of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The normal vector of the terrain surface at the specified location.
+        """
+
+        return jnp.array(self._normal, dtype=float)
+
+    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
+        """
+        Compute the height of the terrain at a specific (x, y) location on a plane.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The height of the terrain at the specified location on the plane.
+        """
+
+        # Equation of the plane:      A x + B y + C z + D = 0
+        # Normal vector coordinates:  (A, B, C)
+        # The height over the origin: -D/C
+
+        # Get the plane equation coefficients from the terrain normal.
+        A, B, C = self._normal
+
+        exceptions.raise_value_error_if(
+            condition=jnp.allclose(C, 0.0),
+            msg="The z component of the normal cannot be zero.",
+        )
+
+        # Compute the final coefficient D considering the terrain height.
+        D = -C * self._height
+
+        # Invert the plane equation to get the height at the given (x, y) coordinates.
+        return jnp.array(-(A * x + B * y + D) / C).astype(float)
+
+    def __hash__(self) -> int:
+
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return hash(
+            (
+                hash(self._height),
+                HashedNumpyArray.hash_of_array(
+                    array=np.array(self._normal, dtype=float)
+                ),
+            )
+        )
+
+    def __eq__(self, other: PlaneTerrain) -> bool:
+
+        if not isinstance(other, PlaneTerrain):
+            return False
+
+        if not (
+            np.allclose(self._height, other._height)
+            and np.allclose(
+                np.array(self._normal, dtype=float),
+                np.array(other._normal, dtype=float),
+            )
+        ):
+            return False
+
+        return True

jaxsim/typing.py

@@ -1,4 +1,5 @@
-from typing import Any, Hashable
+from collections.abc import Hashable
+from typing import Any, TypeVar
 
 import jax
 
@@ -6,34 +7,33 @@ import jax
 # JAX types
 # =========
 
-ScalarJax = jax.Array
-IntJax = ScalarJax
-BoolJax = ScalarJax
-FloatJax = ScalarJax
-
-ArrayJax = jax.Array
-VectorJax = ArrayJax
-MatrixJax = ArrayJax
+Array = jax.Array
+Scalar = Array
+Vector = Array
+Matrix = Array
 
-PyTree = (
-    dict[Hashable, "PyTree"] | list["PyTree"] | tuple["PyTree"] | None | jax.Array | Any
+Int = Scalar
+Bool = Scalar
+Float = Scalar
+
+PyTree: object = (
+    dict[Hashable, TypeVar("PyTree")]
+    | list[TypeVar("PyTree")]
+    | tuple[TypeVar("PyTree")]
+    | jax.Array
+    | Any
+    | None
 )
 
 # =======================
 # Mixed JAX / NumPy types
 # =======================
 
-Array = jax.typing.ArrayLike
-Vector = Array
-Matrix = Array
-
-Int = int | IntJax
-Bool = bool | ArrayJax
-Float = float | FloatJax
+ArrayLike = jax.typing.ArrayLike | tuple
+ScalarLike = int | float | Scalar | ArrayLike
+VectorLike = Vector | ArrayLike | tuple
+MatrixLike = Matrix | ArrayLike
 
-ArrayLike = Array
-VectorLike = Vector
-MatrixLike = Matrix
-IntLike = Int
-BoolLike = Bool
-FloatLike = Float
+IntLike = int | Int | jax.typing.ArrayLike
+BoolLike = bool | Bool | jax.typing.ArrayLike
+FloatLike = float | Float | jax.typing.ArrayLike

jaxsim/utils/__init__.py

@@ -2,7 +2,4 @@ from jax_dataclasses._copy_and_mutate import _Mutability as Mutability
 
 from .jaxsim_dataclass import JaxsimDataclass
 from .tracing import not_tracing, tracing
-from .vmappable import Vmappable
-
-# Leave this below the others to prevent circular imports
-from .oop import jax_tf  # isort: skip
+from .wrappers import HashedNumpyArray, HashlessObject