jcm/train.py [99:120]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    train_step_fn = losses.get_step_fn(
        train_loss_fn,
        train=True,
        optimize_fn=optimize_fn,
        ema_scales_fn=ema_scale_fn,
    )
    # Pmap (and jit-compile) multiple training steps together for faster running
    p_train_step = jax.pmap(
        functools.partial(jax.lax.scan, train_step_fn),
        axis_name="batch",
    )
    eval_step_fn = losses.get_step_fn(
        eval_loss_fn,
        train=False,
        optimize_fn=optimize_fn,
        ema_scales_fn=ema_scale_fn,
    )
    # Pmap (and jit-compile) multiple evaluation steps together for faster running
    p_eval_step = jax.pmap(
        functools.partial(jax.lax.scan, eval_step_fn),
        axis_name="batch",
    )
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



jcm/train.py [263:284]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
                    train_step_fn = losses.get_step_fn(
                        train_loss_fn,
                        train=True,
                        optimize_fn=optimize_fn,
                        ema_scales_fn=ema_scale_fn,
                    )
                    # Pmap (and jit-compile) multiple training steps together for faster running
                    p_train_step = jax.pmap(
                        functools.partial(jax.lax.scan, train_step_fn),
                        axis_name="batch",
                    )
                    eval_step_fn = losses.get_step_fn(
                        eval_loss_fn,
                        train=False,
                        optimize_fn=optimize_fn,
                        ema_scales_fn=ema_scale_fn,
                    )
                    # Pmap (and jit-compile) multiple evaluation steps together for faster running
                    p_eval_step = jax.pmap(
                        functools.partial(jax.lax.scan, eval_step_fn),
                        axis_name="batch",
                    )
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -