code: def main(argv):
      tf.config.experimental.set_visible_devices([], "GPU")
      del argv
      logging.info("*** Starting experiment")
      gin_configs = FLAGS.gin_configs
  
      logging.info("*** Loading Gin configs from: %s", str(gin_configs))
      gin.parse_config_files_and_bindings(
          config_files=gin_configs,
          bindings=FLAGS.gin_bindings,
          skip_unknown=True,
      )
  
      # Load configurations.
      exp_config = configs.ExperimentConfig()
      model_config = configs.ModelConfig(use_stratified_sampling=False)
      train_config = configs.TrainConfig()
      eval_config = configs.EvalConfig()
  
      # Get directory information.
      exp_dir = gpath.GPath(FLAGS.base_folder)
      if exp_config.subname:
          exp_dir = exp_dir / exp_config.subname
      logging.info("\texp_dir = %s", exp_dir)
      if not exp_dir.exists():
          exp_dir.mkdir(parents=True, exist_ok=True)
  
      summary_dir = exp_dir / "summaries" / "eval"
      logging.info("\tsummary_dir = %s", summary_dir)
      if not summary_dir.exists():
          summary_dir.mkdir(parents=True, exist_ok=True)
  
      renders_dir = exp_dir / "renders"
      logging.info("\trenders_dir = %s", renders_dir)
      if not renders_dir.exists():
          renders_dir.mkdir(parents=True, exist_ok=True)
  
      checkpoint_dir = exp_dir / "checkpoints"
      logging.info("\tcheckpoint_dir = %s", checkpoint_dir)
  
      logging.info(
          "Starting host %d. There are %d hosts : %s",
          jax.process_index(),
          jax.process_count(),
          str(jax.process_indexs()),
      )
      logging.info(
          "Found %d accelerator devices: %s.",
          jax.local_device_count(),
          str(jax.local_devices()),
      )
      logging.info(
          "Found %d total devices: %s.",
          jax.device_count(),
          str(jax.devices()),
      )
  
      rng = random.PRNGKey(20200823)
  
      devices_to_use = jax.local_devices()
      n_devices = (
          len(devices_to_use)
          if devices_to_use
          else jax.local_device_count()
      )
  
      datasource_spec = exp_config.datasource_spec
      if datasource_spec is None:
          datasource_spec = {
              "type": exp_config.datasource_type,
              "data_dir": FLAGS.data_dir,
          }
      logging.info("Creating datasource: %s", datasource_spec)
      datasource = datasets.from_config(
          datasource_spec,
          image_scale=exp_config.image_scale,
          use_appearance_id=model_config.use_appearance_metadata,
          use_camera_id=model_config.use_camera_metadata,
          use_warp_id=model_config.use_warp,
          use_time=model_config.warp_metadata_encoder_type == "time",
          random_seed=exp_config.random_seed,
          **exp_config.datasource_kwargs,
      )
  
      # Get training IDs to evaluate.
      train_eval_ids = utils.strided_subset(
          datasource.train_ids, eval_config.num_train_eval
      )
      train_eval_iter = datasource.create_iterator(
          train_eval_ids, batch_size=0
      )
      val_eval_ids = utils.strided_subset(
          datasource.val_ids, eval_config.num_val_eval
      )
      val_eval_iter = datasource.create_iterator(
          val_eval_ids, batch_size=0
      )
  
      test_cameras = datasource.load_test_cameras(
          count=eval_config.num_test_eval
      )
      if test_cameras:
          test_dataset = datasource.create_cameras_dataset(test_cameras)
          test_eval_ids = [f"{x:03d}" for x in range(len(test_cameras))]
          test_eval_iter = datasets.iterator_from_dataset(
              test_dataset, batch_size=0
          )
      else:
          test_eval_ids = None
          test_eval_iter = None
  
      rng, key = random.split(rng)
      params = {}
      model, params["model"] = models.construct_nerf(
          key,
          model_config,
          batch_size=eval_config.chunk,
          appearance_ids=datasource.appearance_ids,
          camera_ids=datasource.camera_ids,
          warp_ids=datasource.warp_ids,
          near=datasource.near,
          far=datasource.far,
          use_warp_jacobian=False,
          use_weights=False,
      )
  
      optimizer_def = optim.Adam(0.0)
      optimizer = optimizer_def.create(params)
      init_state = model_utils.TrainState(optimizer=optimizer)
      del params
  
      def _model_fn(key_0, key_1, params, rays_dict, warp_extra):
          out = model.apply(
              {"params": params},
              rays_dict,
              warp_extra=warp_extra,
              rngs={"coarse": key_0, "fine": key_1},
              mutable=False,
          )
          return jax.lax.all_gather(out, axis_name="batch")
  
      pmodel_fn = jax.pmap(
          # Note rng_keys are useless in eval mode since there's no randomness.
          _model_fn,
          in_axes=(0, 0, 0, 0, 0),  # Only distribute the data input.
          devices=devices_to_use,
          donate_argnums=(3,),  # Donate the 'rays' argument.
          axis_name="batch",
      )
  
      render_fn = functools.partial(
          evaluation.render_image,
          model_fn=pmodel_fn,
          device_count=n_devices,
          chunk=eval_config.chunk,
      )
  
      last_step = 0
      summary_writer = tensorboard.SummaryWriter(str(summary_dir))
  
      while True:
          if not checkpoint_dir.exists():
              logging.info("No checkpoints yet.")
              time.sleep(10)
              continue
  
          state = checkpoints.restore_checkpoint(
              checkpoint_dir, init_state
          )
          state = jax_utils.replicate(state, devices=devices_to_use)
          step = int(state.optimizer.state.step[0])
          if step <= last_step:
              logging.info(
                  "No new checkpoints (%d <= %d).", step, last_step
              )
              time.sleep(10)
              continue
  
          save_dir = renders_dir if eval_config.save_output else None
          process_iterator(
              tag="val",
              item_ids=val_eval_ids,
              iterator=val_eval_iter,
              state=state,
              rng=rng,
              step=step,
              render_fn=render_fn,
              summary_writer=summary_writer,
              save_dir=save_dir,
              datasource=datasource,
          )
  
          process_iterator(
              tag="train",
              item_ids=train_eval_ids,
              iterator=train_eval_iter,
              state=state,
              rng=rng,
              step=step,
              render_fn=render_fn,
              summary_writer=summary_writer,
              save_dir=save_dir,
              datasource=datasource,
          )
  
          if test_eval_iter:
              process_iterator(
                  tag="test",
                  item_ids=test_eval_ids,
                  iterator=test_eval_iter,
                  state=state,
                  rng=rng,
                  step=step,
                  render_fn=render_fn,
                  summary_writer=summary_writer,
                  save_dir=save_dir,
                  datasource=datasource,
              )
  
          if save_dir:
              delete_old_renders(
                  renders_dir, eval_config.max_render_checkpoints
              )
  
          if eval_config.eval_once:
              break
          if step >= train_config.max_steps:
              break
          last_step = step
  
imports: import collections
  import functools
  import time
  from typing import Any, Dict, Optional, Sequence
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  from flax import optim
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  import gin
  import jax
  from jax import numpy
  from jax import random
  import tensorflow
  from nerfies import configs
  from nerfies import datasets
  from nerfies import evaluation
  from nerfies import gpath
  from nerfies import image_utils
  from nerfies import model_utils
  from nerfies import models
  from nerfies import types
  from nerfies import utils
  from nerfies import visualization
label_desc: Use image node_id 5 for calling a global function with example usage: main(argv) and returns: jax