spekk.transformations.wrap

Wrap is a simple Transformation that wraps a function with another function, such as jax.jit().

Classes

Wrap(f, *args, **kwargs)

Simply wraps a function with another function, but useful for keeping information about the spec in a chain of Transformation.