probly.representation.sample.jax_axis_tracking

Axis-tracking for PyTorch tensors.