spekk.transformations.base.Transformation#

class spekk.transformations.base.Transformation[source]#

Bases: ABC

A Transformation takes a function and transforms/modifies it. It also keeps track of changes to the spec when transforming the function. See the schematic in spekk.transformations’s docstring for a visualization of how specs are processed.

A Transformation acts as a higher-order function that takes in a function and returns an instance of TransformedFunction. It gets information about the dimensions of the data (through instances of Spec) that the function operates on and it also defines what happens to the spec when the function is transformed.

transform_function[source]#

The main transformation function that returns a TransformedFunction.

transform_input_spec[source]#

A function that returns the spec that is passed down into the wrapped function after the transformation has been applied.

transform_output_spec[source]#

A function that returns the final returned spec of the transformed function.

__init__()#

Methods

__init__()

transform_function(wrapped_fn, input_spec, ...)

Transform the wrapped function given the spec of the input arguments and the spec of the returned value of the wrapped function.

transform_input_spec(spec)

Return a new spec that represent the input arguments that are passed down to the wrapped function after the transformation has been applied.

transform_output_spec(spec)

Return a new spec that represent the returned value of the final transformed function.

abstract transform_function(wrapped_fn: callable, input_spec: Spec, output_spec: Spec) TransformedFunction[source]#

Transform the wrapped function given the spec of the input arguments and the spec of the returned value of the wrapped function.

abstract transform_input_spec(spec: Spec) Spec[source]#

Return a new spec that represent the input arguments that are passed down to the wrapped function after the transformation has been applied.

For example, if the transformation vectorizes the wrapped function over a dimension, the wrapped function would only see single items of the dimension at a time. Therefore, the input spec will have one less dimension when passed down to the wrapped function.

abstract transform_output_spec(spec: Spec) Spec[source]#

Return a new spec that represent the returned value of the final transformed function.

For example, if the transformation sums over a dimension of the result of calling the wrapped function, the output spec will have one less dimension.