# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# python3
"""Train seq-to-seq model on random supervised training tasks."""

# pytype: disable=wrong-arg-count
# pytype: disable=attribute-error

import collections
import functools
import json
import os
import random
import sys
import time

from absl import app
from absl import flags
from absl import logging
from flax import jax_utils
from flax import linen as nn
from flax import optim
from flax.metrics import tensorboard
from flax.training import checkpoints
from flax.training import common_utils
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf

from latent_programmer import decode
from latent_programmer import models as base_models
from latent_programmer.decomposition_transformer_attention import decomposition_models as models
from latent_programmer.decomposition_transformer_attention import input_pipeline
from latent_programmer.tasks.robust_fill import dsl
from latent_programmer.tasks.robust_fill import tokens as dsl_tokens

sys.path.append('../../')
gfile = tf.io.gfile

FLAGS = flags.FLAGS

flags.DEFINE_integer('seed', 0, 'Fixed random seed for training.')
flags.DEFINE_float('lr', 1e-3, 'Learning rate.')
flags.DEFINE_float('weight_decay', 1e-1,
                   'Decay factor for AdamW-style weight decay.')
flags.DEFINE_integer('embedding_dim', 256, 'Embedding dimension.')
flags.DEFINE_integer('hidden_dim', 512, 'Hidden dimension.')
flags.DEFINE_integer('num_heads', 4, 'Number of layers.')
flags.DEFINE_integer('num_layers', 3, 'Number of Transformer heads.')
flags.DEFINE_boolean('slow_decode', True, 'Use slow decoding for prediction?')

flags.DEFINE_string('dataset_filepattern', None,
                    'Filepattern for TFRecord dataset.')
flags.DEFINE_integer('per_device_batch_size', 16,
                     'Number of program tasks in a batch.')
flags.DEFINE_integer('num_strings_per_task', 4,
                     'Number of input/output strings per task.')
flags.DEFINE_integer('max_program_length', 100,
                     'Maximum number of tokens in program.')
flags.DEFINE_integer('max_characters', 120,
                     'Maximum number of characters in input/output strings.')

flags.DEFINE_string('save_dir', None, 'Directory to save results to.')
flags.DEFINE_integer('num_train_steps', 2000000, 'Number of training steps.')
flags.DEFINE_integer('num_eval_steps', 10, 'Number of evaluation steps.')
flags.DEFINE_integer('log_freq', 1000, 'Number of steps between training logs.')
flags.DEFINE_integer('eval_freq', 2000, 'Number of steps between eval.')
flags.DEFINE_integer('predict_freq', 50000,
                     'Number of steps between prediction (beam search).')
flags.DEFINE_integer('checkpoint_freq', 50000,
                     'Number of steps between checkpoint saves.')
flags.DEFINE_integer('finetune_start_step', -1,
                     'Step the initial checkpoint should start at for '
                     'finetuning, or -1 if not finetuning.')
flags.DEFINE_bool('restore_checkpoints', True,
                  'Whether to restore from existing model checkpoints.')

flags.DEFINE_string('attention_mask_type', 'bos_full_attention',
                    'The kind of attention mask to use. Options are: baseline, '
                    'bos_to_bos, bos_full_attention')

flags.DEFINE_bool('use_relative_attention', True,
                  'Whether to use relative positonal embeddings.')
flags.DEFINE_bool('bos_special_attention', False,
                  'Whether to use special relative attention computation for '
                  'BOS tokens.')


_internal = False
if not _internal:
  flags.DEFINE_string('xm_parameters', None,
                      'String specifying hyperparamter search.')


def create_learning_rate_scheduler(
    base_learning_rate=0.5,
    factors='constant * linear_warmup * rsqrt_normalized_decay',
    warmup_steps=16000,
    decay_factor=0.5,
    steps_per_decay=50000,
    steps_per_cycle=100000):
  """Creates learning rate schedule.

  Interprets factors in the factors string which can consist of:
  * constant: interpreted as the constant value,
  * linear_warmup: interpreted as linear warmup until warmup_steps,
  * rsqrt_decay: divide by square root of max(step, warmup_steps)
  * decay_every: Every k steps decay the learning rate by decay_factor.
  * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.

  Args:
    base_learning_rate: float, the starting constant for the lr schedule.
    factors: a string with factors separated by '*' that defines the schedule.
    warmup_steps: how many steps to warm up for in the warmup schedule.
    decay_factor: The amount to decay the learning rate by.
    steps_per_decay: How often to decay the learning rate.
    steps_per_cycle: Steps per cycle when using cosine decay.

  Returns:
    A function learning_rate(step): float -> {'learning_rate': float}, the
    step-dependent lr.
  """
  factors = [n.strip() for n in factors.split('*')]

  def step_fn(step):
    """Step to learning rate function."""
    ret = 1.0
    for name in factors:
      if name == 'constant':
        ret *= base_learning_rate
      elif name == 'linear_warmup':
        ret *= jnp.minimum(1.0, step / warmup_steps)
      elif name == 'rsqrt_decay':
        ret /= jnp.sqrt(jnp.maximum(1.0, step - warmup_steps))
      elif name == 'rsqrt_normalized_decay':
        ret *= jnp.sqrt(warmup_steps)
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'decay_every':
        ret *= (decay_factor**(step // steps_per_decay))
      elif name == 'cosine_decay':
        progress = jnp.maximum(0.0,
                               (step - warmup_steps) / float(steps_per_cycle))
        ret *= jnp.maximum(0.0,
                           0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
      else:
        raise ValueError('Unknown factor %s.' % name)
    return jnp.asarray(ret, dtype=jnp.float32)

  return step_fn


def compute_weighted_cross_entropy(logits, targets, weights=None):
  """Compute weighted cross entropy and entropy for log probs and targets.

  Args:
   logits: `[batch, length, num_classes]` float array.
   targets: categorical targets `[batch, length]` int array.
   weights: None or array of shape [batch, length, 1]

  Returns:
    Tuple of scalar loss and batch normalizing factor.
  """
  if logits.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))

  onehot_targets = common_utils.onehot(targets, logits.shape[-1])
  loss = -jnp.sum(onehot_targets * nn.log_softmax(logits), axis=-1)
  normalizing_factor = jnp.prod(jnp.asarray(targets.shape))
  if weights is not None:
    loss = loss * weights
    normalizing_factor = weights.sum()

  return loss.sum(), normalizing_factor


def compute_weighted_accuracy(logits, targets, weights=None):
  """Compute weighted accuracy for log probs and targets.

  Args:
   logits: `[batch, length, num_classes]` float array.
   targets: categorical targets `[batch, length]` int array.
   weights: None or array of shape [batch, length, 1]

  Returns:
    Tuple of scalar accuracy and batch normalizing factor.
  """
  if logits.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))
  acc = jnp.equal(jnp.argmax(logits, axis=-1), targets)
  normalizing_factor = jnp.prod(jnp.asarray(targets.shape))
  if weights is not None:
    acc = acc * weights
    normalizing_factor = weights.sum()

  return acc.sum(), normalizing_factor


def compute_metrics(logits, targets, weights):
  """Compute summary metrics."""
  loss, weight_sum = compute_weighted_cross_entropy(logits, targets, weights)
  acc, _ = compute_weighted_accuracy(logits, targets, weights)
  metrics = {
      'loss': loss,
      'accuracy': acc,
      'denominator': weight_sum,
  }
  metrics = jax.lax.psum(metrics, 'batch')
  return metrics


# Train / eval / decode step functions.
# -----------------------------------------------------------------------------


def train_step(optimizer,
               inputs,
               outputs,
               programs,
               learning_rate_fn,
               config,
               dropout_rng):
  """Train on batch of program tasks."""
  # We handle PRNG splitting inside the top pmap, rather
  # than handling it outside in the training loop - doing the
  # latter can add some stalls to the devices.
  dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)

  weights = jnp.where(programs > 0, 1, 0).astype(jnp.float32)

  def loss_fn(params):
    """Loss function used for training."""
    logits = models.DecomposeAttentionTransformer(config).apply(
        {'params': params},
        inputs,
        outputs,
        programs,
        rngs={'dropout': dropout_rng})
    loss, weight_sum = compute_weighted_cross_entropy(logits, programs, weights)
    mean_loss = loss / weight_sum
    return mean_loss, logits

  step = optimizer.state.step
  lr = learning_rate_fn(step)
  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (_, logits), grad = grad_fn(optimizer.target)
  grad = jax.lax.pmean(grad, 'batch')
  new_optimizer = optimizer.apply_gradient(grad, learning_rate=lr)

  # Get metrics.
  metrics = compute_metrics(logits, programs, weights)
  metrics['learning_rate'] = lr
  return new_optimizer, metrics, new_dropout_rng


def eval_step(params, inputs, outputs, programs, eos_token, config):
  """Collect metrics for evaluation during training."""
  weights = jnp.where(
      jnp.logical_and(programs > 0,
                      jnp.logical_and(programs != config.base_config.bos_token,
                                      programs != eos_token)),
      1, 0).astype(jnp.float32)
  logits = models.DecomposeAttentionTransformer(config).apply(
      {'params': params}, inputs, outputs, programs)

  return compute_metrics(logits, programs, weights)


def initialize_cache(inputs, outputs, programs, max_decode_len, config):
  """Initialize a cache for a given input shape and max decode length."""
  target_shape = (programs.shape[0], max_decode_len)
  dtype = config.base_config.dtype
  initial_variables = models.DecomposeAttentionTransformer(config).init(
      jax.random.PRNGKey(0),
      jnp.ones(inputs.shape, dtype),
      jnp.ones(outputs.shape, dtype),
      jnp.ones(target_shape, dtype))
  return initial_variables['cache']


def predict_step(params,
                 inputs,
                 outputs,
                 cache,
                 beam_size,
                 eos_token,
                 max_decode_len,
                 config,
                 slow_decode=True):
  """Predict translation with fast decoding beam search on a batch."""
  # Prepare transformer fast-decoder call for beam search: for beam search, we
  # need to set up our decoder model to handle a batch size equal to
  # batch_size * beam_size, where each batch item's data is expanded in-place
  # rather than tiled.
  flat_encoded = decode.flat_batch_beam_expand(
      models.DecomposeAttentionTransformer(config).apply(
          {'params': params},
          inputs,
          outputs,
          method=models.DecomposeAttentionTransformer.encode),
      beam_size)

  encoded_padding_mask = jnp.where(outputs > 0, 1, 0).astype(jnp.float32)
  flat_encoded_padding_mask = decode.flat_batch_beam_expand(
      encoded_padding_mask, beam_size)

  if slow_decode:
    def tokens_ids_to_logits(flat_ids):
      """Token slice to logits from decoder model."""
      # --> [batch * beam, 1, vocab]
      flat_logits = models.DecomposeAttentionTransformer(config=config).apply(
          {'params': params},
          flat_ids,
          flat_encoded,
          flat_encoded_padding_mask,
          method=models.DecomposeAttentionTransformer.decode)
      return flat_logits
  else:
    def tokens_ids_to_logits(flat_ids, flat_cache):
      """Token slice to logits from decoder model."""
      # --> [batch * beam, 1, vocab]
      flat_logits, new_vars = models.DecomposeAttentionTransformer(
          config=config).apply(
              {'params': params, 'cache': flat_cache},
              flat_ids,
              flat_encoded,
              flat_encoded_padding_mask,
              mutable=['cache'],
              method=models.DecomposeAttentionTransformer.decode)
      new_flat_cache = new_vars['cache']
      # Remove singleton sequence-length dimension:
      # [batch * beam, 1, vocab] --> [batch * beam, vocab]
      flat_logits = flat_logits.squeeze(axis=1)
      return flat_logits, new_flat_cache

  # Using the above-defined single-step decoder function, run a
  # beam search over possible sequences given input encoding.
  beam_seqs, _ = decode.beam_search(
      inputs,
      cache,
      tokens_ids_to_logits,
      beam_size=beam_size,
      alpha=0.6,
      bos_token=config.base_config.bos_token,
      eos_token=eos_token,
      max_decode_len=max_decode_len,
      slow_decode=slow_decode)

  # Beam search returns [n_batch, n_beam, n_length] with beam dimension
  # sorted in increasing order of log-probability.
  return beam_seqs


# Util functions for prediction
# -----------------------------------------------------------------------------


def pad_examples(x, desired_batch_size):
  """Expand batch to desired size by repeating last slice."""
  batch_pad = desired_batch_size - x.shape[0]
  tile_dims = [1] * len(x.shape)
  tile_dims[0] = batch_pad
  return np.concatenate([x, np.tile(x[-1], tile_dims)], axis=0)


def tohost(x):
  """Collect batches from all devices to host and flatten batch dimensions."""
  n_device, n_batch, *remaining_dims = x.shape
  return x.reshape((n_device * n_batch,) + tuple(remaining_dims))


def per_host_sum_pmap(in_tree):
  """Execute psum on in_tree's leaves over one device per host."""
  host2devices = collections.defaultdict(list)
  for d in jax.devices():
    host2devices[d.host_id].append(d)
  devices = [host2devices[k][0] for k in host2devices]
  host_psum = jax.pmap(lambda x: jax.lax.psum(x, 'i'), 'i', devices=devices)
  def pre_pmap(xs):
    return jax.tree_map(lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs)
  def post_pmap(xs):
    return jax.tree_map(lambda x: x[0], xs)
  return post_pmap(host_psum(pre_pmap(in_tree)))


def eval_predicted(predicted, inputs, outputs, parse_beam_fn):
  """Evaluate predicted program beams."""
  best_p, best_score = None, -1

  # predicted shape [beam_size, length]
  for beam in predicted[::-1]:
    try:
      p = parse_beam_fn(beam)
      p_outs = [p(inp) for inp in inputs]
      score = np.sum([p_out == out for p_out, out in zip(p_outs, outputs)])
      if score > best_score:
        best_p, best_score = p, score
    except:  # pylint: disable=bare-except
      pass
    if best_score >= len(inputs):  # Found solution.
      break
  return best_p, best_score


def shorten(key):
  splits = key.split('_')
  return ''.join(s[0] for s in splits)


def main(_):
  tf.enable_v2_behavior()

  tf.random.set_seed(FLAGS.seed)
  np.random.seed(FLAGS.seed)
  random.seed(FLAGS.seed)

  # BOS special attention only makes sense if we are using relative attention
  # and it's not the baseline.
  if FLAGS.bos_special_attention and (not FLAGS.use_relative_attention or
                                      FLAGS.attention_mask_type == 'baseline'):
    raise ValueError(
        "bos_special_attention doesn't work when use_relative_attention={} and "
        'attention_mask_type={}'.format(FLAGS.use_relative_attention,
                                        FLAGS.attention_mask_type))

  if not gfile.isdir(FLAGS.save_dir):
    gfile.makedirs(FLAGS.save_dir)

  hparam_str_dict = dict(seed=FLAGS.seed, lr=FLAGS.lr)
  # Get hyperparmaters
  if FLAGS.xm_parameters:
    for key, value in json.loads(FLAGS.xm_parameters).items():
      if key not in hparam_str_dict:
        hparam_str_dict[key] = value

  hparam_str = ','.join(['%s=%s' % (shorten(k), str(hparam_str_dict[k]))
                         for k in sorted(hparam_str_dict.keys())])

  # Number of local devices for this host.
  n_devices = jax.local_device_count()

  if jax.host_id() == 0:
    summary_writer = tensorboard.SummaryWriter(
        os.path.join(FLAGS.save_dir, 'tb', hparam_str))

  batch_size = FLAGS.per_device_batch_size * n_devices
  io_shape = (FLAGS.per_device_batch_size,
              FLAGS.num_strings_per_task,
              FLAGS.max_characters)
  program_shape = (FLAGS.per_device_batch_size, FLAGS.max_program_length)

  # Setup DSL
  # ---------------------------------------------------------------------------

  # Build token tables.
  id_char_table = {i+1: char for (i, char) in enumerate(dsl.CHARACTER)}
  char_id_table = {char: id for id, char in id_char_table.items()}
  id_token_table, token_id_table = dsl_tokens.build_token_tables()
  io_vocab_size = len(char_id_table) + 1  # For padding.
  program_vocab_size = len(token_id_table) + 1

  bos_token = token_id_table[dsl.BOS]
  eos_token = token_id_table[dsl.EOS]

  # Parse io and program token sequences (for eval).
  def decode_io(inputs, outputs):
    """Decode io examples tokens."""
    def decode_str(s):
      """Decode string tokens."""
      return ''.join([id_char_table[c_id] for c_id in s if c_id > 0])

    inps, outs = [], []
    for inp, out in zip(inputs, outputs):
      inps.append(decode_str(inp))
      outs.append(decode_str(out))
    return inps, outs

  def decode_program(program):
    """Decode program tokens."""
    program = program[:np.argmax(program == eos_token) + 1].astype(np.int32)
    program = program[program != bos_token]

    try:
      return dsl.decode_program(program.tolist(), id_token_table)
    except:  # pylint: disable=bare-except
      return None  # Program does not compile.

  # Load Dataset
  # ---------------------------------------------------------------------------
  logging.info('Initializing dataset.')
  if not FLAGS.dataset_filepattern:
    raise ValueError('Must specify filepattern to dataset.')

  # Training dataset.
  logging.info('Loading dataset from %s', FLAGS.dataset_filepattern)
  padded_shapes = (io_shape[1:], io_shape[1:], program_shape[1:])
  logging.info('padded_shapes: %s', padded_shapes)
  dataset = input_pipeline.create_dataset_from_tf_record(
      FLAGS.dataset_filepattern, token_id_table, char_id_table)
  dataset = dataset.padded_batch(
      batch_size,
      padded_shapes=padded_shapes,
      drop_remainder=True)
  # Split evaluation and training.
  eval_ds = dataset.take(FLAGS.num_eval_steps)
  # Decrease batch of predict dataset to handle beam search.
  predict_ds = eval_ds.unbatch().padded_batch(
      int(np.ceil(batch_size / 10)),
      padded_shapes=padded_shapes)
  train_ds = dataset.skip(FLAGS.num_eval_steps).repeat()
  train_iter = train_ds.as_numpy_iterator()

  # Build Model and Optimizer
  # ---------------------------------------------------------------------------
  use_dropout = False
  base_config = base_models.TransformerConfig(
      vocab_size=io_vocab_size,
      output_vocab_size=program_vocab_size,
      shift=True,
      emb_dim=FLAGS.embedding_dim,
      num_heads=FLAGS.num_heads,
      num_layers=FLAGS.num_layers,
      qkv_dim=FLAGS.embedding_dim,
      mlp_dim=FLAGS.hidden_dim,
      max_len=max(FLAGS.max_characters, FLAGS.max_program_length),
      use_relative_attention=FLAGS.use_relative_attention,
      deterministic=not use_dropout,
      decode=False,
      bos_token=bos_token)
  train_config = models.DecomposeAttentionTransformerConfig(
      base_config=base_config,
      attention_mask_type=FLAGS.attention_mask_type,
      bos_special_attention=FLAGS.bos_special_attention)
  eval_config = models.DecomposeAttentionTransformerConfig(
      base_config=base_config.replace(deterministic=not use_dropout),
      attention_mask_type=FLAGS.attention_mask_type,
      bos_special_attention=FLAGS.bos_special_attention)
  predict_config = models.DecomposeAttentionTransformerConfig(
      base_config=base_config.replace(
          shift=False, deterministic=not use_dropout,
          decode=not FLAGS.slow_decode),
      attention_mask_type=FLAGS.attention_mask_type,
      bos_special_attention=FLAGS.bos_special_attention)

  rng = jax.random.PRNGKey(FLAGS.seed)
  rng = jax.random.fold_in(rng, jax.host_id())
  rng, init_rng = jax.random.split(rng)

  m = models.DecomposeAttentionTransformer(eval_config)
  initial_variables = jax.jit(m.init)(
      {'params': init_rng, 'dropout': init_rng},
      jnp.ones(io_shape, jnp.float32),
      jnp.ones(io_shape, jnp.float32),
      jnp.ones(program_shape, jnp.float32))

  optimizer_def = optim.Adam(
      FLAGS.lr,
      beta1=0.9,
      beta2=0.98,
      eps=1e-9,
      weight_decay=FLAGS.weight_decay)
  optimizer = optimizer_def.create(initial_variables['params'])

  del initial_variables  # Don't keep a copy of the initial model.

  start_step = 0
  if FLAGS.restore_checkpoints:
    # Restore unreplicated optimizer + model state from last checkpoint.
    optimizer = checkpoints.restore_checkpoint(
        os.path.join(FLAGS.save_dir, 'checkpoints', hparam_str), optimizer)
    # Grab last step.
    start_step = int(optimizer.state.step)
    logging.info('Found model checkpointed at step %d.', start_step)
    if FLAGS.finetune_start_step > 0:
      logging.info('Checking that start_step (%s) == finetune_start_step (%s)',
                   start_step, FLAGS.finetune_start_step)
      assert start_step == FLAGS.finetune_start_step

  # Replicate optimizer.
  optimizer = jax_utils.replicate(optimizer)

  # TODO(jxihong): Implement fast decoding.
  assert FLAGS.slow_decode, 'Fast decoding is not implemented yet.'

  if FLAGS.finetune_start_step <= 0:
    learning_rate_fn = create_learning_rate_scheduler(
        base_learning_rate=FLAGS.lr)
  else:
    # Constant LR for finetuning.
    learning_rate_fn = create_learning_rate_scheduler(
        base_learning_rate=FLAGS.lr,
        factors='constant')
  p_train_step = jax.pmap(
      functools.partial(
          train_step,
          learning_rate_fn=learning_rate_fn,
          config=train_config),
      axis_name='batch')
  p_eval_step = jax.pmap(
      functools.partial(eval_step,
                        eos_token=eos_token,
                        config=eval_config),
      axis_name='batch')
  p_init_cache = jax.pmap(
      functools.partial(
          initialize_cache,
          max_decode_len=FLAGS.max_program_length,
          config=predict_config),
      axis_name='batch')
  p_pred_step = jax.pmap(
      functools.partial(
          predict_step,
          eos_token=eos_token,
          max_decode_len=FLAGS.max_program_length,
          config=predict_config,
          slow_decode=FLAGS.slow_decode),
      axis_name='batch',
      static_broadcasted_argnums=(4,))

  # Main Train Loop
  # ---------------------------------------------------------------------------
  dropout_rng = jax.random.split(rng, jax.local_device_count())
  del rng

  metrics_all = []
  tick = time.time()
  for step in range(start_step, FLAGS.num_train_steps):
    inputs, outputs, programs = common_utils.shard(next(train_iter))

    optimizer, metrics, dropout_rng = p_train_step(
        optimizer, inputs, outputs, programs, dropout_rng=dropout_rng)
    metrics_all.append(metrics)
    is_last_step = step == FLAGS.num_train_steps - 1

    # Save a Checkpoint
    if (step % FLAGS.checkpoint_freq == 0 and step > 0) or is_last_step:
      if jax.host_id() == 0:
        # Save unreplicated optimizer + model state.
        checkpoints.save_checkpoint(
            os.path.join(FLAGS.save_dir, 'checkpoints', hparam_str),
            jax_utils.unreplicate(optimizer),
            step)

    # Periodic metric handling.

    # Training Metrics
    if (step and step % FLAGS.log_freq == 0) or is_last_step:
      logging.info('Gathering training metrics.')
      metrics_all = common_utils.get_metrics(metrics_all)
      lr = metrics_all.pop('learning_rate').mean()
      metrics_sums = jax.tree_map(jnp.sum, metrics_all)
      denominator = metrics_sums.pop('denominator')
      summary = jax.tree_map(
          lambda x: x / denominator,  # pylint: disable=cell-var-from-loop
          metrics_sums)
      summary['learning_rate'] = lr
      # Calculate (clipped) perplexity after averaging log-perplexities:
      summary['perplexity'] = jnp.clip(jnp.exp(summary['loss']), a_max=1.0e4)

      if jax.host_id() == 0:
        logging.info('Train in step: %d, loss: %.4f', step, summary['loss'])
        tock = time.time()
        steps_per_sec = FLAGS.log_freq / (tock - tick)
        tick = tock
        summary_writer.scalar('train/steps per second', steps_per_sec, step)
        for key, val in summary.items():
          summary_writer.scalar('train/' + key, val, step)
        summary_writer.flush()
      # Reset metric accumulation for next evaluation cycle.
      metrics_all = []

    # Evaluation Metrics
    if (step and step % FLAGS.eval_freq == 0) or is_last_step:
      logging.info('Gathering evaluation metrics.')
      t_evaluation_start = time.time()
      eval_metrics = []
      for batches in eval_ds.as_numpy_iterator():
        inputs, outputs, programs = common_utils.shard(batches)

        metrics = p_eval_step(optimizer.target, inputs, outputs, programs)
        eval_metrics.append(metrics)

      eval_metrics = common_utils.get_metrics(eval_metrics)
      eval_metrics_sums = jax.tree_map(jnp.sum, eval_metrics)
      eval_denominator = eval_metrics_sums.pop('denominator')
      eval_summary = jax.tree_map(
          lambda x: x / eval_denominator,  # pylint: disable=cell-var-from-loop
          eval_metrics_sums)

      if jax.host_id() == 0:
        logging.info('Evaluation time: %.4f s step %d, loss: %.4f.',
                     time.time()-t_evaluation_start, step, eval_summary['loss'])
        for key, val in eval_summary.items():
          summary_writer.scalar('eval/' + key, val, step)
        summary_writer.flush()

    # Beam search metrics.
    if (step and step % FLAGS.predict_freq == 0) or is_last_step:
      logging.info('Gathering beam search metrics.')
      for beam_size in [1, 5, 10, 20, 50]:
        t_inference_start = time.time()
        pred_acc = 0
        pred_denominator = 0

        ios, targets, predictions, top_of_beams = [], [], [], []
        for batches in predict_ds.as_numpy_iterator():
          pred_batch = batches
          # Handle final odd-sized batch by padding instead of dropping it.
          cur_pred_batch_size = pred_batch[0].shape[0]
          if cur_pred_batch_size % n_devices:
            padded_size = int(
                np.ceil(cur_pred_batch_size / n_devices) * n_devices)
            # pylint: disable=cell-var-from-loop
            pred_batch = jax.tree_map(
                lambda x: pad_examples(x, padded_size), pred_batch)
          inputs, outputs, programs = common_utils.shard(pred_batch)

          cache = (p_init_cache(inputs, outputs, programs)
                   if not FLAGS.slow_decode else None)
          predicted = p_pred_step(optimizer.target, inputs, outputs, cache,
                                  beam_size)
          predicted = tohost(predicted)
          inputs, outputs, programs = map(tohost, (inputs, outputs, programs))

          pred_denominator += programs.shape[0]
          for i, beams in enumerate(predicted):
            inps, outs = decode_io(inputs[i], outputs[i])
            p, p_score = eval_predicted(
                beams, inps, outs, parse_beam_fn=decode_program)
            if p_score >= len(inps):
              pred_acc += 1
            ios.append(' ; '.join(map(str, zip(inps, outs))))
            targets.append(decode_program(programs[i]).to_string())
            try:
              predictions.append(p.to_string())
            except:  # pylint: disable=bare-except
              predictions.append('Did not compile')
            logging.info('ios: %s', ios[-1])
            logging.info('target: %s', targets[-1])
            beams_log = []
            for beam in beams:
              try:
                beams_log.append(decode_program(beam).to_string())
              except:  # pylint: disable=bare-except
                beams_log.append('Did not compile')
            logging.info('predicted beam: %s', '\n'.join(beams_log))

            top_of_beam = []
            for index, beam in enumerate(beams[:-5:-1]):
              try:
                decoded_program = decode_program(beam).to_string()
              except:  # pylint: disable=bare-except
                decoded_program = 'Did not compile'
              top_of_beam.append('index: {}, decoded: {}, tokens: {}'.format(
                  index, decoded_program, beam))
            top_of_beams.append('\n\n'.join(top_of_beam))

        all_pred_acc, all_pred_denominator = per_host_sum_pmap(
            jax.tree_map(np.array, (pred_acc, pred_denominator)))

        # Record beam search results as text summaries.
        message = []
        for n in np.random.choice(np.arange(len(predictions)), 8):
          text = (f'ios: {ios[n]}\n\ntarget: {targets[n]}\n\n'
                  f'predicted: {predictions[n]}\n\n'
                  f'top of beam:\n\n{top_of_beams[n]}\n\n')
          message.append(text)

        # Write to tensorboard.
        if jax.host_id() == 0:
          slow_or_fast = 'slow' if FLAGS.slow_decode else 'fast'
          logging.info(
              'Prediction time, %s (beam %d): %.4f s, step %d, score %.4f',
              slow_or_fast, beam_size, time.time() - t_inference_start, step,
              all_pred_acc / all_pred_denominator)
          summary_writer.scalar(
              'predict-{}/score-{}'.format(slow_or_fast, beam_size),
              all_pred_acc / all_pred_denominator, step)
          summary_writer.text('samples-{}'.format(beam_size),
                              '\n------\n'.join(message), step)
          summary_writer.flush()


if __name__ == '__main__':
  app.run(main)

[((51, 0, 51, 25), 'sys.path.append', 'sys.path.append', ({(51, 16, 51, 24): '"""../../"""'}, {}), "('../../')", False, 'import sys\n'), ((56, 0, 56, 66), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(56, 21, 56, 27): '"""seed"""', (56, 29, 56, 30): '(0)', (56, 32, 56, 65): '"""Fixed random seed for training."""'}, {}), "('seed', 0, 'Fixed random seed for training.')", False, 'from absl import flags\n'), ((57, 0, 57, 48), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(57, 19, 57, 23): '"""lr"""', (57, 25, 57, 29): '(0.001)', (57, 31, 57, 47): '"""Learning rate."""'}, {}), "('lr', 0.001, 'Learning rate.')", False, 'from absl import flags\n'), ((58, 0, 59, 64), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(58, 19, 58, 33): '"""weight_decay"""', (58, 35, 58, 39): '(0.1)', (59, 19, 59, 63): '"""Decay factor for AdamW-style weight decay."""'}, {}), "('weight_decay', 0.1,\n    'Decay factor for AdamW-style weight decay.')", False, 'from absl import flags\n'), ((60, 0, 60, 66), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(60, 21, 60, 36): '"""embedding_dim"""', (60, 38, 60, 41): '(256)', (60, 43, 60, 65): '"""Embedding dimension."""'}, {}), "('embedding_dim', 256, 'Embedding dimension.')", False, 'from absl import flags\n'), ((61, 0, 61, 60), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(61, 21, 61, 33): '"""hidden_dim"""', (61, 35, 61, 38): '(512)', (61, 40, 61, 59): '"""Hidden dimension."""'}, {}), "('hidden_dim', 512, 'Hidden dimension.')", False, 'from absl import flags\n'), ((62, 0, 62, 57), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(62, 21, 62, 32): '"""num_heads"""', (62, 34, 62, 35): '(4)', (62, 37, 62, 56): '"""Number of layers."""'}, {}), "('num_heads', 4, 'Number of layers.')", False, 'from absl import flags\n'), ((63, 0, 63, 69), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(63, 21, 63, 33): '"""num_layers"""', (63, 35, 63, 36): '(3)', (63, 38, 63, 68): '"""Number of Transformer heads."""'}, {}), "('num_layers', 3, 'Number of Transformer heads.')", False, 'from absl import flags\n'), ((64, 0, 64, 78), 'absl.flags.DEFINE_boolean', 'flags.DEFINE_boolean', ({(64, 21, 64, 34): '"""slow_decode"""', (64, 36, 64, 40): '(True)', (64, 42, 64, 77): '"""Use slow decoding for prediction?"""'}, {}), "('slow_decode', True, 'Use slow decoding for prediction?')", False, 'from absl import flags\n'), ((66, 0, 67, 56), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(66, 20, 66, 41): '"""dataset_filepattern"""', (66, 43, 66, 47): 'None', (67, 20, 67, 55): '"""Filepattern for TFRecord dataset."""'}, {}), "('dataset_filepattern', None,\n    'Filepattern for TFRecord dataset.')", False, 'from absl import flags\n'), ((68, 0, 69, 59), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(68, 21, 68, 44): '"""per_device_batch_size"""', (68, 46, 68, 48): '(16)', (69, 21, 69, 58): '"""Number of program tasks in a batch."""'}, {}), "('per_device_batch_size', 16,\n    'Number of program tasks in a batch.')", False, 'from absl import flags\n'), ((70, 0, 71, 64), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(70, 21, 70, 43): '"""num_strings_per_task"""', (70, 45, 70, 46): '(4)', (71, 21, 71, 63): '"""Number of input/output strings per task."""'}, {}), "('num_strings_per_task', 4,\n    'Number of input/output strings per task.')", False, 'from absl import flags\n'), ((72, 0, 73, 60), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(72, 21, 72, 41): '"""max_program_length"""', (72, 43, 72, 46): '(100)', (73, 21, 73, 59): '"""Maximum number of tokens in program."""'}, {}), "('max_program_length', 100,\n    'Maximum number of tokens in program.')", False, 'from absl import flags\n'), ((74, 0, 75, 77), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(74, 21, 74, 37): '"""max_characters"""', (74, 39, 74, 42): '(120)', (75, 21, 75, 76): '"""Maximum number of characters in input/output strings."""'}, {}), "('max_characters', 120,\n    'Maximum number of characters in input/output strings.')", False, 'from absl import flags\n'), ((77, 0, 77, 70), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(77, 20, 77, 30): '"""save_dir"""', (77, 32, 77, 36): 'None', (77, 38, 77, 69): '"""Directory to save results to."""'}, {}), "('save_dir', None, 'Directory to save results to.')", False, 'from absl import flags\n'), ((78, 0, 78, 77), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(78, 21, 78, 38): '"""num_train_steps"""', (78, 40, 78, 47): '(2000000)', (78, 49, 78, 76): '"""Number of training steps."""'}, {}), "('num_train_steps', 2000000, 'Number of training steps.')", False, 'from absl import flags\n'), ((79, 0, 79, 73), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(79, 21, 79, 37): '"""num_eval_steps"""', (79, 39, 79, 41): '(10)', (79, 43, 79, 72): '"""Number of evaluation steps."""'}, {}), "('num_eval_steps', 10, 'Number of evaluation steps.')", False, 'from absl import flags\n'), ((80, 0, 80, 80), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(80, 21, 80, 31): '"""log_freq"""', (80, 33, 80, 37): '(1000)', (80, 39, 80, 79): '"""Number of steps between training logs."""'}, {}), "('log_freq', 1000, 'Number of steps between training logs.'\n    )", False, 'from absl import flags\n'), ((81, 0, 81, 72), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(81, 21, 81, 32): '"""eval_freq"""', (81, 34, 81, 38): '(2000)', (81, 40, 81, 71): '"""Number of steps between eval."""'}, {}), "('eval_freq', 2000, 'Number of steps between eval.')", False, 'from absl import flags\n'), ((82, 0, 83, 73), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(82, 21, 82, 35): '"""predict_freq"""', (82, 37, 82, 42): '(50000)', (83, 21, 83, 72): '"""Number of steps between prediction (beam search)."""'}, {}), "('predict_freq', 50000,\n    'Number of steps between prediction (beam search).')", False, 'from absl import flags\n'), ((84, 0, 85, 65), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(84, 21, 84, 38): '"""checkpoint_freq"""', (84, 40, 84, 45): '(50000)', (85, 21, 85, 64): '"""Number of steps between checkpoint saves."""'}, {}), "('checkpoint_freq', 50000,\n    'Number of steps between checkpoint saves.')", False, 'from absl import flags\n'), ((86, 0, 88, 60), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(86, 21, 86, 42): '"""finetune_start_step"""', (86, 44, 86, 46): '(-1)', (87, 21, 88, 59): '"""Step the initial checkpoint should start at for finetuning, or -1 if not finetuning."""'}, {}), "('finetune_start_step', -1,\n    'Step the initial checkpoint should start at for finetuning, or -1 if not finetuning.'\n    )", False, 'from absl import flags\n'), ((89, 0, 90, 72), 'absl.flags.DEFINE_bool', 'flags.DEFINE_bool', ({(89, 18, 89, 39): '"""restore_checkpoints"""', (89, 41, 89, 45): '(True)', (90, 18, 90, 71): '"""Whether to restore from existing model checkpoints."""'}, {}), "('restore_checkpoints', True,\n    'Whether to restore from existing model checkpoints.')", False, 'from absl import flags\n'), ((92, 0, 94, 53), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(92, 20, 92, 41): '"""attention_mask_type"""', (92, 43, 92, 63): '"""bos_full_attention"""', (93, 20, 94, 52): '"""The kind of attention mask to use. Options are: baseline, bos_to_bos, bos_full_attention"""'}, {}), "('attention_mask_type', 'bos_full_attention',\n    'The kind of attention mask to use. Options are: baseline, bos_to_bos, bos_full_attention'\n    )", False, 'from absl import flags\n'), ((96, 0, 97, 66), 'absl.flags.DEFINE_bool', 'flags.DEFINE_bool', ({(96, 18, 96, 42): '"""use_relative_attention"""', (96, 44, 96, 48): '(True)', (97, 18, 97, 65): '"""Whether to use relative positonal embeddings."""'}, {}), "('use_relative_attention', True,\n    'Whether to use relative positonal embeddings.')", False, 'from absl import flags\n'), ((98, 0, 100, 32), 'absl.flags.DEFINE_bool', 'flags.DEFINE_bool', ({(98, 18, 98, 41): '"""bos_special_attention"""', (98, 43, 98, 48): '(False)', (99, 18, 100, 31): '"""Whether to use special relative attention computation for BOS tokens."""'}, {}), "('bos_special_attention', False,\n    'Whether to use special relative attention computation for BOS tokens.')", False, 'from absl import flags\n'), ((105, 2, 106, 64), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(105, 22, 105, 37): '"""xm_parameters"""', (105, 39, 105, 43): 'None', (106, 22, 106, 63): '"""String specifying hyperparamter search."""'}, {}), "('xm_parameters', None,\n    'String specifying hyperparamter search.')", False, 'from absl import flags\n'), ((181, 19, 181, 65), 'flax.training.common_utils.onehot', 'common_utils.onehot', ({(181, 39, 181, 46): 'targets', (181, 48, 181, 64): 'logits.shape[-1]'}, {}), '(targets, logits.shape[-1])', False, 'from flax.training import common_utils\n'), ((223, 12, 223, 42), 'jax.lax.psum', 'jax.lax.psum', ({(223, 25, 223, 32): 'metrics', (223, 34, 223, 41): '"""batch"""'}, {}), "(metrics, 'batch')", False, 'import jax\n'), ((242, 33, 242, 62), 'jax.random.split', 'jax.random.split', ({(242, 50, 242, 61): 'dropout_rng'}, {}), '(dropout_rng)', False, 'import jax\n'), ((260, 12, 260, 53), 'jax.value_and_grad', 'jax.value_and_grad', (), '', False, 'import jax\n'), ((262, 9, 262, 37), 'jax.lax.pmean', 'jax.lax.pmean', ({(262, 23, 262, 27): 'grad', (262, 29, 262, 36): '"""batch"""'}, {}), "(grad, 'batch')", False, 'import jax\n'), ((319, 30, 320, 38), 'latent_programmer.decode.flat_batch_beam_expand', 'decode.flat_batch_beam_expand', ({(320, 6, 320, 26): 'encoded_padding_mask', (320, 28, 320, 37): 'beam_size'}, {}), '(encoded_padding_mask, beam_size)', False, 'from latent_programmer import decode\n'), ((353, 17, 362, 30), 'latent_programmer.decode.beam_search', 'decode.beam_search', (), '', False, 'from latent_programmer import decode\n'), ((389, 17, 389, 46), 'collections.defaultdict', 'collections.defaultdict', ({(389, 41, 389, 45): 'list'}, {}), '(list)', False, 'import collections\n'), ((390, 11, 390, 24), 'jax.devices', 'jax.devices', ({}, {}), '()', False, 'import jax\n'), ((426, 2, 426, 25), 'tensorflow.compat.v2.enable_v2_behavior', 'tf.enable_v2_behavior', ({}, {}), '()', True, 'import tensorflow.compat.v2 as tf\n'), ((428, 2, 428, 32), 'tensorflow.compat.v2.random.set_seed', 'tf.random.set_seed', ({(428, 21, 428, 31): 'FLAGS.seed'}, {}), '(FLAGS.seed)', True, 'import tensorflow.compat.v2 as tf\n'), ((429, 2, 429, 28), 'numpy.random.seed', 'np.random.seed', ({(429, 17, 429, 27): 'FLAGS.seed'}, {}), '(FLAGS.seed)', True, 'import numpy as np\n'), ((430, 2, 430, 25), 'random.seed', 'random.seed', ({(430, 14, 430, 24): 'FLAGS.seed'}, {}), '(FLAGS.seed)', False, 'import random\n'), ((455, 14, 455, 38), 'jax.local_device_count', 'jax.local_device_count', ({}, {}), '()', False, 'import jax\n'), ((473, 35, 473, 66), 'latent_programmer.tasks.robust_fill.tokens.build_token_tables', 'dsl_tokens.build_token_tables', ({}, {}), '()', True, 'from latent_programmer.tasks.robust_fill import tokens as dsl_tokens\n'), ((505, 2, 505, 39), 'absl.logging.info', 'logging.info', ({(505, 15, 505, 38): '"""Initializing dataset."""'}, {}), "('Initializing dataset.')", False, 'from absl import logging\n'), ((510, 2, 510, 68), 'absl.logging.info', 'logging.info', ({(510, 15, 510, 40): '"""Loading dataset from %s"""', (510, 42, 510, 67): 'FLAGS.dataset_filepattern'}, {}), "('Loading dataset from %s', FLAGS.dataset_filepattern)", False, 'from absl import logging\n'), ((512, 2, 512, 50), 'absl.logging.info', 'logging.info', ({(512, 15, 512, 34): '"""padded_shapes: %s"""', (512, 36, 512, 49): 'padded_shapes'}, {}), "('padded_shapes: %s', padded_shapes)", False, 'from absl import logging\n'), ((513, 12, 514, 63), 'latent_programmer.decomposition_transformer_attention.input_pipeline.create_dataset_from_tf_record', 'input_pipeline.create_dataset_from_tf_record', ({(514, 6, 514, 31): 'FLAGS.dataset_filepattern', (514, 33, 514, 47): 'token_id_table', (514, 49, 514, 62): 'char_id_table'}, {}), '(FLAGS.dataset_filepattern,\n    token_id_table, char_id_table)', False, 'from latent_programmer.decomposition_transformer_attention import input_pipeline\n'), ((545, 17, 548, 56), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformerConfig', 'models.DecomposeAttentionTransformerConfig', (), '', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((560, 8, 560, 38), 'jax.random.PRNGKey', 'jax.random.PRNGKey', ({(560, 27, 560, 37): 'FLAGS.seed'}, {}), '(FLAGS.seed)', False, 'import jax\n'), ((562, 18, 562, 39), 'jax.random.split', 'jax.random.split', ({(562, 35, 562, 38): 'rng'}, {}), '(rng)', False, 'import jax\n'), ((564, 6, 564, 55), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', ({(564, 43, 564, 54): 'eval_config'}, {}), '(eval_config)', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((571, 18, 576, 38), 'flax.optim.Adam', 'optim.Adam', (), '', False, 'from flax import optim\n'), ((595, 14, 595, 44), 'flax.jax_utils.replicate', 'jax_utils.replicate', ({(595, 34, 595, 43): 'optimizer'}, {}), '(optimizer)', False, 'from flax import jax_utils\n'), ((641, 9, 641, 20), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((800, 2, 800, 15), 'absl.app.run', 'app.run', ({(800, 10, 800, 14): 'main'}, {}), '(main)', False, 'from absl import app\n'), ((161, 11, 161, 46), 'jax.numpy.asarray', 'jnp.asarray', (), '', True, 'import jax.numpy as jnp\n'), ((183, 32, 183, 58), 'jax.numpy.asarray', 'jnp.asarray', ({(183, 44, 183, 57): 'targets.shape'}, {}), '(targets.shape)', True, 'import jax.numpy as jnp\n'), ((205, 18, 205, 45), 'jax.numpy.argmax', 'jnp.argmax', (), '', True, 'import jax.numpy as jnp\n'), ((206, 32, 206, 58), 'jax.numpy.asarray', 'jnp.asarray', ({(206, 44, 206, 57): 'targets.shape'}, {}), '(targets.shape)', True, 'import jax.numpy as jnp\n'), ((289, 6, 289, 27), 'jax.random.PRNGKey', 'jax.random.PRNGKey', ({(289, 25, 289, 26): '0'}, {}), '(0)', False, 'import jax\n'), ((290, 6, 290, 35), 'jax.numpy.ones', 'jnp.ones', ({(290, 15, 290, 27): 'inputs.shape', (290, 29, 290, 34): 'dtype'}, {}), '(inputs.shape, dtype)', True, 'import jax.numpy as jnp\n'), ((291, 6, 291, 36), 'jax.numpy.ones', 'jnp.ones', ({(291, 15, 291, 28): 'outputs.shape', (291, 30, 291, 35): 'dtype'}, {}), '(outputs.shape, dtype)', True, 'import jax.numpy as jnp\n'), ((292, 6, 292, 35), 'jax.numpy.ones', 'jnp.ones', ({(292, 15, 292, 27): 'target_shape', (292, 29, 292, 34): 'dtype'}, {}), '(target_shape, dtype)', True, 'import jax.numpy as jnp\n'), ((397, 11, 397, 43), 'jax.tree_map', 'jax.tree_map', ({(397, 24, 397, 38): '(lambda x: x[0])', (397, 40, 397, 42): 'xs'}, {}), '(lambda x: x[0], xs)', False, 'import jax\n'), ((457, 5, 457, 18), 'jax.host_id', 'jax.host_id', ({}, {}), '()', False, 'import jax\n'), ((561, 32, 561, 45), 'jax.host_id', 'jax.host_id', ({}, {}), '()', False, 'import jax\n'), ((565, 22, 565, 37), 'jax.jit', 'jax.jit', ({(565, 30, 565, 36): 'm.init'}, {}), '(m.init)', False, 'import jax\n'), ((567, 6, 567, 37), 'jax.numpy.ones', 'jnp.ones', ({(567, 15, 567, 23): 'io_shape', (567, 25, 567, 36): 'jnp.float32'}, {}), '(io_shape, jnp.float32)', True, 'import jax.numpy as jnp\n'), ((568, 6, 568, 37), 'jax.numpy.ones', 'jnp.ones', ({(568, 15, 568, 23): 'io_shape', (568, 25, 568, 36): 'jnp.float32'}, {}), '(io_shape, jnp.float32)', True, 'import jax.numpy as jnp\n'), ((569, 6, 569, 42), 'jax.numpy.ones', 'jnp.ones', ({(569, 15, 569, 28): 'program_shape', (569, 30, 569, 41): 'jnp.float32'}, {}), '(program_shape, jnp.float32)', True, 'import jax.numpy as jnp\n'), ((588, 4, 588, 68), 'absl.logging.info', 'logging.info', ({(588, 17, 588, 55): '"""Found model checkpointed at step %d."""', (588, 57, 588, 67): 'start_step'}, {}), "('Found model checkpointed at step %d.', start_step)", False, 'from absl import logging\n'), ((609, 6, 612, 30), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((615, 6, 617, 43), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((620, 6, 623, 32), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((626, 6, 631, 40), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((637, 38, 637, 62), 'jax.local_device_count', 'jax.local_device_count', ({}, {}), '()', False, 'import jax\n'), ((244, 12, 244, 41), 'jax.numpy.where', 'jnp.where', ({(244, 22, 244, 34): 'programs > 0', (244, 36, 244, 37): '1', (244, 39, 244, 40): '0'}, {}), '(programs > 0, 1, 0)', True, 'import jax.numpy as jnp\n'), ((278, 11, 278, 55), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', ({(278, 48, 278, 54): 'config'}, {}), '(config)', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((288, 22, 288, 66), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', ({(288, 59, 288, 65): 'config'}, {}), '(config)', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((318, 25, 318, 53), 'jax.numpy.where', 'jnp.where', ({(318, 35, 318, 46): 'outputs > 0', (318, 48, 318, 49): '1', (318, 51, 318, 52): '0'}, {}), '(outputs > 0, 1, 0)', True, 'import jax.numpy as jnp\n'), ((378, 28, 378, 53), 'numpy.tile', 'np.tile', ({(378, 36, 378, 41): 'x[-1]', (378, 43, 378, 52): 'tile_dims'}, {}), '(x[-1], tile_dims)', True, 'import numpy as np\n'), ((393, 33, 393, 53), 'jax.lax.psum', 'jax.lax.psum', ({(393, 46, 393, 47): 'x', (393, 49, 393, 52): '"""i"""'}, {}), "(x, 'i')", False, 'import jax\n'), ((459, 8, 459, 54), 'os.path.join', 'os.path.join', ({(459, 21, 459, 35): 'FLAGS.save_dir', (459, 37, 459, 41): '"""tb"""', (459, 43, 459, 53): 'hparam_str'}, {}), "(FLAGS.save_dir, 'tb', hparam_str)", False, 'import os\n'), ((523, 10, 523, 34), 'numpy.ceil', 'np.ceil', ({(523, 18, 523, 33): 'batch_size / 10'}, {}), '(batch_size / 10)', True, 'import numpy as np\n'), ((585, 8, 585, 63), 'os.path.join', 'os.path.join', ({(585, 21, 585, 35): 'FLAGS.save_dir', (585, 37, 585, 50): '"""checkpoints"""', (585, 52, 585, 62): 'hparam_str'}, {}), "(FLAGS.save_dir, 'checkpoints', hparam_str)", False, 'import os\n'), ((590, 6, 591, 57), 'absl.logging.info', 'logging.info', ({(590, 19, 590, 78): '"""Checking that start_step (%s) == finetune_start_step (%s)"""', (591, 19, 591, 29): 'start_step', (591, 31, 591, 56): 'FLAGS.finetune_start_step'}, {}), "('Checking that start_step (%s) == finetune_start_step (%s)',\n    start_step, FLAGS.finetune_start_step)", False, 'from absl import logging\n'), ((663, 6, 663, 49), 'absl.logging.info', 'logging.info', ({(663, 19, 663, 48): '"""Gathering training metrics."""'}, {}), "('Gathering training metrics.')", False, 'from absl import logging\n'), ((664, 20, 664, 57), 'flax.training.common_utils.get_metrics', 'common_utils.get_metrics', ({(664, 45, 664, 56): 'metrics_all'}, {}), '(metrics_all)', False, 'from flax.training import common_utils\n'), ((666, 21, 666, 55), 'jax.tree_map', 'jax.tree_map', ({(666, 34, 666, 41): 'jnp.sum', (666, 43, 666, 54): 'metrics_all'}, {}), '(jnp.sum, metrics_all)', False, 'import jax\n'), ((668, 16, 670, 23), 'jax.tree_map', 'jax.tree_map', ({(669, 10, 669, 35): 'lambda x: x / denominator', (670, 10, 670, 22): 'metrics_sums'}, {}), '(lambda x: x / denominator, metrics_sums)', False, 'import jax\n'), ((689, 6, 689, 51), 'absl.logging.info', 'logging.info', ({(689, 19, 689, 50): '"""Gathering evaluation metrics."""'}, {}), "('Gathering evaluation metrics.')", False, 'from absl import logging\n'), ((690, 27, 690, 38), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((698, 21, 698, 59), 'flax.training.common_utils.get_metrics', 'common_utils.get_metrics', ({(698, 46, 698, 58): 'eval_metrics'}, {}), '(eval_metrics)', False, 'from flax.training import common_utils\n'), ((699, 26, 699, 61), 'jax.tree_map', 'jax.tree_map', ({(699, 39, 699, 46): 'jnp.sum', (699, 48, 699, 60): 'eval_metrics'}, {}), '(jnp.sum, eval_metrics)', False, 'import jax\n'), ((701, 21, 703, 28), 'jax.tree_map', 'jax.tree_map', ({(702, 10, 702, 40): 'lambda x: x / eval_denominator', (703, 10, 703, 27): 'eval_metrics_sums'}, {}), '(lambda x: x / eval_denominator, eval_metrics_sums)', False, 'import jax\n'), ((714, 6, 714, 52), 'absl.logging.info', 'logging.info', ({(714, 19, 714, 51): '"""Gathering beam search metrics."""'}, {}), "('Gathering beam search metrics.')", False, 'from absl import logging\n'), ((182, 35, 182, 57), 'flax.linen.log_softmax', 'nn.log_softmax', ({(182, 50, 182, 56): 'logits'}, {}), '(logits)', True, 'from flax import linen as nn\n'), ((248, 13, 248, 57), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', ({(248, 50, 248, 56): 'config'}, {}), '(config)', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((311, 6, 311, 50), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', ({(311, 43, 311, 49): 'config'}, {}), '(config)', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((395, 34, 395, 69), 'jax.numpy.broadcast_to', 'jnp.broadcast_to', ({(395, 51, 395, 52): 'x', (395, 54, 395, 68): '((1,) + x.shape)'}, {}), '(x, (1,) + x.shape)', True, 'import jax.numpy as jnp\n'), ((447, 22, 447, 53), 'json.loads', 'json.loads', ({(447, 33, 447, 52): 'FLAGS.xm_parameters'}, {}), '(FLAGS.xm_parameters)', False, 'import json\n'), ((652, 9, 652, 22), 'jax.host_id', 'jax.host_id', ({}, {}), '()', False, 'import jax\n'), ((673, 39, 673, 63), 'jax.numpy.exp', 'jnp.exp', ({(673, 47, 673, 62): "summary['loss']"}, {}), "(summary['loss'])", True, 'import jax.numpy as jnp\n'), ((675, 9, 675, 22), 'jax.host_id', 'jax.host_id', ({}, {}), '()', False, 'import jax\n'), ((676, 8, 676, 76), 'absl.logging.info', 'logging.info', ({(676, 21, 676, 52): '"""Train in step: %d, loss: %.4f"""', (676, 54, 676, 58): 'step', (676, 60, 676, 75): "summary['loss']"}, {}), "('Train in step: %d, loss: %.4f', step, summary['loss'])", False, 'from absl import logging\n'), ((677, 15, 677, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((693, 36, 693, 63), 'flax.training.common_utils.shard', 'common_utils.shard', ({(693, 55, 693, 62): 'batches'}, {}), '(batches)', False, 'from flax.training import common_utils\n'), ((705, 9, 705, 22), 'jax.host_id', 'jax.host_id', ({}, {}), '()', False, 'import jax\n'), ((716, 28, 716, 39), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((146, 15, 146, 52), 'jax.numpy.minimum', 'jnp.minimum', ({(146, 27, 146, 30): '(1.0)', (146, 32, 146, 51): '(step / warmup_steps)'}, {}), '(1.0, step / warmup_steps)', True, 'import jax.numpy as jnp\n'), ((275, 22, 276, 60), 'jax.numpy.logical_and', 'jnp.logical_and', ({(275, 38, 275, 78): 'programs != config.base_config.bos_token', (276, 38, 276, 59): 'programs != eos_token'}, {}), '(programs != config.base_config.bos_token, programs != eos_token\n    )', True, 'import jax.numpy as jnp\n'), ((326, 20, 326, 71), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', (), '', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((337, 30, 338, 24), 'latent_programmer.decomposition_transformer_attention.decomposition_models.DecomposeAttentionTransformer', 'models.DecomposeAttentionTransformer', (), '', True, 'from latent_programmer.decomposition_transformer_attention import decomposition_models as models\n'), ((655, 12, 655, 67), 'os.path.join', 'os.path.join', ({(655, 25, 655, 39): 'FLAGS.save_dir', (655, 41, 655, 54): '"""checkpoints"""', (655, 56, 655, 66): 'hparam_str'}, {}), "(FLAGS.save_dir, 'checkpoints', hparam_str)", False, 'import os\n'), ((656, 12, 656, 44), 'flax.jax_utils.unreplicate', 'jax_utils.unreplicate', ({(656, 34, 656, 43): 'optimizer'}, {}), '(optimizer)', False, 'from flax import jax_utils\n'), ((731, 38, 731, 68), 'flax.training.common_utils.shard', 'common_utils.shard', ({(731, 57, 731, 67): 'pred_batch'}, {}), '(pred_batch)', False, 'from flax.training import common_utils\n'), ((774, 12, 774, 64), 'jax.tree_map', 'jax.tree_map', ({(774, 25, 774, 33): 'np.array', (774, 35, 774, 63): '(pred_acc, pred_denominator)'}, {}), '(np.array, (pred_acc, pred_denominator))', False, 'import jax\n'), ((785, 11, 785, 24), 'jax.host_id', 'jax.host_id', ({}, {}), '()', False, 'import jax\n'), ((707, 21, 707, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((753, 12, 753, 44), 'absl.logging.info', 'logging.info', ({(753, 25, 753, 34): '"""ios: %s"""', (753, 36, 753, 43): 'ios[-1]'}, {}), "('ios: %s', ios[-1])", False, 'from absl import logging\n'), ((754, 12, 754, 51), 'absl.logging.info', 'logging.info', ({(754, 25, 754, 37): '"""target: %s"""', (754, 39, 754, 50): 'targets[-1]'}, {}), "('target: %s', targets[-1])", False, 'from absl import logging\n'), ((148, 24, 148, 61), 'jax.numpy.maximum', 'jnp.maximum', ({(148, 36, 148, 39): '(1.0)', (148, 41, 148, 60): '(step - warmup_steps)'}, {}), '(1.0, step - warmup_steps)', True, 'import jax.numpy as jnp\n'), ((150, 15, 150, 37), 'jax.numpy.sqrt', 'jnp.sqrt', ({(150, 24, 150, 36): 'warmup_steps'}, {}), '(warmup_steps)', True, 'import jax.numpy as jnp\n'), ((495, 23, 495, 54), 'numpy.argmax', 'np.argmax', ({(495, 33, 495, 53): 'program == eos_token'}, {}), '(program == eos_token)', True, 'import numpy as np\n'), ((789, 39, 789, 50), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((151, 24, 151, 55), 'jax.numpy.maximum', 'jnp.maximum', ({(151, 36, 151, 40): 'step', (151, 42, 151, 54): 'warmup_steps'}, {}), '(step, warmup_steps)', True, 'import jax.numpy as jnp\n'), ((727, 16, 727, 56), 'numpy.ceil', 'np.ceil', ({(727, 24, 727, 55): 'cur_pred_batch_size / n_devices'}, {}), '(cur_pred_batch_size / n_devices)', True, 'import numpy as np\n'), ((158, 40, 158, 74), 'jax.numpy.cos', 'jnp.cos', ({(158, 48, 158, 73): '(jnp.pi * (progress % 1.0))'}, {}), '(jnp.pi * (progress % 1.0))', True, 'import jax.numpy as jnp\n')]