jaxsim 0.2.dev188__py3-none-any.whl → 0.2.dev364__py3-none-any.whl

jaxsim/utils/vmappable.py

@@ -1,117 +0,0 @@
-import dataclasses
-from typing import Type
-
-import jax
-import jax.numpy as jnp
-import jax_dataclasses
-
-from . import JaxsimDataclass, Mutability
-
-try:
-    from typing import Self
-except ImportError:
-    from typing_extensions import Self
-
-
-@jax_dataclasses.pytree_dataclass
-class Vmappable(JaxsimDataclass):
-    """Abstract class with utilities for vmappable pytrees."""
-
-    batch_size: jax_dataclasses.Static[int] = dataclasses.field(
-        default=int(0), repr=False, compare=False, hash=False, kw_only=True
-    )
-
-    @property
-    def vectorized(self) -> bool:
-        """Marks this pytree as vectorized."""
-
-        return self.batch_size > 0
-
-    @classmethod
-    def build_from_list(cls: Type[Self], list_of_obj: list[Self]) -> Self:
-        """
-        Build a vectorized pytree from a list of pytree of the same type.
-
-        Args:
-            list_of_obj: The list of pytrees to vectorize.
-
-        Returns:
-            The vectorized pytree having as leaves the stacked leaves of the input list.
-        """
-
-        if set(type(el) for el in list_of_obj) != {cls}:
-            msg = "The input list must contain only objects of type '{}'"
-            raise ValueError(msg.format(cls.__name__))
-
-        # Create a pytree by stacking all the leafs of the input list
-        data_vec: Vmappable = jax.tree_map(
-            lambda *leafs: jnp.array(leafs), *list_of_obj
-        )
-
-        # Store the batch dimension
-        with data_vec.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
-            data_vec.batch_size = len(list_of_obj)
-
-        # Detect the most common mutability in the input list
-        mutabilities = [e._mutability() for e in list_of_obj]
-        mutability = max(set(mutabilities), key=mutabilities.count)
-
-        # Update the mutability of the vectorized pytree
-        data_vec._set_mutability(mutability)
-
-        return data_vec
-
-    def vectorize(self: Self, batch_size: int) -> Self:
-        """
-        Return a vectorized version of this pytree.
-
-        Args:
-            batch_size: The batch size.
-
-        Returns:
-            A vectorized version of this pytree obtained by stacking the leaves of the
-            original pytree along a new batch dimension (the first one).
-        """
-
-        if self.vectorized:
-            raise RuntimeError("Cannot vectorize an already vectorized object")
-
-        if batch_size == 0:
-            return self.copy()
-
-        # TODO validate if mutability is maintained
-
-        return self.__class__.build_from_list(list_of_obj=[self] * batch_size)
-
-    def extract_element(self: Self, index: int) -> Self:
-        """
-        Extract the i-th element from a vectorized pytree.
-
-        Args:
-            index: The index of the element to extract.
-
-        Returns:
-            A non vectorized pytree obtained by extracting the i-th element from the
-            vectorized pytree.
-        """
-
-        if index < 0:
-            raise ValueError("The index of the desired element cannot be negative")
-
-        if index == 0 and self.batch_size == 0:
-            return self.copy()
-
-        if not self.vectorized:
-            raise RuntimeError("Cannot extract elements from a non-vectorized object")
-
-        if index >= self.batch_size:
-            raise ValueError("The index must be smaller than the batch size")
-
-        # Get the i-th pytree by extracting the i-th element from the vectorized pytree
-        data = jax.tree_map(lambda leaf: leaf[index], self)
-
-        # Update the batch size of the extracted scalar pytree
-        with data.mutable_context(mutability=Mutability.MUTABLE):
-            data.batch_size = 0
-
-        return data

jaxsim-0.2.dev188.dist-info/RECORD

@@ -1,81 +0,0 @@
-jaxsim/__init__.py,sha256=acqCVEg71ekNOwazlwA06dXeaPJh3Sr_e5JB8pcI5lo,1862
-jaxsim/_version.py,sha256=25Ei2WVWsVwEC3k_WuiVBPfYOrVetMFq0zkbn_gOAyY,423
-jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
-jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
-jaxsim/api/__init__.py,sha256=UehgJtcV9BYlX6vec-f_IgnMcq-F9BU_STsR0RDjEjA,53
-jaxsim/api/common.py,sha256=I8YgQwsfLcuWXtyp_2IVIyAnWMsfcf4ICmgkMfaGyRs,6475
-jaxsim/api/contact.py,sha256=9LX9FNXSXjMonJgg6svyrFcRfhptfBMywbswBUebHeE,5945
-jaxsim/api/data.py,sha256=mKgyYEJ6JpuHcXe_ddzw13Y553SV_92AzH3mT4FaxBo,25588
-jaxsim/api/joint.py,sha256=J65MgjMzp2O980wIEf2TxG_VCgczUsexGgd4YJ6rxiE,3615
-jaxsim/api/link.py,sha256=Fpbtr2JF9fAouOX-lGAmajF2ZEz6q8ttDYsCiqpLxEk,6939
-jaxsim/api/model.py,sha256=sxFJADqHXXm8nCvwYiorOgl9nzGhA_QvI1dOSOv9its,35206
-jaxsim/api/ode.py,sha256=cbWrXe7s5O3UCpWpI77Lg8Pujr_SiaQfLc7ELBqnJQM,9492
-jaxsim/api/references.py,sha256=hiVrDfXF0jOYrtLpb2_-nmunz0p2maIDFx_NZYPOCks,15419
-jaxsim/high_level/__init__.py,sha256=aWYBCsYmEO76Qt4GEi91Hye_ifGFLvc_bpy9OQplz2o,69
-jaxsim/high_level/common.py,sha256=6nyRlFsNOLEy5JvLH70VPWeHGSL_ZKNxX3Q62ccqSuY,196
-jaxsim/high_level/joint.py,sha256=0WF0QWkZzP0SXw0QYpn3PAwdZq0_uXFr2_f1OATiOBA,4089
-jaxsim/high_level/link.py,sha256=4kcBMh-3w9c-fkTYm3_sXfdwd3NwUm7jKf5BjwEge94,8010
-jaxsim/high_level/model.py,sha256=q29hSs0gH_UQrUyVzuCyF9YRB79WEWyHcVvV_1B9LrA,57054
-jaxsim/integrators/__init__.py,sha256=x4kYr7k-nC0X7HH2bTXzGOl_3B9TBN9blw5BBxQtyQ0,72
-jaxsim/integrators/common.py,sha256=kL4EqUjWZUNPtDrCcrStx-ZTxEYDS-W2WPQ0g-QqHqs,17567
-jaxsim/integrators/fixed_step.py,sha256=_gChRm0RHjDtst-aWjyQKUys6TyYhZqdWap7vLf5LaE,4291
-jaxsim/math/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/math/adjoint.py,sha256=ImkOkWHQKMukBprLTsOPpSuqb1NNPA3_t447zRVo79s,3779
-jaxsim/math/conv.py,sha256=jbr9MU_vGtTpLTQpbqwAT4huF51skCfORFUBmKSdlaI,3138
-jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
-jaxsim/math/inertia.py,sha256=UAB7ym4gXFanejcs_ovZMpteHCc6poWYmt-mLmd5hhk,1640
-jaxsim/math/joint.py,sha256=vNgxPfn6rkHHrIt_2vYwS832mDMYmza_7cnle2Jh8sA,3012
-jaxsim/math/plucker.py,sha256=44NvKVbcZoG8ivFN1BeXxDpuSFdEre1Q6ZXvhnmIiPY,2282
-jaxsim/math/quaternion.py,sha256=ToyRnAWU0JvKSSSX2vaJeSw2lMa5BGU72DjtonOUw0k,3685
-jaxsim/math/rotation.py,sha256=MHOnrpS5Sf4rszhOpZ8w7qXFkEl7UMltYimqqsuYuuU,2187
-jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
-jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
-jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=8sXc_tsDFWBYl8nesgFarYd3hA-PESLMrXsnR3Siz1Y,16400
-jaxsim/mujoco/model.py,sha256=0kG2GERxjVFqWZ1K3352rgUNfchB4kRtIrsvv4pS4oc,10766
-jaxsim/mujoco/visualizer.py,sha256=-qg26t5tleTva6zzQmc5SdnlC8XZ1ZAwZ_lDjdwHJ0A,4400
-jaxsim/parsers/__init__.py,sha256=sonYi-bBWAoB04kp1mxT4uIORxjb7SdZ0ukGPmVx98Y,44
-jaxsim/parsers/kinematic_graph.py,sha256=5wQnbzu8JE0bbnLRxK4ZsD_gQ9kbBpYbhUSzCMiNWko,23610
-jaxsim/parsers/descriptions/__init__.py,sha256=EbTfnrK3oCxA3pNv--YUwllJ6uICENvFgAdRbYtS9ts,238
-jaxsim/parsers/descriptions/collision.py,sha256=p8v24UyWNQFd3n5t3Z25IamfTqw19V9VlF5oyxeQXtU,2935
-jaxsim/parsers/descriptions/joint.py,sha256=hpH0ANvIhbEQk-NGRmWIvPv3lXW385TBIMWNgz5rzM4,4106
-jaxsim/parsers/descriptions/link.py,sha256=SONPaSwtNhIX93RIVng8Fb_Y7I5h3sk-5rsqd5U7Fmw,2493
-jaxsim/parsers/descriptions/model.py,sha256=wenuDrjoBf6prkzm9WyYT0nFWc0l6WBpKNjLoRUDPxo,8937
-jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=mFi1baSJte6EMmWLpVjVuCpicfAAF48aFUzoKzYzPpo,12555
-jaxsim/parsers/rod/utils.py,sha256=C3RfLSnHHR7rgZhnM15QjM_tpJFHKf1Jf2mOsdep3LM,6492
-jaxsim/physics/__init__.py,sha256=TKH7IqQi39eobWT03b820ky6tWVWFFWcO8YPayNpnZc,216
-jaxsim/physics/algos/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/physics/algos/aba.py,sha256=nv7ezClJT6ygx-pui3oZzUrOc-eMlvRYUzDNwCLOUwE,7547
-jaxsim/physics/algos/aba_motors.py,sha256=gGdFrFvImTwYs2V2lONHh85DObppee1MkDUoUGGoc24,8362
-jaxsim/physics/algos/crba.py,sha256=5OcPNaeGMZ3IcemhX4cdXvUUtysJv1SB92QkZTVtizc,4478
-jaxsim/physics/algos/forward_kinematics.py,sha256=myiqBGk5Nv-OIJpkNFVOby2FdHEJM3teEOQrL77dExQ,2513
-jaxsim/physics/algos/jacobian.py,sha256=HQk7UBYwWr2_qIYZBil_t9nqoL-TQB8hi-JHdPcSA0g,3166
-jaxsim/physics/algos/rnea.py,sha256=bIv6TTwdCZVj_s9Fd5i_ZTMib4W2_sn8gxCbMRrzRcg,5322
-jaxsim/physics/algos/rnea_motors.py,sha256=gaiEAOfpw9kZBVImfNKKwpgIgM8eygVu--wszh3npJc,5684
-jaxsim/physics/algos/soft_contacts.py,sha256=mPNs5a3JAZ5dRTWbNMhZNV57UyYxQTqjRQXGA5pk3II,18424
-jaxsim/physics/algos/terrain.py,sha256=Gw9-1AjU4c4Yd2xzo0i-fgWwYlroj03TjScJsz_2m9o,2228
-jaxsim/physics/algos/utils.py,sha256=0OiELbXtv5Jink3l-vMK_OGHgGkZ_wTAAclcd7vDKoc,2230
-jaxsim/physics/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/physics/model/ground_contact.py,sha256=mva-yDzYHREmgUu8jGJmIAsf66_SF6ZISmN-XQQ9Ktw,1924
-jaxsim/physics/model/physics_model.py,sha256=izK5EaKFwNZVqD2hG69yZ5XOW5MmOFyP8_KojaWxBMk,14529
-jaxsim/physics/model/physics_model_state.py,sha256=Av42rJWgicqxDmZZu1T_clpnmeaX_daASYhA7cxHYxQ,9256
-jaxsim/simulation/__init__.py,sha256=WOWkzq7rMGa4xWvjNqTYtD0Nl4yLQtULGW1xU7hD9m0,182
-jaxsim/simulation/integrators.py,sha256=WIlL7xi4UocSlWg4Qms8-6puqRYnK5A4r7TJUNPg5g0,13022
-jaxsim/simulation/ode.py,sha256=ntq_iQPIw3SHj64CZWD2mHAKmt05ZgRpw2UwyTxHDOQ,10380
-jaxsim/simulation/ode_data.py,sha256=5nPynNSC5tRnhAVwOCkdv6FrDn148bC7zZaqVCKgowM,2952
-jaxsim/simulation/ode_integration.py,sha256=VDprQYoHEE_iI7ia1Mm3RyYl-LRvHU8dJEvRoGA4TFA,1947
-jaxsim/simulation/simulator.py,sha256=qCI5QG0WKkBC5GNqauSvI7rSlGD7CLttTzCgLED7iJM,18123
-jaxsim/simulation/simulator_callbacks.py,sha256=QWdY7dilmjrxeieWCB6RQ-cWpwLuUOK8fYWXpnnBcyU,2217
-jaxsim/simulation/utils.py,sha256=YdNA1mYGBAE7xVA-Dw7_OoBEuh0J8RS2X0RPQZf4c5E,329
-jaxsim/sixd/__init__.py,sha256=3tbynXQjvJ6X1IRcDH5eQBgBL0ilTSerDkS8SEF7a8A,62
-jaxsim/utils/__init__.py,sha256=UQsjrWMrhSQUfUXIIQhzd0kEioLHR3U0OhB-sIQqOd4,291
-jaxsim/utils/jaxsim_dataclass.py,sha256=FbjfEoCoYC_F-M3wUggXiEhQ7MMS-V_ciYQca-uSiMQ,3272
-jaxsim/utils/oop.py,sha256=LQhBXkSOD0zgYNJLO7Bl0FPRg-LvtvPzxyQa1WFP0rM,22616
-jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim/utils/vmappable.py,sha256=NqGL9nGFRI5OorCfnjXsjR_yXigzDxL0lW1YhQ_nMTY,3655
-jaxsim-0.2.dev188.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
-jaxsim-0.2.dev188.dist-info/METADATA,sha256=flkZoYKaLN8C6F_IlOOyfucfXnzwHxBN8Kgi9NXNJww,7734
-jaxsim-0.2.dev188.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-jaxsim-0.2.dev188.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.dev188.dist-info/RECORD,,