mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

comment-out new pytree check with a TODO

Browse Source
This commit is contained in:
Matthew Johnson 2019-03-22 07:38:40 -07:00
parent f94c20633a
commit 50f52070fc
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/tree_util.py
View File

@ -67,8 +67,9 @@ def tree_multimap(f, tree, *rest):
all_children = [children]
for other_tree in rest:
other_node_type = node_types.get(type(other_tree))
if node_type != other_node_type:
raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))
# TODO(mattjj): enable this check
# if node_type != other_node_type:
# raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))
other_children, other_node_data = node_type.to_iterable(other_tree)
if other_node_data != node_spec:
raise TypeError('Mismatch: {} != {}'.format(other_node_data, node_spec))