code: def encode(
      codec: BitsBackCodec,
      dataloader: Iterator[Tuple[torch.Tensor, dict]],
      num_samples: int,
      save_interval: int,
  ):
      """
      Encode data using Bits-Back coding.
  
      Args:
          codec: Codec to use for encoding.
          dataloader: Iterator yielding the data to encode.
          num_samples: Number of samples to encode.
          save_interval: Interval after which to save a checkpoint.
      """
      # raise errors leading to wrong encoding
      np.seterr(all="raise")
  
      rate_effective: Dict[str, List[np.ndarray]] = {
          "push": [],
          "pop": [],
      }
      num_complete = (
          codec.data_shape[0] * codec.data_count * dist.get_world_size()
      )
      while num_complete < num_samples:
          batch, _ = next(dataloader)
  
          # discretize batch to [0, 2 ** codec.data_prec - 1]
          discr_batch = codec.quantize(batch.numpy())
  
          # push the discretized batch to the message, and compute rates
          codec.encode(discr_batch)
          stats = codec.statistics()
  
          if codec.track:
              # check rates (for each diffusion step) and append
              assert stats.rate_effective_pop is not None
              assert stats.rate_effective_push is not None
              assert np.allclose(
                  stats.rate_effective_push.sum()
                  - stats.rate_effective_pop.sum(),
                  stats.rate_effective,
              )
              rate_effective["push"].append(stats.rate_effective_push)
              rate_effective["pop"].append(stats.rate_effective_pop)
  
          # log
          stats_dict = stats._asdict()
          stats_dict.pop("rate_effective_push")
          stats_dict.pop("rate_effective_pop")
          logger.logkvs(stats_dict)
          logger.dumpkvs()
  
          # checkpoint
          num_complete = stats.samples * dist.get_world_size()
          if (
              num_complete >= num_samples
              or stats.samples % save_interval < codec.data_shape[0]
          ):
              ckpt_path = (
                  Path(logger.get_dir())
                  / f"encode-ckpt{stats.samples:06d}.npz"
              )
              np.savez(ckpt_path, **codec.state_dict())
  
      if codec.track:
          # log rates
          plots = {}
          for name, rate in rate_effective.items():
              out_path = Path(logger.get_dir()) / f"rates-{name}.npz"
              logger.log(f"saving rates-{name} to {out_path}")
              rate_arr = np.stack(rate)
              np.savez(
                  out_path,
                  mean=rate_arr.mean(axis=0),
                  stdd=rate_arr.std(axis=0),
                  min=rate_arr.min(axis=0),
                  max=rate_arr.max(axis=0),
              )
              plots[name] = plot_evolution(
                  rate_arr, xaxis_title="t", yaxis_title=name
              )
          wandb.log(plots)
  
file_path: NeuralCompression/projects/bits_back_diffusion/encode.py
imports: from itertools import islice
  from pathlib import Path
  from typing import Dict, Iterator, List, Tuple, Union
  import hydra
  import numpy
  import torch
  import torch.distributed
  import wandb
  from bits_back_diffusion.codec import BitsBackCodec
  from bits_back_diffusion.script_util import plot_evolution, setup, update_checkpoint
  from improved_diffusion import logger
  from omegaconf import DictConfig
label_desc: Use image node_id 2 for calling a global function with example usage: encode(codec, dataloader, num_samples, save_interval) without return types