Document New freezing API about pytreeclass

[ASEM000]

Comparison between PyTreeClass and data class field metadata-based approach for freezing tree leaves [Draft]

PyTreeClass	Flax	Equinox
import jax import pytreeclass as pytc @pytc.treeclass class Test: a: tuple[int, ...] = (pytc.freeze(1.), 2., 3.) def __call__(self, x): return sum(self.a) @jax.grad def loss_func(NN:Test,x : float): return (NN(x)-x)**2 @jax.jit def step_func(NN:Test, x:float): dNN = loss_func(NN, x) NN -= 1e-3* dNN return NN def train(NN:Test, epochs:int= 10_000): NN = Test() for _ in range(epochs): NN = step_func(NN, 10.) return NN print(train(NN)) # Test(a=(#1.0, 3.9999967, 4.9999437))	import jax import jax.tree_util as jtu from flax import struct @struct.dataclass class Test: a_0: float = struct.field(pytree_node=False, default=1.) a_1: float = 2. a_2: float = 3. def __call__(self, x): return self.a_0 + self.a_1 + self.a_2 @jax.grad def loss_func(NN:Test,x : float): return (NN(x)-x)**2 @jax.jit def step_func(NN:Test, x:float): dNN = loss_func(NN, x) NN = jtu.tree_map(lambda x,y: x-y*1e-3, NN,dNN) return NN def train(NN:Test, epochs:int= 10_000): NN = Test() for _ in range(epochs): NN = step_func(NN, 10.) return NN print(train(NN)) # Test(a_0=1.0, a_1=Array(3.9999967, dtype=float32, weak_type=True), a_2=Array(4.9999437, dtype=float32, weak_type=True))	import jax import jax.tree_util as jtu import equinox as eqx class Test(eqx.Module): a_0: float = eqx.static_field(default=1.) a_1: float = 2. a_2: float = 3. def __call__(self, x): return self.a_0 + self.a_1 + self.a_2 @jax.grad def loss_func(NN:Test,x : float): return (NN(x)-x)**2 @jax.jit def step_func(NN:Test, x:float): dNN = loss_func(NN, x) NN = jtu.tree_map(lambda x,y: x-y*1e-3, NN,dNN) return NN def train(NN:Test, epochs:int= 10_000): NN = Test() for _ in range(epochs): NN = step_func(NN, 10.) return NN print(train(NN)) # Test(a_0=1.0, a_1=f32[], a_2=f32[]) print(train(NN).a_1) # 3.9999967