Utilities

numerax.utils

Utility functions for the numerax package.

preserve_metadata

preserve_metadata(decorator)

Wrapper that ensures a decorator preserves function metadata for documentation tools.

Overview

This is particularly useful for JAX decorators like @custom_jvp that create special objects which may not preserve __doc__ and other metadata properly for documentation generators like pdoc.

Args

• decorator: The decorator function to wrap

Returns

A new decorator that preserves metadata

Example

import jax
from numerax.utils import preserve_metadata

@preserve_metadata(jax.custom_jvp)
def my_function(x):
    """This docstring will be preserved for automatic
    documentation generation."""
    return x