Python tensorflow.compat.v2.reshape() Examples

def postprocessing(self, env_output):
    observation = env_output.observation
    # [time_step, 1]
    is_start = observation[constants.IS_START].numpy()
    cnt = 0
    mask = []
    for i in range(is_start.shape[0]):
      cnt += is_start[i]
      if cnt == 1:
        mask.append(True)
      else:
        mask.append(False)
    mask = tf.reshape(tf.convert_to_tensor(mask), is_start.shape)
    observation[constants.DISC_MASK] = mask
    env_output = env_output._replace(observation=observation)
    return env_output 

def test_default_construction_2d(self):
    """Tests the default parameters for 2 dimensional Brownian Motion."""
    process = BrownianMotion(dim=2)
    self.assertEqual(process.dim(), 2)
    drift_fn = process.total_drift_fn()
    # Drifts should be zero.
    t0 = np.array([0.2, 0.7, 0.9])
    delta_t = np.array([0.1, 0.8, 0.3])
    t1 = t0 + delta_t
    drifts = self.evaluate(drift_fn(t0, t1))
    self.assertEqual(drifts.shape, (3, 2))
    self.assertArrayNear(drifts.reshape([-1]), np.zeros([3 * 2]), 1e-10)
    variances = self.evaluate(process.total_covariance_fn()(t0, t1))
    self.assertEqual(variances.shape, (3, 2, 2))
    expected_variances = np.eye(2) * delta_t.reshape([-1, 1, 1])
    print(variances, expected_variances)

    self.assertArrayNear(
        variances.reshape([-1]), expected_variances.reshape([-1]), 1e-10) 

def test_time_dependent_construction(self):
    """Tests with time dependent drift and variance."""

    def vol_fn(t):
      return tf.expand_dims(0.2 - 0.1 * tf.exp(-t), axis=-1)

    def variance_fn(t0, t1):
      # The instantaneous volatility is 0.2 - 0.1 e^(-t).
      tot_var = (t1 - t0) * 0.04 - (tf.exp(-2 * t1) - tf.exp(-2 * t0)) * 0.005
      tot_var += 0.04 * (tf.exp(-t1) - tf.exp(-t0))
      return tf.reshape(tot_var, [-1, 1, 1])

    process = BrownianMotion(
        dim=1, drift=0.1, volatility=vol_fn, total_covariance_fn=variance_fn)
    t0 = np.array([0.2, 0.7, 0.9])
    delta_t = np.array([0.1, 0.8, 0.3])
    t1 = t0 + delta_t
    drifts = self.evaluate(process.total_drift_fn()(t0, t1))
    self.assertArrayNear(drifts, 0.1 * delta_t, 1e-10)
    variances = self.evaluate(process.total_covariance_fn()(t0, t1))
    self.assertArrayNear(
        variances.reshape([-1]), [0.00149104, 0.02204584, 0.00815789], 1e-8) 

def outer_multiply(x, y):
  """Performs an outer multiplication of two tensors.

  Given two `Tensor`s, `S` and `T` of shape `s` and `t` respectively, the outer
  product `P` is a `Tensor` of shape `s + t` whose components are given by:

  ```none
  P_{i1,...ik, j1, ... , jm} = S_{i1...ik} T_{j1, ... jm}
  ```

  Args:
    x: A `Tensor` of any shape and numeric dtype.
    y: A `Tensor` of any shape and the same dtype as `x`.

  Returns:
    outer_product: A `Tensor` of shape Shape[x] + Shape[y] and the same dtype
      as `x`.
  """
  x_shape = tf.shape(x)
  padded_shape = tf.concat(
      [x_shape, tf.ones(tf.rank(y), dtype=x_shape.dtype)], axis=0)
  return tf.reshape(x, padded_shape) * y 

def hessian_as_matrix(function: Callable[[Parameters], tf.Tensor],
                      parameters: Parameters) -> tf.Tensor:
  """Computes the Hessian of a given function.

  Same as `hessian`, although return a matrix of size [w_dim, w_dim], where
  `w_dim` is the number of parameters, which makes it easier to work with.

  Args:
    function: A function for which we want to compute the Hessian.
    parameters: Parameters with respect to the Hessian should be computed.

  Returns:
    A tensor of size [w_dim, w_dim] representing the Hessian.
  """
  hessian_as_tensor_list = hessian(function, parameters)
  hessian_as_tensor_list = [
      tf.reshape(e, [e.shape[0], -1]) for e in hessian_as_tensor_list]
  return tf.concat(hessian_as_tensor_list, axis=1) 

def _reshape_inner_dims(
    tensor: tf.Tensor,
    shape: tf.TensorShape,
    new_shape: tf.TensorShape) -> tf.Tensor:
  """Reshapes tensor to: shape(tensor)[:-len(shape)] + new_shape."""
  tensor_shape = tf.shape(tensor)
  ndims = shape.rank
  tensor.shape[-ndims:].assert_is_compatible_with(shape)
  new_shape_inner_tensor = tf.cast(
      [-1 if d is None else d for d in new_shape.as_list()], tf.int64)
  new_shape_outer_tensor = tf.cast(
      tensor_shape[:-ndims], tf.int64)
  full_new_shape = tf.concat(
      (new_shape_outer_tensor, new_shape_inner_tensor), axis=0)
  new_tensor = tf.reshape(tensor, full_new_shape)
  new_tensor.set_shape(tensor.shape[:-ndims] + new_shape)
  return new_tensor 

def argsort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missing-docstring
  # TODO(nareshmodi): make string tensors also work.
  if kind not in ('quicksort', 'stable'):
    raise ValueError("Only 'quicksort' and 'stable' arguments are supported.")
  if order is not None:
    raise ValueError("'order' argument to sort is not supported.")
  stable = (kind == 'stable')

  a = array_ops.array(a).data

  def _argsort(a, axis, stable):
    if axis is None:
      a = tf.reshape(a, [-1])
      axis = 0

    return tf.argsort(a, axis, stable=stable)

  tf_ans = tf.cond(
      tf.rank(a) == 0, lambda: tf.constant([0]),
      lambda: _argsort(a, axis, stable))

  return array_ops.array(tf_ans, dtype=np.intp) 

def _prob(self, y):
    """Called by the base class to compute likelihoods."""
    # Convert to (channels, 1, batch) format by collapsing dimensions and then
    # commuting channels to front.
    y = tf.broadcast_to(
        y, tf.broadcast_dynamic_shape(tf.shape(y), self.batch_shape_tensor()))
    shape = tf.shape(y)
    y = tf.reshape(y, (-1, 1, self.batch_shape.num_elements()))
    y = tf.transpose(y, (2, 1, 0))

    # Evaluate densities.
    # We can use the special rule below to only compute differences in the left
    # tail of the sigmoid. This increases numerical stability: sigmoid(x) is 1
    # for large x, 0 for small x. Subtracting two numbers close to 0 can be done
    # with much higher precision than subtracting two numbers close to 1.
    lower = self._logits_cumulative(y - .5)
    upper = self._logits_cumulative(y + .5)
    # Flip signs if we can move more towards the left tail of the sigmoid.
    sign = tf.stop_gradient(-tf.math.sign(lower + upper))
    p = abs(tf.sigmoid(sign * upper) - tf.sigmoid(sign * lower))

    # Convert back to (broadcasted) input tensor shape.
    p = tf.transpose(p, (2, 1, 0))
    p = tf.reshape(p, shape)
    return p 

def distance_metric(preds, targets):
  """Calculate distances between model predictions and targets within a batch."""
  batch_size = preds.shape[0]
  preds = tf.reshape(
      preds, [batch_size, DOWNSCALED_PANO_HEIGHT, DOWNSCALED_PANO_WIDTH])
  targets = tf.reshape(
      targets, [batch_size, DOWNSCALED_PANO_HEIGHT, DOWNSCALED_PANO_WIDTH])
  distances = []
  for pred, target in zip(preds, targets):
    pred_coord = np.unravel_index(np.argmax(pred), pred.shape)
    target_coord = np.unravel_index(np.argmax(target), target.shape)
    dist = np.sqrt((target_coord[0] - pred_coord[0])**2 +
                   (target_coord[1] - pred_coord[1])**2)
    dist = dist * RESOLUTION_MULTIPLIER
    distances.append(dist)
  return distances 

def take(a, indices, axis=None, out=None, mode='clip'):
  """out argument is not supported, and default mode is clip."""
  if out is not None:
    raise ValueError('out argument is not supported in take.')

  if mode not in {'raise', 'clip', 'wrap'}:
    raise ValueError("Invalid mode '{}' for take".format(mode))

  a = asarray(a).data
  indices = asarray(indices).data

  if axis is None:
    a = tf.reshape(a, [-1])
    axis = 0

  axis_size = tf.shape(a, indices.dtype)[axis]
  if mode == 'clip':
    indices = tf.clip_by_value(indices, 0, axis_size-1)
  elif mode == 'wrap':
    indices = tf.math.floormod(indices, axis_size)
  else:
    raise ValueError("The 'raise' mode to take is not supported.")

  return utils.tensor_to_ndarray(tf.gather(a, indices, axis=axis)) 

def reshape(a, newshape, order='C'):
  """order argument can only b 'C' or 'F'."""
  if order not in {'C', 'F'}:
    raise ValueError('Unsupported order argument {}'.format(order))

  a = asarray(a)
  if isinstance(newshape, arrays_lib.ndarray):
    newshape = newshape.data
  if isinstance(newshape, int):
    newshape = [newshape]

  if order == 'F':
    r = tf.transpose(tf.reshape(tf.transpose(a.data), newshape[::-1]))
  else:
    r = tf.reshape(a.data, newshape)

  return utils.tensor_to_ndarray(r) 

def diagflat(v, k=0):
  """Returns a 2-d array with flattened `v` as diagonal.

  Args:
    v: array_like of any rank. Gets flattened when setting as diagonal. Could be
      an ndarray, a Tensor or any object that can be converted to a Tensor using
      `tf.convert_to_tensor`.
    k: Position of the diagonal. Defaults to 0, the main diagonal. Positive
      values refer to diagonals shifted right, negative values refer to
      diagonals shifted left.

  Returns:
    2-d ndarray.
  """
  v = asarray(v)
  return diag(tf.reshape(v.data, [-1]), k) 

def reshape(self, shape):
    """See tf.reshape."""
    return self._apply_op(lambda t: tf.reshape(t, shape)) 

def _setup_tensors(self):
    """Sets up tensors for efficient computations."""
    date_schedule = dates.PeriodicSchedule(
        start_date=self._start_date,
        end_date=self._maturity_date,
        tenor=self._reset_frequency).dates()

    # rates reset at the begining of coupon period
    reset_dates = date_schedule[:, :-1]
    # payments occur at the end of the coupon period
    payment_dates = date_schedule[:, 1:]
    daycount_fractions = rc.get_daycount_fraction(
        date_schedule[:, :-1],
        date_schedule[:, 1:],
        self._daycount_convention,
        dtype=self._dtype)
    contract_index = tf.repeat(
        tf.range(0, self._batch_size),
        payment_dates.shape.as_list()[-1])

    self._num_caplets = daycount_fractions.shape.as_list()[-1]
    # TODO(b/152164086): Use the functionality from dates library
    self._rate_term = tf.repeat(tf.cast(reset_dates[:, 0].days_until(
        payment_dates[:, 0]), dtype=self._dtype) / 365.0, self._num_caplets)
    self._reset_dates = dates.DateTensor.reshape(reset_dates, [-1])
    self._payment_dates = dates.DateTensor.reshape(payment_dates, [-1])
    self._accrual_start_dates = dates.DateTensor.reshape(reset_dates, [-1])
    self._accrual_end_dates = dates.DateTensor.reshape(payment_dates, [-1])
    self._daycount_fractions = tf.reshape(daycount_fractions, [-1])
    self._contract_index = contract_index
    self._strike = tf.repeat(self._strike, self._num_caplets)
    self._is_cap = tf.repeat(self._is_cap, self._num_caplets) 

def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
  """Returns an array with given shape and dtype filled with `fill_value`.

  Args:
    shape: A valid shape object. Could be a native python object or an object
       of type ndarray, numpy.ndarray or tf.TensorShape.
    fill_value: array_like. Could be an ndarray, a Tensor or any object that can
      be converted to a Tensor using `tf.convert_to_tensor`.
    dtype: Optional, defaults to dtype of the `fill_value`. The type of the
      resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
      `DType`.

  Returns:
    An ndarray.

  Raises:
    ValueError: if `fill_value` can not be broadcast to shape `shape`.
  """
  fill_value = asarray(fill_value, dtype=dtype)
  if utils.isscalar(shape):
    shape = tf.reshape(shape, [1])
  return arrays_lib.tensor_to_ndarray(tf.broadcast_to(fill_value.data, shape))


# Using doc only here since np full_like signature doesn't seem to have the
# shape argument (even though it exists in the documentation online). 

def predict_class(self, text_token_ids, action_panos):
    """Takes in an instruction and action and returns classifier outputs.

    Args:
      text_token_ids: Tensor of token indices for the input instruction.
      action_panos: Tensor of concatenated image panoramas.

    Returns:
      (class_outputs, predictions): Output of last layer of MLP and prediction.
    """
    text_enc_outputs, current_lstm_state = self.encode_instruction(
        text_token_ids)
    lstm_output, next_lstm_state = self.encode_action(current_lstm_state,
                                                      action_panos)
    lstm_output = tf.expand_dims(lstm_output, axis=1)
    batch_size = text_enc_outputs.shape[0]

    # c_text has shape [batch_size, 1, self._text_attention_size]
    c_text = self._text_attention([
        self._text_attention_project_hidden(lstm_output),
        self._text_attention_project_text(text_enc_outputs)
    ])
    # convert ListWrapper output of next_lstm_state to tuples
    result_state = []
    for one_state in next_lstm_state:
      result_state.append((one_state[0], one_state[1]))

    hidden_state = lstm_output
    c_visual = self._visual_attention([
        self._visual_attention_project_ctext(c_text),
        self._visual_attention_project_feature(action_panos),
    ])

    input_feature = tf.concat([hidden_state, c_text, c_visual], axis=2)
    class_outputs = self._mlp_layer(input_feature)
    class_outputs = tf.reshape(class_outputs, (batch_size, 2))
    predictions = tf.argmax(class_outputs, axis=-1)
    return (class_outputs, predictions) 

def test_paths_time_dependent(self):
    """Tests path properties with time dependent drift and variance."""

    def vol_fn(t):
      return tf.expand_dims(0.2 - 0.1 * tf.exp(-t), axis=-1)

    def variance_fn(t0, t1):
      # The instantaneous volatility is 0.2 - 0.1 e^(-t).
      tot_var = (t1 - t0) * 0.04 - (tf.exp(-2 * t1) - tf.exp(-2 * t0)) * 0.005
      tot_var += 0.04 * (tf.exp(-t1) - tf.exp(-t0))
      return tf.reshape(tot_var, [-1, 1, 1])

    process = BrownianMotion(
        dim=1, drift=0.1, volatility=vol_fn, total_covariance_fn=variance_fn)
    times = np.array([0.2, 0.33, 0.7, 0.9, 1.88])
    num_samples = 10000
    paths = self.evaluate(
        process.sample_paths(
            times,
            num_samples=num_samples,
            initial_state=np.array(0.1),
            seed=12134))

    self.assertArrayEqual(paths.shape, (num_samples, 5, 1))
    self.assertArrayNear(
        np.mean(paths, axis=0).reshape([-1]), 0.1 + times * 0.1, 0.05)

    covars = np.cov(paths.reshape([num_samples, 5]), rowvar=False)
    # Expected covariances are: cov_{ij} = variance_fn(0, min(t_i, t_j))
    min_times = np.minimum(times.reshape([-1, 1]),
                           times.reshape([1, -1])).reshape([-1])
    expected_covars = self.evaluate(
        variance_fn(tf.zeros_like(min_times), min_times))
    self.assertArrayNear(covars.reshape([-1]), expected_covars, 0.005) 

def test_sample_paths_wiener(self, use_xla):
    """Tests paths properties for Wiener process (dX = dW)."""

    def drift_fn(_, x):
      return tf.zeros_like(x)

    def vol_fn(_, x):
      return tf.expand_dims(tf.ones_like(x), -1)

    process = GenericItoProcess(dim=1, drift_fn=drift_fn, volatility_fn=vol_fn)
    times = np.array([0.1, 0.2, 0.3])
    num_samples = 10000

    @tf.function
    def fn():
      return process.sample_paths(
          times=times, num_samples=num_samples, seed=42, time_step=0.01)

    if use_xla:
      paths = self.evaluate(tf.xla.experimental.compile(fn))[0]
    else:
      paths = self.evaluate(fn())

    means = np.mean(paths, axis=0).reshape([-1])
    covars = np.cov(paths.reshape([num_samples, -1]), rowvar=False)
    expected_means = np.zeros((3,))
    expected_covars = np.minimum(times.reshape([-1, 1]), times.reshape([1, -1]))
    with self.subTest(name="Means"):
      self.assertAllClose(means, expected_means, rtol=1e-2, atol=1e-2)
    with self.subTest(name="Covar"):
      self.assertAllClose(covars, expected_covars, rtol=1e-2, atol=1e-2) 

def test_sample_paths_2d(self):
    """Tests path properties for 2-dimentional Ito process.

    We construct the following Ito processes.

    dX_1 = mu_1 sqrt(t) dt + s11 dW_1 + s12 dW_2
    dX_2 = mu_2 sqrt(t) dt + s21 dW_1 + s22 dW_2

    mu_1, mu_2 are constants.
    s_ij = a_ij t + b_ij

    For this process expected value at time t is (x_0)_i + 2/3 * mu_i * t^1.5.
    """
    mu = np.array([0.2, 0.7])
    a = np.array([[0.4, 0.1], [0.3, 0.2]])
    b = np.array([[0.33, -0.03], [0.21, 0.5]])

    def drift_fn(t, x):
      return mu * tf.sqrt(t) * tf.ones_like(x, dtype=t.dtype)

    def vol_fn(t, x):
      del x
      return (a * t + b) * tf.ones([2, 2], dtype=t.dtype)

    num_samples = 10000
    process = GenericItoProcess(dim=2, drift_fn=drift_fn, volatility_fn=vol_fn)
    times = np.array([0.1, 0.21, 0.32, 0.43, 0.55])
    x0 = np.array([0.1, -1.1])
    paths = self.evaluate(
        process.sample_paths(
            times,
            num_samples=num_samples,
            initial_state=x0,
            time_step=0.01,
            seed=12134))

    self.assertAllClose(paths.shape, (num_samples, 5, 2), atol=0)
    means = np.mean(paths, axis=0)
    times = np.reshape(times, [-1, 1])
    expected_means = x0 + (2.0 / 3.0) * mu * np.power(times, 1.5)
    self.assertAllClose(means, expected_means, rtol=1e-2, atol=1e-2) 

def _analytic_valuation(discount_rate_fn, model, strikes, expiries, maturities,
                        dim, is_call_options):
  """Performs analytic valuation."""

  discount_factor_expiry = tf.math.exp(
      -discount_rate_fn(expiries) * expiries)
  input_shape = expiries.shape
  variance = _bond_option_variance(
      model, tf.reshape(expiries, shape=[-1]), tf.reshape(maturities, [-1]),
      dim)
  variance = tf.reshape(variance, [dim] + input_shape)
  discount_factor_maturity = tf.math.exp(-discount_rate_fn(maturities) *
                                         maturities)
  forward_bond_price = discount_factor_maturity / discount_factor_expiry
  sqrt_variance = tf.math.sqrt(variance)
  d1 = (tf.expand_dims(tf.math.log(forward_bond_price / strikes), axis=0) +
        0.5 * variance) / sqrt_variance
  d2 = d1 - tf.math.sqrt(variance)
  option_value_call = (
      tf.expand_dims(discount_factor_maturity, axis=0) * _ncdf(d1) -
      tf.expand_dims(strikes * discount_factor_expiry, axis=0) * _ncdf(d2))
  option_value_put = (
      tf.expand_dims(strikes * discount_factor_expiry, axis=0) * _ncdf(-d2)
      - tf.expand_dims(discount_factor_maturity, axis=0) * _ncdf(-d1))

  is_call_options = tf.broadcast_to(is_call_options, [dim] + strikes.shape)
  option_value = tf.where(is_call_options, option_value_call,
                          option_value_put)

  # Make `dim` as the last dimension and return.
  return tf.transpose(
      option_value,
      perm=[i for i in range(1, len(option_value.shape.as_list()))] + [0])


# TODO(b/158501671): Clean-up this implementation. 

def _prepare_grid(times, *params):
  """Prepares grid of times for path generation.

  Args:
    times:  Rank 1 `Tensor` of increasing positive real values. The times at
      which the path points are to be evaluated.
    *params: Parameters of the Heston model. Either scalar `Tensor`s of the
      same `dtype` or instances of `PiecewiseConstantFunc`.

  Returns:
    Tuple `(all_times, mask)`.
    `all_times` is a 1-D real `Tensor` containing all points from 'times`, the
    uniform grid of points between `[0, times[-1]]` with grid size equal to
    `time_step`, and jump locations of piecewise constant parameters The
    `Tensor` is sorted in ascending order and may contain duplicates.
    `mask` is a boolean 1-D `Tensor` of the same shape as 'all_times', showing
    which elements of 'all_times' correspond to THE values from `times`.
    Guarantees that times[0]=0 and mask[0]=False.
  """
  additional_times = []
  for param in params:
    if hasattr(param, 'is_piecewise_constant'):
      if param.is_piecewise_constant:
        # Flatten all jump locations
        additional_times.append(tf.reshape(param.jump_locations(), [-1]))
  zeros = tf.constant([0], dtype=times.dtype)
  all_times = tf.concat([zeros] + [times] + additional_times, axis=0)
  additional_times_mask = [
      tf.zeros_like(times, dtype=tf.bool) for times in additional_times]
  mask = tf.concat([
      tf.cast(zeros, dtype=tf.bool),
      tf.ones_like(times, dtype=tf.bool)
  ] + additional_times_mask, axis=0)
  perm = tf.argsort(all_times, stable=True)
  all_times = tf.gather(all_times, perm)
  mask = tf.gather(mask, perm)
  return all_times, mask 

def _normalize_indexes(self, indexes):
    indexes = math_ops.lower_bound(indexes, 0)
    if isinstance(self.index_ranges, int):
      indexes = math_ops.upper_bound(indexes, self.index_ranges - 1)
    else:
      axes = [1] * indexes.shape.rank
      axes[self.channel_axis] = len(self.index_ranges)
      bounds = tf.reshape([s - 1 for s in self.index_ranges], axes)
      indexes = math_ops.upper_bound(indexes, bounds)
    return indexes 

def encode_pano(self, pano_name, pano_path, num_views):
    """Takes in a pano name and returns image features.

    Args:
      pano_name: String specifying the file name of the pano.
      pano_path: Path to directory that contains panoramic images.
      num_views: Number of viewpoints.

    Returns:
      image_features: Image features for the input panoramic image.
    """
    image_features = self._get_image_features(
        os.path.join(pano_path, '{}_viewpoints_proto'.format(pano_name)))

    return np.reshape(image_features.value, (1, num_views, -1)) 

def setUp(self):
    super(ReferenceTest, self).setUp()
    cpus = tf.config.experimental.list_physical_devices('CPU')
    tf.config.experimental.set_virtual_device_configuration(
        cpus[0], [tf.config.experimental.VirtualDeviceConfiguration()] * 2)

    strategy = tf.distribute.MirroredStrategy()
    dataset = tf.data.Dataset.from_tensor_slices(
        tf.reshape(tf.range(40), (10, 4)))

    self.ds = strategy
    self.dds = strategy.experimental_distribute_dataset(dataset) 

def _lower_tail(self, tail_mass):
    tail = helpers.estimate_tails(
        self._logits_cumulative, -tf.math.log(2 / tail_mass - 1),
        tf.constant([self.batch_shape.num_elements(), 1, 1], tf.int32),
        self.dtype)
    return tf.reshape(tail, self.batch_shape_tensor()) 

def _upper_tail(self, tail_mass):
    tail = helpers.estimate_tails(
        self._logits_cumulative, tf.math.log(2 / tail_mass - 1),
        tf.constant([self.batch_shape.num_elements(), 1, 1], tf.int32),
        self.dtype)
    return tf.reshape(tail, self.batch_shape_tensor()) 

def _compute_indexes(self, broadcast_shape):
    # TODO(jonycgn, ssjhv): Investigate broadcasting in range coding op.
    prior_size = functools.reduce(lambda x, y: x * y, self.prior_shape, 1)
    indexes = tf.range(prior_size, dtype=tf.int32)
    indexes = tf.reshape(indexes, self.prior_shape)
    indexes = tf.broadcast_to(
        indexes, tf.concat([broadcast_shape, self.prior_shape], 0))
    return indexes 

def tile(a, reps):
  a = array_ops.array(a).data
  reps = array_ops.array(reps, dtype=tf.int32).reshape([-1]).data

  a_rank = tf.rank(a)
  reps_size = tf.size(reps)
  reps = tf.pad(
      reps, [[tf.math.maximum(a_rank - reps_size, 0), 0]],
      constant_values=1)
  a_shape = tf.pad(
      tf.shape(a), [[tf.math.maximum(reps_size - a_rank, 0), 0]],
      constant_values=1)
  a = tf.reshape(a, a_shape)

  return arrays.tensor_to_ndarray(tf.tile(a, reps)) 

def _generate_examples(self, path):
        """Yields examples."""

        clean_key = "clean"
        adversarial_key = "adversarial"

        def _parse(serialized_example):
            ds_features = {
                "height": tf.io.FixedLenFeature([], tf.int64),
                "width": tf.io.FixedLenFeature([], tf.int64),
                "label": tf.io.FixedLenFeature([], tf.int64),
                "adv-image": tf.io.FixedLenFeature([], tf.string),
                "clean-image": tf.io.FixedLenFeature([], tf.string),
            }
            example = tf.io.parse_single_example(serialized_example, ds_features)

            img_clean = tf.io.decode_raw(example["clean-image"], tf.float32)
            img_adv = tf.io.decode_raw(example["adv-image"], tf.float32)
            # float values are integers in [0.0, 255.0] for clean and adversarial
            img_clean = tf.cast(img_clean, tf.uint8)
            img_clean = tf.reshape(img_clean, (example["height"], example["width"], 3))
            img_adv = tf.cast(img_adv, tf.uint8)
            img_adv = tf.reshape(img_adv, (example["height"], example["width"], 3))
            return {clean_key: img_clean, adversarial_key: img_adv}, example["label"]

        ds = tf.data.TFRecordDataset(filenames=[path])
        ds = ds.map(lambda x: _parse(x))
        default_graph = tf.compat.v1.keras.backend.get_session().graph
        ds = tfds.as_numpy(ds, graph=default_graph)

        for i, (img, label) in enumerate(ds):
            yield str(i), {
                "images": img,
                "label": label,
            } 

def ravel(a):  # pylint: disable=missing-docstring
  a = asarray(a)
  if a.ndim == 1:
    return a
  return utils.tensor_to_ndarray(tf.reshape(a.data, [-1]))