jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def input_fn_builder(
    tfrecord_file, seq_length, is_training, batch_size, drop_remainder=True
):
    """Creates an `input_fn` closure to be passed to Estimator."""
    logging.info(
        "Creating input fun with batch_size: {} and drop remainder: {}".format(
            str(batch_size), str(drop_remainder)
        )
    )
    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "score": tf.FixedLenFeature([], tf.float32),
    }

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)
        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.cast(t, tf.int32)
            example[name] = t
        return example

    def input_fn(params):  # pylint: disable=unused-argument
        """Acutal data generator."""
        d = tf.data.TFRecordDataset(tfrecord_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=FLAGS.shuffle_buffer_size)
        d = d.map(lambda record: _decode_record(record, name_to_features))
        d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
        return d

    return input_fn
//  [python]
def get_masked_mention_word_scores(self, encoded_doc, span_starts, span_ends):
      num_words = util.shape(encoded_doc, 0) # T
      num_c = util.shape(span_starts, 0) # NC
      doc_range = tf.tile(tf.expand_dims(tf.range(0, num_words), 0), [num_c, 1]) # [K, T]
      mention_mask = tf.logical_and(doc_range >= tf.expand_dims(span_starts, 1), doc_range <= tf.expand_dims(span_ends, 1)) #[K, T]
      with tf.variable_scope("mention_word_attn", reuse=tf.AUTO_REUSE):
        word_attn = tf.squeeze(util.projection(encoded_doc, 1, initializer=tf.truncated_normal_initializer(stddev=0.02)), 1)
      mention_word_attn = tf.nn.softmax(tf.log(tf.to_float(mention_mask)) + tf.expand_dims(word_attn, 0))
      return mention_word_attn
//  [python]
def get_masked_mention_word_scores(self, encoded_doc, span_starts, span_ends):
      num_words = util.shape(encoded_doc, 0) # T
      num_c = util.shape(span_starts, 0) # NC
      doc_range = tf.tile(tf.expand_dims(tf.range(0, num_words), 0), [num_c, 1]) # [K, T]
      mention_mask = tf.logical_and(doc_range >= tf.expand_dims(span_starts, 1), doc_range <= tf.expand_dims(span_ends, 1)) #[K, T]
      with tf.variable_scope("mention_word_attn", reuse=tf.AUTO_REUSE):
        word_attn = tf.squeeze(util.projection(encoded_doc, 1, initializer=tf.truncated_normal_initializer(stddev=0.02)), 1)
      mention_word_attn = tf.nn.softmax(tf.log(tf.cast(mention_mask, tf.float32)) + tf.expand_dims(word_attn, 0))
      return mention_word_attn
//  [python]
def _interactive_input_tensor_to_features_dict(feature_map, hparams):
  """Convert the interactive input format (see above) to a dictionary.

  Args:
    feature_map: dict with inputs.
    hparams: model hyperparameters

  Returns:
    a features dictionary, as expected by the decoder.
  """
  inputs = tf.convert_to_tensor(feature_map["inputs"])
  input_is_image = False if len(inputs.get_shape()) < 3 else True

  x = inputs
  if input_is_image:
    x = tf.image.resize_images(x, [299, 299])
    x = tf.reshape(x, [1, 299, 299, -1])
    x = tf.to_int32(x)
  else:
    # Remove the batch dimension.
    num_samples = x[0]
    length = x[2]
    x = tf.slice(x, [3], tf.to_int32([length]))
    x = tf.reshape(x, [1, -1, 1, 1])
    # Transform into a batch of size num_samples to get that many random
    # decodes.
    x = tf.tile(x, tf.to_int32([num_samples, 1, 1, 1]))

  p_hparams = hparams.problem_hparams
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = (
      IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
  features["inputs"] = x
  # Save inputs to "partial_targets" when prepending inputs to targets. Also
  # keep "inputs" as some models crash if they don't exist.
  if getattr(hparams, "prepend_mode", "none") != "none":
    shape = tf.shape(x)
    partial_targets = tf.reshape(x, [shape[0], shape[1]])
    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
    features["partial_targets"] = partial_targets
  return features
//  [python]
def _interactive_input_tensor_to_features_dict(feature_map, hparams):
  """Convert the interactive input format (see above) to a dictionary.

  Args:
    feature_map: dict with inputs.
    hparams: model hyperparameters

  Returns:
    a features dictionary, as expected by the decoder.
  """
  inputs = tf.convert_to_tensor(feature_map["inputs"])
  input_is_image = False if len(inputs.get_shape()) < 3 else True

  x = inputs
  if input_is_image:
    x = tf.image.resize_images(x, [299, 299])
    x = tf.reshape(x, [1, 299, 299, -1])
    x = tf.cast(x, tf.int32)
  else:
    # Remove the batch dimension.
    num_samples = x[0]
    length = x[2]
    x = tf.slice(x, [3], tf.cast([length], tf.int32))
    x = tf.reshape(x, [1, -1, 1, 1])
    # Transform into a batch of size num_samples to get that many random
    # decodes.
    x = tf.tile(x, tf.cast([num_samples, 1, 1, 1], tf.int32))

  p_hparams = hparams.problem_hparams
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = (
      IMAGE_DECODE_LENGTH if input_is_image else inputs[1])
  features["inputs"] = x
  # Save inputs to "partial_targets" when prepending inputs to targets. Also
  # keep "inputs" as some models crash if they don't exist.
  if getattr(hparams, "prepend_mode", "none") != "none":
    shape = tf.shape(x)
    partial_targets = tf.reshape(x, [shape[0], shape[1]])
    partial_targets = tf.pad(partial_targets, [[0, 0], [0, 1]])
    features["partial_targets"] = partial_targets
  return features
//  [python]
def get_cosine_learning_rate_fn(
    training_steps, learning_rate, min_learning_rate_mult,
    constant_fraction, warmup_fraction):
  """Get a function that does cosine learning rate decay with warmup.

  The learning rate starts at zero, is "warmed up" linearly over
  `warmup_fraction * training_steps` iterations to achieve a final value of
  `learning_rate`. A constant learning rate of `learning_rate` is held up until
  `training_steps*constant_fraction` at which point a cosine decay is started
  to a final learning rate of `min_learning_rate_mult * learning_rate`.

  The cosine decay sets the learning rate using a monotomically decreasing
  section of the cosine function from 0 to pi/2. It has been proven to be useful
  in large large language modeling (gpt, megatron-lm) and image classification.
  See https://arxiv.org/abs/1608.03983 for more information on the cosine decay.


  Args:
    training_steps: number of training steps the schedule should be run for.
    learning_rate: base learning rate. This is the learning rate used just after
      warmup and where the decay starts from.
    min_learning_rate_mult: a multiplicative factor to control how low the
      learning rate should be decayed to.
    constant_fraction: the fraction of training steps number of steps to take
      before starting the decay. This includes the time spent warming up the
      learning rate.
    warmup_fraction: the fraction of training steps to use for a learning rate
      warmup.

  Returns:
    A function that takes as input a training iteration and returns the learning
    rate from the specified schedule.
  """

  def fn(global_step):
    """Returns a learning rate given the current training iteration."""

    float_training_steps = tf.to_float(training_steps)
    global_step = tf.to_float(global_step)

    # ensure we don't train longer than training steps
    global_step = tf.minimum(global_step, float_training_steps)

    constant_steps = float_training_steps * constant_fraction
    x = tf.maximum(tf.to_float(global_step), tf.to_float(constant_steps))

    min_learning_rate = min_learning_rate_mult * learning_rate

    if warmup_fraction:
      min_warmup_fraction = tf.minimum(warmup_fraction, constant_fraction)
      warmup_steps = float_training_steps * min_warmup_fraction
      is_warmup = tf.to_float(
          tf.greater(tf.to_float(warmup_steps), tf.to_float(global_step)))
      warmup_lr = (global_step / warmup_steps) * learning_rate
    else:
      warmup_lr = learning_rate
      is_warmup = 0.0

    step = x - constant_steps

    constant_and_decay = (learning_rate - min_learning_rate) * (
        tf.math.cos(step * np.pi /
                    (float_training_steps - constant_steps)) / 2.0 +
        0.5) + min_learning_rate

    new_learning_rate = constant_and_decay * (1.0 - is_warmup) + is_warmup * (
        warmup_lr)
    return new_learning_rate

  return fn
//  [python]
def get_cosine_learning_rate_fn(
    training_steps, learning_rate, min_learning_rate_mult,
    constant_fraction, warmup_fraction):
  """Get a function that does cosine learning rate decay with warmup.

  The learning rate starts at zero, is "warmed up" linearly over
  `warmup_fraction * training_steps` iterations to achieve a final value of
  `learning_rate`. A constant learning rate of `learning_rate` is held up until
  `training_steps*constant_fraction` at which point a cosine decay is started
  to a final learning rate of `min_learning_rate_mult * learning_rate`.

  The cosine decay sets the learning rate using a monotomically decreasing
  section of the cosine function from 0 to pi/2. It has been proven to be useful
  in large large language modeling (gpt, megatron-lm) and image classification.
  See https://arxiv.org/abs/1608.03983 for more information on the cosine decay.


  Args:
    training_steps: number of training steps the schedule should be run for.
    learning_rate: base learning rate. This is the learning rate used just after
      warmup and where the decay starts from.
    min_learning_rate_mult: a multiplicative factor to control how low the
      learning rate should be decayed to.
    constant_fraction: the fraction of training steps number of steps to take
      before starting the decay. This includes the time spent warming up the
      learning rate.
    warmup_fraction: the fraction of training steps to use for a learning rate
      warmup.

  Returns:
    A function that takes as input a training iteration and returns the learning
    rate from the specified schedule.
  """

  def fn(global_step):
    """Returns a learning rate given the current training iteration."""

    float_training_steps = tf.cast(training_steps, tf.float32)
    global_step = tf.cast(global_step, tf.float32)

    # ensure we don't train longer than training steps
    global_step = tf.minimum(global_step, float_training_steps)

    constant_steps = float_training_steps * constant_fraction
    x = tf.maximum(tf.cast(global_step, tf.float32), tf.cast(constant_steps, tf.float32))

    min_learning_rate = min_learning_rate_mult * learning_rate

    if warmup_fraction:
      min_warmup_fraction = tf.minimum(warmup_fraction, constant_fraction)
      warmup_steps = float_training_steps * min_warmup_fraction
      is_warmup = tf.cast(
          tf.greater(tf.cast(warmup_steps, tf.float32), tf.cast(global_step, tf.float32)), tf.float32)
      warmup_lr = (global_step / warmup_steps) * learning_rate
    else:
      warmup_lr = learning_rate
      is_warmup = 0.0

    step = x - constant_steps

    constant_and_decay = (learning_rate - min_learning_rate) * (
        tf.math.cos(step * np.pi /
                    (float_training_steps - constant_steps)) / 2.0 +
        0.5) + min_learning_rate

    new_learning_rate = constant_and_decay * (1.0 - is_warmup) + is_warmup * (
        warmup_lr)
    return new_learning_rate

  return fn
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, c):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k) # [k]
    antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0) # [k, k]
    antecedents_mask = antecedent_offsets >= 1 # [k, k]
    fast_antecedent_scores = self.get_fast_antecedent_scores(top_span_emb) # [k, k]
    fast_antecedent_scores += tf.log(tf.to_float(antecedents_mask)) # [k, k]
    if self.config['use_prior']:
      antecedent_distance_buckets = self.bucket_distance(antecedent_offsets) # [k, c]
      with tf.variable_scope("antecedent_distance", reuse=tf.AUTO_REUSE):
        distance_scores = util.projection(tf.nn.dropout(tf.get_variable("antecedent_distance_emb", [10, self.config["feature_size"]], initializer=tf.truncated_normal_initializer(stddev=0.02)), self.dropout), 1, initializer=tf.truncated_normal_initializer(stddev=0.02)) #[10, 1]
      antecedent_distance_scores = tf.gather(tf.squeeze(distance_scores, 1), antecedent_distance_buckets) # [k, c]
      fast_antecedent_scores += antecedent_distance_scores

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def coarse_to_fine_pruning(self, top_span_emb, c):
    k = util.shape(top_span_emb, 0)
    top_span_range = tf.range(k) # [k]
    antecedent_offsets = tf.expand_dims(top_span_range, 1) - tf.expand_dims(top_span_range, 0) # [k, k]
    antecedents_mask = antecedent_offsets >= 1 # [k, k]
    fast_antecedent_scores = self.get_fast_antecedent_scores(top_span_emb) # [k, k]
    fast_antecedent_scores += tf.log(tf.cast(antecedents_mask, tf.float32)) # [k, k]
    if self.config['use_prior']:
      antecedent_distance_buckets = self.bucket_distance(antecedent_offsets) # [k, c]
      with tf.variable_scope("antecedent_distance", reuse=tf.AUTO_REUSE):
        distance_scores = util.projection(tf.nn.dropout(tf.get_variable("antecedent_distance_emb", [10, self.config["feature_size"]], initializer=tf.truncated_normal_initializer(stddev=0.02)), self.dropout), 1, initializer=tf.truncated_normal_initializer(stddev=0.02)) #[10, 1]
      antecedent_distance_scores = tf.gather(tf.squeeze(distance_scores, 1), antecedent_distance_buckets) # [k, c]
      fast_antecedent_scores += antecedent_distance_scores

    _, top_antecedents = tf.nn.top_k(fast_antecedent_scores, c, sorted=False) # [k, c]
    top_antecedents_mask = util.batch_gather(antecedents_mask, top_antecedents) # [k, c]
    top_fast_antecedent_scores = util.batch_gather(fast_antecedent_scores, top_antecedents) # [k, c]
    top_antecedent_offsets = util.batch_gather(antecedent_offsets, top_antecedents) # [k, c]
    return top_antecedents, top_antecedents_mask, top_fast_antecedent_scores, top_antecedent_offsets
//  [python]
def test_function_integration(self):
        t_max = 1
        ode = NeuralODE(SineDumpingModel(), t=np.linspace(0, t_max, 1000))
        x0 = tf.to_float([0])
        xN_euler = ode.forward(x0)

        ode = NeuralODE(
            SineDumpingModel(), t=np.linspace(0, t_max, 100),
            solver=neural_ode.rk2_step
        )

        xN_rk2 = ode.forward(x0)

        ode = NeuralODE(
            SineDumpingModel(), t=np.linspace(0, t_max, 50),
            solver=neural_ode.rk4_step
        )

        xN_rk4 = ode.forward(x0)
        xN_exact = [np.log(2 - np.cos(t_max))]

        self.assertAllClose(xN_euler.numpy(), xN_exact, atol=1e-4)
        self.assertAllClose(xN_rk2.numpy(), xN_exact, atol=1e-4)
        self.assertAllClose(xN_rk4.numpy(), xN_exact, atol=1e-4)
//  [python]
def test_function_integration(self):
        t_max = 1
        ode = NeuralODE(SineDumpingModel(), t=np.linspace(0, t_max, 1000))
        x0 = tf.cast([0], tf.float32)
        xN_euler = ode.forward(x0)

        ode = NeuralODE(
            SineDumpingModel(), t=np.linspace(0, t_max, 100),
            solver=neural_ode.rk2_step
        )

        xN_rk2 = ode.forward(x0)

        ode = NeuralODE(
            SineDumpingModel(), t=np.linspace(0, t_max, 50),
            solver=neural_ode.rk4_step
        )

        xN_rk4 = ode.forward(x0)
        xN_exact = [np.log(2 - np.cos(t_max))]

        self.assertAllClose(xN_euler.numpy(), xN_exact, atol=1e-4)
        self.assertAllClose(xN_rk2.numpy(), xN_exact, atol=1e-4)
        self.assertAllClose(xN_rk4.numpy(), xN_exact, atol=1e-4)
//  [python]
def test_backprop(self):
        t_max = 1
        t_grid = np.linspace(0, t_max, 40)

        ode = NeuralODE(SineDumpingModel(), t=t_grid,
                        solver=neural_ode.rk4_step)
        x0 = tf.to_float([0])
        xN = ode.forward(x0)
        with tf.GradientTape() as g:
            g.watch(xN)
            loss = xN ** 2

        dLoss = g.gradient(loss, xN)
        x0_rec, dLdx0, dLdW = ode.backward(xN, dLoss)
        self.assertAllClose(x0_rec.numpy(), x0)
        self.assertEqual(dLdW, [])

        with tf.GradientTape() as g:
            g.watch(x0)
            xN = ode.forward(x0)
            loss = xN ** 2

        dLdx0_exact = g.gradient(loss, x0)
        self.assertAllClose(dLdx0_exact, dLdx0)
//  [python]
def test_backprop(self):
        t_max = 1
        t_grid = np.linspace(0, t_max, 40)

        ode = NeuralODE(SineDumpingModel(), t=t_grid,
                        solver=neural_ode.rk4_step)
        x0 = tf.cast([0], tf.float32)
        xN = ode.forward(x0)
        with tf.GradientTape() as g:
            g.watch(xN)
            loss = xN ** 2

        dLoss = g.gradient(loss, xN)
        x0_rec, dLdx0, dLdW = ode.backward(xN, dLoss)
        self.assertAllClose(x0_rec.numpy(), x0)
        self.assertEqual(dLdW, [])

        with tf.GradientTape() as g:
            g.watch(x0)
            xN = ode.forward(x0)
            loss = xN ** 2

        dLdx0_exact = g.gradient(loss, x0)
        self.assertAllClose(dLdx0_exact, dLdx0)
//  [python]
def test_multiple_inputs(self):
        t_max = 1
        t_grid = np.linspace(0, t_max, 100)
        ode = NeuralODE(DoubleSineDumpingModel(), t=t_grid)

        xy0 = tf.to_float([0.0, 0.0])
        xyN_euler = ode.forward(xy0)

        ode = NeuralODE(DoubleSineDumpingModel(), t=t_grid,
                        solver=neural_ode.rk2_step)

        xyN_rk2 = ode.forward(xy0)

        ode = NeuralODE(DoubleSineDumpingModel(), t=t_grid,
                        solver=neural_ode.rk4_step)

        xyN_rk4 = ode.forward(xy0)

        xN_exact = np.log(2 - np.cos(t_max))
        yN_exact = np.log((np.sin(t_max) + 2) / 2)
        xyN_exact = tf.to_float([xN_exact, yN_exact])

        self.assertAllClose(xyN_euler, xyN_exact, atol=1e-2)
        self.assertAllClose(xyN_rk2, xyN_exact, atol=1e-5)
        self.assertAllClose(xyN_rk4, xyN_exact, atol=1e-6)
//  [python]
def test_multiple_inputs(self):
        t_max = 1
        t_grid = np.linspace(0, t_max, 100)
        ode = NeuralODE(DoubleSineDumpingModel(), t=t_grid)

        xy0 = tf.cast([0.0, 0.0], tf.float32)
        xyN_euler = ode.forward(xy0)

        ode = NeuralODE(DoubleSineDumpingModel(), t=t_grid,
                        solver=neural_ode.rk2_step)

        xyN_rk2 = ode.forward(xy0)

        ode = NeuralODE(DoubleSineDumpingModel(), t=t_grid,
                        solver=neural_ode.rk4_step)

        xyN_rk4 = ode.forward(xy0)

        xN_exact = np.log(2 - np.cos(t_max))
        yN_exact = np.log((np.sin(t_max) + 2) / 2)
        xyN_exact = tf.cast([xN_exact, yN_exact], tf.float32)

        self.assertAllClose(xyN_euler, xyN_exact, atol=1e-2)
        self.assertAllClose(xyN_rk2, xyN_exact, atol=1e-5)
        self.assertAllClose(xyN_rk4, xyN_exact, atol=1e-6)
//  [python]
def test_affine_exact_solution(self):
        tf.set_random_seed(1234)
        t_grid = np.linspace(0, 1.0, 40)

        W = tf.to_float(np.random.randn(8, 8))
        x0 = tf.random_normal(shape=[3, 8])

        class Affine(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                y = tf.matmul(x, W)
                return y

        model = Affine()
        ode = NeuralODE(model, t=t_grid)
        x1 = ode.forward(x0)
        x1_exact = tf.matmul(x0, tf.linalg.expm(W))
        self.assertAllClose(x1, x1_exact, atol=1e-5)
//  [python]
def test_affine_exact_solution(self):
        tf.random.set_seed(1234)
        t_grid = np.linspace(0, 1.0, 40)

        W = tf.cast(np.random.randn(8, 8), dtype=tf.float32)
        x0 = tf.random.normal(shape=[3, 8])

        class Affine(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                y = tf.matmul(x, W)
                return y

        model = Affine()
        ode = NeuralODE(model, t=t_grid)
        x1 = ode.forward(x0)
        x1_exact = tf.matmul(x0, tf.linalg.expm(W))
        self.assertAllClose(x1, x1_exact, atol=1e-5)
//  [python]
def test_jvp_on_simple_problem(self):
        tf.set_random_seed(1234)
        W = np.random.randn(8, 8)
        some_vec = np.random.randn(3, 8)
        x0 = tf.random_normal(shape=[3, 8])

        class Affine(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                y = tf.matmul(x, tf.to_float(W))
                return y

        model = Affine()
        with tf.GradientTape() as g:
            g.watch(x0)
            x1 = model([0.0, x0])

        dLdx_jvp = g.gradient(x1, x0, output_gradients=tf.to_float(some_vec))

        with tf.GradientTape() as g:
            g.watch(x1)
            x2 = model([0.0, x1])
            loss = tf.reduce_sum(x2 * tf.to_float(some_vec))

        dLdx_loss = g.gradient(loss, x1)

        self.assertAllClose(dLdx_jvp, dLdx_loss)
//  [python]
def test_jvp_on_simple_problem(self):
        tf.set_random_seed(1234)
        W = np.random.randn(8, 8)
        some_vec = np.random.randn(3, 8)
        x0 = tf.random_normal(shape=[3, 8])

        class Affine(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                y = tf.matmul(x, tf.cast(W, tf.float32))
                return y

        model = Affine()
        with tf.GradientTape() as g:
            g.watch(x0)
            x1 = model([0.0, x0])

        dLdx_jvp = g.gradient(x1, x0, output_gradients=tf.cast(some_vec, tf.float32))

        with tf.GradientTape() as g:
            g.watch(x1)
            x2 = model([0.0, x1])
            loss = tf.reduce_sum(x2 * tf.cast(some_vec, tf.float32))

        dLdx_loss = g.gradient(loss, x1)

        self.assertAllClose(dLdx_jvp, dLdx_loss)
//  [python]
def test_determinant_estimation(self):
        """Similar to previous one"""
        tf.set_random_seed(1234)
        t_grid = np.linspace(0, 1.0, 5)

        W = tf.to_float(np.random.randn(8, 8))
        logdet0 = tf.zeros(shape=[1])

        class DeterminantEstimation(tf.keras.Model):
            def call(self, inputs, **kwargs):
                trace_dens = - tf.trace(W)
                return trace_dens

        ode = NeuralODE(DeterminantEstimation(), t=t_grid)
        logdet = ode.forward(logdet0)

        W_exact = tf.linalg.expm(W)
        logdet_exact = - tf.log(tf.linalg.det(W_exact))
        self.assertAllClose([logdet_exact], logdet)
//  [python]
def test_determinant_estimation(self):
        """Similar to previous one"""
        tf.set_random_seed(1234)
        t_grid = np.linspace(0, 1.0, 5)

        W = tf.cast(np.random.randn(8, 8), dtype=tf.float32)
        logdet0 = tf.zeros(shape=[1])

        class DeterminantEstimation(tf.keras.Model):
            def call(self, inputs, **kwargs):
                trace_dens = - tf.trace(W)
                return trace_dens

        ode = NeuralODE(DeterminantEstimation(), t=t_grid)
        logdet = ode.forward(logdet0)

        W_exact = tf.linalg.expm(W)
        logdet_exact = - tf.log(tf.linalg.det(W_exact))
        self.assertAllClose([logdet_exact], logdet)
//  [python]
def make_2d_timing_signal(height, width, dim):
    position = tf.to_float(tf.range(width))[None, :]  # [1, width]
    signal = tf.tile(make_sinusoid_signal(position, dim // 2), [height, 1, 1])
    position = tf.to_float(tf.range(height))[:, None]  # [height, 1]
    signal = tf.concat([signal, tf.tile(make_sinusoid_signal(position, dim // 2), [1, width, 1])], axis=-1)
    return signal
//  [python]
def make_2d_timing_signal(height, width, dim):
    position = tf.cast(tf.range(width), tf.float32)[None, :]  # [1, width]
    signal = tf.tile(make_sinusoid_signal(position, dim // 2), [height, 1, 1])
    position = tf.cast(tf.range(height), tf.float32)[:, None]  # [height, 1]
    signal = tf.concat([signal, tf.tile(make_sinusoid_signal(position, dim // 2), [1, width, 1])], axis=-1)
    return signal
//  [python]
def test_affine_with_density_exact_solution(self):
        tf.set_random_seed(1234)
        t_grid = np.linspace(0, 1.0, 50)

        W = tf.to_float(np.random.randn(8, 8))
        x0 = tf.random_normal(shape=[3, 8])
        logdet0 = tf.zeros(shape=[3, 1])
        h0 = tf.concat([x0, logdet0], axis=1)

        class Affine(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                x = x[:, :8]
                y = tf.matmul(x, W)
                trace_dens = - tf.reshape(tf.trace(W), [1, 1])
                trace_dens = tf.tile(trace_dens, [3, 1])
                return tf.concat([y, trace_dens], axis=1)

        model = Affine()
        ode = NeuralODE(model, t=t_grid)
        h1 = ode.forward(h0)
        x1, logdet = h1[:, :8], h1[:, 8]

        W_exact = tf.linalg.expm(W)
        x1_exact = tf.matmul(x0[:, :8], W_exact)
        self.assertAllClose(x1, x1_exact, atol=1e-5)
        logdet_exact = - tf.log(tf.linalg.det(W_exact))
        self.assertAllClose([logdet_exact] * 3, logdet)
//  [python]
def test_affine_with_density_exact_solution(self):
        tf.random.set_seed(1234)
        t_grid = np.linspace(0, 1.0, 50)

        W = tf.cast(np.random.randn(8, 8), dtype=tf.float32)
        x0 = tf.random.normal(shape=[3, 8])
        logdet0 = tf.zeros(shape=[3, 1])
        h0 = tf.concat([x0, logdet0], axis=1)

        class Affine(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                x = x[:, :8]
                y = tf.matmul(x, W)
                trace_dens = - tf.reshape(tf.trace(W), [1, 1])
                trace_dens = tf.tile(trace_dens, [3, 1])
                return tf.concat([y, trace_dens], axis=1)

        model = Affine()
        ode = NeuralODE(model, t=t_grid)
        h1 = ode.forward(h0)
        x1, logdet = h1[:, :8], h1[:, 8]

        W_exact = tf.linalg.expm(W)
        x1_exact = tf.matmul(x0[:, :8], W_exact)
        self.assertAllClose(x1, x1_exact, atol=1e-5)
        logdet_exact = - tf.math.log(tf.linalg.det(W_exact))
        self.assertAllClose([logdet_exact] * 3, logdet)
//  [python]
def embedding_to_padding(emb):
    """Calculates the padding mask based on which embeddings are all zero.

    We have hacked symbol_modality to return all-zero embeddings for padding.

    Parameters
    ----------
    emb:
        A Tensor with shape [..., depth].

    Returns
    -------
    y
        A float Tensor with shape [...]. Each element is 1 if its corresponding
        embedding vector is all zero, and is 0 otherwise.
    """
    emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1)
    return tf.compat.v1.to_float(tf.equal(emb_sum, 0.0))
//  [python]
def embedding_to_padding(emb):
    """Calculates the padding mask based on which embeddings are all zero.

    We have hacked symbol_modality to return all-zero embeddings for padding.

    Parameters
    ----------
    emb:
        A Tensor with shape [..., depth].

    Returns
    -------
    y
        A float Tensor with shape [...]. Each element is 1 if its corresponding
        embedding vector is all zero, and is 0 otherwise.
    """
    emb_sum = tf.reduce_sum(tf.abs(emb), axis=-1)
    return tf.cast(tf.equal(emb_sum, 0.0), tf.float32)
//  [python]
def make_sinusoid_signal(position, num_channels, min_timescale=1.0, max_timescale=1e4):
    """
    Generate sinusoid timing signal
    :param position: tensor of position values
    :return: a tensor of the shape tf.shape(positon) + [num_channels]
    """
    num_timescales = num_channels // 2
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (tf.to_float(num_timescales) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.to_float(tf.range(num_timescales)) * -log_timescale_increment)

    # scaled_time: tf.shape(positon) + [num_timescales]
    scaled_time = tf.expand_dims(position, -1) * tf.reshape(inv_timescales, [1] * position.shape.ndims + [-1])
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=-1)
    signal = tf.pad(signal, [[0, 0]] * position.shape.ndims + [[0, tf.mod(num_channels, 2)]])
    return signal
//  [python]
def make_sinusoid_signal(position, num_channels, min_timescale=1.0, max_timescale=1e4):
    """
    Generate sinusoid timing signal
    :param position: tensor of position values
    :return: a tensor of the shape tf.shape(positon) + [num_channels]
    """
    num_timescales = num_channels // 2
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)

    # scaled_time: tf.shape(positon) + [num_timescales]
    scaled_time = tf.expand_dims(position, -1) * tf.reshape(inv_timescales, [1] * position.shape.ndims + [-1])
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=-1)
    signal = tf.pad(signal, [[0, 0]] * position.shape.ndims + [[0, tf.mod(num_channels, 2)]])
    return signal
//  [python]
def test_grad_no_grad(self):
        class CosGradModel(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                with tf.GradientTape() as g:
                    g.watch(x)
                    y = tf.sin(t * x)
                    dy = g.gradient(y, x)

                return dy

        class CosModel(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                y = t * tf.cos(t * x)
                return y

        t_vec = np.linspace(0, 1.0, 40)
        ode = NeuralODE(CosModel(), t=t_vec, solver=neural_ode.rk4_step)
        x0 = tf.to_float([1.0])

        cos_xN = ode.forward(x0)
        x0_rec, dLdx0, _ = ode.backward(cos_xN, cos_xN)

        ode = NeuralODE(CosGradModel(), t=t_vec, solver=neural_ode.rk4_step)
        cos_grad_xN = ode.forward(x0)
        x0_grad_rec, dLdx0_grad, _ = ode.backward(cos_grad_xN, cos_grad_xN)

        self.assertAllClose(cos_grad_xN, cos_xN)
        self.assertAllClose(x0_grad_rec, x0_rec)
        self.assertAllClose(x0, x0_rec)
        self.assertAllClose(dLdx0, dLdx0_grad)

        ode = NeuralODE(CosModel(), t=t_vec, solver=neural_ode.rk4_step)
        with tf.GradientTape() as g:
            g.watch(x0)
            cos_xN = ode.forward(x0)
            loss = 0.5 * cos_xN ** 2

        dLdx0_exact = g.gradient(loss, [x0])

        self.assertAllClose(dLdx0_exact, dLdx0_grad)

        # build static graph
        ode = NeuralODE(CosGradModel(), t=t_vec, solver=neural_ode.rk4_step)
        ode = neural_ode.defun_neural_ode(ode)
        cos_grad_xN = ode.forward(x0)
        x0_grad_rec, dLdx0_grad, _ = ode.backward(cos_grad_xN, cos_grad_xN)

        self.assertAllClose(dLdx0_grad, dLdx0_exact)
        self.assertAllClose(x0_grad_rec, x0)
//  [python]
def test_grad_no_grad(self):
        class CosGradModel(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                with tf.GradientTape() as g:
                    g.watch(x)
                    y = tf.sin(t * x)
                    dy = g.gradient(y, x)

                return dy

        class CosModel(tf.keras.Model):
            def call(self, inputs, **kwargs):
                t, x = inputs
                y = t * tf.cos(t * x)
                return y

        t_vec = np.linspace(0, 1.0, 40)
        ode = NeuralODE(CosModel(), t=t_vec, solver=neural_ode.rk4_step)
        x0 = tf.cast([1.0], dtype=tf.float32)

        cos_xN = ode.forward(x0)
        x0_rec, dLdx0, _ = ode.backward(cos_xN, cos_xN)

        ode = NeuralODE(CosGradModel(), t=t_vec, solver=neural_ode.rk4_step)
        cos_grad_xN = ode.forward(x0)
        x0_grad_rec, dLdx0_grad, _ = ode.backward(cos_grad_xN, cos_grad_xN)

        self.assertAllClose(cos_grad_xN, cos_xN)
        self.assertAllClose(x0_grad_rec, x0_rec)
        self.assertAllClose(x0, x0_rec)
        self.assertAllClose(dLdx0, dLdx0_grad)

        ode = NeuralODE(CosModel(), t=t_vec, solver=neural_ode.rk4_step)
        with tf.GradientTape() as g:
            g.watch(x0)
            cos_xN = ode.forward(x0)
            loss = 0.5 * cos_xN ** 2

        dLdx0_exact = g.gradient(loss, [x0])

        self.assertAllClose(dLdx0_exact, dLdx0_grad)

        # build static graph
        ode = NeuralODE(CosGradModel(), t=t_vec, solver=neural_ode.rk4_step)
        ode = neural_ode.defun_neural_ode(ode)
        cos_grad_xN = ode.forward(x0)
        x0_grad_rec, dLdx0_grad, _ = ode.backward(cos_grad_xN, cos_grad_xN)

        self.assertAllClose(dLdx0_grad, dLdx0_exact)
        self.assertAllClose(x0_grad_rec, x0)
//  [python]
def _create_loss(features, params, logits, labels):
    # follows the logic of tf.losses.sigmoid_cross_entropy
    labels = tf.to_float(labels)

    # unweighted loss
    unweighted_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)

    # calc weighted loss (assemble cost-sensitive learning)
    weight_column = params.weight_column
    if weight_column:
        weights = _get_weights_and_check_match_logits(features=features, weight_column=weight_column,
                                                      logits=logits)  # get weight col tensor
    else:
        weights = 1.0
    weighted_loss = tf.losses.compute_weighted_loss(unweighted_loss, weights=weights,
                                                    reduction=tf.losses.Reduction.MEAN)
    return weighted_loss, unweighted_loss, weights, labels
//  [python]
def _create_loss(features, params, logits, labels):
    labels = tf.cast(labels, tf.float32)

    unweighted_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)

    weight_column = params.weight_column
    if weight_column:
        weights = _get_weights_and_check_match_logits(features=features, weight_column=weight_column,
                                                      logits=logits)
    else:
        weights = 1.0
    weighted_loss = tf.losses.compute_weighted_loss(unweighted_loss, weights=weights,
                                                    reduction=tf.losses.Reduction.MEAN)
    return weighted_loss, unweighted_loss, weights, labels
//  [python]
def make_transformer_timing_signal(inp, min_timescale=1.0, max_timescale=1e4, offset=0, inp_reverse=None):
    """
    Generate timing signal like described in transformer article
    :param inp: (batch_size * ninp * hid_dim)
    :param offset: add this number to all character positions.
        if offset == 'random', picks this number uniformly from [-32000,32000] integers
    :type offset: number, tf.Tensor or 'random'
    """
    with tf.name_scope("timing_signal"):
        ninp = tf.shape(inp)[1]
        hid_size = tf.shape(inp)[2]

        position = tf.to_float(tf.range(ninp))[None, :]

        if offset == 'random':
            BIG_LEN = 32000
            offset = tf.random_uniform(tf.shape(position), minval=-BIG_LEN, maxval=BIG_LEN, dtype=tf.int32)

        # force broadcasting over batch axis
        if isinstance(offset * 1, tf.Tensor):  # multiply by 1 to also select variables, special generators, etc.
            assert offset.shape.ndims in (0, 1, 2)
            new_shape = [tf.shape(offset)[i] for i in range(offset.shape.ndims)]
            new_shape += [1] * (2 - len(new_shape))
            offset = tf.reshape(offset, new_shape)

        position += tf.to_float(offset)

        if inp_reverse is not None:
            position = tf.multiply(
                position,
                tf.where(
                    tf.equal(inp_reverse, 0),
                    tf.ones_like(inp_reverse, dtype=tf.float32),
                    -1.0 * tf.ones_like(inp_reverse, dtype=tf.float32)
                )[:, None, None]  # (batch_size * ninp * dim)
            )

        return make_sinusoid_signal(position, hid_size, min_timescale=min_timescale, max_timescale=max_timescale)
//  [python]
def make_transformer_timing_signal(inp, min_timescale=1.0, max_timescale=1e4, offset=0, inp_reverse=None):
    """
    Generate timing signal like described in transformer article
    :param inp: (batch_size * ninp * hid_dim)
    :param offset: add this number to all character positions.
        if offset == 'random', picks this number uniformly from [-32000,32000] integers
    :type offset: number, tf.Tensor or 'random'
    """
    with tf.name_scope("timing_signal"):
        ninp = tf.shape(inp)[1]
        hid_size = tf.shape(inp)[2]

        position = tf.cast(tf.range(ninp), tf.float32)[None, :]

        if offset == 'random':
            BIG_LEN = 32000
            offset = tf.random_uniform(tf.shape(position), minval=-BIG_LEN, maxval=BIG_LEN, dtype=tf.int32)

        # force broadcasting over batch axis
        if isinstance(offset * 1, tf.Tensor):  # multiply by 1 to also select variables, special generators, etc.
            assert offset.shape.ndims in (0, 1, 2)
            new_shape = [tf.shape(offset)[i] for i in range(offset.shape.ndims)]
            new_shape += [1] * (2 - len(new_shape))
            offset = tf.reshape(offset, new_shape)

        position += tf.cast(offset, tf.float32)

        if inp_reverse is not None:
            position = tf.multiply(
                position,
                tf.where(
                    tf.equal(inp_reverse, 0),
                    tf.ones_like(inp_reverse, dtype=tf.float32),
                    -1.0 * tf.ones_like(inp_reverse, dtype=tf.float32)
                )[:, None, None]  # (batch_size * ninp * dim)
            )

        return make_sinusoid_signal(position, hid_size, min_timescale=min_timescale, max_timescale=max_timescale)
//  [python]
def set_keep_prob(self, keep_prob=None):
        """This method only supports Eager Execution"""
        if keep_prob is not None:
            self.keep_prob = keep_prob
        w, h = tf.to_float(self.w), tf.to_float(self.h)
        self.gamma = (1. - self.keep_prob) * (w * h) / (self.block_size ** 2) / \
                     ((w - self.block_size + 1) * (h - self.block_size + 1))
//  [python]
def set_keep_prob(self, keep_prob=None):
        """This method only supports Eager Execution"""
        if keep_prob is not None:
            self.keep_prob = keep_prob
        w, h = tf.cast(self.w, tf.float32), tf.cast(self.h, tf.float32)
        self.gamma = (1. - self.keep_prob) * (w * h) / (self.block_size ** 2) / \
                     ((w - self.block_size + 1) * (h - self.block_size + 1))
//  [python]
def call(self, inputs, training=None, **kwargs):
        def drop():
            inp, mask = inputs
            output = inp * mask
            output = tf.cond(self.scale,
                             true_fn=lambda: output * tf.to_float(tf.size(mask)) / tf.reduce_sum(mask),
                             false_fn=lambda: output)
            return output

        if training is None:
            training = K.learning_phase()
        output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)),
                         true_fn=lambda: inputs,
                         false_fn=drop)
        return output
//  [python]
def call(self, inputs, training=None, **kwargs):
        def drop():
            inp, mask = inputs
            output = inp * mask
            output = tf.cond(self.scale,
                             true_fn=lambda: output * tf.cast(tf.size(mask), tf.float32) / tf.reduce_sum(mask),
                             false_fn=lambda: output)
            return output

        if training is None:
            training = K.learning_phase()
        output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)),
                         true_fn=lambda: inputs,
                         false_fn=drop)
        return output
//  [python]
def call(self, inputs, training=None, **kwargs):
        def drop():
            mask = self._create_mask(tf.shape(inputs))
            output = inputs * mask
            output = tf.cond(self.scale,
                             true_fn=lambda: output * tf.to_float(tf.size(mask)) / tf.reduce_sum(mask),
                             false_fn=lambda: output)
            return output

        if training is None:
            training = K.learning_phase()
        output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)),
                         true_fn=lambda: inputs,
                         false_fn=drop)
        return output
//  [python]
def call(self, inputs, training=None, **kwargs):
        def drop():
            mask = self._create_mask(tf.shape(inputs))
            output = inputs * mask
            output = tf.cond(self.scale,
                             true_fn=lambda: output * tf.cast(tf.size(mask), tf.float32) / tf.reduce_sum(mask),
                             false_fn=lambda: output)
            return output

        if training is None:
            training = K.learning_phase()
        output = tf.cond(tf.logical_or(tf.logical_not(training), tf.equal(self.keep_prob, 1.0)),
                         true_fn=lambda: inputs,
                         false_fn=drop)
        return output
//  [python]
def set_keep_prob(self, keep_prob=None):
        """This method only supports Eager Execution"""
        if keep_prob is not None:
            self.keep_prob = keep_prob
        d, w, h = tf.to_float(self.d), tf.to_float(self.w), tf.to_float(self.h)
        self.gamma = ((1. - self.keep_prob) * (d * w * h) / (self.block_size ** 3) /
                      ((d - self.block_size + 1) * (w - self.block_size + 1) * (h - self.block_size + 1)))
//  [python]
def set_keep_prob(self, keep_prob=None):
        """This method only supports Eager Execution"""
        if keep_prob is not None:
            self.keep_prob = keep_prob
        d, w, h = tf.cast(self.d, tf.float32), tf.cast(self.w, tf.float32), tf.cast(self.h, tf.float32)
        self.gamma = ((1. - self.keep_prob) * (d * w * h) / (self.block_size ** 3) /
                      ((d - self.block_size + 1) * (w - self.block_size + 1) * (h - self.block_size + 1)))
//  [python]
def _summarize_input(self, groundtruth_boxes_list, match_list):
    """Creates tensorflow summaries for the input boxes and anchors.

    This function creates four summaries corresponding to the average
    number (over images in a batch) of (1) groundtruth boxes, (2) anchors
    marked as positive, (3) anchors marked as negative, and (4) anchors marked
    as ignored.

    Args:
      groundtruth_boxes_list: a list of 2-D tensors of shape [num_boxes, 4]
        containing corners of the groundtruth boxes.
      match_list: a list of matcher.Match objects encoding the match between
        anchors and groundtruth boxes for each image of the batch,
        with rows of the Match objects corresponding to groundtruth boxes
        and columns corresponding to anchors.
    """
    num_boxes_per_image = tf.stack(
        [tf.shape(x)[0] for x in groundtruth_boxes_list])
    pos_anchors_per_image = tf.stack(
        [match.num_matched_columns() for match in match_list])
    neg_anchors_per_image = tf.stack(
        [match.num_unmatched_columns() for match in match_list])
    ignored_anchors_per_image = tf.stack(
        [match.num_ignored_columns() for match in match_list])
    tf.summary.scalar('Input/AvgNumGroundtruthBoxesPerImage',
                      tf.reduce_mean(tf.to_float(num_boxes_per_image)))
    tf.summary.scalar('Input/AvgNumPositiveAnchorsPerImage',
                      tf.reduce_mean(tf.to_float(pos_anchors_per_image)))
    tf.summary.scalar('Input/AvgNumNegativeAnchorsPerImage',
                      tf.reduce_mean(tf.to_float(neg_anchors_per_image)))
    tf.summary.scalar('Input/AvgNumIgnoredAnchorsPerImage',
                      tf.reduce_mean(tf.to_float(ignored_anchors_per_image)))
//  [python]
def _summarize_input(self, groundtruth_boxes_list, match_list):
    """Creates tensorflow summaries for the input boxes and anchors.

    This function creates four summaries corresponding to the average
    number (over images in a batch) of (1) groundtruth boxes, (2) anchors
    marked as positive, (3) anchors marked as negative, and (4) anchors marked
    as ignored.

    Args:
      groundtruth_boxes_list: a list of 2-D tensors of shape [num_boxes, 4]
        containing corners of the groundtruth boxes.
      match_list: a list of matcher.Match objects encoding the match between
        anchors and groundtruth boxes for each image of the batch,
        with rows of the Match objects corresponding to groundtruth boxes
        and columns corresponding to anchors.
    """
    num_boxes_per_image = tf.stack(
        [tf.shape(x)[0] for x in groundtruth_boxes_list])
    pos_anchors_per_image = tf.stack(
        [match.num_matched_columns() for match in match_list])
    neg_anchors_per_image = tf.stack(
        [match.num_unmatched_columns() for match in match_list])
    ignored_anchors_per_image = tf.stack(
        [match.num_ignored_columns() for match in match_list])
    tf.summary.scalar('Input/AvgNumGroundtruthBoxesPerImage',
                      tf.reduce_mean(tf.cast(num_boxes_per_image, tf.float32)))
    tf.summary.scalar('Input/AvgNumPositiveAnchorsPerImage',
                      tf.reduce_mean(tf.cast(pos_anchors_per_image, tf.float32)))
    tf.summary.scalar('Input/AvgNumNegativeAnchorsPerImage',
                      tf.reduce_mean(tf.cast(neg_anchors_per_image, tf.float32)))
    tf.summary.scalar('Input/AvgNumIgnoredAnchorsPerImage',
                      tf.reduce_mean(tf.cast(ignored_anchors_per_image, tf.float32)))