File size: 397 Bytes
0a937d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from typing import Any
import jax
def tree_chunk(tree: Any, n_chunk: int, axis: int = 0) -> Any:
return jax.tree_map(
lambda v: v.reshape(v.shape[:axis] + (n_chunk, -1) + v.shape[axis + 1:]),
tree
)
def tree_unchunk(tree: Any, axis: int = 0) -> Any:
return jax.tree_map(
lambda x: x.reshape(x.shape[:axis] + (-1,) + x.shape[axis + 2:]),
tree
)
|