spekk.transformations.for_all.ForAll

class spekk.transformations.for_all.ForAll(dimension: str, *additional_dimensions: str, vmap_impl: Callable[[callable, Sequence[int | None]], callable] | None = None)

Bases: Transformation

Vectorize/”make looped” a function such that it works on arrays instead of scalars.

>>> from spekk.transformations import ForAll, compose
>>> f = lambda x, y: x + y
>>> data =       {"x": range(2), "y": range(3)}
>>> spec = Spec( {"x": ["dim1"], "y": ["dim2"]} )

Transform f to run on a grid defined by the “dim1” and “dim2” dimensions:

>>> tf = compose(f, ForAll("dim2"), ForAll("dim1")).build(spec)
>>> tf.output_spec
Spec(['dim1', 'dim2'])
>>> result = tf(**data)
>>> util.shape(result)
(2, 3)
>>> result
[[0, 1, 2], [1, 2, 3]]

You can also use vmap over multiple dimensions at once. Note that the order of the dimensions in ForAll("dim2"), ForAll("dim1") and ForAll("dim1", "dim2") is reversed:

>>> tf = compose(f, ForAll("dim1", "dim2")).build(spec)
>>> tf(**data)  # This results in the same as in the previous example
[[0, 1, 2], [1, 2, 3]]

__init__(dimension: str, *additional_dimensions: str, vmap_impl: Callable[[callable, Sequence[int | None]], callable] | None = None)

Methods

__init__(dimension, *additional_dimensions)
transform_function(to_be_transformed, ...)	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.

transform_function(to_be_transformed: callable, input_spec: Spec, output_spec: Spec) → callable

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: 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.

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.

transform_output_spec(spec: Spec) → Spec

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.