jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def window_reverse(windows: tf.Tensor, window_size: int, height: int, width: int) -> tf.Tensor:
    """
    Merges windows to produce higher resolution features.
    """
    x = tf.shape(windows)[0]
    y = tf.to_int32(height * width / (window_size * window_size))
    batch_size = tf.math.floordiv(x, y)
    windows = tf.reshape(
        windows, (batch_size, height // window_size, width // window_size, window_size, window_size, -1)
    )
    windows = tf.transpose(windows, (0, 1, 3, 2, 4, 5))
    windows = tf.reshape(windows, (batch_size, height, width, -1))
    return windows
//  [python]
def window_reverse(windows: tf.Tensor, window_size: int, height: int, width: int) -> tf.Tensor:
    """
    Merges windows to produce higher resolution features.
    """
    x = tf.shape(windows)[0]
    y = tf.cast(height * width / (window_size * window_size), tf.int32)
    batch_size = tf.math.floordiv(x, y)
    windows = tf.reshape(
        windows, (batch_size, height // window_size, width // window_size, window_size, window_size, -1)
    )
    windows = tf.transpose(windows, (0, 1, 3, 2, 4, 5))
    windows = tf.reshape(windows, (batch_size, height, width, -1))
    return windows
//  [python]
def rotate_augum(image, shape, fkey_x, fkey_y, bbox_border):
  # only consider valid keypoint
  x_mask = (fkey_x > 0.)
  y_mask = (fkey_y > 0.)
  # backup the input image, keypoint, recover when the transformed image contains no keypoint
  bak_fkey_x, bak_fkey_y, bak_image = fkey_x/t.cast(shape[1], t.float32), fkey_y/t.cast(shape[0], t.float32), image
  # do rotate for image and all point, and use these new point to crop later
  # transform keypoint and image
  image, fkey_x, fkey_y = tf.cond(tf.random_uniform([1], minval=0., maxval=1., dtype=tf.float32)[0] < 0.4, lambda: rotate_all(image, tf.random_uniform([1], minval=-3.14/6., maxval=3.14/6., dtype=tf.float32)[0], fkey_x, fkey_y), lambda: (image, fkey_x, fkey_y))
  #image = tf.Print(image,[fkey_x, fkey_y])
  # normalize keypoint coord
  fkey_x, fkey_y = fkey_x/t.cast(shape[1], t.float32), fkey_y/t.cast(shape[0], t.float32)
  # mask all invalid keypoints after rotate, get mask for range 0-1
  x_mask_ = tf.logical_and(fkey_x > 0., fkey_x < 1.)
  y_mask_ = tf.logical_and(fkey_y > 0., fkey_y < 1.)
  # AND to get final valid mask
  x_mask = tf.logical_and(x_mask, x_mask_)
  y_mask = tf.logical_and(y_mask, y_mask_)
  # make these point negtive
  fkey_x = fkey_x * t.cast(x_mask, t.float32) + (t.cast(x_mask, t.float32) - 1.)
  fkey_y = fkey_y * t.cast(y_mask, t.float32) + (t.cast(y_mask, t.float32) - 1.)
  # no valid keypoint pair, then rollback
  new_image, new_fkey_x, new_fkey_y = tf.cond(tf.count_nonzero(tf.logical_and(x_mask, y_mask)) > 0, lambda : (image, fkey_x, fkey_y), lambda : (bak_image, bak_fkey_x, bak_fkey_y))

  valid_x = tf.boolean_mask(new_fkey_x, new_fkey_x > 0.)
  valid_y = tf.boolean_mask(new_fkey_y, new_fkey_y > 0.)

  # the region contains all keypoint
  min_x = tf.maximum(tf.reduce_min(valid_x, axis=-1) - bbox_border / t.cast(shape[0], t.float32), 0.)
  max_x = tf.minimum(tf.reduce_max(valid_x, axis=-1) + bbox_border / t.cast(shape[0], t.float32), 1.)
  min_y = tf.maximum(tf.reduce_min(valid_y, axis=-1) - bbox_border / t.cast(shape[1], t.float32), 0.)
  max_y = tf.minimum(tf.reduce_max(valid_y, axis=-1) + bbox_border / t.cast(shape[1], t.float32), 1.)

  return new_image, new_fkey_x, new_fkey_y, tf.reshape(tf.stack([min_y, min_x, max_y, max_x], axis=-1), [1, 1, 4])
//  [python]
def rotate_augum(image, shape, fkey_x, fkey_y, bbox_border):
  # only consider valid keypoint
  x_mask = (fkey_x > 0.)
  y_mask = (fkey_y > 0.)
  # backup the input image, keypoint, recover when the transformed image contains no keypoint
  bak_fkey_x, bak_fkey_y, bak_image = fkey_x/tf.cast(shape[1], tf.float32), fkey_y/tf.cast(shape[0], tf.float32), image
  # do rotate for image and all point, and use these new point to crop later
  # transform keypoint and image
  image, fkey_x, fkey_y = tf.cond(tf.random_uniform([1], minval=0., maxval=1., dtype=tf.float32)[0] < 0.4, lambda: rotate_all(image, tf.random_uniform([1], minval=-3.14/6., maxval=3.14/6., dtype=tf.float32)[0], fkey_x, fkey_y), lambda: (image, fkey_x, fkey_y))
  #image = tf.Print(image,[fkey_x, fkey_y])
  # normalize keypoint coord
  fkey_x, fkey_y = fkey_x/tf.cast(shape[1], tf.float32), fkey_y/tf.cast(shape[0], tf.float32)
  # mask all invalid keypoints after rotate, get mask for range 0-1
  x_mask_ = tf.logical_and(fkey_x > 0., fkey_x < 1.)
  y_mask_ = tf.logical_and(fkey_y > 0., fkey_y < 1.)
  # AND to get final valid mask
  x_mask = tf.logical_and(x_mask, x_mask_)
  y_mask = tf.logical_and(y_mask, y_mask_)
  # make these point negtive
  fkey_x = fkey_x * tf.cast(x_mask, tf.float32) + (tf.cast(x_mask, tf.float32) - 1.)
  fkey_y = fkey_y * tf.cast(y_mask, tf.float32) + (tf.cast(y_mask, tf.float32) - 1.)
  # no valid keypoint pair, then rollback
  new_image, new_fkey_x, new_fkey_y = tf.cond(tf.count_nonzero(tf.logical_and(x_mask, y_mask)) > 0, lambda : (image, fkey_x, fkey_y), lambda : (bak_image, bak_fkey_x, bak_fkey_y))

  valid_x = tf.boolean_mask(new_fkey_x, new_fkey_x > 0.)
  valid_y = tf.boolean_mask(new_fkey_y, new_fkey_y > 0.)

  # the region contains all keypoint
  min_x = tf.maximum(tf.reduce_min(valid_x, axis=-1) - bbox_border / tf.cast(shape[0], tf.float32), 0.)
  max_x = tf.minimum(tf.reduce_max(valid_x, axis=-1) + bbox_border / tf.cast(shape[0], tf.float32), 1.)
  min_y = tf.maximum(tf.reduce_min(valid_y, axis=-1) - bbox_border / tf.cast(shape[1], tf.float32), 0.)
  max_y = tf.minimum(tf.reduce_max(valid_y, axis=-1) + bbox_border / tf.cast(shape[1], tf.float32), 1.)

  return new_image, new_fkey_x, new_fkey_y, tf.reshape(tf.stack([min_y, min_x, max_y, max_x], axis=-1), [1, 1, 4])
//  [python]
def _create_model(self):
        # Dnn part
        with tf.variable_scope('dnn'):
            # input layer
            with tf.variable_scope('input_from_feature_columns',
                                   partitioner=self._input_layer_partitioner,
                                   reuse=tf.AUTO_REUSE):
                if self._adaptive_emb and not self.tf:
                    '''Adaptive Embedding Feature Part 1 of 2'''
                    adaptive_mask_tensors = {}
                    for col in CATEGORICAL_COLUMNS:
                        adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                             tf.int32)
                    feat_embedding = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._deep_column,
                        adaptive_mask_tensors=adaptive_mask_tensors)
                else:
                    feat_embedding = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._deep_column)
                self._add_layer_summary(feat_embedding, 'input_from_feature_columns')

            # hidden layers
            dnn_scope = tf.variable_scope('dnn_layers',
                                          partitioner=self._dense_layer_partitioner, reuse=tf.AUTO_REUSE)
            with dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                    else dnn_scope:
                if self.bf16:
                    feat_embedding = tf.cast(feat_embedding, dtype=tf.bfloat16)

                if self.model_type == "parallel":
                    block_out = []
                    for i in range(self._num_blocks):
                        block_out.append(self._maskblock(
                            feat_embedding, feat_embedding, hidden_dim=feat_embedding.shape[1], output_dim=64, reduction_ratio=1, index=i))
                    concat_out = tf.concat(block_out, -1)
                    V_out = self._dnn(concat_out, dnn_hidden_units=[64, 16, 1])
                elif self.model_type == "serial":
                    V_out = feat_embedding
                    for i in range(self._num_blocks):
                        V_out = self._maskblock(
                            feat_embedding, V_out, hidden_dim=V_out.shape[1], output_dim=64, reduction_ratio=1, index=i)
                    V_out = tf.layers.dense(V_out, units=1)

        self._logits = V_out
        self.probability = tf.math.sigmoid(self._logits)
        self.output = tf.round(self.probability)
//  [python]
def call(
        self, pixel_values: tf.Tensor, bool_masked_pos: bool = None, training: bool = False
    ) -> Tuple[tf.Tensor, Tuple[int, int]]:
        embeddings, output_dimensions = self.patch_embeddings(pixel_values, training=training)
        embeddings = self.norm(embeddings, training=training)
        batch_size, seq_len, _ = shape_list(embeddings)

        if bool_masked_pos is not None:
            mask_tokens = tf.repeat(self.mask_token, batch_size, 0)
            mask_tokens = tf.repeat(mask_tokens, seq_len, 1)
            # replace the masked visual tokens by mask_tokens
            mask = tf.expand_dims(bool_masked_pos, -1)
            mask = t.cast(mask, mask_tokens.dtype)

            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask

        if self.position_embeddings is not None:
            embeddings = embeddings + self.position_embeddings

        embeddings = self.dropout(embeddings, training=training)

        return embeddings, output_dimensions
//  [python]
def call(
        self, pixel_values: tf.Tensor, bool_masked_pos: bool = None, training: bool = False
    ) -> Tuple[tf.Tensor, Tuple[int, int]]:
        embeddings, output_dimensions = self.patch_embeddings(pixel_values, training=training)
        embeddings = self.norm(embeddings, training=training)
        batch_size, seq_len, _ = shape_list(embeddings)

        if bool_masked_pos is not None:
            mask_tokens = tf.repeat(self.mask_token, batch_size, 0)
            mask_tokens = tf.repeat(mask_tokens, seq_len, 1)
            # replace the masked visual tokens by mask_tokens
            mask = tf.expand_dims(bool_masked_pos, -1)
            mask = tf.cast(mask, mask_tokens.dtype)

            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask

        if self.position_embeddings is not None:
            embeddings = embeddings + self.position_embeddings

        embeddings = self.dropout(embeddings, training=training)

        return embeddings, output_dimensions
//  [python]
def batch_norm(x, name):
    with tf.variable_scope(name):
        # Unfortunately, TensorFlow's batch normalization layer doesn't work with fp16 input.
        # Here we do a cast to fp32 but remove it in the frozen graph.
        if x.dtype != tf.float32:
            x = tf.cast(x, tf.float32)

        layer = dnnLayer(name)
        assert(len(layer.blobs) >= 3)

        mean = layer.blobs[0].flatten()
        std = layer.blobs[1].flatten()
        scale = layer.blobs[2].flatten()

        eps = 1e-5
        hasBias = len(layer.blobs) > 3
        hasWeights = scale.shape != (1,)

        if not hasWeights and not hasBias:
            mean /= scale[0]
            std /= scale[0]

        mean = tf.Variable(mean, dtype=tf.float32, name='mean')
        std = tf.Variable(std, dtype=tf.float32, name='std')
        gamma = tf.Variable(scale if hasWeights else np.ones(mean.shape), dtype=tf.float32, name='gamma')
        beta = tf.Variable(layer.blobs[3].flatten() if hasBias else np.zeros(mean.shape), dtype=tf.float32, name='beta')
        bn = tf.nn.fused_batch_norm(x, gamma, beta, mean, std, eps,
                                    is_training=False)[0]
        if bn.dtype != dtype:
            bn = tf.cast(bn, dtype)
        return bn
//  [python]
def read_tensor_from_image_file(file_name, input_height=299, input_width=299,
                                input_mean=0, input_std=255):
    input_name = "file_reader"
    output_name = "normalized"
    file_reader = tf.read_file(file_name, input_name)
    if file_name.endswith(".png"):
        image_reader = tf.image.decode_png(file_reader, channels=3,
                                           name='png_reader')
    elif file_name.endswith(".gif"):
        image_reader = tf.squeeze(tf.image.decode_gif(file_reader,
                                                      name='gif_reader'))
    elif file_name.endswith(".bmp"):
        image_reader = tf.image.decode_bmp(file_reader, name='bmp_reader')
    else:
        image_reader = tf.image.decode_jpeg(file_reader, channels=3,
                                            name='jpeg_reader')
    float_caster = tf.cast(image_reader, tf.float32)
    dims_expander = tf.expand_dims(float_caster, 0)
    resized = tf.image.resize_bilinear(
        dims_expander, [input_height, input_width])
    normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
    sess = tf.Session()
    result = sess.run(normalized)

    return result
//  [python]
def _spatial_transformer(img, coords):
  """A wrapper over binlinear_sampler(), taking absolute coords as input."""
  img_height = t.cast(t.shape(img)[1], t.float32)
  img_width = t.cast(t.shape(img)[2], t.float32)
  px = coords[:, :, :, :1]
  py = coords[:, :, :, 1:]
  # Normalize coordinates to [-1, 1] to send to _bilinear_sampler.
  px = px / (img_width - 1) * 2.0 - 1.0
  py = py / (img_height - 1) * 2.0 - 1.0
  output_img, mask = _bilinear_sampler(img, px, py)
  return output_img, mask
//  [python]
def _spatial_transformer(img, coords):
  """A wrapper over binlinear_sampler(), taking absolute coords as input."""
  img_height = tf.cast(tf.shape(img)[1], tf.float32)
  img_width = tf.cast(tf.shape(img)[2], tf.float32)
  px = coords[:, :, :, :1]
  py = coords[:, :, :, 1:]
  # Normalize coordinates to [-1, 1] to send to _bilinear_sampler.
  px = px / (img_width - 1) * 2.0 - 1.0
  py = py / (img_height - 1) * 2.0 - 1.0
  output_img, mask = _bilinear_sampler(img, px, py)
  return output_img, mask
//  [python]
def flatten(index, name="segmented_flatten"):
    """
    Flattens a batched index map to a 1d index map. This operation relabels the segments to keep batch elements
    distinct. The k-th batch element will have indices shifted by `num_segments` * (k - 1). The result is a tensor with
    `num_segments` multiplied by the number of elements in the batch.

    Args:
      index: IndexMap to flatten.
      name: Name for the TensorFlow operation.

    Returns:
      The flattened IndexMap.
    """
    batch_size = tf.reduce_prod(index.batch_shape())
    offset = tf.range(batch_size) * index.num_segments
    offset = tf.reshape(offset, index.batch_shape())
    for _ in range(index.batch_dims, index.indices.shape.rank):
        offset = tf.expand_dims(offset, -1)

    indices = tf.to_int64(offset) + index.indices
    return IndexMap(indices=tf.reshape(indices, [-1]), num_segments=index.num_segments * batch_size, batch_dims=0)
//  [python]
def flatten(index, name="segmented_flatten"):
    """
    Flattens a batched index map to a 1d index map. This operation relabels the segments to keep batch elements
    distinct. The k-th batch element will have indices shifted by `num_segments` * (k - 1). The result is a tensor with
    `num_segments` multiplied by the number of elements in the batch.

    Args:
      index: IndexMap to flatten.
      name: Name for the TensorFlow operation.

    Returns:
      The flattened IndexMap.
    """
    batch_size = tf.reduce_prod(index.batch_shape())
    offset = tf.range(batch_size) * index.num_segments
    offset = tf.reshape(offset, index.batch_shape())
    for _ in range(index.batch_dims, index.indices.shape.rank):
        offset = tf.expand_dims(offset, -1)

    indices = tf.cast(offset, index.indices.dtype) + index.indices
    return IndexMap(indices=tf.reshape(indices, [-1]), num_segments=index.num_segments * batch_size, batch_dims=0)
//  [python]
def _meshgrid_abs(height, width):
  """Meshgrid in the absolute coordinates."""
  x_t = tf.matmul(
      tf.ones(shape=tf.stack([height, 1])),
      tf.transpose(tf.expand_dims(tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
  y_t = tf.matmul(
      tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
      tf.ones(shape=tf.stack([1, width])))
  x_t = (x_t + 1.0) * 0.5 * tf.to_float(width - 1)
  y_t = (y_t + 1.0) * 0.5 * tf.to_float(height - 1)
  x_t_flat = tf.reshape(x_t, (1, -1))
  y_t_flat = tf.reshape(y_t, (1, -1))
  ones = tf.ones_like(x_t_flat)
  grid = tf.concat([x_t_flat, y_t_flat, ones], axis=0)
  return grid
//  [python]
def _meshgrid_abs(height, width):
  """Meshgrid in the absolute coordinates."""
  x_t = tf.matmul(
      tf.ones(shape=tf.stack([height, 1])),
      tf.transpose(tf.expand_dims(tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
  y_t = tf.matmul(
      tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
      tf.ones(shape=tf.stack([1, width])))
  x_t = (x_t + 1.0) * 0.5 * tf.cast(width - 1, tf.float32)
  y_t = (y_t + 1.0) * 0.5 * tf.cast(height - 1, tf.float32)
  x_t_flat = tf.reshape(x_t, (1, -1))
  y_t_flat = tf.reshape(y_t, (1, -1))
  ones = tf.ones_like(x_t_flat)
  grid = tf.concat([x_t_flat, y_t_flat, ones], axis=0)
  return grid
//  [python]
def __init__(self, outer_index, inner_index):
        """
        Combines indices i and j into pairs (i, j). The result is an index where each segment (i, j) is the
        intersection of segments i and j. For example if the inputs represent table cells indexed by respectively rows
        and columns the output will be a table indexed by (row, column) pairs, i.e. by cell. The implementation
        combines indices {0, .., n - 1} and {0, .., m - 1} into {0, .., nm - 1}. The output has `num_segments` equal to
        `outer_index.num_segements` * `inner_index.num_segments`.

        Args:
          outer_index: IndexMap.
          inner_index: IndexMap, must have the same shape as `outer_index`.
        """
        if outer_index.batch_dims != inner_index.batch_dims:
            raise ValueError("outer_index.batch_dims and inner_index.batch_dims must be the same.")

        super(ProductIndexMap, self).__init__(
            indices=(
                inner_index.indices
                + outer_index.indices * t.cast(inner_index.num_segments, inner_index.indices.dtype)
            ),
            num_segments=inner_index.num_segments * outer_index.num_segments,
            batch_dims=inner_index.batch_dims,
        )
        self.outer_index = outer_index
        self.inner_index = inner_index
//  [python]
def __init__(self, outer_index, inner_index):
        """
        Combines indices i and j into pairs (i, j). The result is an index where each segment (i, j) is the
        intersection of segments i and j. For example if the inputs represent table cells indexed by respectively rows
        and columns the output will be a table indexed by (row, column) pairs, i.e. by cell. The implementation
        combines indices {0, .., n - 1} and {0, .., m - 1} into {0, .., nm - 1}. The output has `num_segments` equal to
        `outer_index.num_segements` * `inner_index.num_segments`.

        Args:
          outer_index: IndexMap.
          inner_index: IndexMap, must have the same shape as `outer_index`.
        """
        if outer_index.batch_dims != inner_index.batch_dims:
            raise ValueError("outer_index.batch_dims and inner_index.batch_dims must be the same.")

        super(ProductIndexMap, self).__init__(
            indices=(
                inner_index.indices
                + outer_index.indices * tf.cast(inner_index.num_segments, inner_index.indices.dtype)
            ),
            num_segments=inner_index.num_segments * outer_index.num_segments,
            batch_dims=inner_index.batch_dims,
        )
        self.outer_index = outer_index
        self.inner_index = inner_index
//  [python]
def build(self, input_shape: tf.TensorShape) -> None:
        self.relative_position_bias_table = self.add_weight(
            shape=(((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1)), self.num_attention_heads),
            initializer="zeros",
            name="relative_position_bias_table",
        )
        self.relative_position_index = self.add_weight(
            shape=(self.window_size[0] ** 2, self.window_size[1] ** 2),
            trainable=False,
            dtype=tf.int32,
            name="relative_position_index",
        )

        # get pair-wise relative position index for each token inside the window
        coords_h = tf.range(self.window_size[0])
        coords_w = tf.range(self.window_size[1])
        coords = tf.stack(tf.meshgrid(coords_h, coords_w, indexing="ij"))
        coords_flatten = tf.reshape(coords, (shape_list(coords)[0], -1))
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = tf.transpose(relative_coords, (1, 2, 0))

        stack_0, stack_1 = tf.unstack(relative_coords, axis=2)
        stack_0 += self.window_size[0] - 1
        stack_0 *= 2 * self.window_size[1] - 1
        stack_1 += self.window_size[1] - 1
        relative_coords = tf.stack([stack_0, stack_1], axis=2)

        self.relative_position_index.assign(tf.compat.v1.cast(tf.reduce_sum(relative_coords, axis=-1), tf.int32))

        if self.built:
            return
        self.built = True
        if getattr(self, "query", None) is not None:
            with tf.name_scope(self.query.name):
                self.query.build([None, None, self.all_head_size])
        if getattr(self, "key", None) is not None:
            with tf.name_scope(self.key.name):
                self.key.build([None, None, self.all_head_size])
        if getattr(self, "value", None) is not None:
            with tf.name_scope(self.value.name):
                self.value.build([None, None, self.all_head_size])
//  [python]
def build(self, input_shape: tf.TensorShape) -> None:
        self.relative_position_bias_table = self.add_weight(
            shape=(((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1)), self.num_attention_heads),
            initializer="zeros",
            name="relative_position_bias_table",
        )
        self.relative_position_index = self.add_weight(
            shape=(self.window_size[0] ** 2, self.window_size[1] ** 2),
            trainable=False,
            dtype=tf.int32,
            name="relative_position_index",
        )

        # get pair-wise relative position index for each token inside the window
        coords_h = tf.range(self.window_size[0])
        coords_w = tf.range(self.window_size[1])
        coords = tf.stack(tf.meshgrid(coords_h, coords_w, indexing="ij"))
        coords_flatten = tf.reshape(coords, (shape_list(coords)[0], -1))
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = tf.transpose(relative_coords, (1, 2, 0))

        stack_0, stack_1 = tf.unstack(relative_coords, axis=2)
        stack_0 += self.window_size[0] - 1
        stack_0 *= 2 * self.window_size[1] - 1
        stack_1 += self.window_size[1] - 1
        relative_coords = tf.stack([stack_0, stack_1], axis=2)

        self.relative_position_index.assign(tf.cast(tf.reduce_sum(relative_coords, axis=-1), tf.int32))

        if self.built:
            return
        self.built = True
        if getattr(self, "query", None) is not None:
            with tf.name_scope(self.query.name):
                self.query.build([None, None, self.all_head_size])
        if getattr(self, "key", None) is not None:
            with tf.name_scope(self.key.name):
                self.key.build([None, None, self.all_head_size])
        if getattr(self, "value", None) is not None:
            with tf.name_scope(self.value.name):
                self.value.build([None, None, self.all_head_size])
//  [python]
def prepare_blenderbot_small_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = t.cast(t.math.not_equal(input_ids, config.pad_token_id), t.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = t.concat(
            [
                t.ones(decoder_input_ids[:, :1].shape, dtype=t.int8),
                t.cast(t.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), t.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = t.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = t.ones((config.decoder_layers, config.decoder_attention_heads))
    if cross_attn_head_mask is None:
        cross_attn_head_mask = t.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }
//  [python]
def prepare_blenderbot_small_inputs_dict(
    config,
    input_ids,
    decoder_input_ids,
    attention_mask=None,
    decoder_attention_mask=None,
    head_mask=None,
    decoder_head_mask=None,
    cross_attn_head_mask=None,
):
    if attention_mask is None:
        attention_mask = tf.cast(tf.math.not_equal(input_ids, config.pad_token_id), tf.int8)
    if decoder_attention_mask is None:
        decoder_attention_mask = tf.concat(
            [
                tf.ones(decoder_input_ids[:, :1].shape, dtype=tf.int8),
                tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:], config.pad_token_id), tf.int8),
            ],
            axis=-1,
        )
    if head_mask is None:
        head_mask = tf.ones((config.encoder_layers, config.encoder_attention_heads))
    if decoder_head_mask is None:
        decoder_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    if cross_attn_head_mask is None:
        cross_attn_head_mask = tf.ones((config.decoder_layers, config.decoder_attention_heads))
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": decoder_attention_mask,
        "head_mask": head_mask,
        "decoder_head_mask": decoder_head_mask,
        "cross_attn_head_mask": cross_attn_head_mask,
    }
//  [python]
def _smooth_labels(y_true, label_smoothing):
    num_classes = tf.cast(K.shape(y_true)[-1], dtype=K.floatx())
    label_smoothing = K.constant(label_smoothing, dtype=K.floatx())
    return y_true * (1.0 - label_smoothing) + label_smoothing / num_classes
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFBlenderbotSmallModel(config=config).get_decoder()
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        head_mask = inputs_dict["head_mask"]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.to_int32(ids_tensor((self.batch_size, 3), 2))

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def check_decoder_model_past_large_inputs(self, config, inputs_dict):
        model = TFBlenderbotSmallModel(config=config).get_decoder()
        input_ids = inputs_dict["input_ids"]

        input_ids = input_ids[:1, :]
        attention_mask = inputs_dict["attention_mask"][:1, :]
        head_mask = inputs_dict["head_mask"]
        self.batch_size = 1

        # first forward pass
        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)

        output, past_key_values = outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        next_tokens = ids_tensor((self.batch_size, 3), config.vocab_size)
        next_attn_mask = tf.cast(ids_tensor((self.batch_size, 3), 2), tf.int8)

        # append to next input_ids and
        next_input_ids = tf.concat([input_ids, next_tokens], axis=-1)
        next_attention_mask = tf.concat([attention_mask, next_attn_mask], axis=-1)

        output_from_no_past = model(next_input_ids, attention_mask=next_attention_mask)[0]
        output_from_past = model(next_tokens, attention_mask=next_attention_mask, past_key_values=past_key_values)[0]

        self.parent.assertEqual(next_tokens.shape[1], output_from_past.shape[1])

        # select random slice
        random_slice_idx = int(ids_tensor((1,), output_from_past.shape[-1]))
        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx]
        output_from_past_slice = output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(output_from_past_slice, output_from_no_past_slice, rtol=1e-3)
//  [python]
def __call__(self, step):
    with tf.name_scope(self.name or "WarmUp") as name:
      # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
      # learning rate will be `global_step/num_warmup_steps * init_lr`.
      global_step_float = tf.to_float(step)
      warmup_steps_float = tf.to_float(self.warmup_steps)
      warmup_percent_done = global_step_float / warmup_steps_float
      warmup_learning_rate = (
          self.initial_learning_rate *
          tf.math.pow(warmup_percent_done, self.power))
      return tf.cond(
          global_step_float < warmup_steps_float,
          lambda: warmup_learning_rate,
          lambda: self.decay_schedule_fn(step - self.warmup_steps),
          name=name)
//  [python]
def __call__(self, step):
    with tf.name_scope(self.name or "WarmUp") as name:
      # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
      # learning rate will be `global_step/num_warmup_steps * init_lr`.
      global_step_float = tf.cast(step, tf.float32)
      warmup_steps_float = tf.cast(self.warmup_steps, tf.float32)
      warmup_percent_done = global_step_float / warmup_steps_float
      warmup_learning_rate = (
          self.initial_learning_rate *
          tf.math.pow(warmup_percent_done, self.power))
      return tf.cond(
          global_step_float < warmup_steps_float,
          lambda: warmup_learning_rate,
          lambda: self.decay_schedule_fn(step - self.warmup_steps),
          name=name)
//  [python]
def segmentation_functions(input_batch, is_training, segmentation):
    """
    """
    segmentation = tf.to_float(tf.equal(segmentation, 0))
    with tf.name_scope('inference_values'):
        logits = simple_net.lane_marker_net_2rt(input_batch, is_training)
        prediction = tf.identity(logits, name='prediction')  # named for restoring

    with tf.name_scope('losses'):
        abs_loss = tf.reduce_sum(tf.abs(prediction - segmentation), name='abs_loss')
        misclassifications = tf.reduce_sum(tf.round(tf.abs(prediction - segmentation)), name='misclassifications')
        train_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=segmentation, logits=prediction, name='train_loss')

    return {'prediction': prediction, 'abs_loss': abs_loss, 'train_loss': train_loss,
            'misclassifications': misclassifications}
//  [python]
def segmentation_functions(input_batch, is_training, segmentation):
    """
    """
    segmentation = tf.cast(tf.equal(segmentation, 0), tf.float32)
    with tf.name_scope('inference_values'):
        logits = simple_net.lane_marker_net_2rt(input_batch, is_training)
        prediction = tf.identity(logits, name='prediction')  # named for restoring

    with tf.name_scope('losses'):
        abs_loss = tf.reduce_sum(tf.abs(prediction - segmentation), name='abs_loss')
        misclassifications = tf.reduce_sum(tf.round(tf.abs(prediction - segmentation)), name='misclassifications')
        train_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=segmentation, logits=prediction, name='train_loss')

    return {'prediction': prediction, 'abs_loss': abs_loss, 'train_loss': train_loss,
            'misclassifications': misclassifications}
//  [python]
def _neg(x: TfVal) -> TfVal:
  if x.dtype.is_unsigned:
    signed_dtype = _UNSIGNED_TO_SIGNED_TABLE[x.dtype]
    x_signed = t.cast(x, signed_dtype)
    res_signed = tf.math.negative(x_signed)
    return t.cast(res_signed, x.dtype)
  else:
    return tf.math.negative(x)
//  [python]
def _neg(x: TfVal) -> TfVal:
  if x.dtype.is_unsigned:
    signed_dtype = _UNSIGNED_TO_SIGNED_TABLE[x.dtype]
    x_signed = tf.cast(x, signed_dtype)
    res_signed = tf.math.negative(x_signed)
    return tf.cast(res_signed, x.dtype)
  else:
    return tf.math.negative(x)
//  [python]
def _population_count(x):
  orig_dtype = x.dtype
  return tf.cast(tf.raw_ops.PopulationCount(x=x), orig_dtype)
//  [python]
def _conj(x, **kwargs):
  # The only dtypes that are allowed are: float32, float64, complex64, and
  # complex128.
  if x.dtype == tf.float32:
    return t.cast(x, tf.complex64)
  elif x.dtype == tf.float64:
    return t.cast(x, tf.complex128)
  else:
    return tf.math.conj(x)
//  [python]
def _conj(x, **kwargs):
  # The only dtypes that are allowed are: float32, float64, complex64, and
  # complex128.
  if x.dtype == tf.float32:
    return tf.cast(x, tf.complex64)
  elif x.dtype == tf.float64:
    return tf.cast(x, tf.complex128)
  else:
    return tf.math.conj(x)
//  [python]
def train_step(
      self,
      inputs: Tuple[Any, Any],
      model: tf_keras.Model,
      optimizer: tf_keras.optimizers.Optimizer,
      metrics: Optional[List[Any]] = None) -> Dict[str, Any]:
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync

    with tf.GradientTape() as tape:
      outputs = model(
          inputs=images,
          image_info=labels['image_info'],
          training=True)
      outputs = tf.nest.map_structure(
          lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          labels=labels,
          model_outputs=outputs,
          aux_losses=model.losses)
      scaled_loss = losses['total_loss'] / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])

    if self.task_config.evaluation.report_train_mean_iou:
      segmentation_labels = {
          'masks': labels['category_mask'],
          'valid_masks': labels['valid_mask'],
          'image_info': labels['image_info']
      }
      self.process_metrics(
          metrics=[self.train_mean_iou],
          labels=segmentation_labels,
          model_outputs=outputs['segmentation_outputs'])
      logs.update({
          self.train_mean_iou.name:
              self.train_mean_iou.result()
      })

    return logs
//  [python]
def build_losses(self,
                   labels: Mapping[str, tf.Tensor],
                   model_outputs: Mapping[str, tf.Tensor],
                   aux_losses: Optional[Any] = None):
    """Panoptic deeplab losses.

    Args:
      labels: labels.
      model_outputs: Output logits from panoptic deeplab.
      aux_losses: auxiliarly loss tensors, i.e. `losses` in keras.Model.

    Returns:
      The total loss tensor.
    """
    loss_config = self._task_config.losses
    segmentation_loss_fn = (
        panoptic_deeplab_losses.WeightedBootstrappedCrossEntropyLoss(
            loss_config.label_smoothing,
            loss_config.class_weights,
            loss_config.ignore_label,
            top_k_percent_pixels=loss_config.top_k_percent_pixels))
    instance_center_heatmap_loss_fn = panoptic_deeplab_losses.CenterHeatmapLoss(
    )
    instance_center_offset_loss_fn = panoptic_deeplab_losses.CenterOffsetLoss()

    semantic_weights = tf.to_float(
        labels['semantic_weights'],
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    things_mask = tf.to_float(
        tf.squeeze(labels['things_mask'], axis=3),
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    valid_mask = tf.to_float(
        tf.squeeze(labels['valid_mask'], axis=3),
        dtype=model_outputs['instance_centers_heatmap'].dtype)

    segmentation_loss = segmentation_loss_fn(
        model_outputs['segmentation_outputs'],
        labels['category_mask'],
        sample_weight=semantic_weights)
    instance_center_heatmap_loss = instance_center_heatmap_loss_fn(
        model_outputs['instance_centers_heatmap'],
        labels['instance_centers_heatmap'],
        sample_weight=valid_mask)
    instance_center_offset_loss = instance_center_offset_loss_fn(
        model_outputs['instance_centers_offset'],
        labels['instance_centers_offset'],
        sample_weight=things_mask)

    model_loss = (
        loss_config.segmentation_loss_weight * segmentation_loss +
        loss_config.center_heatmap_loss_weight * instance_center_heatmap_loss +
        loss_config.center_offset_loss_weight * instance_center_offset_loss)

    total_loss = model_loss
    if aux_losses:
      total_loss += tf.add_n(aux_losses)

    losses = {
        'total_loss': total_loss,
        'model_loss': model_loss,
        'segmentation_loss': segmentation_loss,
        'instance_center_heatmap_loss': instance_center_heatmap_loss,
        'instance_center_offset_loss': instance_center_offset_loss
    }

    return losses
//  [python]
def build_losses(self,
                   labels: Mapping[str, tf.Tensor],
                   model_outputs: Mapping[str, tf.Tensor],
                   aux_losses: Optional[Any] = None):
    """Panoptic deeplab losses.

    Args:
      labels: labels.
      model_outputs: Output logits from panoptic deeplab.
      aux_losses: auxiliarly loss tensors, i.e. `losses` in keras.Model.

    Returns:
      The total loss tensor.
    """
    loss_config = self._task_config.losses
    segmentation_loss_fn = (
        panoptic_deeplab_losses.WeightedBootstrappedCrossEntropyLoss(
            loss_config.label_smoothing,
            loss_config.class_weights,
            loss_config.ignore_label,
            top_k_percent_pixels=loss_config.top_k_percent_pixels))
    instance_center_heatmap_loss_fn = panoptic_deeplab_losses.CenterHeatmapLoss(
    )
    instance_center_offset_loss_fn = panoptic_deeplab_losses.CenterOffsetLoss()

    semantic_weights = tf.cast(
        labels['semantic_weights'],
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    things_mask = tf.cast(
        tf.squeeze(labels['things_mask'], axis=3),
        dtype=model_outputs['instance_centers_heatmap'].dtype)
    valid_mask = tf.cast(
        tf.squeeze(labels['valid_mask'], axis=3),
        dtype=model_outputs['instance_centers_heatmap'].dtype)

    segmentation_loss = segmentation_loss_fn(
        model_outputs['segmentation_outputs'],
        labels['category_mask'],
        sample_weight=semantic_weights)
    instance_center_heatmap_loss = instance_center_heatmap_loss_fn(
        model_outputs['instance_centers_heatmap'],
        labels['instance_centers_heatmap'],
        sample_weight=valid_mask)
    instance_center_offset_loss = instance_center_offset_loss_fn(
        model_outputs['instance_centers_offset'],
        labels['instance_centers_offset'],
        sample_weight=things_mask)

    model_loss = (
        loss_config.segmentation_loss_weight * segmentation_loss +
        loss_config.center_heatmap_loss_weight * instance_center_heatmap_loss +
        loss_config.center_offset_loss_weight * instance_center_offset_loss)

    total_loss = model_loss
    if aux_losses:
      total_loss += tf.add_n(aux_losses)

    losses = {
        'total_loss': total_loss,
        'model_loss': model_loss,
        'segmentation_loss': segmentation_loss,
        'instance_center_heatmap_loss': instance_center_heatmap_loss,
        'instance_center_offset_loss': instance_center_offset_loss
    }

    return losses
//  [python]
def _iota(*, dtype, shape, dimension):
  dtype = _to_tf_dtype(dtype)
  # Some dtypes are unsupported, like uint32, so we just fall back to int32.
  # TODO(mattjj, necula): improve tf.range dtype handling
  shape_tf = _eval_shape(shape)
  vec = tf.range(tf.cast(shape_tf[dimension], tf.int32), dtype=tf.int32)
  vec_shape = [-1 if i == dimension else 1 for i in range(len(shape))]
  return t.cast(tf.broadcast_to(tf.reshape(vec, vec_shape), shape_tf), dtype)
//  [python]
def _iota(*, dtype, shape, dimension):
  dtype = _to_tf_dtype(dtype)
  # Some dtypes are unsupported, like uint32, so we just fall back to int32.
  # TODO(mattjj, necula): improve tf.range dtype handling
  shape_tf = _eval_shape(shape)
  vec = tf.range(tf.cast(shape_tf[dimension], tf.int32), dtype=tf.int32)
  vec_shape = [-1 if i == dimension else 1 for i in range(len(shape))]
  return tf.cast(tf.broadcast_to(tf.reshape(vec, vec_shape), shape_tf), dtype)
//  [python]
def _concatenate(*operands, dimension):
  return tf.concat(operands, axis=t.cast(dimension, tf.int32))
//  [python]
def _concatenate(*operands, dimension):
  return tf.concat(operands, axis=tf.cast(dimension, tf.int32))
//  [python]
def _shift_right_arithmetic_raw(x, y):
  if x.dtype.is_unsigned:
    assert x.dtype == y.dtype
    orig_dtype = x.dtype
    signed_dtype = _UNSIGNED_TO_SIGNED_TABLE[orig_dtype]
    x = t.cast(x, signed_dtype)
    y = t.cast(y, signed_dtype)
    res = tf.bitwise.right_shift(x, y)
    return t.cast(res, orig_dtype)
  else:
    return tf.bitwise.right_shift(x, y)
//  [python]
def _shift_right_arithmetic_raw(x, y):
  if x.dtype.is_unsigned:
    assert x.dtype == y.dtype
    orig_dtype = x.dtype
    signed_dtype = _UNSIGNED_TO_SIGNED_TABLE[orig_dtype]
    x = tf.cast(x, signed_dtype)
    y = tf.cast(y, signed_dtype)
    res = tf.bitwise.right_shift(x, y)
    return tf.cast(res, orig_dtype)
  else:
    return tf.bitwise.right_shift(x, y)
//  [python]
def _shift_right_logical_raw(x, y):
  if x.dtype.is_unsigned:
    return tf.bitwise.right_shift(x, y)
  else:
    assert x.dtype == y.dtype
    orig_dtype = x.dtype
    unsigned_dtype = _SIGNED_TO_UNSIGNED_TABLE[orig_dtype]
    x = t.cast(x, unsigned_dtype)
    y = t.cast(y, unsigned_dtype)
    res = tf.bitwise.right_shift(x, y)
    return t.cast(res, orig_dtype)
//  [python]
def _shift_right_logical_raw(x, y):
  if x.dtype.is_unsigned:
    return tf.bitwise.right_shift(x, y)
  else:
    assert x.dtype == y.dtype
    orig_dtype = x.dtype
    unsigned_dtype = _SIGNED_TO_UNSIGNED_TABLE[orig_dtype]
    x = tf.cast(x, unsigned_dtype)
    y = tf.cast(y, unsigned_dtype)
    res = tf.bitwise.right_shift(x, y)
    return tf.cast(res, orig_dtype)
//  [python]
def _shift_in_bounds(x: TfVal, y: TfVal) -> TfVal:
  # Return the TF expression for when y is within bounds (0 <= y < |x|)
  x_bits = 8 * x.dtype.size
  # TF does not have comparisons for uint16 and uint32 (despite what the
  # documentation says)
  y_comp = t.cast(
      y, _UNSIGNED_TO_SIGNED_TABLE[y.dtype]) if y.dtype.is_unsigned else y
  y_lt_x_bits = tf.math.less(y_comp, x_bits)
  y_ge_0 = tf.math.greater_equal(y_comp, 0)
  return tf.logical_and(y_lt_x_bits, y_ge_0)
//  [python]
def _shift_in_bounds(x: TfVal, y: TfVal) -> TfVal:
  # Return the TF expression for when y is within bounds (0 <= y < |x|)
  x_bits = 8 * x.dtype.size
  # TF does not have comparisons for uint16 and uint32 (despite what the
  # documentation says)
  y_comp = tf.cast(
      y, _UNSIGNED_TO_SIGNED_TABLE[y.dtype]) if y.dtype.is_unsigned else y
  y_lt_x_bits = tf.math.less(y_comp, x_bits)
  y_ge_0 = tf.math.greater_equal(y_comp, 0)
  return tf.logical_and(y_lt_x_bits, y_ge_0)
//  [python]
def _maybe_cast_to_int64(x: TfVal) -> TfVal:
  if x.dtype != tf.int32 and x.dtype != tf.int64:
    return t.cast(x, tf.int64)
  return x
//  [python]
def _maybe_cast_to_int64(x: TfVal) -> TfVal:
  if x.dtype != tf.int32 and x.dtype != tf.int64:
    return tf.cast(x, tf.int64)
  return x
//  [python]
def _iota_2x32_shape(*, shape):
  def _add(x, y):
    return x + y
  def _mul(x, y):
    if not core.is_constant_dim(x):
      x = tf.cast(_eval_shape((x,))[0], y.dtype)
      x = tf.broadcast_to(x, tf.shape(y))
    return x * y
  def _cast32(xs): return t.cast(xs, _to_tf_dtype(jnp.uint32))
  iotas = [_iota(dtype=jnp.uint64, shape=shape, dimension=dimension)
           for dimension in range(len(shape))]
  counts = prng.bcast_iotas_to_reshaped_iota(_add, _mul, shape, iotas)
  counts_lo = _cast32(counts)
  counts_hi = _cast32(t.bitwise.right_shift(counts, 32))
  return counts_hi, counts_lo
//  [python]
def _iota_2x32_shape(*, shape):
  def _add(x, y):
    return x + y
  def _mul(x, y):
    if not core.is_constant_dim(x):
      x = tf.cast(_eval_shape((x,))[0], y.dtype)
      x = tf.broadcast_to(x, tf.shape(y))
    return x * y
  def _cast32(xs): return tf.cast(xs, _to_tf_dtype(jnp.uint32))
  iotas = [_iota(dtype=jnp.uint64, shape=shape, dimension=dimension)
           for dimension in range(len(shape))]
  counts = prng.bcast_iotas_to_reshaped_iota(_add, _mul, shape, iotas)
  counts_lo = _cast32(counts)
  counts_hi = _cast32(tf.bitwise.right_shift(counts, 32))
  return counts_hi, counts_lo