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

jaxsim/utils/jaxsim_dataclass.py

@@ -1,8 +1,9 @@
 import abc
 import contextlib
-import copy
 import dataclasses
-from typing import ClassVar, Generator
+import functools
+from collections.abc import Callable, Iterator, Sequence
+from typing import Any, ClassVar
 
 import jax.flatten_util
 import jax_dataclasses
@@ -19,51 +20,224 @@ except ImportError:
 
 @jax_dataclasses.pytree_dataclass
 class JaxsimDataclass(abc.ABC):
-    """"""
+    """Class extending `jax_dataclasses.pytree_dataclass` instances with utilities."""
 
     # This attribute is set by jax_dataclasses
     __mutability__: ClassVar[Mutability] = Mutability.FROZEN
 
     @contextlib.contextmanager
-    def editable(self: Self, validate: bool = True) -> Generator[Self, None, None]:
-        """"""
+    def editable(self: Self, validate: bool = True) -> Iterator[Self]:
+        """
+        Context manager to operate on a mutable copy of the object.
+
+        Args:
+            validate: Whether to validate the output PyTree upon exiting the context.
+
+        Yields:
+            A mutable copy of the object.
+
+        Note:
+            This context manager is useful to operate on an r/w copy of a PyTree making
+            sure that the output object does not trigger JIT recompilations.
+        """
 
         mutability = (
             Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
         )
 
-        with JaxsimDataclass.mutable_context(self.copy(), mutability=mutability) as obj:
+        with self.copy().mutable_context(mutability=mutability) as obj:
             yield obj
 
     @contextlib.contextmanager
     def mutable_context(
-        self: Self, mutability: Mutability, restore_after_exception: bool = True
-    ) -> Generator[Self, None, None]:
-        """"""
+        self: Self,
+        mutability: Mutability = Mutability.MUTABLE,
+        restore_after_exception: bool = True,
+    ) -> Iterator[Self]:
+        """
+        Context manager to temporarily change the mutability of the object.
+
+        Args:
+            mutability: The mutability to set.
+            restore_after_exception:
+                Whether to restore the original object in case of an exception
+                occurring within the context.
+
+        Yields:
+            The object with the new mutability.
+
+        Note:
+            This context manager is useful to operate in place on a PyTree without
+            the need to make a copy while optionally keeping active the checks on
+            the PyTree structure, shapes, and dtypes.
+        """
 
         if restore_after_exception:
             self_copy = self.copy()
 
-        original_mutability = self._mutability()
+        original_mutability = self.mutability()
 
-        def restore_self():
-            self._set_mutability(mutability=Mutability.MUTABLE)
+        original_dtypes = JaxsimDataclass.get_leaf_dtypes(tree=self)
+        original_shapes = JaxsimDataclass.get_leaf_shapes(tree=self)
+        original_weak_types = JaxsimDataclass.get_leaf_weak_types(tree=self)
+        original_structure = jax.tree_util.tree_structure(tree=self)
+
+        def restore_self() -> None:
+            self.set_mutability(mutability=Mutability.MUTABLE_NO_VALIDATION)
             for f in dataclasses.fields(self_copy):
                 setattr(self, f.name, getattr(self_copy, f.name))
 
         try:
-            self._set_mutability(mutability)
+            self.set_mutability(mutability=mutability)
             yield self
+
+            if mutability is not Mutability.MUTABLE_NO_VALIDATION:
+                new_structure = jax.tree_util.tree_structure(tree=self)
+                if original_structure != new_structure:
+                    msg = "Pytree structure has changed from {} to {}"
+                    raise ValueError(msg.format(original_structure, new_structure))
+
+                new_shapes = JaxsimDataclass.get_leaf_shapes(tree=self)
+                if original_shapes != new_shapes:
+                    msg = "Leaves shapes have changed from {} to {}"
+                    raise ValueError(msg.format(original_shapes, new_shapes))
+
+                new_dtypes = JaxsimDataclass.get_leaf_dtypes(tree=self)
+                if original_dtypes != new_dtypes:
+                    msg = "Leaves dtypes have changed from {} to {}"
+                    raise ValueError(msg.format(original_dtypes, new_dtypes))
+
+                new_weak_types = JaxsimDataclass.get_leaf_weak_types(tree=self)
+                if original_weak_types != new_weak_types:
+                    msg = "Leaves weak types have changed from {} to {}"
+                    raise ValueError(msg.format(original_weak_types, new_weak_types))
+
         except Exception as e:
             if restore_after_exception:
                 restore_self()
-            self._set_mutability(original_mutability)
+            self.set_mutability(original_mutability)
             raise e
+
         finally:
-            self._set_mutability(original_mutability)
+            self.set_mutability(original_mutability)
+
+    @staticmethod
+    def get_leaf_shapes(tree: jtp.PyTree) -> tuple[tuple[int, ...] | None]:
+        """
+        Get the leaf shapes of a PyTree.
+
+        Args:
+            tree: The PyTree to consider.
+
+        Returns:
+            A tuple containing the leaf shapes of the PyTree or `None` is the leaf is
+            not a numpy-like array.
+        """
+
+        return tuple(
+            map(
+                lambda leaf: getattr(leaf, "shape", None),
+                jax.tree_util.tree_leaves(tree),
+            )
+        )
+
+    @staticmethod
+    def get_leaf_dtypes(tree: jtp.PyTree) -> tuple:
+        """
+        Get the leaf dtypes of a PyTree.
+
+        Args:
+            tree: The PyTree to consider.
+
+        Returns:
+            A tuple containing the leaf dtypes of the PyTree or `None` is the leaf is
+            not a numpy-like array.
+        """
+
+        return tuple(
+            map(
+                lambda leaf: getattr(leaf, "dtype", None),
+                jax.tree_util.tree_leaves(tree),
+            )
+        )
+
+    @staticmethod
+    def get_leaf_weak_types(tree: jtp.PyTree) -> tuple[bool, ...]:
+        """
+        Get the leaf weak types of a PyTree.
+
+        Args:
+            tree: The PyTree to consider.
+
+        Returns:
+            A tuple marking whether the leaf contains a JAX array with weak type.
+        """
+
+        return tuple(
+            map(
+                lambda leaf: getattr(leaf, "weak_type", None),
+                jax.tree_util.tree_leaves(tree),
+            )
+        )
+
+    @staticmethod
+    def check_compatibility(*trees: Sequence[Any]) -> None:
+        """
+        Check whether the PyTrees are compatible in structure, shape, and dtype.
+
+        Args:
+            *trees: The PyTrees to compare.
+
+        Raises:
+            ValueError: If the PyTrees have incompatible structures, shapes, or dtypes.
+        """
+
+        target_structure = jax.tree_util.tree_structure(trees[0])
+
+        compatible_structure = functools.reduce(
+            lambda compatible, tree: compatible
+            and jax.tree_util.tree_structure(tree) == target_structure,
+            trees[1:],
+            True,
+        )
+
+        if not compatible_structure:
+            raise ValueError("Pytrees have incompatible structures.")
+
+        target_shapes = JaxsimDataclass.get_leaf_shapes(trees[0])
+
+        compatible_shapes = functools.reduce(
+            lambda compatible, tree: compatible
+            and JaxsimDataclass.get_leaf_shapes(tree) == target_shapes,
+            trees[1:],
+            True,
+        )
+
+        if not compatible_shapes:
+            raise ValueError("Pytrees have incompatible shapes.")
+
+        target_dtypes = JaxsimDataclass.get_leaf_dtypes(trees[0])
+
+        compatible_dtypes = functools.reduce(
+            lambda compatible, tree: compatible
+            and JaxsimDataclass.get_leaf_dtypes(tree) == target_dtypes,
+            trees[1:],
+            True,
+        )
+
+        if not compatible_dtypes:
+            raise ValueError("Pytrees have incompatible dtypes.")
 
     def is_mutable(self, validate: bool = False) -> bool:
-        """"""
+        """
+        Check whether the object is mutable.
+
+        Args:
+            validate: Additionally checks if the object also has validation enabled.
+
+        Returns:
+            True if the object is mutable, False otherwise.
+        """
 
         return (
             self.__mutability__ is Mutability.MUTABLE
@@ -71,39 +245,120 @@ class JaxsimDataclass(abc.ABC):
             else self.__mutability__ is Mutability.MUTABLE_NO_VALIDATION
         )
 
-    def set_mutability(self, mutable: bool = True, validate: bool = False) -> None:
-        if not mutable:
-            mutability = Mutability.FROZEN
-        else:
-            mutability = (
-                Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
-            )
+    def mutability(self) -> Mutability:
+        """
+        Get the mutability type of the object.
 
-        self._set_mutability(mutability=mutability)
+        Returns:
+            The mutability type of the object.
+        """
 
-    def _mutability(self) -> Mutability:
         return self.__mutability__
 
-    def _set_mutability(self, mutability: Mutability) -> None:
+    def set_mutability(self, mutability: Mutability) -> None:
+        """
+        Set the mutability of the object in-place.
+
+        Args:
+            mutability: The desired mutability type.
+        """
+
         jax_dataclasses._copy_and_mutate._mark_mutable(
             self, mutable=mutability, visited=set()
         )
 
     def mutable(self: Self, mutable: bool = True, validate: bool = False) -> Self:
-        self.set_mutability(mutable=mutable, validate=validate)
+        """
+        Return a mutable reference of the object.
+
+        Args:
+            mutable: Whether to make the object mutable.
+            validate: Whether to enable validation on the object.
+
+        Returns:
+            A mutable reference of the object.
+        """
+
+        if mutable:
+            mutability = (
+                Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
+            )
+        else:
+            mutability = Mutability.FROZEN
+
+        self.set_mutability(mutability=mutability)
         return self
 
     def copy(self: Self) -> Self:
-        obj = jax.tree_util.tree_map(lambda leaf: leaf, self)
-        obj._set_mutability(mutability=self._mutability())
+        """
+        Return a copy of the object.
+
+        Returns:
+            A copy of the object.
+        """
+
+        # Make a copy calling tree_map.
+        obj = jax.tree.map(lambda leaf: leaf, self)
+
+        # Make sure that the copied object and all the copied leaves have the same
+        # mutability of the original object.
+        obj.set_mutability(mutability=self.mutability())
+
         return obj
 
     def replace(self: Self, validate: bool = True, **kwargs) -> Self:
-        with self.editable(validate=validate) as obj:
-            _ = [obj.__setattr__(k, copy.copy(v)) for k, v in kwargs.items()]
+        """
+        Return a new object replacing in-place the specified fields with new values.
+
+        Args:
+            validate: Whether to validate that the new fields do not alter the PyTree.
+            **kwargs: The fields to replace.
+
+        Returns:
+            A reference of the object with the specified fields replaced.
+        """
+
+        # Use the dataclasses replace method.
+        obj = dataclasses.replace(self, **kwargs)
+
+        if validate:
+            JaxsimDataclass.check_compatibility(self, obj)
+
+        # Make sure that all the new leaves have the same mutability of the object.
+        obj.set_mutability(mutability=self.mutability())
 
-        obj._set_mutability(mutability=self._mutability())
         return obj
 
-    def flatten(self) -> jtp.VectorJax:
-        return jax.flatten_util.ravel_pytree(self)[0]
+    def flatten(self) -> jtp.Vector:
+        """
+        Flatten the object into a 1D vector.
+
+        Returns:
+            A 1D vector containing the flattened object.
+        """
+
+        return self.flatten_fn()(self)
+
+    @classmethod
+    def flatten_fn(cls: type[Self]) -> Callable[[Self], jtp.Vector]:
+        """
+        Return a function to flatten the object into a 1D vector.
+
+        Returns:
+            A function to flatten the object into a 1D vector.
+        """
+
+        return lambda pytree: jax.flatten_util.ravel_pytree(pytree)[0]
+
+    def unflatten_fn(self: Self) -> Callable[[jtp.Vector], Self]:
+        """
+        Return a function to unflatten a 1D vector into the object.
+
+        Returns:
+            A function to unflatten a 1D vector into the object.
+
+        Notes:
+            Due to JAX internals, the function to unflatten a PyTree needs to be
+            created from an existing instance of the PyTree.
+        """
+        return jax.flatten_util.ravel_pytree(self)[1]

jaxsim/utils/tracing.py

@@ -6,20 +6,14 @@ import jax.interpreters.partial_eval
 
 
 def tracing(var: Any) -> bool | jax.Array:
-    """Returns True if the variable is being traced by JAX, False otherwise."""
+    """Return True if the variable is being traced by JAX, False otherwise."""
 
-    return jax.numpy.array(
-        [
-            isinstance(var, t)
-            for t in (
-                jax._src.core.Tracer,
-                jax.interpreters.partial_eval.DynamicJaxprTracer,
-            )
-        ]
-    ).any()
+    return isinstance(
+        var, jax._src.core.Tracer | jax.interpreters.partial_eval.DynamicJaxprTracer
+    )
 
 
 def not_tracing(var: Any) -> bool | jax.Array:
-    """Returns True if the variable is not being traced by JAX, False otherwise."""
+    """Return True if the variable is not being traced by JAX, False otherwise."""
 
     return True if tracing(var) is False else False

jaxsim/utils/wrappers.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable
+from typing import Generic, TypeVar
+
+import jax
+import jax_dataclasses
+import numpy as np
+import numpy.typing as npt
+
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class HashlessObject(Generic[T]):
+    """
+    A class that wraps an object and makes it hashless.
+
+    This is useful for creating particular JAX pytrees.
+    For example, to create a pytree with a static leaf that is ignored
+    by JAX when it compares two instances to trigger a JIT recompilation.
+    """
+
+    obj: T
+
+    def get(self: HashlessObject[T]) -> T:
+        """
+        Get the wrapped object.
+        """
+        return self.obj
+
+    def __hash__(self) -> int:
+
+        return 0
+
+    def __eq__(self, other: HashlessObject[T]) -> bool:
+
+        if not isinstance(other, HashlessObject) and isinstance(
+            other.get(), type(self.get())
+        ):
+            return False
+
+        return hash(self) == hash(other)
+
+
+@dataclasses.dataclass
+class CustomHashedObject(Generic[T]):
+    """
+    A class that wraps an object and computes its hash with a custom hash function.
+    """
+
+    obj: T
+
+    hash_function: Callable[[T], int] = hash
+
+    def get(self: CustomHashedObject[T]) -> T:
+        """
+        Get the wrapped object.
+        """
+        return self.obj
+
+    def __hash__(self) -> int:
+
+        return self.hash_function(self.obj)
+
+    def __eq__(self, other: CustomHashedObject[T]) -> bool:
+
+        if not isinstance(other, CustomHashedObject) and isinstance(
+            other.get(), type(self.get())
+        ):
+            return False
+
+        return hash(self) == hash(other)
+
+
+@jax_dataclasses.pytree_dataclass
+class HashedNumpyArray:
+    """
+    A class that wraps a numpy array and makes it hashable.
+
+    This is useful for creating particular JAX pytrees.
+    For example, to create a pytree with a plain NumPy or JAX NumPy array as static leaf.
+
+    Note:
+        Calculating with the wrapper class the hash of a very large array can be
+        very expensive. If the array is large and only the equality operator is needed,
+        set `large_array=True` to use a faster comparison method.
+    """
+
+    array: jax.Array | npt.NDArray
+
+    precision: float | None = dataclasses.field(
+        default=1e-9, repr=False, compare=False, hash=False
+    )
+
+    large_array: jax_dataclasses.Static[bool] = dataclasses.field(
+        default=False, repr=False, compare=False, hash=False
+    )
+
+    def get(self) -> jax.Array | npt.NDArray:
+        """
+        Get the wrapped array.
+        """
+        return self.array
+
+    def __hash__(self) -> int:
+
+        return HashedNumpyArray.hash_of_array(
+            array=self.array, precision=self.precision
+        )
+
+    def __eq__(self, other: HashedNumpyArray) -> bool:
+
+        if not isinstance(other, HashedNumpyArray):
+            return False
+
+        if self.large_array:
+            return np.allclose(
+                self.array,
+                other.array,
+                **(dict(atol=self.precision) if self.precision is not None else {}),
+            )
+
+        return hash(self) == hash(other)
+
+    @staticmethod
+    def hash_of_array(
+        array: jax.Array | npt.NDArray, precision: float | None = 1e-9
+    ) -> int:
+        """
+        Calculate the hash of a NumPy array.
+
+        Args:
+            array: The array to hash.
+            precision: Optionally limit the precision over which the hash is computed.
+
+        Returns:
+            The hash of the array.
+        """
+
+        array = np.array(array).flatten()
+
+        array = np.where(array == np.nan, hash(np.nan), array)
+        array = np.where(array == np.inf, hash(np.inf), array)
+        array = np.where(array == -np.inf, hash(-np.inf), array)
+
+        if precision is not None:
+
+            integer1 = (array * precision).astype(int)
+            integer2 = (array - integer1 / precision).astype(int)
+
+            decimal_array = ((array - integer1 * 1e9 - integer2) / precision).astype(
+                int
+            )
+
+            array = np.hstack([integer1, integer2, decimal_array]).astype(int)
+
+        return hash(tuple(array.tolist()))

{jaxsim-0.2.dev191.dist-info → jaxsim-0.6.1.dev2.dist-info}/LICENSE

@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2022, Artificial and Mechanical Intelligence 
+Copyright (c) 2022, Artificial and Mechanical Intelligence
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without