Document New freezing API about pytreeclass HOT 1 CLOSED

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