Object that is jax valid type

[NITHISHM2410]

I'm trying to create a class to track my model losses. I'm unable to use a normal python object inside flax.train_state.TrainState.
When I do that, I get this error:
Value <__main__.F1Scoe object at 0x7a7481686b00> with type <class '__main__.F1Scoe'> is not a valid JAX type.

I went through the Jax documentation and I added these below methods to my class.

    def flatten(self):
        children = (self.tp, self.tn, self.fp, self.fn, self.count)
        aux_data = {'threshold': self.threshold, 'num_classes': self.num_classes, 'average': self.average}
        return children, aux_data

    @classmethod
    def unflatten(cls, aux_data, children):
        obj = cls(**aux_data)
        obj.tp, obj.tn, obj.fp, obj.fn, obj.count = children
        return obj

register_pytree_node(F1Scoe, F1Scoe.flatten, F1Scoe.unflatten)

When I did the above I found that when my data members like tp, fp, etc, are jnp.arrays the time taken to update these data members increases a lot and when I leave them as python scalar , the time taken to update them is very less even if the data member holds a object with same size.
For ex. When data members take the form self.fp = jnp.array(0) , time taken to update them is very high compared to the form self.fp=0.

All the updates to these tracker's data members like 'fp' are called from jax.pmap function which automatically 'jit'.

So my question is, is there a way to solve this time taken issue or is there an alternate way to create loss and metric trackers instead of the above method?

[jakevdp]

Hi - indeed the pytree registry is the best way to create a custom type that is compatible with JAX transformations.

Regarding your question, I don't entirely understand when you say "the time taken to update these data members increases a lot". What is the context of these updates? Can you show an example of the issue?

[NITHISHM2410]

Thq for the reply. I hope you were able to understand my problem as shown in the linked notebook.

[jakevdp]

Thanks - looks like maybe a 10% difference? I suspect it comes from the fact that creating JAX arrays takes a few hundred microseconds longer than creating integers, so when you do it 3000 times, together it will take on order a second longer.

For example:

import jax.numpy as jnp

%timeit x = jnp.array(0)
212 µs ± 86.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit x = 0
30.7 ns ± 3.78 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

[NITHISHM2410]

Ok, will try a work around.

[NITHISHM2410]

Initializing those data members as np.array instead of jnp.array solved my issue. Still don't understand the reason.

[jakevdp]

Initializing a size-1 numpy array is about 200x faster than initializing a size-1 JAX array:

import numpy as np

%timeit x = np.array(0)
# 1.04 µs ± 385 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

When you do this 3000 times, using JAX it will take seconds, using numpy it will take milliseconds, and using python ints it will take microseconds.