Tentative solution to pytree

[mavenlin]

We may move the discussions on pytree here, there are some future PRs that would depend on a systematic solution of this.

[samuela]

ah this is a cute idea. i like it if it works! so iiuc the idea is that we replace

torch2jax/torch2jax/__init__.py

Lines 1174 to 1175 in bee0974

	# use the torch's tree_map, because out is generated from torch code
	return torch.utils._pytree.tree_map(lambda x: x.value if isinstance(x, Torchish) else x, out)

with _tree_coerce and then it automatically handles pytorch pytree types as we encounter them? i wonder if this will work inside of jax.jit?

[mavenlin]

This PR is the prerequisite for functions like topk, sort, max to be properly handled as they have custom return types under torch.return_types., which needs to be systematically handled with the tree node registration.

[samuela]

this contains a lot of formatting changes that have no impact on functionality. let's remove those before merging

[samuela]

this is a cleaner solution than i thought would be possible, thanks @mavenlin !