def _free_accelerators(self, exceptions=(), keep_constants=True):
    """Deletes all live buffers from accelerator with no safety guarantees."""
    live_buffers = jax.live_arrays()
    logging.info('Deleting %d live buffers.', len(live_buffers))
    exceptions_buffers = []
    for x in fastmath.tree_flatten(exceptions):
      if hasattr(x, 'device_buffer'):  # DeviceArray
        exceptions_buffers.append(x.device_buffer)
      if hasattr(x, 'device_buffers'):  # ShardedDeviceArray
        exceptions_buffers.extend(x.device_buffers)
    for b in live_buffers:
      should_delete = True
      for e in exceptions_buffers:
        if b is e:
          should_delete = False
      if keep_constants and not b.shape:
        should_delete = False
      if should_delete:
        b.delete()
Deletes all live buffers from accelerator with no safety guarantees.
trax/optimizers/trainer.py