spekk.transformations.for_all.specced_vmap

Contents

spekk.transformations.for_all.specced_vmap#

spekk.transformations.for_all.specced_vmap(f: callable, spec: Spec, dimension: str, vmap_impl: Callable[[callable, Sequence[int | None]], callable] | None = None)[source]#

Similar to vmap, but flattens/decomposes the kwargs to a list that is supported by vmap.