jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def tower_loss(self, x, y_, z_):
        y_conv, z_conv = self.construct_net(x)

        # Cast the nn result back to fp32 to avoid loss overflow/underflow
        if self.model_dtype != tf.float32:
            y_conv = tf.cast(y_conv, tf.float32)
            z_conv = tf.cast(z_conv, tf.float32)

        # Calculate loss on policy head
        cross_entropy = \
            tf.nn.softmax_cross_entropy_with_logits(labels=y_,
                                                    logits=y_conv)
        policy_loss = tf.reduce_mean(cross_entropy)

        # Loss on value head
        mse_loss = \
            tf.reduce_mean(tf.squared_difference(z_, z_conv))

        # Regularizer
        reg_variables = tf.get_collection(tf.GraphKeys.WEIGHTS)
        reg_term = self.l2_scale * tf.add_n(
            [tf.cast(tf.nn.l2_loss(v), tf.float32) for v in reg_variables])

        # For training from a (smaller) dataset of strong players, you will
        # want to reduce the factor in front of self.mse_loss here.
        loss = 1.0 * policy_loss + 1.0 * mse_loss + reg_term

        return loss, policy_loss, mse_loss, reg_term, y_conv
//  [python]
def init(self, batch_size, macrobatch=1, gpus_num=None, logbase='leelalogs'):
        self.batch_size = batch_size
        self.macrobatch = macrobatch
        self.logbase = logbase
        # Input batch placeholders
        self.planes = tf.placeholder(tf.string, name='in_planes')
        self.probs = tf.placeholder(tf.string, name='in_probs')
        self.winner = tf.placeholder(tf.string, name='in_winner')

        # Mini-batches come as raw packed strings. Decode
        # into tensors to feed into network.
        planes = tf.decode_raw(self.planes, tf.uint8)
        probs = tf.decode_raw(self.probs, tf.float32)
        winner = tf.decode_raw(self.winner, tf.float32)

        planes = tf.to_int32(planes)

        planes = tf.reshape(planes, (batch_size, 18, 19*19))
        probs = tf.reshape(probs, (batch_size, 19*19 + 1))
        winner = tf.reshape(winner, (batch_size, 1))

        if gpus_num is None:
            gpus_num = self.gpus_num
        self.init_net(planes, probs, winner, gpus_num)
//  [python]
def init(self, batch_size, macrobatch=1, gpus_num=None, logbase='leelalogs'):
        self.batch_size = batch_size
        self.macrobatch = macrobatch
        self.logbase = logbase
        # Input batch placeholders
        self.planes = tf.placeholder(tf.string, name='in_planes')
        self.probs = tf.placeholder(tf.string, name='in_probs')
        self.winner = tf.placeholder(tf.string, name='in_winner')

        # Mini-batches come as raw packed strings. Decode
        # into tensors to feed into network.
        planes = tf.decode_raw(self.planes, tf.uint8)
        probs = tf.decode_raw(self.probs, tf.float32)
        winner = tf.decode_raw(self.winner, tf.float32)

        planes = tf.cast(planes, self.model_dtype)

        planes = tf.reshape(planes, (batch_size, 18, 19*19))
        probs = tf.reshape(probs, (batch_size, 19*19 + 1))
        winner = tf.reshape(winner, (batch_size, 1))

        if gpus_num is None:
            gpus_num = self.gpus_num
        self.init_net(planes, probs, winner, gpus_num)
//  [python]
def infer_step(self, batch_x):
        if batch_x.dtype != self.precision:
            batch_x = tf.to_float(batch_x)
        output = self.graph_func(batch_x)
        return itemgetter(self.output_tensor_name)(output)
//  [python]
def infer_step(self, batch_x):
        if batch_x.dtype != self.precision:
            batch_x = tf.cast(batch_x, self.precision)
        output = self.graph_func(batch_x)
        return itemgetter(self.output_tensor_name)(output)
//  [python]
def _generalized_inverse_pth_root(
      self,
      input_t: TfValue,
      exponent: float,
      epsilon: float = 1e-12) -> Tuple[float, float]:
    """Compute inverse of a square matrix."""
    input_t_f64 = t.cast(input_t, t.float64)
    epsilon_f64 = t.cast(epsilon, t.float64)
    s, u, v = tf.linalg.svd(
        input_t_f64 +
        tf.eye(tf.shape(input_t_f64)[0], dtype=tf.float64) * epsilon_f64,
        full_matrices=True)
    inv_s = tf.reshape(
        tf.pow(tf.maximum(s, epsilon_f64), t.cast(exponent, t.float64)),
        [1, -1])
    val = tf.matmul(u * inv_s, v, adjoint_b=True)
    return t.cast(val, t.float32), tf.reduce_max(tf.abs(u - v))
//  [python]
def _generalized_inverse_pth_root(
      self,
      input_t: TfValue,
      exponent: float,
      epsilon: float = 1e-12) -> Tuple[float, float]:
    """Compute inverse of a square matrix."""
    input_t_f64 = tf.cast(input_t, tf.float64)
    epsilon_f64 = tf.cast(epsilon, tf.float64)
    s, u, v = tf.linalg.svd(
        input_t_f64 +
        tf.eye(tf.shape(input_t_f64)[0], dtype=tf.float64) * epsilon_f64,
        full_matrices=True)
    inv_s = tf.reshape(
        tf.pow(tf.maximum(s, epsilon_f64), tf.cast(exponent, tf.float64)),
        [1, -1])
    val = tf.matmul(u * inv_s, v, adjoint_b=True)
    return tf.cast(val, tf.float32), tf.reduce_max(tf.abs(u - v))
//  [python]
def padded_one_hot_encoding(indices, depth, left_pad):
  """Returns a zero padded one-hot tensor.

  This function converts a sparse representation of indices (e.g., [4]) to a
  zero padded one-hot representation (e.g., [0, 0, 0, 0, 1] with depth = 4 and
  left_pad = 1). If `indices` is empty, the result will simply be a tensor of
  shape (0, depth + left_pad). If depth = 0, then this function just returns
  `None`.

  Args:
    indices: an integer tensor of shape [num_indices].
    depth: depth for the one-hot tensor (integer).
    left_pad: number of zeros to left pad the one-hot tensor with (integer).

  Returns:
    padded_onehot: a tensor with shape (num_indices, depth + left_pad). Returns
      `None` if the depth is zero.

  Raises:
    ValueError: if `indices` does not have rank 1 or if `left_pad` or `depth are
      either negative or non-integers.

  TODO: add runtime checks for depth and indices.
  """
  if depth < 0 or not isinstance(depth, (int, long) if six.PY2 else int):
    raise ValueError('`depth` must be a non-negative integer.')
  if left_pad < 0 or not isinstance(left_pad, (int, long) if six.PY2 else int):
    raise ValueError('`left_pad` must be a non-negative integer.')
  if depth == 0:
    return None
  if len(indices.get_shape().as_list()) != 1:
    raise ValueError('`indices` must have rank 1')

  def one_hot_and_pad():
    one_hot = tf.to_float(tf.one_hot(tf.to_int64(indices), depth,
                                 on_value=1, off_value=0))
    return tf.pad(one_hot, [[0, 0], [left_pad, 0]], mode='CONSTANT')
  result = tf.cond(tf.greater(tf.size(indices), 0), one_hot_and_pad,
                   lambda: tf.zeros((depth + left_pad, 0)))
  return tf.reshape(result, [-1, depth + left_pad])
//  [python]
def padded_one_hot_encoding(indices, depth, left_pad):
  """Returns a zero padded one-hot tensor.

  This function converts a sparse representation of indices (e.g., [4]) to a
  zero padded one-hot representation (e.g., [0, 0, 0, 0, 1] with depth = 4 and
  left_pad = 1). If `indices` is empty, the result will simply be a tensor of
  shape (0, depth + left_pad). If depth = 0, then this function just returns
  `None`.

  Args:
    indices: an integer tensor of shape [num_indices].
    depth: depth for the one-hot tensor (integer).
    left_pad: number of zeros to left pad the one-hot tensor with (integer).

  Returns:
    padded_onehot: a tensor with shape (num_indices, depth + left_pad). Returns
      `None` if the depth is zero.

  Raises:
    ValueError: if `indices` does not have rank 1 or if `left_pad` or `depth are
      either negative or non-integers.

  TODO: add runtime checks for depth and indices.
  """
  if depth < 0 or not isinstance(depth, (int, long) if six.PY2 else int):
    raise ValueError('`depth` must be a non-negative integer.')
  if left_pad < 0 or not isinstance(left_pad, (int, long) if six.PY2 else int):
    raise ValueError('`left_pad` must be a non-negative integer.')
  if depth == 0:
    return None
  if len(indices.get_shape().as_list()) != 1:
    raise ValueError('`indices` must have rank 1')

  def one_hot_and_pad():
    one_hot = tf.cast(tf.one_hot(tf.cast(indices, tf.int64), depth,
                                 on_value=1, off_value=0), tf.float32)
    return tf.pad(one_hot, [[0, 0], [left_pad, 0]], mode='CONSTANT')
  result = tf.cond(tf.greater(tf.size(indices), 0), one_hot_and_pad,
                   lambda: tf.zeros((depth + left_pad, 0)))
  return tf.reshape(result, [-1, depth + left_pad])
//  [python]
def compute_squared_error_label_binning_tensorflow(logits,
                                                   y,
                                                   m,
                                                   temperature,
                                                   name_scope='sq_err_loss',
                                                   precision=32):
  """Computes and returns the squared soft-binned ECE (label-binned) tensor.

  Soft-binned ECE (label-binned, l2-norm) is defined in equation (12) in this
  paper: https://arxiv.org/abs/2108.00106. It is a softened version of ECE
  (label-binned) which is defined in equation (7).

  Args:
    logits: tensor of predicted logits of (batch-size, num-classes) shape
    y: tensor of labels in [0,num-classes) of (batch-size,) shape
    m: number of bins
    temperature: soft binning temperature
    name_scope: name
    precision: precision for tf.int and tf.float dtypes

  Returns:
    A tensor of () shape containing a single value: squared soft-binned ECE.
  """

  if precision == 32:
    int_dtype = tf.int32
    float_dtype = tf.float32
  elif precision == 64:
    int_dtype = tf.int64
    float_dtype = tf.float64
  else:
    raise NotImplementedError

  with tf.name_scope(name_scope) as scope:  # pylint: disable=unused-variable

    # <float_dtype>[N, n]
    logits = tf.to_float(logits, name='logits')

    # <int_dtype>[N]
    y = tf.to_int32(y, name='y')

    # <float_dtype>[m+1]
    bins = tf.linspace(
        tf.constant(1. / logits.shape[1], dtype=float_dtype),
        1,
        m + 1,
        name='bins')

    # <float_dtype>[m]
    b = tf.identity(bins[1:] * .5 + bins[:-1] * .5, name='b')

    # <float_dtype>[m]
    q_hat = tf.nn.softmax(logits, axis=1, name='q_hat')

    # <int_dtype>[N]
    y_hat = tf.math.argmax(q_hat, axis=1, output_type=int_dtype, name='y_hat')

    # <float_dtype>[N]
    p_hat = tf.reduce_max(q_hat, axis=1, name='y_hat')

    # <float_dtype>[N]
    a = tf.cast(tf.math.equal(y_hat, y), dtype=float_dtype, name='a')

    # <float_dtype>[N, m]
    c_numerator = tf.math.exp(
        -tf.math.pow(
            tf.expand_dims(b, 0) - tf.expand_dims(p_hat, 1),
            2,
        ) / temperature,
        name='c_numerator')

    # <float_dtype>[N]
    c_denominator = tf.einsum('ij->i', c_numerator, name='c_denominator')

    # <float_dtype>[N, m]
    c_pre = tf.identity(
        c_numerator / tf.expand_dims(c_denominator, 1), name='c_pre')

    # <float_dtype>[N, m]
    c = tf.identity((1 - EPS) * c_pre + EPS * (1 / m), name='c')

    # <float_dtype>[m]
    a_bar_numerator = tf.einsum('ij,i->j', c, a, name='a_bar_numerator')

    # <float_dtype>[m]
    a_bar_denominator = tf.einsum('ij->j', c, name='a_bar_denominator')

    # <float_dtype>[m]
    a_bar = tf.identity(a_bar_numerator / a_bar_denominator, name='a_bar')

    # <float_dtype>[1]
    squared_error = tf.reduce_sum(
        c * tf.math.pow(tf.expand_dims(a_bar, 0) - tf.expand_dims(p_hat, 1), 2),
        name='squared_error')

    # <float_dtype>[1]
    squared_error_scaled = squared_error / logits.shape[0]

    return squared_error_scaled
//  [python]
def compute_squared_error_label_binning_tensorflow(logits,
                                                   y,
                                                   m,
                                                   temperature,
                                                   name_scope='sq_err_loss',
                                                   precision=32):
  """Computes and returns the squared soft-binned ECE (label-binned) tensor.

  Soft-binned ECE (label-binned, l2-norm) is defined in equation (12) in this
  paper: https://arxiv.org/abs/2108.00106. It is a softened version of ECE
  (label-binned) which is defined in equation (7).

  Args:
    logits: tensor of predicted logits of (batch-size, num-classes) shape
    y: tensor of labels in [0,num-classes) of (batch-size,) shape
    m: number of bins
    temperature: soft binning temperature
    name_scope: name
    precision: precision for tf.int and tf.float dtypes

  Returns:
    A tensor of () shape containing a single value: squared soft-binned ECE.
  """

  if precision == 32:
    int_dtype = tf.int32
    float_dtype = tf.float32
  elif precision == 64:
    int_dtype = tf.int64
    float_dtype = tf.float64
  else:
    raise NotImplementedError

  with tf.name_scope(name_scope) as scope:  # pylint: disable=unused-variable

    # <float_dtype>[N, n]
    logits = tf.cast(logits, float_dtype, name='logits')

    # <int_dtype>[N]
    y = tf.cast(y, int_dtype, name='y')

    # <float_dtype>[m+1]
    bins = tf.linspace(
        tf.constant(1. / logits.shape[1], dtype=float_dtype),
        1,
        m + 1,
        name='bins')

    # <float_dtype>[m]
    b = tf.identity(bins[1:] * .5 + bins[:-1] * .5, name='b')

    # <float_dtype>[m]
    q_hat = tf.nn.softmax(logits, axis=1, name='q_hat')

    # <int_dtype>[N]
    y_hat = tf.math.argmax(q_hat, axis=1, output_type=int_dtype, name='y_hat')

    # <float_dtype>[N]
    p_hat = tf.reduce_max(q_hat, axis=1, name='y_hat')

    # <float_dtype>[N]
    a = tf.cast(tf.math.equal(y_hat, y), dtype=float_dtype, name='a')

    # <float_dtype>[N, m]
    c_numerator = tf.math.exp(
        -tf.math.pow(
            tf.expand_dims(b, 0) - tf.expand_dims(p_hat, 1),
            2,
        ) / temperature,
        name='c_numerator')

    # <float_dtype>[N]
    c_denominator = tf.einsum('ij->i', c_numerator, name='c_denominator')

    # <float_dtype>[N, m]
    c_pre = tf.identity(
        c_numerator / tf.expand_dims(c_denominator, 1), name='c_pre')

    # <float_dtype>[N, m]
    c = tf.identity((1 - EPS) * c_pre + EPS * (1 / m), name='c')

    # <float_dtype>[m]
    a_bar_numerator = tf.einsum('ij,i->j', c, a, name='a_bar_numerator')

    # <float_dtype>[m]
    a_bar_denominator = tf.einsum('ij->j', c, name='a_bar_denominator')

    # <float_dtype>[m]
    a_bar = tf.identity(a_bar_numerator / a_bar_denominator, name='a_bar')

    # <float_dtype>[1]
    squared_error = tf.reduce_sum(
        c * tf.math.pow(tf.expand_dims(a_bar, 0) - tf.expand_dims(p_hat, 1), 2),
        name='squared_error')

    # <float_dtype>[1]
    squared_error_scaled = squared_error / logits.shape[0]

    return squared_error_scaled
//  [python]
def _preconditioned_update(self, var: TfValue,
                             partitioned_grads: Sequence[TfValue],
                             diagonal_grad_update: TfValue) -> TfValue:
    """Computes the matrix preconditioned update.

    Args:
      var: Variable for which we are computing the preconditioned gradient.
      partitioned_grads: Partitioned gradients.
      diagonal_grad_update: Update as given by diagonal adagrad.

    Returns:
      scaled preconditioned gradient.
    """

    var_dtype = var.dtype.base_dtype
    first_moment_averaging = self._get_hyper('first_moment_averaging',
                                             var_dtype)
    local_steps = t.cast(self.iterations, var_dtype)
    first_moment_averaging_t = t.pow(first_moment_averaging, local_steps + 1)
    precond_grad = self._compute_preconditioned_raw_grad(var, partitioned_grads)
    if self._first_moment_averaging > 0.0:
      gbar = self.get_slot(var, 'precond_grad_momentum')
      matrix_preconditioned_grad = gbar.assign(
          gbar * first_moment_averaging_t + precond_grad *
          (1.0 - first_moment_averaging_t),
          use_locking=self._use_locking)
    else:
      matrix_preconditioned_grad = precond_grad

    # We use the direction from Shampoo while using the step size scale from
    # adaptive optimizers.
    precond_l2_norm = t.norm(matrix_preconditioned_grad, ord=2)
    diagonal_l2_norm = t.norm(diagonal_grad_update, ord=2)
    multiplier = t.where(
        t.greater(precond_l2_norm, 0.0),
        t.maximum(diagonal_l2_norm, 1e-30) /
        (t.maximum(precond_l2_norm, 1e-30)), 1.0)
    return matrix_preconditioned_grad * multiplier
//  [python]
def _preconditioned_update(self, var: TfValue,
                             partitioned_grads: Sequence[TfValue],
                             diagonal_grad_update: TfValue) -> TfValue:
    """Computes the matrix preconditioned update.

    Args:
      var: Variable for which we are computing the preconditioned gradient.
      partitioned_grads: Partitioned gradients.
      diagonal_grad_update: Update as given by diagonal adagrad.

    Returns:
      scaled preconditioned gradient.
    """

    var_dtype = var.dtype.base_dtype
    first_moment_averaging = self._get_hyper('first_moment_averaging',
                                             var_dtype)
    local_steps = tf.cast(self.iterations, var_dtype)
    first_moment_averaging_t = tf.pow(first_moment_averaging, local_steps + 1)
    precond_grad = self._compute_preconditioned_raw_grad(var, partitioned_grads)
    if self._first_moment_averaging > 0.0:
      gbar = self.get_slot(var, 'precond_grad_momentum')
      matrix_preconditioned_grad = gbar.assign(
          gbar * first_moment_averaging_t + precond_grad *
          (1.0 - first_moment_averaging_t),
          use_locking=self._use_locking)
    else:
      matrix_preconditioned_grad = precond_grad

    # We use the direction from Shampoo while using the step size scale from
    # adaptive optimizers.
    precond_l2_norm = tf.norm(matrix_preconditioned_grad, ord=2)
    diagonal_l2_norm = tf.norm(diagonal_grad_update, ord=2)
    multiplier = tf.where(
        tf.greater(precond_l2_norm, 0.0),
        tf.maximum(diagonal_l2_norm, 1e-30) /
        (tf.maximum(precond_l2_norm, 1e-30)), 1.0)
    return matrix_preconditioned_grad * multiplier
//  [python]
def _test_rel_op(data, func):
    with tf.Graph().as_default():
        in1 = tf.placeholder(shape=data[0].shape, dtype=data[0].dtype, name="in1")
        in2 = tf.placeholder(shape=data[1].shape, dtype=data[1].dtype, name="in2")
        op = func(in1, in2, name="op")
        _ = tf.to_int32(op, name="out1")
        io_info = {
            "in_data": [data[0], data[1]],
            "in_name": ["in1:0", "in2:0"],
            "out_name": "out1:0",
        }
        graph_def, golden = get_graph_def(**io_info)
    verify_model(graph_def, golden, **io_info)
//  [python]
def _test_rel_op(data, func):
    with tf.Graph().as_default():
        in1 = tf.placeholder(shape=data[0].shape, dtype=data[0].dtype, name="in1")
        in2 = tf.placeholder(shape=data[1].shape, dtype=data[1].dtype, name="in2")
        op = func(in1, in2, name="op")
        _ = tf.cast(op, tf.int32, name="out1")
        io_info = {
            "in_data": [data[0], data[1]],
            "in_name": ["in1:0", "in2:0"],
            "out_name": "out1:0",
        }
        graph_def, golden = get_graph_def(**io_info)
    verify_model(graph_def, golden, **io_info)
//  [python]
def bincount(x, weights=None, minlength=None):
    if minlength is not None:
        x = t.cast(x, tf.int32)
    if isinstance(x, tf.SparseTensor):
        result = tf.sparse.bincount(
            x,
            weights=weights,
            minlength=minlength,
            axis=-1,
        )
        if x.shape.rank == 1:
            output_shape = (minlength,)
        else:
            batch_size = tf.shape(result)[0]
            output_shape = (batch_size, minlength)
        return tf.SparseTensor(
            indices=result.indices,
            values=result.values,
            dense_shape=output_shape,
        )
    return tf.math.bincount(x, weights=weights, minlength=minlength, axis=-1)
//  [python]
def bincount(x, weights=None, minlength=None):
    if minlength is not None:
        x = tf.cast(x, tf.int32)
    if isinstance(x, tf.SparseTensor):
        result = tf.sparse.bincount(
            x,
            weights=weights,
            minlength=minlength,
            axis=-1,
        )
        if x.shape.rank == 1:
            output_shape = (minlength,)
        else:
            batch_size = tf.shape(result)[0]
            output_shape = (batch_size, minlength)
        return tf.SparseTensor(
            indices=result.indices,
            values=result.values,
            dense_shape=output_shape,
        )
    return tf.math.bincount(x, weights=weights, minlength=minlength, axis=-1)
//  [python]
def sqrt(x):
    x = convert_to_tensor(x)
    if tf.as_dtype(x.dtype).is_integer:
        x = t.cast(x, dtype=config.floatx())
    return tfnp.sqrt(x)
//  [python]
def sqrt(x):
    x = convert_to_tensor(x)
    if tf.as_dtype(x.dtype).is_integer:
        x = tf.cast(x, dtype=config.floatx())
    return tfnp.sqrt(x)
//  [python]
def testDenseDVIMoments(self):
    """Verifies DenseDVI's moments empirically with samples."""
    tf.random.set_seed(377269)
    batch_size = 3
    num_features = 5
    units = 128
    num_samples = 50000
    inputs = t.cast(np.random.rand(batch_size, num_features), dtype=tf.float32)
    layer = ed.layers.DenseDVI(units, activation=tf.nn.relu)

    outputs1 = layer(inputs)
    mean1 = outputs1.distribution.mean()
    covariance1 = outputs1.distribution.covariance()

    kernel_samples = layer.kernel.distribution.sample(num_samples)
    outputs2 = layer.activation(
        tf.einsum("bd,sdu->sbu", inputs, kernel_samples) +
        tf.reshape(layer.bias, [1, 1, units]))
    mean2 = tf.reduce_mean(outputs2, axis=0)
    centered_outputs2 = tf.transpose(a=outputs2 - mean2, perm=[1, 2, 0])
    covariance2 = tf.matmul(centered_outputs2,
                            centered_outputs2,
                            transpose_b=True) / float(num_samples)

    # Check % of mismatches is not too high according to heuristic thresholds.
    num_mismatches = np.sum(np.abs(mean1 - mean2) > 5e-3)
    percent_mismatches = num_mismatches / float(batch_size * units)
    self.assertLessEqual(percent_mismatches, 0.05)
    num_mismatches = np.sum(np.abs(covariance1 - covariance2) > 5e-3)
    percent_mismatches = num_mismatches / float(batch_size * units * units)
    self.assertLessEqual(percent_mismatches, 0.05)
//  [python]
def testDenseDVIMoments(self):
    """Verifies DenseDVI's moments empirically with samples."""
    tf.random.set_seed(377269)
    batch_size = 3
    num_features = 5
    units = 128
    num_samples = 50000
    inputs = tf.cast(np.random.rand(batch_size, num_features), dtype=tf.float32)
    layer = ed.layers.DenseDVI(units, activation=tf.nn.relu)

    outputs1 = layer(inputs)
    mean1 = outputs1.distribution.mean()
    covariance1 = outputs1.distribution.covariance()

    kernel_samples = layer.kernel.distribution.sample(num_samples)
    outputs2 = layer.activation(
        tf.einsum("bd,sdu->sbu", inputs, kernel_samples) +
        tf.reshape(layer.bias, [1, 1, units]))
    mean2 = tf.reduce_mean(outputs2, axis=0)
    centered_outputs2 = tf.transpose(a=outputs2 - mean2, perm=[1, 2, 0])
    covariance2 = tf.matmul(centered_outputs2,
                            centered_outputs2,
                            transpose_b=True) / float(num_samples)

    # Check % of mismatches is not too high according to heuristic thresholds.
    num_mismatches = np.sum(np.abs(mean1 - mean2) > 5e-3)
    percent_mismatches = num_mismatches / float(batch_size * units)
    self.assertLessEqual(percent_mismatches, 0.05)
    num_mismatches = np.sum(np.abs(covariance1 - covariance2) > 5e-3)
    percent_mismatches = num_mismatches / float(batch_size * units * units)
    self.assertLessEqual(percent_mismatches, 0.05)
//  [python]
def call(
      self,
      inputs,
      inputs_embeddings=None,
      attention_mask=None,
      training=False):

    if inputs is not None and inputs_embeddings is not None:
      raise ValueError("You cannot specify both inputs and "
                       "inputs_embeddings at the same time")
    elif inputs is not None:
      input_shape = get_shape(inputs)
      inputs = tf.reshape(inputs, (-1, input_shape[-1]))
    elif inputs_embeddings is not None:
      input_shape = get_shape(inputs_embeddings)[:-1]
    else:
      raise ValueError("You have to specify either inputs or inputs_embeddings")

    if inputs_embeddings is None:
      err_msg = "You have to intialize the model with valid token embeddings"
      assert self.embeddings is not None, err_msg
      inputs_embeddings = self.embeddings(inputs)

    batch_size, seq_length = input_shape
    mask_seq_length = seq_length
    if attention_mask is None:
      attention_mask = tf.fill((batch_size, mask_seq_length), 1)

    # Provided a padding mask of dimensions [batch_size, mask_seq_length]
    # make the mask broadcastable to
    #         [batch_size, num_heads, mask_seq_length, mask_seq_length]
    attention_mask = tf.to_float(attention_mask)
    attention_mask = attention_mask[:, None, None, :]
    attention_mask = (1.0 - attention_mask) * -1e9

    all_hidden_states = ()
    all_attentions = ()
    position_bias = None

    hidden_states = self.dropout(inputs_embeddings, training=training)

    for n, layer in enumerate(self.layers):
      # temporary --- to reduce memory consumption
      # all_hidden_states = all_hidden_states + (hidden_states,)

      layer_outputs = layer(inputs=hidden_states,
                            attention_mask=attention_mask,
                            position_bias=position_bias,
                            training=training)

      # layer_outputs is a dictionary with the following keys:
      # hidden_states, key_value_state,
      #                           self_attention_weights, self_position_bias

      hidden_states = layer_outputs["hidden_states"]
      if n == 0:
        position_bias = layer_outputs["self_position_bias"]

      all_attentions = all_attentions + (
          layer_outputs["self_attention_weights"],
          )

    hidden_states = self.final_layer_norm(hidden_states)
    hidden_states = self.dropout(hidden_states, training=training)

    # Add last layer
    all_hidden_states = all_hidden_states + (hidden_states,)

    outputs = {
        "hidden_states": all_hidden_states,
        "attention_weights": all_attentions,
    }

    return outputs
//  [python]
def call(
      self,
      inputs,
      inputs_embeddings=None,
      attention_mask=None,
      training=False):

    if inputs is not None and inputs_embeddings is not None:
      raise ValueError("You cannot specify both inputs and "
                       "inputs_embeddings at the same time")
    elif inputs is not None:
      input_shape = get_shape(inputs)
      inputs = tf.reshape(inputs, (-1, input_shape[-1]))
    elif inputs_embeddings is not None:
      input_shape = get_shape(inputs_embeddings)[:-1]
    else:
      raise ValueError("You have to specify either inputs or inputs_embeddings")

    if inputs_embeddings is None:
      err_msg = "You have to intialize the model with valid token embeddings"
      assert self.embeddings is not None, err_msg
      inputs_embeddings = self.embeddings(inputs)

    batch_size, seq_length = input_shape
    mask_seq_length = seq_length
    if attention_mask is None:
      attention_mask = tf.fill((batch_size, mask_seq_length), 1)

    # Provided a padding mask of dimensions [batch_size, mask_seq_length]
    # make the mask broadcastable to
    #         [batch_size, num_heads, mask_seq_length, mask_seq_length]
    attention_mask = tf.cast(attention_mask, dtype=tf.float32)
    attention_mask = attention_mask[:, None, None, :]
    attention_mask = (1.0 - attention_mask) * -1e9

    all_hidden_states = ()
    all_attentions = ()
    position_bias = None

    hidden_states = self.dropout(inputs_embeddings, training=training)

    for n, layer in enumerate(self.layers):
      # temporary --- to reduce memory consumption
      # all_hidden_states = all_hidden_states + (hidden_states,)

      layer_outputs = layer(inputs=hidden_states,
                            attention_mask=attention_mask,
                            position_bias=position_bias,
                            training=training)

      # layer_outputs is a dictionary with the following keys:
      # hidden_states, key_value_state,
      #                           self_attention_weights, self_position_bias

      hidden_states = layer_outputs["hidden_states"]
      if n == 0:
        position_bias = layer_outputs["self_position_bias"]

      all_attentions = all_attentions + (
          layer_outputs["self_attention_weights"],
          )

    hidden_states = self.final_layer_norm(hidden_states)
    hidden_states = self.dropout(hidden_states, training=training)

    # Add last layer
    all_hidden_states = all_hidden_states + (hidden_states,)

    outputs = {
        "hidden_states": all_hidden_states,
        "attention_weights": all_attentions,
    }

    return outputs
//  [python]
def _apply_dense(self, grad: TfValue, var: TfValue):
    """Shampoo Computes preconditioner periodically and updates model."""
    partitioned_grads = partition_tensor(grad, self._partition_info)
    shape = var.get_shape()
    fallback_to_diagonal = self._fallback_to_diagonal_for_shape(shape)

    precond_statistics_update = []
    if not fallback_to_diagonal:
      precond_statistics_update = self._updated_statistics(
          var, partitioned_grads)

    var_dtype = var.dtype.base_dtype
    lr_t = self._get_hyper('learning_rate', var_dtype)
    beta1_t = self._get_hyper('first_moment_averaging', var_dtype)
    epsilon_t = self._get_hyper('epsilon', var_dtype)
    start_preconditioning_steps = self._get_hyper('start_preconditioning_steps',
                                                  var_dtype)
    local_step = t.cast(self.iterations + 1, var_dtype)

    lr = lr_t
    if self._gradient_norm_adjuster is None:
      per_coord_lr = 1.0
    elif self._gradient_norm_adjuster == 'adagrad':
      v = self.get_slot(var, 'accumulator')
      v_t = v.assign_add(grad * grad)
      v_sqrt = t.sqrt(v_t)
      per_coord_lr = 1.0 / (v_sqrt + epsilon_t)
    else:
      raise NotImplementedError('Gradient norm adjuster %s is not supported!' %
                                self._gradient_norm_adjuster)

    update_vs = []
    if self._first_moment_averaging > 0.0:
      # m_t = beta1 * m + (1 - beta1) * g_t
      scaled_g = (1.0 - beta1_t) * (grad * per_coord_lr)
      m = self.get_slot(var, 'momentum')
      m_t = m.assign(m * beta1_t + scaled_g, use_locking=self._use_locking)
      gbar_updated = m_t
    else:
      gbar_updated = per_coord_lr * grad
      m_t = tf.no_op()

    if not fallback_to_diagonal:
      # Update the preconditioner statistics followed by computing the
      # preconditioned gradient.
      with tf.control_dependencies(precond_statistics_update):
        s = t.cast(
            t.greater_equal(
                t.cast(local_step, t.int32),
                t.cast(start_preconditioning_steps, t.int32)), t.float32)
        preconditioned_grad = self._preconditioned_update(
            var, partitioned_grads, gbar_updated)
        # slowly adapt from diagonal to preconditioned gradient.
        w = t.minimum(
            1.0,
            t.maximum((local_step - start_preconditioning_steps) /
                       start_preconditioning_steps, 0.0))
        warmup_update = s * lr * (
            w * preconditioned_grad + (1.0 - w) * gbar_updated)
        fallback_update = (1 - s) * (lr * gbar_updated)
        var_update = var.assign_sub(
            warmup_update + fallback_update, use_locking=self._use_locking)
    else:
      var_update = var.assign_sub(
          lr * gbar_updated, use_locking=self._use_locking)

    update_vs.append(var_update)
    update_vs.append(m_t)
    update_vs.append(v_t)

    # Create an op that groups all the above operations
    return tf.group(*update_vs)
//  [python]
def _apply_dense(self, grad: TfValue, var: TfValue):
    """Shampoo Computes preconditioner periodically and updates model."""
    partitioned_grads = partition_tensor(grad, self._partition_info)
    shape = var.get_shape()
    fallback_to_diagonal = self._fallback_to_diagonal_for_shape(shape)

    precond_statistics_update = []
    if not fallback_to_diagonal:
      precond_statistics_update = self._updated_statistics(
          var, partitioned_grads)

    var_dtype = var.dtype.base_dtype
    lr_t = self._get_hyper('learning_rate', var_dtype)
    beta1_t = self._get_hyper('first_moment_averaging', var_dtype)
    epsilon_t = self._get_hyper('epsilon', var_dtype)
    start_preconditioning_steps = self._get_hyper('start_preconditioning_steps',
                                                  var_dtype)
    local_step = tf.cast(self.iterations + 1, var_dtype)

    lr = lr_t
    if self._gradient_norm_adjuster is None:
      per_coord_lr = 1.0
    elif self._gradient_norm_adjuster == 'adagrad':
      v = self.get_slot(var, 'accumulator')
      v_t = v.assign_add(grad * grad)
      v_sqrt = tf.sqrt(v_t)
      per_coord_lr = 1.0 / (v_sqrt + epsilon_t)
    else:
      raise NotImplementedError('Gradient norm adjuster %s is not supported!' %
                                self._gradient_norm_adjuster)

    update_vs = []
    if self._first_moment_averaging > 0.0:
      # m_t = beta1 * m + (1 - beta1) * g_t
      scaled_g = (1.0 - beta1_t) * (grad * per_coord_lr)
      m = self.get_slot(var, 'momentum')
      m_t = m.assign(m * beta1_t + scaled_g, use_locking=self._use_locking)
      gbar_updated = m_t
    else:
      gbar_updated = per_coord_lr * grad
      m_t = tf.no_op()

    if not fallback_to_diagonal:
      # Update the preconditioner statistics followed by computing the
      # preconditioned gradient.
      with tf.control_dependencies(precond_statistics_update):
        s = tf.cast(
            tf.greater_equal(
                tf.cast(local_step, tf.int32),
                tf.cast(start_preconditioning_steps, tf.int32)), tf.float32)
        preconditioned_grad = self._preconditioned_update(
            var, partitioned_grads, gbar_updated)
        # slowly adapt from diagonal to preconditioned gradient.
        w = tf.minimum(
            1.0,
            tf.maximum((local_step - start_preconditioning_steps) /
                       start_preconditioning_steps, 0.0))
        warmup_update = s * lr * (
            w * preconditioned_grad + (1.0 - w) * gbar_updated)
        fallback_update = (1 - s) * (lr * gbar_updated)
        var_update = var.assign_sub(
            warmup_update + fallback_update, use_locking=self._use_locking)
    else:
      var_update = var.assign_sub(
          lr * gbar_updated, use_locking=self._use_locking)

    update_vs.append(var_update)
    update_vs.append(m_t)
    update_vs.append(v_t)

    # Create an op that groups all the above operations
    return tf.group(*update_vs)
//  [python]
def multiply(x1, x2):
    if isinstance(x1, tf.SparseTensor):
        if isinstance(x2, tf.SparseTensor):
            ones_like_int8 = functools.partial(tf.ones_like, dtype=tf.int8)
            zeros_like_int8 = functools.partial(tf.zeros_like, dtype=tf.int8)

            # compute the intersection of indices in the form of a sparse tensor
            # containing ones as values
            ones1 = tf.sparse.map_values(ones_like_int8, x1)
            ones2 = tf.sparse.map_values(ones_like_int8, x2)
            # tf.sets.intersection ignores the last dimension when comparing,
            # so we need to add a dummy extra dimension and then remove it
            intersection = tf.sparse.reshape(
                tf.sets.intersection(
                    tf.sparse.expand_dims(ones1, axis=-1),
                    tf.sparse.expand_dims(ones2, axis=-1),
                ),
                x1.dense_shape,
            )

            # compute the masks to remove indices in x1 and x2 that are not part
            # of the intersection, then trim x1 and x2
            zeros1 = tf.sparse.map_values(zeros_like_int8, x1)
            zeros2 = tf.sparse.map_values(zeros_like_int8, x2)
            mask1 = tf.sparse.add(zeros1, intersection)
            mask2 = tf.sparse.add(zeros2, intersection)
            x1_trimmed = tf.sparse.retain(x1, tf.cast(mask1.values, tf.bool))
            x2_trimmed = tf.sparse.retain(x2, tf.cast(mask2.values, tf.bool))

            # now it is an element-wise multiplication on the values
            return tf.sparse.map_values(tf.multiply, x1_trimmed, x2_trimmed)
        else:
            return x1 * x2
    elif isinstance(x2, tf.SparseTensor):
        return x2 * x1
    return tfnp.multiply(x1, x2)
//  [python]
def test_readme_gradient_int(self):
    x = np.array(2, dtype=np.int16)

    def f_jax(x):  # x: int16
      return x.astype(np.float32) * 2.

    print(jax.grad(f_jax, allow_int=True)(x))
    # returns a special `float0`: array((b'',), dtype=[('float0', 'V')])

    print(jax2tf.convert(jax.grad(f_jax, allow_int=True))(x))
    # returns a 0 with same shape as x, but with dtype int32

    def f_tf(x):  # x: int16
      return tf.cast(x, tf.float32) * 2.

    xv = tf.Variable(x)
    with tf.GradientTape(persistent=True) as tape:
      print(tape.gradient(f_tf(xv), xv))
      # returns None
      print(tape.gradient(f_tf(xv), xv,
                          unconnected_gradients=tf.UnconnectedGradients.ZERO))
//  [python]
def testPerArmPolicyDistribution(self, policy_class):
    tf.compat.v1.set_random_seed(3000)
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(2, 3, 4)
    time_step_spec = ts.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
    reward_network = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
    )

    policy = policy_class(
        time_step_spec,
        action_spec,
        reward_network=reward_network,
        accepts_per_arm_features=True,
    )
    action_feature = tf.to_float(
        tf.reshape(tf.random.shuffle(tf.range(24)), shape=[2, 4, 3])
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: action_feature,
    }
    time_step = ts.restart(observations, batch_size=2)
    distribution_step = policy.distribution(time_step)
    # Initialize all variables
    self.evaluate(tf.compat.v1.global_variables_initializer())
    info = self.evaluate(distribution_step.info)
    self.assertAllEqual(info.chosen_arm_features.shape, [2, 3])
//  [python]
def testPerArmPolicyDistribution(self, policy_class):
    tf.compat.v1.set_random_seed(3000)
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(2, 3, 4)
    time_step_spec = ts.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
    reward_network = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
    )

    policy = policy_class(
        time_step_spec,
        action_spec,
        reward_network=reward_network,
        accepts_per_arm_features=True,
    )
    action_feature = tf.cast(
        tf.reshape(tf.random.shuffle(tf.range(24)), shape=[2, 4, 3]),
        dtype=tf.float32,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: action_feature,
    }
    time_step = ts.restart(observations, batch_size=2)
    distribution_step = policy.distribution(time_step)
    # Initialize all variables
    self.evaluate(tf.compat.v1.global_variables_initializer())
    info = self.evaluate(distribution_step.info)
    self.assertAllEqual(info.chosen_arm_features.shape, [2, 3])
//  [python]
def construct_padding_mask(self, inp):
        """
        Args: Original input of word ids, shape [batch, seq_len]
        Returns: a mask of shape [batch, seq_len, seq_len], where <pad> is 0 and others are 1s.
        """
        seq_len = inp.shape.as_list()[1]
        mask = tf.to_float(tf.not_equal(inp, self._pad_id))  # mask '<pad>'
        mask = tf.tile(tf.expand_dims(mask, 1), [1, seq_len, 1])
        return mask
//  [python]
def construct_padding_mask(self, inp):
        """
        Args: Original input of word ids, shape [batch, seq_len]
        Returns: a mask of shape [batch, seq_len, seq_len], where <pad> is 0 and others are 1s.
        """
        seq_len = inp.shape.as_list()[1]
        mask = tf.cast(tf.not_equal(inp, self._pad_id), tf.float32)  # mask '<pad>'
        mask = tf.tile(tf.expand_dims(mask, 1), [1, seq_len, 1])
        return mask
//  [python]
def testPerArmRewardsVariableNumActions(self, policy_class):
    tf.compat.v1.set_random_seed(3000)
    obs_spec = bandit_spec_utils.create_per_arm_observation_spec(
        2, 3, 4, add_num_actions_feature=True
    )
    time_step_spec = ts.time_step_spec(obs_spec)
    action_spec = tensor_spec.BoundedTensorSpec((), tf.int32, 0, 3)
    reward_network = (
        global_and_arm_feature_network.create_feed_forward_common_tower_network(
            obs_spec, (4, 3), (3, 4), (4, 2)
        )
    )

    policy = policy_class(
        time_step_spec,
        action_spec,
        reward_network=reward_network,
        accepts_per_arm_features=True,
        emit_policy_info=('predicted_rewards_mean',),
    )
    action_feature = tf.cast(
        tf.reshape(tf.random.shuffle(tf.range(24)), shape=[2, 4, 3]),
        dtype=tf.float32,
    )
    observations = {
        bandit_spec_utils.GLOBAL_FEATURE_KEY: tf.constant(
            [[1, 2], [3, 4]], dtype=tf.float32
        ),
        bandit_spec_utils.PER_ARM_FEATURE_KEY: action_feature,
        bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY: tf.constant(
            [2, 3], dtype=tf.int32
        ),
    }
    time_step = ts.restart(observations, batch_size=2)
    action_step = policy.action(time_step, seed=1)
    self.assertEqual(action_step.action.shape.as_list(), [2])
    self.assertEqual(action_step.action.dtype, tf.int32)
    # Initialize all variables
    self.evaluate(tf.compat.v1.global_variables_initializer())
    action, p_info, first_arm_features = self.evaluate([
        action_step.action,
        action_step.info,
        observations[bandit_spec_utils.PER_ARM_FEATURE_KEY][0],
    ])
    self.assertAllEqual(action.shape, [2])
    self.assertAllEqual(p_info.predicted_rewards_mean.shape, [2, 4])
    self.assertAllEqual(p_info.chosen_arm_features.shape, [2, 3])
    first_action = action[0]
    self.assertAllEqual(
        p_info.chosen_arm_features[0], first_arm_features[first_action]
    )
//  [python]
def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        Args:
            Q (tf.tensor): of shape (h * batch, q_size, d_model)
            K (tf.tensor): of shape (h * batch, k_size, d_model)
            V (tf.tensor): of shape (h * batch, k_size, d_model)
            mask (tf.tensor): of shape (h * batch, q_size, k_size)
        """

        d = self.d_model // self.h
        assert d == Q.shape[-1] == K.shape[-1] == V.shape[-1]

        out = tf.matmul(Q, tf.transpose(K, [0, 2, 1]))  # [h*batch, q_size, k_size]
        out = out / tf.sqrt(tf.to_float(d))  # scaled by sqrt(d_k)

        if mask is not None:
            # masking out (0.0) => setting to -inf.
            out = tf.multiply(out, mask) + (1.0 - mask) * (-1e10)

        out = tf.nn.softmax(out)  # [h * batch, q_size, k_size]
        out = tf.layers.dropout(out, training=self._is_training)
        out = tf.matmul(out, V)  # [h * batch, q_size, d_model]

        return out
//  [python]
def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        Args:
            Q (tf.tensor): of shape (h * batch, q_size, d_model)
            K (tf.tensor): of shape (h * batch, k_size, d_model)
            V (tf.tensor): of shape (h * batch, k_size, d_model)
            mask (tf.tensor): of shape (h * batch, q_size, k_size)
        """

        d = self.d_model // self.h
        assert d == Q.shape[-1] == K.shape[-1] == V.shape[-1]

        out = tf.matmul(Q, tf.transpose(K, [0, 2, 1]))  # [h*batch, q_size, k_size]
        out = out / tf.sqrt(tf.cast(d, tf.float32))  # scaled by sqrt(d_k)

        if mask is not None:
            # masking out (0.0) => setting to -inf.
            out = tf.multiply(out, mask) + (1.0 - mask) * (-1e10)

        out = tf.nn.softmax(out)  # [h * batch, q_size, k_size]
        out = tf.layers.dropout(out, training=self._is_training)
        out = tf.matmul(out, V)  # [h * batch, q_size, d_model]

        return out
//  [python]
def timeWarpDataElement(dat, timeScalingRange):
    warpDat = {}
    warpDat['seqClassIDs'] = dat['seqClassIDs']
    warpDat['nSeqElements'] = dat['nSeqElements']
    warpDat['transcription'] = dat['transcription']

    #nTimeSteps, inputFeatures need to be modified
    globalTimeFactor = 1 + \
        (tf.random.uniform(shape=[], dtype=tf.float32)-0.5)*timeScalingRange
    warpDat['nTimeSteps'] = t.cast(
        t.cast(dat['nTimeSteps'], dtype=tf.float32)*globalTimeFactor, dtype=tf.int64)

    b = tf.shape(dat['inputFeatures'])[0]
    t = tf.cast(tf.shape(dat['inputFeatures'])[1], dtype=tf.int32)
    warppedT = t.cast(t.cast(t, dtype=tf.float32) * globalTimeFactor, dtype=tf.int32)
    newIdx = tf.linspace(tf.zeros_like(dat['nTimeSteps'], dtype=tf.int32),
                         tf.ones_like(dat['nTimeSteps'], dtype=tf.int32) * (t - 1),
                         warppedT,
                         axis=1)
    newIdx = t.cast(newIdx, dtype=tf.int32)
    batchIdx = tf.tile(tf.range(b)[:, None, None], [1, warppedT, 1])
    newIdx = tf.concat([batchIdx, newIdx[..., None]], axis=-1)
    warpDat['inputFeatures'] = tf.gather_nd(dat['inputFeatures'], newIdx)
    #warpDat['classLabelsOneHot'] = tf.gather(
    #    dat['classLabelsOneHot'], newIdx, axis=0)
    warpDat['newClassSignal'] = tf.gather_nd(
        dat['newClassSignal'], newIdx)
    warpDat['ceMask'] = tf.gather_nd(dat['ceMask'], newIdx)

    return warpDat
//  [python]
def timeWarpDataElement(dat, timeScalingRange):
    warpDat = {}
    warpDat['seqClassIDs'] = dat['seqClassIDs']
    warpDat['nSeqElements'] = dat['nSeqElements']
    warpDat['transcription'] = dat['transcription']

    #nTimeSteps, inputFeatures need to be modified
    globalTimeFactor = 1 + \
        (tf.random.uniform(shape=[], dtype=tf.float32)-0.5)*timeScalingRange
    warpDat['nTimeSteps'] = tf.cast(
        tf.cast(dat['nTimeSteps'], dtype=tf.float32)*globalTimeFactor, dtype=tf.int64)

    b = tf.shape(dat['inputFeatures'])[0]
    t = tf.cast(tf.shape(dat['inputFeatures'])[1], dtype=tf.int32)
    warppedT = tf.cast(tf.cast(t, dtype=tf.float32) * globalTimeFactor, dtype=tf.int32)
    newIdx = tf.linspace(tf.zeros_like(dat['nTimeSteps'], dtype=tf.int32),
                         tf.ones_like(dat['nTimeSteps'], dtype=tf.int32) * (t - 1),
                         warppedT,
                         axis=1)
    newIdx = tf.cast(newIdx, dtype=tf.int32)
    batchIdx = tf.tile(tf.range(b)[:, None, None], [1, warppedT, 1])
    newIdx = tf.concat([batchIdx, newIdx[..., None]], axis=-1)
    warpDat['inputFeatures'] = tf.gather_nd(dat['inputFeatures'], newIdx)
    #warpDat['classLabelsOneHot'] = tf.gather(
    #    dat['classLabelsOneHot'], newIdx, axis=0)
    warpDat['newClassSignal'] = tf.gather_nd(
        dat['newClassSignal'], newIdx)
    warpDat['ceMask'] = tf.gather_nd(dat['ceMask'], newIdx)

    return warpDat