jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def validation_step(self, inputs, model, metrics=None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors. output_dict = { "video_ids":
        batch_video_ids, "video_matrix": batch_video_matrix, "labels":
        batch_labels, "num_frames": batch_frames}.
      model: the model, forward definition.
      metrics: a nested structure of metrics objects.

    Returns:
      a dictionary of logs.
    """
    # Will require `num_frames` if `num_sample_frames` is None since
    # video_matrix is padded to max_frames in this case.
    require_num_frames = (
        self.task_config.validation_data.num_sample_frames is None
    )
    outputs = self.inference_step(
        model, inputs, require_num_frames=require_num_frames
    )['predictions']
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    labels, label_weights = self._preprocess_labels(inputs, training=False)
    outputs, labels, label_weights = self._postprocess_outputs(
        outputs, labels, label_weights, training=False
    )

    all_losses = self.build_losses(
        labels=labels,
        model_outputs=outputs,
        label_weights=label_weights,
        aux_losses=model.losses,
    )

    logs = {self.loss: all_losses['total_loss']}
    logs.update(
        self.process_metrics(
            metrics,
            labels=labels,
            outputs=outputs,
            model_losses=all_losses,
            label_weights=inputs.get('label_weights', None),
            training=False,
        )
    )

    return logs
//  [python]
def step(self, mode, agent=None, action_fn=None, **kwargs):
    """Returns [next_contexts..., next_timer] list of ops.

    Args:
      mode: a string representing the mode=[train, explore, eval].
      **kwargs: kwargs for context_transition_fn.
    Returns:
      a list of ops that set the context.
    """
    if agent is None:
      ops = []
      if self._context_transition_fn is not None:
        def sampler_fn():
          samples = self.sample_contexts(mode, 1)[0]
          return [s[0] for s in samples]
        values = self._context_transition_fn(self.vars, self.t, sampler_fn, **kwargs)
        ops += [tf.assign(var, value) for var, value in zip(self.vars, values)]
      ops.append(tf.assign_add(self.t, 1))  # increment timer
      return ops
    else:
      ops = agent.tf_context.step(mode, **kwargs)
      state = kwargs['state']
      next_state = kwargs['next_state']
      state_repr = kwargs['state_repr']
      next_state_repr = kwargs['next_state_repr']
      with tf.control_dependencies(ops):  # Step high level context before computing low level one.
        # Get the context transition function output.
        values = self._context_transition_fn(self.vars, self.t, None,
                                             state=state_repr,
                                             next_state=next_state_repr)
        # Select a new goal every C steps, otherwise use context transition.
        low_level_context = [
            tf.cond(tf.equal(self.t % self.meta_action_every_n, 0),
                    lambda: tf.to_float(action_fn(next_state, context=None)),
                    lambda: values)]
        ops = [tf.assign(var, value)
               for var, value in zip(self.vars, low_level_context)]
        with tf.control_dependencies(ops):
          return [tf.assign_add(self.t, 1)]  # increment timer
        return ops
//  [python]
def step(self, mode, agent=None, action_fn=None, **kwargs):
    """Returns [next_contexts..., next_timer] list of ops.

    Args:
      mode: a string representing the mode=[train, explore, eval].
      **kwargs: kwargs for context_transition_fn.
    Returns:
      a list of ops that set the context.
    """
    if agent is None:
      ops = []
      if self._context_transition_fn is not None:
        def sampler_fn():
          samples = self.sample_contexts(mode, 1)[0]
          return [s[0] for s in samples]
        values = self._context_transition_fn(self.vars, self.t, sampler_fn, **kwargs)
        ops += [tf.assign(var, value) for var, value in zip(self.vars, values)]
      ops.append(tf.assign_add(self.t, 1))  # increment timer
      return ops
    else:
      ops = agent.tf_context.step(mode, **kwargs)
      state = kwargs['state']
      next_state = kwargs['next_state']
      state_repr = kwargs['state_repr']
      next_state_repr = kwargs['next_state_repr']
      with tf.control_dependencies(ops):  # Step high level context before computing low level one.
        # Get the context transition function output.
        values = self._context_transition_fn(self.vars, self.t, None,
                                             state=state_repr,
                                             next_state=next_state_repr)
        # Select a new goal every C steps, otherwise use context transition.
        low_level_context = [
            tf.cond(tf.equal(self.t % self.meta_action_every_n, 0),
                    lambda: tf.cast(action_fn(next_state, context=None), tf.float32),
                    lambda: values)]
        ops = [tf.assign(var, value)
               for var, value in zip(self.vars, low_level_context)]
        with tf.control_dependencies(ops):
          return [tf.assign_add(self.t, 1)]  # increment timer
        return ops
//  [python]
def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):
    # calculate attention
    matmul_qk = tf.matmul(q, k, transpose_b=True)
    
    dk = tf.to_float(shape_list(k)[-1])
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e4)

    if attention_mask is not None:
        # Apply the attention mask
        scaled_attention_logits = scaled_attention_logits + attention_mask

    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1) 

    # Mask heads if we want to
    if head_mask is not None:
        attention_weights = attention_weights * head_mask

    output = tf.matmul(attention_weights, v) 

    return output, attention_weights
//  [python]
def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=None):
    # calculate attention
    matmul_qk = tf.matmul(q, k, transpose_b=True)
    
    dk = tf.cast(shape_list(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e4)

    if attention_mask is not None:
        # Apply the attention mask
        scaled_attention_logits = scaled_attention_logits + attention_mask

    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1) 

    # Mask heads if we want to
    if head_mask is not None:
        attention_weights = attention_weights * head_mask

    output = tf.matmul(attention_weights, v) 

    return output, attention_weights
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors. output_dict = { "video_ids":
        batch_video_ids, "video_matrix": batch_video_matrix, "labels":
        batch_labels, "num_frames": batch_frames, }
      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.
    """
    # Will require `num_frames` if `num_sample_frames` is None since
    # video_matrix is padded to max_frames in this case.
    require_num_frames = self.task_config.train_data.num_sample_frames is None
    inputs_tensor, extra_inputs = self._preprocess_model_inputs(
        inputs,
        require_num_frames=require_num_frames,
        training=True,
    )
    labels, label_weights = self._preprocess_labels(inputs, training=True)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(inputs_tensor, **extra_inputs)['predictions']
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)
      # Post-process model / label outputs.
      outputs, labels, label_weights = self._postprocess_outputs(
          outputs, labels, label_weights, training=True
      )

      # Computes per-replica loss
      all_losses = self.build_losses(
          model_outputs=outputs,
          labels=labels,
          label_weights=label_weights,
          aux_losses=model.losses,
      )

      loss = all_losses['total_loss']
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = 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 before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)

    # Apply gradient clipping.
    if self.task_config.gradient_clip_norm > 0:
      grads, _ = tf.clip_by_global_norm(
          grads, self.task_config.gradient_clip_norm
      )
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}
    logs.update(
        self.process_metrics(
            metrics,
            labels=labels,
            outputs=outputs,
            model_losses=all_losses,
            label_weights=label_weights,
            training=True,
        )
    )
    return logs
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors. output_dict = { "video_ids":
        batch_video_ids, "video_matrix": batch_video_matrix, "labels":
        batch_labels, "num_frames": batch_frames, }
      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.
    """
    # Will require `num_frames` if `num_sample_frames` is None since
    # video_matrix is padded to max_frames in this case.
    require_num_frames = self.task_config.train_data.num_sample_frames is None
    inputs_tensor, extra_inputs = self._preprocess_model_inputs(
        inputs,
        require_num_frames=require_num_frames,
        training=True,
    )
    labels, label_weights = self._preprocess_labels(inputs, training=True)

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(inputs_tensor, **extra_inputs)['predictions']
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
      # Post-process model / label outputs.
      outputs, labels, label_weights = self._postprocess_outputs(
          outputs, labels, label_weights, training=True
      )

      # Computes per-replica loss
      all_losses = self.build_losses(
          model_outputs=outputs,
          labels=labels,
          label_weights=label_weights,
          aux_losses=model.losses,
      )

      loss = all_losses['total_loss']
      # Scales loss as the default gradients allreduce performs sum inside the
      # optimizer.
      scaled_loss = 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 before apply_gradients when LossScaleOptimizer is
    # used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)

    # Apply gradient clipping.
    if self.task_config.gradient_clip_norm > 0:
      grads, _ = tf.clip_by_global_norm(
          grads, self.task_config.gradient_clip_norm
      )
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: loss}
    logs.update(
        self.process_metrics(
            metrics,
            labels=labels,
            outputs=outputs,
            model_losses=all_losses,
            label_weights=label_weights,
            training=True,
        )
    )
    return logs
//  [python]
def all(x, axis=None, keepdims=False):
    x = t.cast(x, "bool")
    return tf.reduce_all(x, axis=axis, keepdims=keepdims)
//  [python]
def all(x, axis=None, keepdims=False):
    x = tf.cast(x, "bool")
    return tf.reduce_all(x, axis=axis, keepdims=keepdims)
//  [python]
def any(x, axis=None, keepdims=False):
    x = t.cast(x, "bool")
    return tf.reduce_any(x, axis=axis, keepdims=keepdims)
//  [python]
def any(x, axis=None, keepdims=False):
    x = tf.cast(x, "bool")
    return tf.reduce_any(x, axis=axis, keepdims=keepdims)
//  [python]
def to_image_frame(camera_frame_points,
                   image_height,
                   image_width,
                   camera_intrinsics,
                   forward_axis=2):
  """Project points from camera frame (3D) into the image (2D).

  Args:
    camera_frame_points: A tf.Tensor of shape (N, 3) of 3D locations
      we want projected into the image.
    image_height: The image height as an int.
    image_width: The image width as an int.
    camera_intrinsics: The camera intrinsics to project from camera
      frame to the image plane. Should be a (3, 3) tensor.
    forward_axis: Which axis point forwards from the camera.

  Returns:
    image_points: A tf.Tensor of shape (N, 2) with pixel locations of dtype
      tf.float32. Each row has the format of (x, y).
    within_image: A boolean tf.Tensor of shape (N,) indicating whether
      each point falls within the image or not.

  Raises:
    ValueError: If shape of camera_frame_points is not (N, 3).
  """
  if camera_frame_points.get_shape().as_list()[1] != 3:
    raise ValueError('Expecting camera frame points to be shape (N, 3)')

  within_image = camera_frame_points[:, forward_axis] >= 0
  image_frame_points = tf.tensordot(camera_frame_points, camera_intrinsics,
                                    axes=(1, 1))
  image_frame_points = image_frame_points[:, :2] / image_frame_points[:, 2:3]

  # Compute image bounds check
  rounded = tf.to_int32(tf.round(image_frame_points))
  within_image &= (rounded[:, 0] >= 0) & (rounded[:, 0] < image_width)
  within_image &= (rounded[:, 1] >= 0) & (rounded[:, 1] < image_height)

  return image_frame_points, within_image
//  [python]
def to_image_frame(camera_frame_points,
                   image_height,
                   image_width,
                   camera_intrinsics,
                   forward_axis=2):
  """Project points from camera frame (3D) into the image (2D).

  Args:
    camera_frame_points: A tf.Tensor of shape (N, 3) of 3D locations
      we want projected into the image.
    image_height: The image height as an int.
    image_width: The image width as an int.
    camera_intrinsics: The camera intrinsics to project from camera
      frame to the image plane. Should be a (3, 3) tensor.
    forward_axis: Which axis point forwards from the camera.

  Returns:
    image_points: A tf.Tensor of shape (N, 2) with pixel locations of dtype
      tf.float32. Each row has the format of (x, y).
    within_image: A boolean tf.Tensor of shape (N,) indicating whether
      each point falls within the image or not.

  Raises:
    ValueError: If shape of camera_frame_points is not (N, 3).
  """
  if camera_frame_points.get_shape().as_list()[1] != 3:
    raise ValueError('Expecting camera frame points to be shape (N, 3)')

  within_image = camera_frame_points[:, forward_axis] >= 0
  image_frame_points = tf.tensordot(camera_frame_points, camera_intrinsics,
                                    axes=(1, 1))
  image_frame_points = image_frame_points[:, :2] / image_frame_points[:, 2:3]

  # Compute image bounds check
  rounded = tf.cast(tf.round(image_frame_points), dtype=tf.int32)
  within_image &= (rounded[:, 0] >= 0) & (rounded[:, 0] < image_width)
  within_image &= (rounded[:, 1] >= 0) & (rounded[:, 1] < image_height)

  return image_frame_points, within_image
//  [python]
def bincount(x, weights=None, minlength=0, sparse=False):
    x = convert_to_tensor(x)
    dtypes_to_resolve = [x.dtype]
    if standardize_dtype(x.dtype) not in ["int32", "int64"]:
        x = t.cast(x, t.int32)
    if weights is not None:
        weights = convert_to_tensor(weights)
        dtypes_to_resolve.append(weights.dtype)
        dtype = dtypes.result_type(*dtypes_to_resolve)
        if standardize_dtype(weights.dtype) not in [
            "int32",
            "int64",
            "float32",
            "float64",
        ]:
            if "int" in standardize_dtype(weights.dtype):
                weights = t.cast(weights, t.int32)
            else:
                weights = t.cast(weights, t.float32)
    else:
        dtype = "int32"
    if sparse or isinstance(x, t.SparseTensor):
        output = t.sparse.bincount(
            x,
            weights=weights,
            minlength=minlength,
            axis=-1,
        )
        actual_length = output.shape[-1]
        if actual_length is None:
            actual_length = t.shape(output)[-1]
        output = cast(output, dtype)
        if x.shape.rank == 1:
            output_shape = (actual_length,)
        else:
            batch_size = output.shape[0]
            if batch_size is None:
                batch_size = t.shape(output)[0]
            output_shape = (batch_size, actual_length)
        return t.SparseTensor(
            indices=output.indices,
            values=output.values,
            dense_shape=output_shape,
        )
    return t.cast(
        t.math.bincount(x, weights=weights, minlength=minlength, axis=-1),
        dtype,
    )
//  [python]
def bincount(x, weights=None, minlength=0, sparse=False):
    x = convert_to_tensor(x)
    dtypes_to_resolve = [x.dtype]
    if standardize_dtype(x.dtype) not in ["int32", "int64"]:
        x = tf.cast(x, tf.int32)
    if weights is not None:
        weights = convert_to_tensor(weights)
        dtypes_to_resolve.append(weights.dtype)
        dtype = dtypes.result_type(*dtypes_to_resolve)
        if standardize_dtype(weights.dtype) not in [
            "int32",
            "int64",
            "float32",
            "float64",
        ]:
            if "int" in standardize_dtype(weights.dtype):
                weights = tf.cast(weights, tf.int32)
            else:
                weights = tf.cast(weights, tf.float32)
    else:
        dtype = "int32"
    if sparse or isinstance(x, tf.SparseTensor):
        output = tf.sparse.bincount(
            x,
            weights=weights,
            minlength=minlength,
            axis=-1,
        )
        actual_length = output.shape[-1]
        if actual_length is None:
            actual_length = tf.shape(output)[-1]
        output = cast(output, dtype)
        if x.shape.rank == 1:
            output_shape = (actual_length,)
        else:
            batch_size = output.shape[0]
            if batch_size is None:
                batch_size = tf.shape(output)[0]
            output_shape = (batch_size, actual_length)
        return tf.SparseTensor(
            indices=output.indices,
            values=output.values,
            dense_shape=output_shape,
        )
    return tf.cast(
        tf.math.bincount(x, weights=weights, minlength=minlength, axis=-1),
        dtype,
    )
//  [python]
def max(x, axis=None, keepdims=False, initial=None):
    x = convert_to_tensor(x)

    # The TensorFlow numpy API implementation doesn't support `initial` so we
    # handle it manually here.
    if initial is not None:
        if standardize_dtype(x.dtype) == "bool":
            x = tf.reduce_any(x, axis=axis, keepdims=keepdims)
            x = tf.math.maximum(t.cast(x, "int32"), t.cast(initial, "int32"))
            return t.cast(x, "bool")
        else:
            x = tf.reduce_max(x, axis=axis, keepdims=keepdims)
            return tf.math.maximum(x, initial)

    # TensorFlow returns -inf by default for an empty list, but for consistency
    # with other backends and the numpy API we want to throw in this case.
    if tf.executing_eagerly():
        size_x = size(x)
        tf.assert_greater(
            size_x,
            tf.constant(0, dtype=size_x.dtype),
            message="Cannot compute the max of an empty tensor.",
        )

    if standardize_dtype(x.dtype) == "bool":
        return tf.reduce_any(x, axis=axis, keepdims=keepdims)
    else:
        return tf.reduce_max(x, axis=axis, keepdims=keepdims)
//  [python]
def max(x, axis=None, keepdims=False, initial=None):
    x = convert_to_tensor(x)

    # The TensorFlow numpy API implementation doesn't support `initial` so we
    # handle it manually here.
    if initial is not None:
        if standardize_dtype(x.dtype) == "bool":
            x = tf.reduce_any(x, axis=axis, keepdims=keepdims)
            x = tf.math.maximum(tf.cast(x, "int32"), tf.cast(initial, "int32"))
            return tf.cast(x, "bool")
        else:
            x = tf.reduce_max(x, axis=axis, keepdims=keepdims)
            return tf.math.maximum(x, initial)

    # TensorFlow returns -inf by default for an empty list, but for consistency
    # with other backends and the numpy API we want to throw in this case.
    if tf.executing_eagerly():
        size_x = size(x)
        tf.assert_greater(
            size_x,
            tf.constant(0, dtype=size_x.dtype),
            message="Cannot compute the max of an empty tensor.",
        )

    if standardize_dtype(x.dtype) == "bool":
        return tf.reduce_any(x, axis=axis, keepdims=keepdims)
    else:
        return tf.reduce_max(x, axis=axis, keepdims=keepdims)
//  [python]
def append(x1, x2, axis=None):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    if axis is None:
        return tf.concat([tf.reshape(x1, [-1]), tf.reshape(x2, [-1])], axis=0)
    else:
        return tf.concat([x1, x2], axis=axis)
//  [python]
def append(x1, x2, axis=None):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    if axis is None:
        return tf.concat([tf.reshape(x1, [-1]), tf.reshape(x2, [-1])], axis=0)
    else:
        return tf.concat([x1, x2], axis=axis)
//  [python]
def arccos(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.acos(x)
//  [python]
def arccos(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.acos(x)
//  [python]
def arccosh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.acosh(x)
//  [python]
def arccosh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.acosh(x)
//  [python]
def arange(start, stop=None, step=1, dtype=None):
    if dtype is None:
        dtypes_to_resolve = [
            getattr(start, "dtype", type(start)),
            getattr(step, "dtype", type(step)),
        ]
        if stop is not None:
            dtypes_to_resolve.append(getattr(stop, "dtype", type(stop)))
        dtype = dtypes.result_type(*dtypes_to_resolve)
    dtype = standardize_dtype(dtype)
    try:
        out = tf.range(start, stop, delta=step, dtype=dtype)
    except tf.errors.NotFoundError:
        # Some dtypes may not work in eager mode on CPU or GPU.
        out = tf.range(start, stop, delta=step, dtype="float32")
        out = t.cast(out, dtype)
    return out
//  [python]
def arange(start, stop=None, step=1, dtype=None):
    if dtype is None:
        dtypes_to_resolve = [
            getattr(start, "dtype", type(start)),
            getattr(step, "dtype", type(step)),
        ]
        if stop is not None:
            dtypes_to_resolve.append(getattr(stop, "dtype", type(stop)))
        dtype = dtypes.result_type(*dtypes_to_resolve)
    dtype = standardize_dtype(dtype)
    try:
        out = tf.range(start, stop, delta=step, dtype=dtype)
    except tf.errors.NotFoundError:
        # Some dtypes may not work in eager mode on CPU or GPU.
        out = tf.range(start, stop, delta=step, dtype="float32")
        out = tf.cast(out, dtype)
    return out
//  [python]
def arcsin(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.asin(x)
//  [python]
def arcsin(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.asin(x)
//  [python]
def arcsinh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.asinh(x)
//  [python]
def arcsinh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.asinh(x)
//  [python]
def arctan(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.atan(x)
//  [python]
def arctan(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.atan(x)
//  [python]
def arctan2(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    return tf.math.atan2(x1, x2)
//  [python]
def arctan2(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.math.atan2(x1, x2)
//  [python]
def arctanh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.atanh(x)
//  [python]
def arctanh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.atanh(x)
//  [python]
def argmax(x, axis=None, keepdims=False):
    _x = x
    if axis is None:
        x = tf.reshape(x, [-1])
    y = t.cast(tf.argmax(x, axis=axis), dtype="int32")
    if keepdims:
        y = _keepdims(_x, y, axis)
    return y
//  [python]
def argmax(x, axis=None, keepdims=False):
    _x = x
    if axis is None:
        x = tf.reshape(x, [-1])
    y = tf.cast(tf.argmax(x, axis=axis), dtype="int32")
    if keepdims:
        y = _keepdims(_x, y, axis)
    return y
//  [python]
def argmin(x, axis=None, keepdims=False):
    _x = x
    if axis is None:
        x = tf.reshape(x, [-1])
    y = t.cast(tf.argmin(x, axis=axis), dtype="int32")
    if keepdims:
        y = _keepdims(_x, y, axis)
    return y
//  [python]
def argmin(x, axis=None, keepdims=False):
    _x = x
    if axis is None:
        x = tf.reshape(x, [-1])
    y = tf.cast(tf.argmin(x, axis=axis), dtype="int32")
    if keepdims:
        y = _keepdims(_x, y, axis)
    return y
//  [python]
def argsort(x, axis=-1):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "bool":
        x = t.cast(x, "uint8")

    x_shape = x.shape
    if x_shape.rank == 0:
        return t.cast([0], "int32")

    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    return tf.argsort(x, axis=axis)
//  [python]
def argsort(x, axis=-1):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "bool":
        x = tf.cast(x, "uint8")

    x_shape = x.shape
    if x_shape.rank == 0:
        return tf.cast([0], "int32")

    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    return tf.argsort(x, axis=axis)
//  [python]
def clip(x, x_min, x_max):
    dtype = standardize_dtype(x.dtype)
    if dtype == "bool":
        x = t.cast(x, "int32")
    return tf.clip_by_value(x, x_min, x_max)
//  [python]
def clip(x, x_min, x_max):
    dtype = standardize_dtype(x.dtype)
    if dtype == "bool":
        x = tf.cast(x, "int32")
    return tf.clip_by_value(x, x_min, x_max)
//  [python]
def ceil(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.ceil(x)
//  [python]
def ceil(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.ceil(x)
//  [python]
def cos(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.cos(x)
//  [python]
def cos(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.cos(x)
//  [python]
def average(x, axis=None, weights=None):
    x = convert_to_tensor(x)

    if weights is None:  # Treat all weights as 1
        dtype = dtypes.result_type(x.dtype, float)
        x = t.cast(x, dtype)
        avg = tf.reduce_mean(x, axis=axis)
    else:
        weights = convert_to_tensor(weights)
        dtype = dtypes.result_type(x.dtype, weights.dtype, float)
        x = t.cast(x, dtype)
        weights = t.cast(weights, dtype)

        def _rank_equal_case():
            weights_sum = tf.reduce_sum(weights, axis=axis)
            return tf.reduce_sum(x * weights, axis=axis) / weights_sum

        def _rank_not_equal_case():
            weights_sum = tf.reduce_sum(weights)
            axes = tf.convert_to_tensor([[axis], [0]])
            return tf.tensordot(x, weights, axes) / weights_sum

        if axis is None:
            avg = _rank_equal_case()
        else:
            if len(x.shape) == len(weights.shape):
                avg = _rank_equal_case()
            else:
                avg = _rank_not_equal_case()
    return avg
//  [python]
def average(x, axis=None, weights=None):
    x = convert_to_tensor(x)

    if weights is None:  # Treat all weights as 1
        dtype = dtypes.result_type(x.dtype, float)
        x = tf.cast(x, dtype)
        avg = tf.reduce_mean(x, axis=axis)
    else:
        weights = convert_to_tensor(weights)
        dtype = dtypes.result_type(x.dtype, weights.dtype, float)
        x = tf.cast(x, dtype)
        weights = tf.cast(weights, dtype)

        def _rank_equal_case():
            weights_sum = tf.reduce_sum(weights, axis=axis)
            return tf.reduce_sum(x * weights, axis=axis) / weights_sum

        def _rank_not_equal_case():
            weights_sum = tf.reduce_sum(weights)
            axes = tf.convert_to_tensor([[axis], [0]])
            return tf.tensordot(x, weights, axes) / weights_sum

        if axis is None:
            avg = _rank_equal_case()
        else:
            if len(x.shape) == len(weights.shape):
                avg = _rank_equal_case()
            else:
                avg = _rank_not_equal_case()
    return avg
//  [python]
def concatenate(xs, axis=0):
    sparse_count = builtins.sum(isinstance(x, tf.SparseTensor) for x in xs)
    if sparse_count:
        if sparse_count == len(xs):
            return tf.sparse.concat(axis=axis, sp_inputs=xs)
        else:
            xs = [
                (
                    convert_to_tensor(x, sparse=False)
                    if isinstance(x, tf.SparseTensor)
                    else x
                )
                for x in xs
            ]
    xs = tree.map_structure(convert_to_tensor, xs)
    dtype_set = set([x.dtype for x in xs])
    if len(dtype_set) > 1:
        dtype = dtypes.result_type(*dtype_set)
        xs = tree.map_structure(lambda x: tf.cast(x, dtype), xs)
    return tf.concat(xs, axis=axis)
//  [python]
def cosh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = t.cast(x, dtype)
    return tf.math.cosh(x)
//  [python]
def cosh(x):
    x = convert_to_tensor(x)
    if standardize_dtype(x.dtype) == "int64":
        dtype = config.floatx()
    else:
        dtype = dtypes.result_type(x.dtype, float)
    x = tf.cast(x, dtype)
    return tf.math.cosh(x)
//  [python]
def cumprod(x, axis=None, dtype=None):
    x = convert_to_tensor(x, dtype=dtype)
    # tf.math.cumprod doesn't support bool
    if standardize_dtype(x.dtype) == "bool":
        x = t.cast(x, "int32")
    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    return tf.math.cumprod(x, axis=axis)
//  [python]
def cumprod(x, axis=None, dtype=None):
    x = convert_to_tensor(x, dtype=dtype)
    # tf.math.cumprod doesn't support bool
    if standardize_dtype(x.dtype) == "bool":
        x = tf.cast(x, "int32")
    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    return tf.math.cumprod(x, axis=axis)
//  [python]
def cumsum(x, axis=None, dtype=None):
    x = convert_to_tensor(x, dtype=dtype)
    # tf.math.cumprod doesn't support bool
    if standardize_dtype(x.dtype) == "bool":
        x = t.cast(x, "int32")
    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    return tf.math.cumsum(x, axis=axis)
//  [python]
def cumsum(x, axis=None, dtype=None):
    x = convert_to_tensor(x, dtype=dtype)
    # tf.math.cumprod doesn't support bool
    if standardize_dtype(x.dtype) == "bool":
        x = tf.cast(x, "int32")
    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    return tf.math.cumsum(x, axis=axis)
//  [python]
def equal(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.equal(x1, x2)
//  [python]
def dot(x, y):
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    result_dtype = dtypes.result_type(x.dtype, y.dtype)
    # GPU only supports float types
    compute_dtype = dtypes.result_type(result_dtype, float)
    x = tf.to_float(x)
    y = tf.to_float(y)

    x_shape = x.shape
    y_shape = y.shape
    if x_shape.rank == 0 or y_shape.rank == 0:
        output = x * y
    elif y_shape.rank == 1:
        output = tf.tensordot(x, y, axes=[[-1], [-1]])
    else:
        output = tf.tensordot(x, y, axes=[[-1], [-2]])
    return tf.cast(output, result_dtype)
//  [python]
def dot(x, y):
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    result_dtype = dtypes.result_type(x.dtype, y.dtype)
    # GPU only supports float types
    compute_dtype = dtypes.result_type(result_dtype, float)
    x = tf.cast(x, compute_dtype)
    y = tf.cast(y, compute_dtype)

    x_shape = x.shape
    y_shape = y.shape
    if x_shape.rank == 0 or y_shape.rank == 0:
        output = x * y
    elif y_shape.rank == 1:
        output = tf.tensordot(x, y, axes=[[-1], [-1]])
    else:
        output = tf.tensordot(x, y, axes=[[-1], [-2]])
    return tf.cast(output, result_dtype)
//  [python]
def exp(x):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    if "int" in ori_dtype or ori_dtype == "bool":
        x = t.cast(x, config.floatx())
    return tf.exp(x)
//  [python]
def exp(x):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    if "int" in ori_dtype or ori_dtype == "bool":
        x = tf.cast(x, config.floatx())
    return tf.exp(x)
//  [python]
def expm1(x):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    if "int" in ori_dtype or ori_dtype == "bool":
        x = t.cast(x, config.floatx())
    return tf.math.expm1(x)
//  [python]
def expm1(x):
    x = convert_to_tensor(x)
    ori_dtype = standardize_dtype(x.dtype)
    if "int" in ori_dtype or ori_dtype == "bool":
        x = tf.cast(x, config.floatx())
    return tf.math.expm1(x)
//  [python]
def floor(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = t.cast(x, dtype)
    return tf.floor(x)
//  [python]
def floor(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = tf.cast(x, dtype)
    return tf.floor(x)
//  [python]
def cross(x1, x2, axisa=-1, axisb=-1, axisc=-1, axis=None):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)

    if axis is not None:
        axisa = axis
        axisb = axis
        axisc = axis
    x1 = moveaxis(x1, axisa, -1)
    x2 = moveaxis(x2, axisb, -1)

    def maybe_pad_zeros(x, size_of_last_dim):
        def pad_zeros(x):
            return tf.pad(
                x,
                tf.concat(
                    [
                        tf.zeros([tf.rank(x) - 1, 2], "int32"),
                        tf.constant([[0, 1]], "int32"),
                    ],
                    axis=0,
                ),
            )

        if isinstance(size_of_last_dim, int):
            if size_of_last_dim == 2:
                return pad_zeros(x)
            return x

        return tf.cond(
            tf.equal(size_of_last_dim, 2), lambda: pad_zeros(x), lambda: x
        )

    x1_dim = shape_op(x1)[-1]
    x2_dim = shape_op(x2)[-1]

    x1 = maybe_pad_zeros(x1, x1_dim)
    x2 = maybe_pad_zeros(x2, x2_dim)

    # Broadcast each other
    shape = shape_op(x1)

    shape = tf.broadcast_dynamic_shape(shape, shape_op(x2))
    x1 = tf.broadcast_to(x1, shape)
    x2 = tf.broadcast_to(x2, shape)

    c = tf.linalg.cross(x1, x2)

    if isinstance(x1_dim, int) and isinstance(x2_dim, int):
        if (x1_dim == 2) & (x2_dim == 2):
            return c[..., 2]
        return moveaxis(c, -1, axisc)

    return tf.cond(
        (x1_dim == 2) & (x2_dim == 2),
        lambda: c[..., 2],
        lambda: moveaxis(c, -1, axisc),
    )
//  [python]
def cross(x1, x2, axisa=-1, axisb=-1, axisc=-1, axis=None):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)

    if axis is not None:
        axisa = axis
        axisb = axis
        axisc = axis
    x1 = moveaxis(x1, axisa, -1)
    x2 = moveaxis(x2, axisb, -1)

    def maybe_pad_zeros(x, size_of_last_dim):
        def pad_zeros(x):
            return tf.pad(
                x,
                tf.concat(
                    [
                        tf.zeros([tf.rank(x) - 1, 2], "int32"),
                        tf.constant([[0, 1]], "int32"),
                    ],
                    axis=0,
                ),
            )

        if isinstance(size_of_last_dim, int):
            if size_of_last_dim == 2:
                return pad_zeros(x)
            return x

        return tf.cond(
            tf.equal(size_of_last_dim, 2), lambda: pad_zeros(x), lambda: x
        )

    x1_dim = shape_op(x1)[-1]
    x2_dim = shape_op(x2)[-1]

    x1 = maybe_pad_zeros(x1, x1_dim)
    x2 = maybe_pad_zeros(x2, x2_dim)

    # Broadcast each other
    shape = shape_op(x1)

    shape = tf.broadcast_dynamic_shape(shape, shape_op(x2))
    x1 = tf.broadcast_to(x1, shape)
    x2 = tf.broadcast_to(x2, shape)

    c = tf.linalg.cross(x1, x2)

    if isinstance(x1_dim, int) and isinstance(x2_dim, int):
        if (x1_dim == 2) & (x2_dim == 2):
            return c[..., 2]
        return moveaxis(c, -1, axisc)

    return tf.cond(
        (x1_dim == 2) & (x2_dim == 2),
        lambda: c[..., 2],
        lambda: moveaxis(c, -1, axisc),
    )
//  [python]
def greater_equal(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    return tf.greater_equal(x1, x2)
//  [python]
def greater_equal(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.greater_equal(x1, x2)
//  [python]
def greater(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    return tf.greater(x1, x2)
//  [python]
def greater(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.greater(x1, x2)
//  [python]
def less(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.less(x1, x2)
//  [python]
def isclose(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    if "float" in dtype:
        # atol defaults to 1e-08
        # rtol defaults to 1e-05
        return tf.abs(x1 - x2) <= (1e-08 + 1e-05 * tf.abs(x2))
    else:
        return tf.equal(x1, x2)
//  [python]
def isclose(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    if "float" in dtype:
        # atol defaults to 1e-08
        # rtol defaults to 1e-05
        return tf.abs(x1 - x2) <= (1e-08 + 1e-05 * tf.abs(x2))
    else:
        return tf.equal(x1, x2)
//  [python]
def less_equal(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    return tf.less_equal(x1, x2)
//  [python]
def less_equal(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    return tf.less_equal(x1, x2)
//  [python]
def log(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = tf.cast(x, dtype)
    return tf.math.log(x)
//  [python]
def log10(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = tf.cast(x, dtype)
    return tf.math.log(x) / tf.math.log(tf.constant(10, x.dtype))
//  [python]
def log1p(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = tf.cast(x, dtype)
    return tf.math.log1p(x)
//  [python]
def log2(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = t.cast(x, dtype)
    return tf.math.log(x) / tf.math.log(tf.constant(2, x.dtype))
//  [python]
def log2(x):
    x = convert_to_tensor(x)
    dtype = (
        config.floatx()
        if standardize_dtype(x.dtype) == "int64"
        else dtypes.result_type(x.dtype, float)
    )
    x = tf.cast(x, dtype)
    return tf.math.log(x) / tf.math.log(tf.constant(2, x.dtype))
//  [python]
def logical_and(x1, x2):
    x1 = t.cast(x1, "bool")
    x2 = t.cast(x2, "bool")
    return tf.logical_and(x1, x2)
//  [python]
def logical_and(x1, x2):
    x1 = tf.cast(x1, "bool")
    x2 = tf.cast(x2, "bool")
    return tf.logical_and(x1, x2)
//  [python]
def logical_not(x):
    x = t.cast(x, "bool")
    return tf.logical_not(x)
//  [python]
def logical_not(x):
    x = tf.cast(x, "bool")
    return tf.logical_not(x)
//  [python]
def logical_or(x1, x2):
    x1 = t.cast(x1, "bool")
    x2 = t.cast(x2, "bool")
    return tf.logical_or(x1, x2)
//  [python]
def logical_or(x1, x2):
    x1 = tf.cast(x1, "bool")
    x2 = tf.cast(x2, "bool")
    return tf.logical_or(x1, x2)
//  [python]
def logaddexp(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
    x1 = t.cast(x1, dtype)
    x2 = t.cast(x2, dtype)
    delta = x1 - x2
    return tf.where(
        tf.math.is_nan(delta),
        x1 + x2,
        tf.maximum(x1, x2) + tf.math.log1p(tf.math.exp(-tf.abs(delta))),
    )
//  [python]
def logaddexp(x1, x2):
    x1 = convert_to_tensor(x1)
    x2 = convert_to_tensor(x2)
    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
    x1 = tf.cast(x1, dtype)
    x2 = tf.cast(x2, dtype)
    delta = x1 - x2
    return tf.where(
        tf.math.is_nan(delta),
        x1 + x2,
        tf.maximum(x1, x2) + tf.math.log1p(tf.math.exp(-tf.abs(delta))),
    )
//  [python]
def linspace(
    start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0
):
    if dtype is None:
        dtypes_to_resolve = [
            getattr(start, "dtype", type(start)),
            getattr(stop, "dtype", type(stop)),
            float,
        ]
        dtype = dtypes.result_type(*dtypes_to_resolve)
    else:
        dtype = standardize_dtype(dtype)
    start = convert_to_tensor(start, dtype=dtype)
    stop = convert_to_tensor(stop, dtype=dtype)
    if num < 0:
        raise ValueError(
            f"`num` must be a non-negative integer. Received: num={num}"
        )
    step = tf.convert_to_tensor(np.nan)
    if endpoint:
        result = tf.linspace(start, stop, num, axis=axis)
        if num > 1:
            step = (stop - start) / (num - 1)
    else:
        # tf.linspace doesn't support endpoint=False, so we manually handle it
        if num > 0:
            step = (stop - start) / num
        if num > 1:
            new_stop = t.cast(stop, step.dtype) - step
            start = t.cast(start, new_stop.dtype)
            result = tf.linspace(start, new_stop, num, axis=axis)
        else:
            result = tf.linspace(start, stop, num, axis=axis)
    if dtype is not None:
        if "int" in dtype:
            result = tf.floor(result)
        result = t.cast(result, dtype)
    if retstep:
        return (result, step)
    else:
        return result