Index

_

__init__() (ArgSpec method) (PytreePolicy method)

A

ArgSpec (class in stringjax_tools.auto_vectorise) auto_vmap() (in module stringjax_tools.auto_vectorise)

auto_vmap_default_shapes() (in module stringjax_tools.auto_vectorise) auto_vmap_defaults() (in module stringjax_tools.auto_vectorise)

C

clear_auto_vmap_caches() (in module stringjax_tools.auto_vectorise)

clear_vmap_caches() (in module stringjax_tools.vmap) configure_compilation_cache() (in module stringjax_tools.cache)

F

flatten() (PytreePolicy method)

flatten_func() (in module stringjax_tools.pytrees)

G

get_auto_vmap_default_shapes() (in module stringjax_tools.auto_vectorise)

get_auto_vmap_defaults() (in module stringjax_tools.auto_vectorise)

I

is_static() (in module stringjax_tools.jit)

J

jit_with_dynamic_static_args() (in module stringjax_tools.jit)

jit_with_static_args() (in module stringjax_tools.jit)

M

make_flatteners() (PytreePolicy method) make_pytree_flatteners() (in module stringjax_tools.pytrees) module stringjax_tools.auto_vectorise stringjax_tools.cache stringjax_tools.jit stringjax_tools.pytrees stringjax_tools.vmap

P

PytreePolicy (class in stringjax_tools.pytrees)

R

register() (PytreePolicy method)

reset_auto_vmap_default_shapes() (in module stringjax_tools.auto_vectorise) reset_auto_vmap_defaults() (in module stringjax_tools.auto_vectorise)

S

set_auto_vmap_default_shapes() (in module stringjax_tools.auto_vectorise) set_auto_vmap_defaults() (in module stringjax_tools.auto_vectorise) stringjax_tools.auto_vectorise module stringjax_tools.cache module

stringjax_tools.jit module stringjax_tools.pytrees module stringjax_tools.vmap module

U

unflatten() (PytreePolicy method)

unflatten_func_class() (in module stringjax_tools.pytrees)

V

vmapping_func() (in module stringjax_tools.vmap)

vmapping_func_cached() (in module stringjax_tools.vmap)