jax.tree_util.keystr

Contents

jax.tree_util.keystr#

jax.tree_util.keystr(keys)[source]#

Helper to pretty-print a tuple of keys.

Parameters:

keys (tuple[TypeVar(KeyEntry, bound= Hashable), ...]) – A tuple of KeyEntry or any class that can be converted to string.

Returns:

A string that joins all string representations of the keys.