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

DOC: fix issue in pytrees.md

Browse Source
This commit is contained in:
Jake VanderPlas 2021-04-15 09:57:53 -07:00
parent 3c003a68fc
commit 89dc691da7
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/pytrees.md
View File

@ -264,12 +264,13 @@ on your class and decorate it with {func}`~jax.tree_util.register_pytree_node_cl
```{code-cell}
from jax.tree_util import register_pytree_node_class
@register_pytree_node_class
class RegisteredSpecial2(Special):
def __repr__(self):
return "RegisteredSpecial2(x={}, y={})".format(self.x, self.y)
def tree_flatten(self):
children = (v.x, v.y)
children = (self.x, self.y)
aux_data = None
return (children, aux_data)