Easily freeze parameters about dynax

[fhchl]

An example:

import equinox as eqx
import jax.tree_util as jtu

# Class of which all parameters are trained
class Class1(eqx.Module): 
  param1: float
  param2: float

# inherits param1, but param2 is set as fixed now
class Class2(Class1): 
  param2: float = eqx.static_field() 

# make two examples
c1 = Class1(1, 2)
c2 = Class2(1, 3) 

# with tree flatten we can check which parameters are part of the pytree
# (they are trainable) and which are part of the structure (they are fixed)
print(jtu.tree_flatten(c1), jtu.tree_flatten(c2)) 

# this function returns a copy of obj, that behaves like cls
def change_class(obj, cls):
  # make a copy 
  flat, treedef = jtu.tree_flatten(obj)
  obj = treedef.unflatten(flat)
  # change class object.__setattr__(obj, "__class__", cls)
  return obj
  
# the value 2 is now part of the strucure. seems to work!
c1 = change_class(c1, Class2)
print(jtu.tree_flatten(c1))

[fhchl]

Fixed by new free_field, static_field functionality.