yardenas / offline-meta-mbrl Goto Github PK

Hey again,

I post this here so that the equinox issue is not clattered.

I did some basic tests to assert whether we get the same output from the same inputs and we do.

I tested the execution time for the two implementations for the following case:

state_size = 50
sequence_size = 500
input_dim = 1

These are the results:

Yours: 2.63 ms ± 58.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Mine: 44.3 µs ± 4.87 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

This time of 2.64 ms is after I jitted your convolve function. Without jit it took:

Yours: 20.3 ms ± 330 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Mine: 2.14 ms ± 91.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

I suspect this has to do with the vmap?? Do you have any suggestions?