Optional imports

Peek at JAX / Torch / TensorFlow (and any other module) without importing them. Used so extras stay unloaded until the caller imports them.

anytensor.optional

Peek at optional libraries that are already imported — never import them.

AnyTensor does not import JAX / Torch / TensorFlow unless the caller already did. Use :func:module_if_loaded instead of try: import … when a missing extra must stay unloaded (broken installs, memory, import order).

A callback can run now if the module is present, or later when it is first imported in this process — the pattern used by TorchScript divert and by JAX pytree registration on structured types. raises=True still registers that callback before raising, so a later import can complete the side effect.

module_if_loaded

module_if_loaded(name, callback=None, *, raises=False)

Return name if it is already imported, else None.

Never imports name. If callback is given and the module is already loaded, it is invoked immediately with the module. If not, callback is registered first and invoked later when that module is imported in this process (including when a submodule import loads the parent).

If raises is true and the module is still absent after that registration, raise :class:RuntimeError. The pending callback is kept, so a later import still runs it.

Parameters:

Name	Type	Description	Default
name	str	Absolute module name ("torch", "jax", "tensorflow").	required
callback	Optional[_Callback]	Optional callback(module) run now or on a future import.	None
raises	bool	If true, raise when name is not already imported.	False

Returns:

Type	Description
ModuleType | None	The loaded module, or None if it is not in sys.modules
ModuleType | None	(only when raises is false).

Examples:

Check without importing::

torch = module_if_loaded("torch")
if torch is None:
    return

Require an extra that must already be imported::

jax = module_if_loaded("jax", raises=True)

Register a side effect for now-or-later (JAX pytree, TorchScript)::

def _register_jax_pytree(jax):
    jax.tree_util.register_pytree_node(Ragged, flatten, unflatten)

module_if_loaded("jax", _register_jax_pytree)
# or module_if_loaded("jax", _register_jax_pytree, raises=True)
# to fail now while still running the callback on a later import.

Source code in anytensor/optional.py

def module_if_loaded(
    name: str, callback: Optional[_Callback] = None, *, raises: bool = False
) -> ModuleType | None:
    """Return ``name`` if it is already imported, else ``None``.

    Never imports ``name``. If ``callback`` is given and the module is
    already loaded, it is invoked immediately with the module. If not,
    ``callback`` is registered first and invoked later when that module is
    imported in this process (including when a submodule import loads the
    parent).

    If ``raises`` is true and the module is still absent after that
    registration, raise :class:`RuntimeError`. The pending callback is kept,
    so a later import still runs it.

    Args:
        name: Absolute module name (``"torch"``, ``"jax"``, ``"tensorflow"``).
        callback: Optional ``callback(module)`` run now or on a future import.
        raises: If true, raise when ``name`` is not already imported.

    Returns:
        The loaded module, or ``None`` if it is not in ``sys.modules``
        (only when ``raises`` is false).

    Examples:
        Check without importing::

            torch = module_if_loaded("torch")
            if torch is None:
                return

        Require an extra that must already be imported::

            jax = module_if_loaded("jax", raises=True)

        Register a side effect for now-or-later (JAX pytree, TorchScript)::

            def _register_jax_pytree(jax):
                jax.tree_util.register_pytree_node(Ragged, flatten, unflatten)

            module_if_loaded("jax", _register_jax_pytree)
            # or module_if_loaded("jax", _register_jax_pytree, raises=True)
            # to fail now while still running the callback on a later import.
    """
    mod = _module(name)
    if callback is not None:
        if mod is not None:
            callback(mod)
        else:
            with _lock:
                _pending.setdefault(name, []).append(callback)
            _install_hook()
            # A concurrent import may have finished while we registered.
            _fire_ready()
            mod = _module(name)
    if mod is None and raises:
        raise RuntimeError(f"{name} is not imported")
    return mod