jaxvacua.util.set_auto_vmap_default_shapes

set_auto_vmap_default_shapes(**sample_shapes)

Update JAXVacua-local exact sample-shape defaults for auto_vmap().

Shape entries may be concrete shapes, attribute names such as "h12", or callables accepted by stringjax_tools.auto_vectorise. Changes affect subsequently decorated functions, while built-in JAXVacua methods use the defaults active when their classes were imported.

Return type:

dict[str, Any]