tree._grow(nd.array([[6, 8]]))
tree._grow(nd.array([[9, 2], [1, 9]]))
tree(nd.array([[1, 1], [2, 2], [-1, -1]]))
# %%
[
node._decision._gate() for node in tree._embeddlayer._children.values()
if hasattr(node, "_decision")
]
# %%
tree._weightlayer(nd.array([[1, 1], [2, 2], [-1, -1]]))
tree._routerlayer(nd.array([[1, 1], [2, 2], [-1, -1]]))
tree._embeddlayer(nd.array([[1, 1], [2, 2], [-1, -1]]))
# %%
root = next(iter(tree._structure.items()))[0]
print(root._box._min_list.data(), root._box._max_list.data(),
root._decision._split.data())
# %%
# for box in tree._weightlayer:
# print(box._min_list.data(), box._max_list.data())
#
(6 - nd.topk(nd.reverse(nd.sum(router_mat, axis=2), axis=1)))[:, 0]
nd.choose_element_0index(nd.sum(router_mat, axis=2), nd.array([4, 6, 1, 0, 6]))
weight
nd.sum(router_mat, axis=2) * weight
nd.sum(nd.sum(router_mat, axis=2) * weight, axis=1)
tree.collect_params()
[x for x in tree._weightlayer]
[key for key, value in tree._weightlayer._children.items()]
tree._weightlayer(nd.array([[1, 3], [5, 6], [6, 2], [20, 20]]))
tree._routerlayer(nd.array([[1, 3], [5, 6], [6, 2], [2, 2]]))
tree(nd.array([[1, 3], [5, 6], [6, 2], [3, 0]]))
nd.sum(router_mat, axis=2)
nd.maximum(nd.sum(router_mat, axis=2) - weight, 0)
-1 * (nd.maximum(nd.sum(router_mat, axis=2) - weight, 0) -
nd.sum(router_mat, axis=2))
nd.sum(nd.expand_dims(weight, axis=2) * router_mat, axis=1)
nd.sum(nd.sum(nd.expand_dims(weight, axis=2) * router_mat, axis=1), axis=1)
weight
router_mat
