jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def grad_planner(
    state,
    num_actions,
    horizon,
    proposals,
    iterations,
    imagine,
    objective,
    kl_scale,
    step_size,
):
    dtype = prec.global_policy().compute_dtype
    B, P = list(state.values())[0].shape[0], proposals
    H, A = horizon, num_actions
    flat_state = {k: tf.repeat(v, P, 0) for k, v in state.items()}
    mean = tf.zeros((B, H, A), dtype)
    rawstd = 0.54 * tf.ones((B, H, A), dtype)
    for _ in range(iterations):
        proposals = tf.random.normal((B, P, H, A), dtype=dtype)
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(mean)
            tape.watch(rawstd)
            std = tf.nn.softplus(rawstd)
            proposals = proposals * std[:, None] + mean[:, None]
            proposals = (
                tf.stop_gradient(tf.clip_by_value(proposals, -1, 1))
                + proposals
                - tf.stop_gradient(proposals)
            )
            flat_proposals = tf.reshape(proposals, (B * P, H, A))
            states = imagine(flat_proposals, flat_state)
            scores = objective(states)
            scores = tf.reshape(tf.reduce_sum(scores, -1), (B, P))
            div = tfd.kl_divergence(
                tfd.Normal(mean, std),
                tfd.Normal(tf.zeros_like(mean), tf.ones_like(std)),
            )
            elbo = tf.reduce_sum(scores) - kl_scale * div
            elbo /= tf.cast(tf.reduce_prod(tf.shape(scores)), dtype)
        grad_mean, grad_rawstd = tape.gradient(elbo, [mean, rawstd])
        e, v = tf.nn.moments(grad_mean, [1, 2], keepdims=True)
        grad_mean /= tf.sqrt(e * e + v + 1e-4)
        e, v = tf.nn.moments(grad_rawstd, [1, 2], keepdims=True)
        grad_rawstd /= tf.sqrt(e * e + v + 1e-4)
        mean = tf.clip_by_value(mean + step_size * grad_mean, -1, 1)
        rawstd = rawstd + step_size * grad_rawstd
    return mean[:, 0, :]
//  [python]
def batch(self, ds, batch_size, drop_remainder=True, batch_constraints=None):
    if batch_constraints is None:
      batch_constraints = []
      for k, v in get_input_modalities().items():
        bound = v.get_constraints()
        if bound is not None:
          def _fn(ex):
            mask = t.cast(ex[f"inputs/{k}/mask"], t.bool)
            return t.reduce_sum(t.cast(t.reduce_any(mask, -1), t.int32))
          batch_bound = int(round(bound * batch_size))
          if self.pack_max_len is None:
            logging.info(f"Adding batch constraint {k}/{bound} => "
                         f"({batch_bound} per batch of {batch_size})")
          else:
            bound *= self.max_to_pack
            logging.info(f"Adding batch constraint {k}/{bound} => "
                         f"({batch_bound} per batch of {batch_size} groups of {self.max_to_pack})")
          batch_constraints.append((_fn, batch_bound))
    if self.pack:
      enc, dec = self.pack_max_len
      if self.max_to_pack == 2:
        ds = pair_examples(ds, enc, dec, self.pack_pool_size)
      else:
        raise NotImplementedError()
      if batch_constraints:
        ds = batch_with_constraints(ds, batch_size, self.constraint_pool_size, batch_constraints)
      else:
        ds = ds.batch(batch_size, drop_remainder=drop_remainder)
      return unfold(ds, batch_size, n=self.max_to_pack)
    else:
      if batch_constraints:
        return batch_with_constraints(ds, batch_size, self.constraint_pool_size, batch_constraints)
      else:
        return ds.batch(batch_size, drop_remainder=drop_remainder)
//  [python]
def batch(self, ds, batch_size, drop_remainder=True, batch_constraints=None):
    if batch_constraints is None:
      batch_constraints = []
      for k, v in get_input_modalities().items():
        bound = v.get_constraints()
        if bound is not None:
          def _fn(ex):
            mask = tf.cast(ex[f"inputs/{k}/mask"], tf.bool)
            return tf.reduce_sum(tf.cast(tf.reduce_any(mask, -1), tf.int32))
          batch_bound = int(round(bound * batch_size))
          if self.pack_max_len is None:
            logging.info(f"Adding batch constraint {k}/{bound} => "
                         f"({batch_bound} per batch of {batch_size})")
          else:
            bound *= self.max_to_pack
            logging.info(f"Adding batch constraint {k}/{bound} => "
                         f"({batch_bound} per batch of {batch_size} groups of {self.max_to_pack})")
          batch_constraints.append((_fn, batch_bound))
    if self.pack:
      enc, dec = self.pack_max_len
      if self.max_to_pack == 2:
        ds = pair_examples(ds, enc, dec, self.pack_pool_size)
      else:
        raise NotImplementedError()
      if batch_constraints:
        ds = batch_with_constraints(ds, batch_size, self.constraint_pool_size, batch_constraints)
      else:
        ds = ds.batch(batch_size, drop_remainder=drop_remainder)
      return unfold(ds, batch_size, n=self.max_to_pack)
    else:
      if batch_constraints:
        return batch_with_constraints(ds, batch_size, self.constraint_pool_size, batch_constraints)
      else:
        return ds.batch(batch_size, drop_remainder=drop_remainder)
//  [python]
def build_network(self):
        layer_configs = self.config
        feature_size = 784

        # define placeholder
        self.x = tf.placeholder(tf.float32, [None, feature_size], name="input_x")
        self.y = tf.placeholder(tf.int32, [None, self.label_size], name="input_y")
        label_number = 10

        # define network
        input_layer = self.x
        in_height = in_width = int(math.sqrt(feature_size))
        for i, layer_config in enumerate(layer_configs):
            input_layer = tf.nn.relu(self.get_layer(layer_config, input_layer, in_height, in_width, i))

        output_layer = input_layer
        if self.is_flatten(output_layer):
            output_layer = tf.contrib.layers.flatten(output_layer)  # flatten
        output_layer = tf.layers.dense(output_layer, label_number)
        child_logit = tf.nn.softmax_cross_entropy_with_logits(logits=output_layer, labels=self.y)
        child_loss = tf.reduce_mean(child_logit)

        self.train_step = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(child_loss)
        child_accuracy = tf.equal(tf.argmax(output_layer, 1), tf.argmax(self.y, 1))
        self.accuracy = tf.reduce_mean(tf.to_float(child_accuracy))
//  [python]
def build_network(self):
        layer_configs = self.config
        feature_size = 784

        # define placeholder
        self.x = tf.placeholder(tf.float32, [None, feature_size], name="input_x")
        self.y = tf.placeholder(tf.int32, [None, self.label_size], name="input_y")
        label_number = 10

        # define network
        input_layer = self.x
        in_height = in_width = int(math.sqrt(feature_size))
        for i, layer_config in enumerate(layer_configs):
            input_layer = tf.nn.relu(self.get_layer(layer_config, input_layer, in_height, in_width, i))

        output_layer = input_layer
        if self.is_flatten(output_layer):
            output_layer = tf.contrib.layers.flatten(output_layer)  # flatten
        output_layer = tf.layers.dense(output_layer, label_number)
        child_logit = tf.nn.softmax_cross_entropy_with_logits(logits=output_layer, labels=self.y)
        child_loss = tf.reduce_mean(child_logit)

        self.train_step = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(child_loss)
        child_accuracy = tf.equal(tf.argmax(output_layer, 1), tf.argmax(self.y, 1))
        self.accuracy = tf.reduce_mean(tf.cast(child_accuracy, "float"))
//  [python]
def get_mask_zero(raw,mask):
    mask=tf.tile(tf.expand_dims(mask,axis=-1),[1,1,tf.shape(raw)[-1]])
    mask=tf.to_float(mask)
    return raw*mask
//  [python]
def get_mask_zero(raw,mask):
    mask=tf.tile(tf.expand_dims(mask,axis=-1),[1,1,tf.shape(raw)[-1]])
    mask=tf.cast(mask,tf.float32)
    return raw*mask
//  [python]
def extract_slice_from_ray_batch(ray_batch, key, use_viewdirs=True, oid=None):
  """Extract a slice of data from a ray batch.

  Args:
    ray_batch: [R, M] tf.float32.
    key: str.
    use_viewdirs: bool.
    oid: int.

  Returns:
    x: [R, K] tf.float32 or tf.int32.
  """
  if key == 'origin':
    start = 0
    end = 3
  elif key == 'direction':
    start = 3
    end = 6
  elif key == 'far':
    start = 7
    end = 8
  else:
    if key == 'example_id':
      start = 11
      end = 12
    elif key == 'metadata':
      start = 12
      end = 13
    elif key == 'object_bounds':
      start = 13 + 2 * oid
      end = start + 2
    else:
      raise ValueError(f'Invalid key: {key}')
    # Adjust index based on whether we are using viewdirs.
    # TODO(guom): Get rid of use_viewdirs dependency everywhere.
    if not use_viewdirs:
      start -= 3
      end -= 3
  x = ray_batch[:, start:end]

  if key == 'metadata':
    x = t.cast(x, dtype=t.int32)
  return x
//  [python]
def extract_slice_from_ray_batch(ray_batch, key, use_viewdirs=True, oid=None):
  """Extract a slice of data from a ray batch.

  Args:
    ray_batch: [R, M] tf.float32.
    key: str.
    use_viewdirs: bool.
    oid: int.

  Returns:
    x: [R, K] tf.float32 or tf.int32.
  """
  if key == 'origin':
    start = 0
    end = 3
  elif key == 'direction':
    start = 3
    end = 6
  elif key == 'far':
    start = 7
    end = 8
  else:
    if key == 'example_id':
      start = 11
      end = 12
    elif key == 'metadata':
      start = 12
      end = 13
    elif key == 'object_bounds':
      start = 13 + 2 * oid
      end = start + 2
    else:
      raise ValueError(f'Invalid key: {key}')
    # Adjust index based on whether we are using viewdirs.
    # TODO(guom): Get rid of use_viewdirs dependency everywhere.
    if not use_viewdirs:
      start -= 3
      end -= 3
  x = ray_batch[:, start:end]

  if key == 'metadata':
    x = tf.cast(x, dtype=tf.int32)
  return x
//  [python]
def batch_with_constraints(
    ds: tf.data.Dataset,
    batch_size: int,
    pool_size: int,
    batch_constraints: List[Tuple[Callable, int]]
):
  batch_fns = [x[0] for x in batch_constraints]
  max_counts = t.cast(t.stack([x[1] for x in batch_constraints], 0), t.int64)
  key = next(iter(ds.element_spec.keys()))
  n_const = len(batch_constraints)

  null_batch = {k: t.zeros([0] + list(v.shape), dtype=v.dtype)
                for k, v in ds.element_spec.items()}

  empty_batch = {k: t.zeros([batch_size]+list(v.shape), dtype=v.dtype)
                 for k, v in ds.element_spec.items()}
  empty_batch["counts"] = t.zeros([batch_size, n_const], dtype=t.int64)

  empty_overflow: Dict[str, t.Tensor] = {
    k: t.zeros([pool_size] + list(v.shape), dtype=v.dtype)
    for k, v in ds.element_spec.items()}
  empty_overflow["counts"] = t.zeros([pool_size, n_const], dtype=t.int64)

  def _revise_batch(state, ex):
    batch, sz, overflow, of_sz = state

    overflow: Dict[str, t.Tensor] = overflow
    batch: Dict[str, t.Tensor] = batch
    ex["counts"] = t.cast(t.stack([fn(ex) for fn in batch_fns], 0), t.int64)
    # t.print("Enter", ex["counts"], "batch:", batch['counts'][:sz], "pool:", overflow["counts"][:of_sz])

    total_counts = t.reduce_sum(batch["counts"][:sz], 0)
    if t.reduce_all(total_counts + ex["counts"] <= max_counts):
      # t.print("Add to batch", sz)
      batch = {k: _set_vector(v, sz, ex[k]) for k, v in batch.items()}
      sz += 1
      return_batch = sz == batch_size
    else:
      # t.print("Add to overflow", of_sz)
      overflow = {k: _set_vector(v, of_sz, ex[k]) for k, v in overflow.items()}
      of_sz += 1
      return_batch = of_sz == pool_size

    if return_batch and of_sz == 0:
      output_batch = {k: v for k, v in batch.items() if k != "counts"}
      return ({k: v*0 for k, v in batch.items()}, t.zeros((), t.int64), overflow, of_sz), output_batch
    elif return_batch:
      output_batch = {k: v for k, v in batch.items() if k != "counts"}
      is_valid = t.reduce_all(t.cumsum(overflow["counts"][:of_sz], axis=0) <= max_counts, -1)
      # t.print("valid", is_valid)
      ix = t.minimum(t.argmax(is_valid) + 1, t.minimum(of_sz, batch_size))
      # t.print("Return batch from", ix)

      batch = {k: pad_to(v[:ix], batch_size) for k, v in overflow.items()}
      overflow = {k: pad_to(v[ix:], pool_size) for k, v in overflow.items()}
      sz = ix
      of_sz = of_sz - ix
      return (batch, sz, overflow, of_sz), output_batch
    else:
      return (batch, sz, overflow, of_sz), null_batch

  ds = ds.scan((empty_batch, t.zeros((), t.int64), empty_overflow, t.zeros((), t.int64)), _revise_batch)
  ds = ds.filter(lambda x: t.shape(x[key])[0] > 0)
  ds = ds.map(lambda batch: {
    k: v if len(v.shape) == 0 else t.ensure_shape(v, [batch_size] + list(v.shape)[1:])
    for k, v in batch.items()})
  return ds
//  [python]
def batch_with_constraints(
    ds: tf.data.Dataset,
    batch_size: int,
    pool_size: int,
    batch_constraints: List[Tuple[Callable, int]]
):
  batch_fns = [x[0] for x in batch_constraints]
  max_counts = tf.cast(tf.stack([x[1] for x in batch_constraints], 0), tf.int64)
  key = next(iter(ds.element_spec.keys()))
  n_const = len(batch_constraints)

  null_batch = {k: tf.zeros([0] + list(v.shape), dtype=v.dtype)
                for k, v in ds.element_spec.items()}

  empty_batch = {k: tf.zeros([batch_size]+list(v.shape), dtype=v.dtype)
                 for k, v in ds.element_spec.items()}
  empty_batch["counts"] = tf.zeros([batch_size, n_const], dtype=tf.int64)

  empty_overflow: Dict[str, tf.Tensor] = {
    k: tf.zeros([pool_size] + list(v.shape), dtype=v.dtype)
    for k, v in ds.element_spec.items()}
  empty_overflow["counts"] = tf.zeros([pool_size, n_const], dtype=tf.int64)

  def _revise_batch(state, ex):
    batch, sz, overflow, of_sz = state

    overflow: Dict[str, tf.Tensor] = overflow
    batch: Dict[str, tf.Tensor] = batch
    ex["counts"] = tf.cast(tf.stack([fn(ex) for fn in batch_fns], 0), tf.int64)
    # tf.print("Enter", ex["counts"], "batch:", batch['counts'][:sz], "pool:", overflow["counts"][:of_sz])

    total_counts = tf.reduce_sum(batch["counts"][:sz], 0)
    if tf.reduce_all(total_counts + ex["counts"] <= max_counts):
      # tf.print("Add to batch", sz)
      batch = {k: _set_vector(v, sz, ex[k]) for k, v in batch.items()}
      sz += 1
      return_batch = sz == batch_size
    else:
      # tf.print("Add to overflow", of_sz)
      overflow = {k: _set_vector(v, of_sz, ex[k]) for k, v in overflow.items()}
      of_sz += 1
      return_batch = of_sz == pool_size

    if return_batch and of_sz == 0:
      output_batch = {k: v for k, v in batch.items() if k != "counts"}
      return ({k: v*0 for k, v in batch.items()}, tf.zeros((), tf.int64), overflow, of_sz), output_batch
    elif return_batch:
      output_batch = {k: v for k, v in batch.items() if k != "counts"}
      is_valid = tf.reduce_all(tf.cumsum(overflow["counts"][:of_sz], axis=0) <= max_counts, -1)
      # tf.print("valid", is_valid)
      ix = tf.minimum(tf.argmax(is_valid) + 1, tf.minimum(of_sz, batch_size))
      # tf.print("Return batch from", ix)

      batch = {k: pad_to(v[:ix], batch_size) for k, v in overflow.items()}
      overflow = {k: pad_to(v[ix:], pool_size) for k, v in overflow.items()}
      sz = ix
      of_sz = of_sz - ix
      return (batch, sz, overflow, of_sz), output_batch
    else:
      return (batch, sz, overflow, of_sz), null_batch

  ds = ds.scan((empty_batch, tf.zeros((), tf.int64), empty_overflow, tf.zeros((), tf.int64)), _revise_batch)
  ds = ds.filter(lambda x: tf.shape(x[key])[0] > 0)
  ds = ds.map(lambda batch: {
    k: v if len(v.shape) == 0 else tf.ensure_shape(v, [batch_size] + list(v.shape)[1:])
    for k, v in batch.items()})
  return ds
//  [python]
def inference_fn(W,b,x_data,y_target):
    prediction = tf.sign(tf.subtract(tf.matmul(x_data, W), b))
    accuracy = tf.reduce_mean(tf.to_float(tf.equal(prediction, y_target)))
    return accuracy
//  [python]
def inference_fn(W,b,x_data,y_target):
    prediction = tf.sign(tf.subtract(tf.matmul(x_data, W), b))
    accuracy = tf.reduce_mean(tf.cast(tf.equal(prediction, y_target), tf.float32))
    return accuracy
//  [python]
def testResourcesWithCond(self):
    def create_v1_saved_model():
      saved_model_dir = os.path.join(self.get_temp_dir(), 'resources_with_cond')
      with tf.Graph().as_default():
        with tf.compat.v1.Session() as sess:
          in_tensor = tf.compat.v1.placeholder(
              shape=[1], dtype=tf.float32, name='input'
          )

          def body(i, arr):
            n = tf.raw_ops.StackPushV2(
                handle=arr, elem=tf.cast(i, dtype=tf.float32)
            )
            return n, arr

          arr = tf.raw_ops.StackV2(max_size=10, elem_type=tf.float32)
          n, result_arr = tf.cond(
              in_tensor < 10, lambda: body(0, arr), lambda: body(1, arr)
          )

          with ops.control_dependencies([result_arr, n]):
            out_tensor = tf.raw_ops.StackPopV2(
                handle=result_arr, elem_type=tf.float32
            )

          inputs = {'x': in_tensor}
          outputs = {'a': out_tensor}
          saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
      return saved_model_dir

    saved_model_dir = create_v1_saved_model()

    converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,
        tf.lite.OpsSet.SELECT_TF_OPS,
    ]
    tflite_model = converter.convert()
    self.assertIsNotNone(tflite_model)

    # Check values from converted model.
    interp = interpreter.Interpreter(model_content=tflite_model)
    input_details = interp.get_input_details()
    output_details = interp.get_output_details()

    interp.allocate_tensors()

    input_data = np.array([1.0], dtype=np.float32)
    interp.set_tensor(input_details[0]['index'], input_data)

    interp.invoke()
    actual_value = interp.get_tensor(output_details[0]['index'])
    self.assertEqual(0.0, actual_value)
//  [python]
def testResourcesWithWhile(self):
    def create_v1_saved_model():
      saved_model_dir = os.path.join(
          self.get_temp_dir(), 'resources_with_while'
      )
      with tf.Graph().as_default():
        with tf.compat.v1.Session() as sess:
          in_tensor = tf.compat.v1.placeholder(
              shape=[1], dtype=tf.float32, name='input'
          )

          def cond(i, arr, m):
            del arr
            del m
            return i < 10

          def body(i, arr, m):
            del m
            n = tf.raw_ops.StackPushV2(
                handle=arr, elem=tf.cast(i, dtype=tf.float32)
            )
            return i + 1, arr, n

          arr = tf.raw_ops.StackV2(max_size=10, elem_type=tf.float32)
          _, result_arr, n = tf.while_loop(cond, body, [0, arr, 0.0])

          with ops.control_dependencies([result_arr, n]):
            out_tensor = tf.raw_ops.StackPopV2(
                handle=result_arr, elem_type=tf.float32
            )

          inputs = {'x': in_tensor}
          outputs = {'a': out_tensor}
          saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
      return saved_model_dir

    saved_model_dir = create_v1_saved_model()

    converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,
        tf.lite.OpsSet.SELECT_TF_OPS,
    ]
    tflite_model = converter.convert()
    self.assertIsNotNone(tflite_model)

    # Check values from converted model.
    interp = interpreter.Interpreter(model_content=tflite_model)
    input_details = interp.get_input_details()
    output_details = interp.get_output_details()

    interp.allocate_tensors()

    input_data = np.array([1.0], dtype=np.float32)
    interp.set_tensor(input_details[0]['index'], input_data)

    interp.invoke()
    actual_value = interp.get_tensor(output_details[0]['index'])
    self.assertEqual(9.0, actual_value)
//  [python]
def merge_labels(labels):
  """Converts the ground-truth labels into loss targets."""
  boxes = box_ops.yxyx_to_xcycwh(labels['bbox'])
  classes = t.cast(labels['classes'], boxes.dtype)
  return tf.concat([classes[..., None], boxes], axis=-1)
//  [python]
def merge_labels(labels):
  """Converts the ground-truth labels into loss targets."""
  boxes = box_ops.yxyx_to_xcycwh(labels['bbox'])
  classes = tf.cast(labels['classes'], boxes.dtype)
  return tf.concat([classes[..., None], boxes], axis=-1)
//  [python]
def inference(self):

		timage = t.cast((self.x + 1) * 127, t.uint8)
		t.summary.image("raw_image_real", timage[:,:,:,:3])
		# t.summary.image("raw_image_real", timage[:,:,:,4:])

		mean, logsigma = self.encoder(self.x)
		recon_x = self.decoder(mean, logsigma)

		timage = t.cast((recon_x + 1) * 127, t.uint8)
		t.summary.image("raw_image_recon", timage[:,:,:,:3])
		# t.summary.image("raw_image_recon", timage[:,:,:,4:])

		
		self.saver = t.train.Saver(t.get_collection(t.GraphKeys.GLOBAL_VARIABLES, scope=self.name))

		return recon_x, mean, logsigma
//  [python]
def inference(self):

		timage = tf.cast((self.x + 1) * 127, tf.uint8)
		tf.summary.image("raw_image_real", timage[:,:,:,:3])
		# tf.summary.image("raw_image_real", timage[:,:,:,4:])

		mean, logsigma = self.encoder(self.x)
		recon_x = self.decoder(mean, logsigma)

		timage = tf.cast((recon_x + 1) * 127, tf.uint8)
		tf.summary.image("raw_image_recon", timage[:,:,:,:3])
		# tf.summary.image("raw_image_recon", timage[:,:,:,4:])

		
		self.saver = tf.train.Saver(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.name))

		return recon_x, mean, logsigma
//  [python]
def test_dropblock(self):
        with self.test_session() as sess:
            data = tf.ones(shape=[1,32,32,1],dtype=tf.float32)
            data =wnnl.dropblock(data,keep_prob=0.8,block_size=4,is_training=True)
            data = tf.reshape(data,shape=[32,32])
            data = tf.to_int32(data)
            wmlu.show_list(sess.run(data).tolist())
//  [python]
def test_dropblock(self):
        with self.test_session() as sess:
            data = tf.ones(shape=[1,32,32,1],dtype=tf.float32)
            data =wnnl.dropblock(data,keep_prob=0.8,block_size=4,is_training=True)
            data = tf.reshape(data,shape=[32,32])
            data = tf.cast(data,tf.int32)
            wmlu.show_list(sess.run(data).tolist())
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    start_positions = labels['start_positions']
    end_positions = labels['end_positions']
    start_logits, end_logits = model_outputs

    start_loss = tf_keras.losses.sparse_categorical_crossentropy(
        start_positions,
        tf.cast(start_logits, dtype=tf.float32),
        from_logits=True)
    end_loss = tf_keras.losses.sparse_categorical_crossentropy(
        end_positions, tf.cast(end_logits, dtype=tf.float32), from_logits=True)

    loss = (tf.reduce_mean(start_loss) + tf.reduce_mean(end_loss)) / 2
    return loss
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    start_positions = labels['start_positions']
    end_positions = labels['end_positions']
    is_impossible = labels['is_impossible']
    is_impossible = tf.to_float(tf.reshape(is_impossible, [-1]))

    start_logits = model_outputs['start_logits']
    end_logits = model_outputs['end_logits']
    class_logits = model_outputs['class_logits']

    start_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        start_positions, start_logits)
    end_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        end_positions, end_logits)
    is_impossible_loss = tf_keras.losses.binary_crossentropy(
        is_impossible, class_logits, from_logits=True)

    loss = (tf.reduce_mean(start_loss) + tf.reduce_mean(end_loss)) / 2
    loss += tf.reduce_mean(is_impossible_loss) / 2
    return loss
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    start_positions = labels['start_positions']
    end_positions = labels['end_positions']
    is_impossible = labels['is_impossible']
    is_impossible = tf.cast(tf.reshape(is_impossible, [-1]), tf.float32)

    start_logits = model_outputs['start_logits']
    end_logits = model_outputs['end_logits']
    class_logits = model_outputs['class_logits']

    start_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        start_positions, start_logits)
    end_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        end_positions, end_logits)
    is_impossible_loss = tf_keras.losses.binary_crossentropy(
        is_impossible, class_logits, from_logits=True)

    loss = (tf.reduce_mean(start_loss) + tf.reduce_mean(end_loss)) / 2
    loss += tf.reduce_mean(is_impossible_loss) / 2
    return loss
//  [python]
def _calculate_aggregate_mask(answer, pooled_output, cell_selection_preference, labels, aggregation_classifier):
    """
    Finds examples where the model should select cells with no aggregation.

    Returns a mask that determines for which examples should the model select answers directly from the table, without
    any aggregation function. If the answer is a piece of text the case is unambiguous as aggregation functions only
    apply to numbers. If the answer is a number but does not appear in the table then we must use some aggregation
    case. The ambiguous case is when the answer is a number that also appears in the table. In this case we use the
    aggregation function probabilities predicted by the model to decide whether to select or aggregate. The threshold
    for this is a hyperparameter *cell_selection_preference*

    Args:
        answer (`tf.Tensor` of shape `(batch_size, )`):
            Answer for every example in the batch. Nan if there is no scalar answer.
        pooled_output (`tf.Tensor` of shape `(batch_size, hidden_size)`):
            Output of the pooler (BertPooler) on top of the encoder layer.
        cell_selection_preference (`float`):
            Preference for cell selection in ambiguous cases.
        labels (`tf.Tensor` of shape `(batch_size, sequence_length)`):
            Labels per token. aggregation_classifier (`torch.nn.Linear`): Aggregation head

    Returns:
        aggregate_mask (`tf.Tensor` of shape `(batch_size,)`): A mask set to 1 for examples that should use aggregation
        functions.
    """
    # tf.Tensor(batch_size,)
    aggregate_mask_init = t.cast(t.logical_not(t.math.is_nan(answer)), t.float32)
    logits_aggregation = aggregation_classifier(pooled_output)
    dist_aggregation = tfp.distributions.Categorical(logits=logits_aggregation)
    # Index 0 corresponds to "no aggregation".
    aggregation_ops_total_mass = tf.reduce_sum(dist_aggregation.probs_parameter()[:, 1:], axis=1)
    # Cell selection examples according to current model.
    is_pred_cell_selection = aggregation_ops_total_mass <= cell_selection_preference
    # Examples with non-empty cell selection supervision.
    is_cell_supervision_available = tf.reduce_sum(labels, axis=1) > 0
    aggregate_mask = tf.where(
        tf.logical_and(is_pred_cell_selection, is_cell_supervision_available),
        tf.zeros_like(aggregate_mask_init, dtype=tf.float32),
        aggregate_mask_init,
    )
    aggregate_mask = tf.stop_gradient(aggregate_mask)
    return aggregate_mask
//  [python]
def _calculate_aggregate_mask(answer, pooled_output, cell_selection_preference, labels, aggregation_classifier):
    """
    Finds examples where the model should select cells with no aggregation.

    Returns a mask that determines for which examples should the model select answers directly from the table, without
    any aggregation function. If the answer is a piece of text the case is unambiguous as aggregation functions only
    apply to numbers. If the answer is a number but does not appear in the table then we must use some aggregation
    case. The ambiguous case is when the answer is a number that also appears in the table. In this case we use the
    aggregation function probabilities predicted by the model to decide whether to select or aggregate. The threshold
    for this is a hyperparameter *cell_selection_preference*

    Args:
        answer (`tf.Tensor` of shape `(batch_size, )`):
            Answer for every example in the batch. Nan if there is no scalar answer.
        pooled_output (`tf.Tensor` of shape `(batch_size, hidden_size)`):
            Output of the pooler (BertPooler) on top of the encoder layer.
        cell_selection_preference (`float`):
            Preference for cell selection in ambiguous cases.
        labels (`tf.Tensor` of shape `(batch_size, sequence_length)`):
            Labels per token. aggregation_classifier (`torch.nn.Linear`): Aggregation head

    Returns:
        aggregate_mask (`tf.Tensor` of shape `(batch_size,)`): A mask set to 1 for examples that should use aggregation
        functions.
    """
    # tf.Tensor(batch_size,)
    aggregate_mask_init = tf.cast(tf.logical_not(tf.math.is_nan(answer)), tf.float32)
    logits_aggregation = aggregation_classifier(pooled_output)
    dist_aggregation = tfp.distributions.Categorical(logits=logits_aggregation)
    # Index 0 corresponds to "no aggregation".
    aggregation_ops_total_mass = tf.reduce_sum(dist_aggregation.probs_parameter()[:, 1:], axis=1)
    # Cell selection examples according to current model.
    is_pred_cell_selection = aggregation_ops_total_mass <= cell_selection_preference
    # Examples with non-empty cell selection supervision.
    is_cell_supervision_available = tf.reduce_sum(labels, axis=1) > 0
    aggregate_mask = tf.where(
        tf.logical_and(is_pred_cell_selection, is_cell_supervision_available),
        tf.zeros_like(aggregate_mask_init, dtype=tf.float32),
        aggregate_mask_init,
    )
    aggregate_mask = tf.stop_gradient(aggregate_mask)
    return aggregate_mask
//  [python]
def call(self, inputs, mask=None, training=None):
        """
        Invoke the query normalization layer for the input feature tensor

        Parameters
        ----------
        inputs: Tensor object
            Input ranking feature tensor
            Shape: [batch_size, sequence_len, num_features]
        mask: Tensor object
            Mask to be used to identify records to ignore in query
            Shape: [batch_size, sequence_len]
        training: bool
            If the layer should be run as training or not

        Returns
        -------
        Tensor object
            Normed input feature tensor
            Shape: [batch_size, sequence_len, encoding_size]
        """
        # NOTE: Tensorflow math ops are not consistent on saving and reloading if they are float32.
        #       Both float16 and float64 give consistent results.
        inputs = t.cast(inputs, t.float64)
        mask = t.expand_dims(t.cast(mask, t.float64), axis=-1)

        inputs = t.multiply(inputs, mask)

        count = t.math.reduce_sum(mask, axis=1)
        mean = t.math.divide(t.math.reduce_sum(inputs, axis=1), count)
        mean = t.multiply(t.expand_dims(mean, axis=1), mask)

        var = t.math.divide(t.math.reduce_sum(t.math.pow(t.math.subtract(inputs, mean), 2), axis=1), count)
        std = t.expand_dims(t.math.sqrt(var), axis=1)

        zscore = t.math.divide_no_nan(t.math.subtract(inputs, mean), std)

        return t.cast(zscore, t.float32)
//  [python]
def call(self, inputs, mask=None, training=None):
        """
        Invoke the query normalization layer for the input feature tensor

        Parameters
        ----------
        inputs: Tensor object
            Input ranking feature tensor
            Shape: [batch_size, sequence_len, num_features]
        mask: Tensor object
            Mask to be used to identify records to ignore in query
            Shape: [batch_size, sequence_len]
        training: bool
            If the layer should be run as training or not

        Returns
        -------
        Tensor object
            Normed input feature tensor
            Shape: [batch_size, sequence_len, encoding_size]
        """
        # NOTE: Tensorflow math ops are not consistent on saving and reloading if they are float32.
        #       Both float16 and float64 give consistent results.
        inputs = tf.cast(inputs, tf.float64)
        mask = tf.expand_dims(tf.cast(mask, tf.float64), axis=-1)

        inputs = tf.multiply(inputs, mask)

        count = tf.math.reduce_sum(mask, axis=1)
        mean = tf.math.divide(tf.math.reduce_sum(inputs, axis=1), count)
        mean = tf.multiply(tf.expand_dims(mean, axis=1), mask)

        var = tf.math.divide(tf.math.reduce_sum(tf.math.pow(tf.math.subtract(inputs, mean), 2), axis=1), count)
        std = tf.expand_dims(tf.math.sqrt(var), axis=1)

        zscore = tf.math.divide_no_nan(tf.math.subtract(inputs, mean), std)

        return tf.cast(zscore, tf.float32)
//  [python]
def get_surf_loss(self, surf, x_offset, y_offset, max_dim_per_batch, batch_size, w, h):
        x_offset = tf.concat([tf.reshape(x_offset, [batch_size, -1]), tf.ones([batch_size, 1]) * -1], axis = 1)
        y_offset = tf.concat([tf.reshape(y_offset, [batch_size, -1]), tf.ones([batch_size, 1]) * -1], axis = 1)

        # stab
        unstab_surf = tf.to_float(surf[:, 0, :, :])
        unstab_surf_x_norm = tf.expand_dims((unstab_surf[:, :, 0] / (w - 1)) * 2 - 1, axis = 2)
        unstab_surf_y_norm = tf.expand_dims((unstab_surf[:, :, 1] / (h - 1)) * 2 - 1, axis = 2)
        unstab_surf_norm = tf.concat([unstab_surf_x_norm, unstab_surf_y_norm], axis = 2)

        # unstab
        stab_surf = surf[:, 1, :, :]
        stab_surf_x = tf.to_float(stab_surf[:, :, 0])
        stab_surf_y = tf.to_float(stab_surf[:, :, 1])
        idx = tf.cast(stab_surf_x + stab_surf_y * w, 'int32')

        stab_surf_transformed_x_norm = tf.expand_dims(tf.batch_gather(x_offset, idx), axis = 2)
        stab_surf_transformed_y_norm = tf.expand_dims(tf.batch_gather(y_offset, idx), axis = 2)
        stab_surf_transformed_norm = tf.concat([stab_surf_transformed_x_norm, stab_surf_transformed_y_norm], axis = 2)

        MSE = tf.reduce_sum(tf.squared_difference(stab_surf_transformed_norm, unstab_surf_norm), axis = [1, 2])
        MSE = tf.div_no_nan(MSE, max_dim_per_batch)

        return tf.reduce_mean(MSE)
//  [python]
def get_surf_loss(self, surf, x_offset, y_offset, max_dim_per_batch, batch_size, w, h):
        x_offset = tf.concat([tf.reshape(x_offset, [batch_size, -1]), tf.ones([batch_size, 1]) * -1], axis = 1)
        y_offset = tf.concat([tf.reshape(y_offset, [batch_size, -1]), tf.ones([batch_size, 1]) * -1], axis = 1)

        # stab
        unstab_surf = tf.cast(surf[:, 0, :, :], 'float32')
        unstab_surf_x_norm = tf.expand_dims((unstab_surf[:, :, 0] / (w - 1)) * 2 - 1, axis = 2)
        unstab_surf_y_norm = tf.expand_dims((unstab_surf[:, :, 1] / (h - 1)) * 2 - 1, axis = 2)
        unstab_surf_norm = tf.concat([unstab_surf_x_norm, unstab_surf_y_norm], axis = 2)

        # unstab
        stab_surf = surf[:, 1, :, :]
        stab_surf_x = tf.cast(stab_surf[:, :, 0], 'float32')
        stab_surf_y = tf.cast(stab_surf[:, :, 1], 'float32')
        idx = tf.cast(stab_surf_x + stab_surf_y * w, 'int32')

        stab_surf_transformed_x_norm = tf.expand_dims(tf.batch_gather(x_offset, idx), axis = 2)
        stab_surf_transformed_y_norm = tf.expand_dims(tf.batch_gather(y_offset, idx), axis = 2)
        stab_surf_transformed_norm = tf.concat([stab_surf_transformed_x_norm, stab_surf_transformed_y_norm], axis = 2)

        MSE = tf.reduce_sum(tf.squared_difference(stab_surf_transformed_norm, unstab_surf_norm), axis = [1, 2])
        MSE = tf.div_no_nan(MSE, max_dim_per_batch)

        return tf.reduce_mean(MSE)
//  [python]
def _parse_data(
      self, decoded_tensors: Mapping[str,
                                     tf.Tensor]) -> Tuple[tf.Tensor, tf.Tensor]:
    label = tf.to_int32(decoded_tensors['image/class/label'])
    image_bytes = decoded_tensors['image/encoded']
    image = tf.io.decode_jpeg(image_bytes, channels=3)
    image = tf.image.resize(
        image, self._output_size, method=tf.image.ResizeMethod.BILINEAR)
    image = tf.ensure_shape(image, self._output_size + [3])

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(
        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)

    image = tf.image.convert_image_dtype(image, self._dtype)
    return image, label
//  [python]
def _parse_data(
      self, decoded_tensors: Mapping[str,
                                     tf.Tensor]) -> Tuple[tf.Tensor, tf.Tensor]:
    label = tf.cast(decoded_tensors['image/class/label'], dtype=tf.int32)
    image_bytes = decoded_tensors['image/encoded']
    image = tf.io.decode_jpeg(image_bytes, channels=3)
    image = tf.image.resize(
        image, self._output_size, method=tf.image.ResizeMethod.BILINEAR)
    image = tf.ensure_shape(image, self._output_size + [3])

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(
        image, offset=preprocess_ops.MEAN_RGB, scale=preprocess_ops.STDDEV_RGB)

    image = tf.image.convert_image_dtype(image, self._dtype)
    return image, label
//  [python]
def optimize(self, loss):
    """Build the graph to optimize the loss function."""

    # Optimizer nodes.
    # Linear learning rate decay.
    opts = self._options
    words_to_train = float(opts.words_per_epoch * opts.epochs_to_train)
    lr = opts.learning_rate * tf.maximum(
        0.0001, 1.0 - tf.cast(self._words, tf.float32) / words_to_train)
    self._lr = lr
    optimizer = tf.train.GradientDescentOptimizer(lr)
    train = optimizer.minimize(loss,
                               global_step=self.global_step,
                               gate_gradients=optimizer.GATE_NONE)
    self._train = train
//  [python]
def top_k_with_unique(inputs, k):
  """Finds the values and indices of the k largests entries.

  Instead of doing sort like tf.nn.top_k, this function finds the max value
  k times. The running time is proportional to k, which is be faster when k
  is small. The current implementation supports only inputs of rank 2.
  In addition, iota is used to replace the lower bits of each element, this
  makes the selection more stable when there are equal elements. The
  overhead is that output values are approximated.

  Args:
    inputs: A tensor with rank of 2. [batch_size, original_size].
    k: An integer, number of top elements to select.

  Returns:
    top_values: A tensor, the k largest elements in sorted order.
      [batch_size, k].
    indices: A tensor, indices of the top_values. [batch_size, k].
  """
  unique_inputs = _create_make_unique(t.cast(inputs, t.float32))
  top_values, indices = _create_topk_unique(unique_inputs, k)
  top_values = t.cast(top_values, inputs.dtype)
  return top_values, indices
//  [python]
def top_k_with_unique(inputs, k):
  """Finds the values and indices of the k largests entries.

  Instead of doing sort like tf.nn.top_k, this function finds the max value
  k times. The running time is proportional to k, which is be faster when k
  is small. The current implementation supports only inputs of rank 2.
  In addition, iota is used to replace the lower bits of each element, this
  makes the selection more stable when there are equal elements. The
  overhead is that output values are approximated.

  Args:
    inputs: A tensor with rank of 2. [batch_size, original_size].
    k: An integer, number of top elements to select.

  Returns:
    top_values: A tensor, the k largest elements in sorted order.
      [batch_size, k].
    indices: A tensor, indices of the top_values. [batch_size, k].
  """
  unique_inputs = _create_make_unique(tf.cast(inputs, tf.float32))
  top_values, indices = _create_topk_unique(unique_inputs, k)
  top_values = tf.cast(top_values, inputs.dtype)
  return top_values, indices
//  [python]
def dc(out, x, mask, lam1):
    """
    build dc layer
    input:
    out: tensor
    x: original tensor
    mask: sampling mask
    lam1:lambda
    renturn:
    dc layer
    """

    # out,x,mask,lam1=input
    def r2c(x):
        """
        r2c
        """
        return tf.complex(x[..., 0], x[..., 1])

    def c2r(x):
        """
        c2r
        """
        return tf.stack([tf.math.real(x), tf.math.imag(x)], axis=-1)

    out = r2c(out)
    x = r2c(x)
    lam1 = tf.to_float(lam1)
    kspace = tf.signal.fft2d(out)
    k0 = tf.signal.fft2d(x)
    temp = kspace * (1 - (lam1 / (0.5 + lam1)) * mask) + \
           (lam1 / (0.5 + lam1)) * k0 * mask
    coilImgs = tf.signal.ifft2d(temp)
    return c2r(coilImgs)
//  [python]
def dc(out, x, mask, lam1):
    """
    build dc layer
    input:
    out: tensor
    x: original tensor
    mask: sampling mask
    lam1:lambda
    renturn:
    dc layer
    """

    # out,x,mask,lam1=input
    def r2c(x):
        """
        r2c
        """
        return tf.complex(x[..., 0], x[..., 1])

    def c2r(x):
        """
        c2r
        """
        return tf.stack([tf.math.real(x), tf.math.imag(x)], axis=-1)

    out = r2c(out)
    x = r2c(x)
    lam1 = tf.cast(lam1, tf.complex64)
    kspace = tf.signal.fft2d(out)
    k0 = tf.signal.fft2d(x)
    temp = kspace * (1 - (lam1 / (0.5 + lam1)) * mask) + \
           (lam1 / (0.5 + lam1)) * k0 * mask
    coilImgs = tf.signal.ifft2d(temp)
    return c2r(coilImgs)
//  [python]
def filter_sentence_by_count(item, minimum_sentences):
    n_sentences = t.cast(item.row_splits[1], t.int32)
    if t.less(n_sentences, minimum_sentences):
        return t.constant(False)
    return t.constant(True)
//  [python]
def filter_sentence_by_count(item, minimum_sentences):
    n_sentences = tf.cast(item.row_splits[1], tf.int32)
    if tf.less(n_sentences, minimum_sentences):
        return tf.constant(False)
    return tf.constant(True)
//  [python]
def makeModel(x, mask, train, nLayers, K):
    """
    This is the main function that creates the model.
    input:
    x: tensor
    mask: sampling mask
    train: bool
    nLayers: res block layer num
    K: dc layer num
    return:
    model
    """
    with tf.compat.forward_compatibility_horizon(2019, 5, 1):
        out = {}
        out['dc0'] = x
        mask = tf.cast(mask, dtype=tf.complex64)
        with tf.name_scope('myModel'):
            with tf.variable_scope('Wts', reuse=tf.AUTO_REUSE):
                for i in range(1, K + 1):
                    j = str(i)
                    # residual block
                    out['dw' +
                        j] = residual_block(out['dc' + str(i - 1)], train, nLayers)
                    lam1 = getLambda()
                    # rhs= x + lam1 * out['dw' + j]
                    # out['dc'+j]=dc(rhs,mask,lam1)
                    out['dc' + j] = dc(out['dw' + j], x, mask, lam1)
    return out
//  [python]
def attention_mask_square(nd):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    dtype = tf_utils.get_dtype()
    ns = nd
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return t.cast(m, dtype)
//  [python]
def attention_mask_square(nd):
    """1's in the lower triangle, counting from the lower right corner.

    Same as tf.matrix_band_part(tf.ones([nd, ns]), -1, ns-nd), but doesn't produce garbage on TPUs.
    """
    dtype = tf_utils.get_dtype()
    ns = nd
    i = tf.range(nd)[:, None]
    j = tf.range(ns)
    m = i >= j - ns + nd
    return tf.cast(m, dtype)
//  [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.to_float(image_reader)
  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 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 forward(self, examples, labels):
    """Build the graph for the forward pass."""
    opts = self._options

    # Declare all variables we need.
    # Embedding: [vocab_size, emb_dim]
    init_width = 0.5 / opts.emb_dim
    emb = tf.Variable(
        tf.random_uniform(
            [opts.vocab_size, opts.emb_dim], -init_width, init_width),
        name="emb")
    self._emb = emb

    # Softmax weight: [vocab_size, emb_dim]. Transposed.
    sm_w_t = tf.Variable(
        tf.zeros([opts.vocab_size, opts.emb_dim]),
        name="sm_w_t")

    # Softmax bias: [vocab_size].
    sm_b = tf.Variable(tf.zeros([opts.vocab_size]), name="sm_b")

    # Global step: scalar, i.e., shape [].
    self.global_step = tf.Variable(0, name="global_step")

    # Nodes to compute the nce loss w/ candidate sampling.
    labels_matrix = tf.reshape(
        tf.to_int64(labels),
        [opts.batch_size, 1])

    # Negative sampling.
    sampled_ids, _, _ = (tf.nn.fixed_unigram_candidate_sampler(
        true_classes=labels_matrix,
        num_true=1,
        num_sampled=opts.num_samples,
        unique=True,
        range_max=opts.vocab_size,
        distortion=0.75,
        unigrams=opts.vocab_counts.tolist()))

    # Embeddings for examples: [batch_size, emb_dim]
    example_emb = tf.nn.embedding_lookup(emb, examples)

    # Weights for labels: [batch_size, emb_dim]
    true_w = tf.nn.embedding_lookup(sm_w_t, labels)
    # Biases for labels: [batch_size, 1]
    true_b = tf.nn.embedding_lookup(sm_b, labels)

    # Weights for sampled ids: [num_sampled, emb_dim]
    sampled_w = tf.nn.embedding_lookup(sm_w_t, sampled_ids)
    # Biases for sampled ids: [num_sampled, 1]
    sampled_b = tf.nn.embedding_lookup(sm_b, sampled_ids)

    # True logits: [batch_size, 1]
    true_logits = tf.reduce_sum(tf.multiply(example_emb, true_w), 1) + true_b

    # Sampled logits: [batch_size, num_sampled]
    # We replicate sampled noise labels for all examples in the batch
    # using the matmul.
    sampled_b_vec = tf.reshape(sampled_b, [opts.num_samples])
    sampled_logits = tf.matmul(example_emb,
                               sampled_w,
                               transpose_b=True) + sampled_b_vec
    return true_logits, sampled_logits