IMLCV.base.datastructures
=========================

.. py:module:: IMLCV.base.datastructures


Attributes
----------

.. autoapisummary::

   IMLCV.base.datastructures.TNode
   IMLCV.base.datastructures.field
   IMLCV.base.datastructures.P
   IMLCV.base.datastructures.P2
   IMLCV.base.datastructures.T


Classes
-------

.. autoapisummary::

   IMLCV.base.datastructures.MyPyTreeNode


Functions
---------

.. autoapisummary::

   IMLCV.base.datastructures.jit_decorator
   IMLCV.base.datastructures.vmap_decorator
   IMLCV.base.datastructures.custom_jvp_decorator
   IMLCV.base.datastructures.Partial_decorator


Module Contents
---------------

.. py:data:: TNode

.. py:class:: MyPyTreeNode(*args, **kwargs)

   Base class for dataclasses that should act like a JAX pytree node.


   .. py:method:: __init_subclass__(**kwargs)
      :classmethod:



   .. py:method:: replace(**overrides) -> TNode
      :abstractmethod:



.. py:data:: field

.. py:data:: P

.. py:data:: P2

.. py:data:: T

.. py:function:: jit_decorator(f: Callable[P, T], static_argnums: int | Sequence[int] | None = None, static_argnames: str | Iterable[str] | None = None)

.. py:function:: vmap_decorator(f: Callable[P, T], in_axes: int | Sequence[Any] | None = 0, out_axes: Any = 0)

.. py:function:: custom_jvp_decorator(f: Callable[P, T], nondiff_argnums: Sequence[int] = ()) -> Callable[P, T]

.. py:function:: Partial_decorator(f: Callable[P, T], *partial_args, **partial_kwargs)