tensorboardX is outdated. torch provides an internal API via from torch.utils.tensorboard import SummaryWriter. Replace this and reduce dependencies. Only remaining required dependency should be tensorboard.
Furthermore, trax appears to provide a jax-compatible version as well: jaxboard with standalone installation. Check if this can simply act as a drop-in replacement with parameter FrozenDict.
Addressed in 28f45f4. Example for JAX tensorboard logging can be found in examples/jax_vae. We don't log gradients since they would have to be passed seperately!
tensorboardX
is outdated.torch
provides an internal API viafrom torch.utils.tensorboard import SummaryWriter
. Replace this and reduce dependencies. Only remaining required dependency should betensorboard
.Furthermore,
trax
appears to provide ajax
-compatible version as well:jaxboard
with standalone installation. Check if this can simply act as a drop-in replacement with parameterFrozenDict
.