Train fit using pytrees

[saisua]

I've taken a liking to Jax as an oportunity to bridge base Python and ml, while keeping performance.

My specific use case would be to use a regular Python class decorated as a pytree, with util functions to allow easy formatting and input of data.

I've created a simple training loop with such class, and I'd like to fit it using the pytrees. However, it is unclear to me what should be the workflow.

If possible, I'd like to get a pytree in and return a modified (copy) version of that same pytree, while being able to use the functions inside the pytree, if possible (all within the purity constraints of Jax).

Trying back and forth I got that I can flatten the pytree inside the function, and unflatten it at the end with the arrays. Although this works, it takes about 20s per sample of compilation the first time, and 0.1s later.
To me it seems I am in damp terrain, so I'd like to ask what is the actual expected workflow.

1. Should I not fit in pytrees, but the flattened versions?
2. Should all class functions be static?
3. Should I do something else?
4. Is flattening triggered automatically at any point during the call of a jittes function?

I know we can use datasets, but I wanted to check if it was possible to use pytrees inside jitted fns to use the pytree class functions to modify the functionality of the training (even if it has to recompile when you change the pytree class)

[ToshiyukiBandai]

Could you give us any example to illustrate your need? I am not sure what you really intend to do, but Equinox can handle Pytrees operations very well.

[jakevdp]

You can use pytrees directly within JIT-compiled functions. You do not need to manually flatten and unflatten them: flattening and unflattening are internal mechanisms used by JAX at the boundaries of transformations like jax.jit.

Beyond that, it's hard to answer your question because the details are very vague. I'd suggest putting together a short example illustrating the kind of operation you have in mind.

[saisua]

Hi, thank you for your fast responses and sorry for my late one.

I've tried to condense my case as much as possible, while keeping the most core stuff I think I could be messing up.

This is my very much simplified code of my pytrees:

@register_pytree_node_class
class Nodeset:
  node_info: List[jnp.ndarray]
  def __init__(self, ...):
    ...
  def __mul__(self, ...):
    #Used during git
  def tree_flatten(self) -> Tuple:
    ...
  @classmethod
  def tree_unflatten(cls, aux_data, children):
    ...

@register_pytree_node_class
class Graph:
  nodes: List[Nodeset]
  def __init__(self, ...):
    ...
  def add_nodeset(self, ...):
    # Not used during jit
    ...
  def to_json(self):
    # Not used during jit
    ...
  def connectivity(self) -> jnp.ndarray:
    # Used during jit
    # Should it be static?
    ...
  def __mul__(self, ...) -> Graph:
    # Used during jit
    ...
  def tree_flatten(self) -> Tuple:
    nodeset_ch, nodeset_aux = zip(*(ns.tree_flatten() for ns in self.nodes))
    ...
  @classmethod
  def tree_unflatten(cls, aux_data, children):
    ...

graphs: List[Graph] = ...

This is an approximation of the transform I want to reach:

@jax.jit
def transform(graphs: List[graphs], th: float) -> List[Graph]:
  # Could not get vmap to work, but it is outside the question
  new_graphs: List[Graph] = []
  for graph in graphs:
    new_graph = graph * (graph.connectivity() > th)
    new_graph.nodes[0].node_info[0] = new_graph.nodes[0].node_info[0] * 2
    new_graphs.append(new_graph)
  return new_graphs

And this is an approximation of the transform I have now:

@jax.jit
def transform(graphs: List[graphs]) -> List[Graph]
  new_graphs = []
  for graph in graphs:
    nodes, aux = graph.tree_flatten()
    ...
    nodes[0][0] = graph.nodes[0].node_info[0] * 2
    new_graph = Graph.tree_unflatten(aux, nodes)
    new_graphs.append(new_graph)
  return new_graphs

[jakevdp]

It looks like the first version of transform should work. What goes wrong when you try it?