text: import os
  import sys
  import tensorflow as tf
  import contextlib
  import json
  import importlib
  from tensorflow.python.eager.context import num_gpus
  
  
  def get_current_epoch(output_dir):
      try:
          with open(os.path.join(output_dir, "stats.json"), "r") as f:
              return json.load(f)["epoch"]
      except:
          return 0
  
  
  class ModelCheckpoint(tf.keras.callbacks.ModelCheckpoint):
      def on_epoch_end(self, epoch, logs=None):
          super().on_epoch_end(epoch, logs=logs)
          with open(os.path.join(os.path.dirname(self.filepath), "stats.json"), "w") as f:
              return json.dump({"epoch": epoch + 1}, f)
  
  
  def get_distribution_scope(batch_size):
      if num_gpus() > 1:
          strategy = tf.distribute.MirroredStrategy()
          assert (
              batch_size % strategy.num_replicas_in_sync == 0
          ), f"Batch size {batch_size} cannot be divided onto {num_gpus()} GPUs"
          distribution_scope = strategy.scope
      else:
          if sys.version_info >= (3, 7):
              distribution_scope = contextlib.nullcontext
          else:
              distribution_scope = contextlib.suppress
  
      return distribution_scope()
  
  
  def import_all_sub_modules(module):
      with os.scandir(os.path.join(os.path.dirname(__file__), module)) as it:
          for entry in it:
              name = entry.name
              if name.endswith(".py") and not name.startswith("_") and entry.is_file():
                  importlib.import_module(f"bnn_optimization.{module}.{name[:-3]}")
  
  
  def prepare_registry():
      for module in ("data", "models"):
          import_all_sub_modules(module)
  