fedjax.tree_util
fedjax.tree_util.tree_mean
fedjax.tree_util.tree_weight
fedjax.tree_util.tree_sum
fedjax.tree_util.tree_add
fedjax.tree_util.tree_zeros_like
fedjax.tree_util.tree_inverse_weight
fedjax.tree_util.tree_l2_norm
fedjax.tree_util.tree_size
fedjax.tree_util.tree_clip_by_global_norm