(qp, box_qp)
| 4 | import numpy as np |
| 5 | |
| 6 | def add_box(qp, box_qp): |
| 7 | return jp.tree_map(lambda x, y: np.concatenate([x,y], axis=-2), qp, box_qp) |
| 8 | |
| 9 | def remove_box(qp): |
| 10 | return jp.tree_map(lambda x: x[..., :19, :], qp), jp.tree_map(lambda x: x[..., 19:, :], qp) |