class_docstr: Base class for dataclasses that should act like a JAX pytree node.
  
  See ``flax.struct.dataclass`` for the ``jax.tree_util`` behavior.
  This base class additionally avoids type checking errors when using PyType.
  
  Example::
  
    >>> from flax import struct
    >>> import jax
    >>> from typing import Any, Callable
  
    >>> class Model(struct.PyTreeNode):
    ...   params: Any
    ...   # use pytree_node=False to indicate an attribute should not be touched
    ...   # by Jax transformations.
    ...   apply_fn: Callable = struct.field(pytree_node=False)
  
    ...   def __apply__(self, *args):
    ...     return self.apply_fn(*args)
  
    >>> params = {}
    >>> params_b = {}
    >>> apply_fn = lambda v, x: x
    >>> model = Model(params, apply_fn)
  
    >>> # model.params = params_b  # Model is immutable. This will raise an error.
    >>> model_b = model.replace(params=params_b)  # Use the replace method instead.
  
    >>> # This class can now be used safely in Jax to compute gradients w.r.t. the
    >>> # parameters.
    >>> model = Model(params, apply_fn)
    >>> loss_fn = lambda model: 3.
    >>> model_grad = jax.grad(loss_fn)(model)
class_docstr_tok: ["Base","class","for","dataclasses","that","should","act","like","a","JAX","pytree","node",".","See","``","flax.struct.dataclass","``","for","the","``","jax.tree_util","``","behavior",".","This","base","class","additionally","avoids","type","checking","errors","when","using","PyType",".","Example",":",":",">",">",">","from","flax","import","struct",">",">",">","import","jax",">",">",">","from","typing","import","Any",",","Callable",">",">",">","class","Model","(","struct.PyTreeNode",")",":","...","params",":","Any","...","#","use","pytree_node=False","to","indicate","an","attribute","should","not","be","touched","...","#","by","Jax","transformations",".","...","apply_fn",":","Callable","=","struct.field","(","pytree_node=False",")","...","def","__apply__","(","self",",","*","args",")",":","...","return","self.apply_fn","(","*","args",")",">",">",">","params","=","{","}",">",">",">","params_b","=","{","}",">",">",">","apply_fn","=","lambda","v",",","x",":","x",">",">",">","model","=","Model","(","params",",","apply_fn",")",">",">",">","#","model.params","=","params_b","#","Model","is","immutable",".","This","will","raise","an","error",".",">",">",">","model_b","=","model.replace","(","params=params_b",")","#","Use","the","replace","method","instead",".",">",">",">","#","This","class","can","now","be","used","safely","in","Jax","to","compute","gradients","w.r.t",".","the",">",">",">","#","parameters",".",">",">",">","model","=","Model","(","params",",","apply_fn",")",">",">",">","loss_fn","=","lambda","model",":","3",".",">",">",">","model_grad","=","jax.grad","(","loss_fn",")","(","model",")"]
code: def __init_subclass__(cls):
      dataclass(cls)
  
file_path: flax/flax/struct.py
imports: import dataclasses
  from typing import TypeVar
  import jax
  from typing_extensions import dataclass_transform
  from .None import serialization
label_desc: Use image node_id 1 for calling the PyTreeNode obj's underlying member method code with example usage: obj.__init_subclass__(cls) without return types