quimb.experimental.autojittn ============================ .. py:module:: quimb.experimental.autojittn .. autoapi-nested-parse:: Decorator for automatically just in time compiling tensor network functions. TODO:: - [ ] go via an intermediate pytree / array function, that could be shared e.g. with the TNOptimizer class. Classes ------- .. autoapisummary:: quimb.experimental.autojittn.AutojittedTN Functions --------- .. autoapisummary:: quimb.experimental.autojittn.autojit_tn Module Contents --------------- .. py:class:: AutojittedTN(fn, decorator=ar.autojit, check_inputs=True, **decorator_opts) Class to hold the ``autojit_tn`` decorated function callable. .. py:attribute:: fn .. py:attribute:: fn_store .. py:attribute:: decorator_opts .. py:attribute:: check_inputs :value: True .. py:attribute:: decorator .. py:method:: setup_fn(tn, *args, **kwargs) .. py:method:: __call__(tn, *args, **kwargs) .. py:function:: autojit_tn(fn=None, decorator=ar.autojit, check_inputs=True, **decorator_opts) Decorate a tensor network function to be just in time compiled / traced. This traces solely array operations resulting in a completely static computational graph with no side-effects. The resulting function can be much faster if called repeatedly with only numeric changes, or hardware accelerated if a library such as ``jax`` is used. :param fn: The function to be decorated. It should take as its first argument a :class:`~quimb.tensor.tensor_core.TensorNetwork` and return either act inplace on it or return a raw scalar or array. :type fn: callable :param decorator: The decorator to use to wrap the underlying array function. For example ``jax.jit``. Defaults to ``autoray.autojit``. :type decorator: callable :param check_inputs: Whether to check the inputs to the function every call to see if a new compiled function needs to be generated. If ``False`` the same compiled function will be used for all inputs which might be incorrect. Defaults to ``True``. :type check_inputs: bool, optional :param decorator_opts: Options to pass to the decorator, e.g. ``backend`` for ``autoray.autojit``.