Have PyTreeDef appear correctly in docs

2026-09-09 11:24:55 +08:00 · 2023-06-14 10:42:14 -07:00
parent 356f5b7f7b
commit 926dc53856
1 changed files with 2 additions and 0 deletions
@@ -160,6 +160,8 @@ elif has_jax:
            from `jax.tree_util.tree_structure(...)`.
            """

+        PyTreeDef.__module__ = "builtins"
+
    else:
        from jax.tree_util import PyTreeDef as PyTreeDef