Add files using upload-large-folder tool

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/losses/losses.py

@@ -0,0 +1,25 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Loss operations for use in neural networks.
+
+Note: All the losses are added to the `GraphKeys.LOSSES` collection by default.
+
+API docstring: tensorflow.losses
+"""
+
+# pylint: disable=wildcard-import
+from tensorflow.python.ops.losses.losses_impl import *
+from tensorflow.python.ops.losses.util import *
+# pylint: enable=wildcard-import

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/losses/losses_impl.py

@@ -0,0 +1,1102 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Implementation of Loss operations for use in neural networks."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import confusion_matrix
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import weights_broadcast_ops
+from tensorflow.python.ops.losses import util
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.deprecation import deprecated_args
+from tensorflow.python.util.deprecation import deprecated_argument_lookup
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["losses.Reduction"])
+class Reduction:
+  """Types of loss reduction.
+
+  Contains the following values:
+
+  * `NONE`: Un-reduced weighted losses with the same shape as input.
+  * `SUM`: Scalar sum of weighted losses.
+  * `MEAN`: Scalar `SUM` divided by sum of weights. DEPRECATED.
+  * `SUM_OVER_BATCH_SIZE`: Scalar `SUM` divided by number of elements in losses.
+  * `SUM_OVER_NONZERO_WEIGHTS`: Scalar `SUM` divided by number of non-zero
+     weights. DEPRECATED.
+  * `SUM_BY_NONZERO_WEIGHTS`: Same as `SUM_OVER_NONZERO_WEIGHTS`. DEPRECATED.
+  """
+
+  NONE = "none"
+  SUM = "weighted_sum"
+  SUM_OVER_BATCH_SIZE = "weighted_sum_over_batch_size"
+  MEAN = "weighted_mean"
+  SUM_BY_NONZERO_WEIGHTS = "weighted_sum_by_nonzero_weights"
+  SUM_OVER_NONZERO_WEIGHTS = SUM_BY_NONZERO_WEIGHTS
+
+  @classmethod
+  def all(cls):
+    return (
+        cls.NONE,
+        cls.SUM,
+        cls.MEAN,
+        cls.SUM_OVER_BATCH_SIZE,
+        cls.SUM_OVER_NONZERO_WEIGHTS,
+        cls.SUM_BY_NONZERO_WEIGHTS)
+
+  @classmethod
+  def validate(cls, key):
+    if key not in cls.all():
+      raise ValueError(f"Invalid Reduction Key {key}. Key should be one of "
+                       f"{cls.all()}.")
+
+
+def _safe_mean(losses, num_present):
+  """Computes a safe mean of the losses.
+
+  Args:
+    losses: `Tensor` whose elements contain individual loss measurements.
+    num_present: The number of measurable elements in `losses`.
+
+  Returns:
+    A scalar representing the mean of `losses`. If `num_present` is zero,
+      then zero is returned.
+  """
+  total_loss = math_ops.reduce_sum(losses)
+  return math_ops.div_no_nan(total_loss, num_present, name="value")
+
+
+def _num_present(losses, weights, per_batch=False):
+  """Computes the number of elements in the loss function induced by `weights`.
+
+  A given weights tensor induces different numbers of usable elements in the
+  `losses` tensor. The `weights` tensor is broadcast across `losses` for all
+  possible dimensions. For example, if `losses` is a tensor of dimension
+  `[4, 5, 6, 3]` and `weights` is a tensor of shape `[4, 5]`, then `weights` is,
+  in effect, tiled to match the shape of `losses`. Following this effective
+  tile, the total number of present elements is the number of non-zero weights.
+
+  Args:
+    losses: `Tensor` of shape `[batch_size, d1, ... dN]`.
+    weights: `Tensor` of shape `[]`, `[batch_size]` or
+      `[batch_size, d1, ... dK]`, where K < N.
+    per_batch: Whether to return the number of elements per batch or as a sum
+      total.
+
+  Returns:
+    The number of present (non-zero) elements in the losses tensor. If
+      `per_batch` is `True`, the value is returned as a tensor of size
+      `[batch_size]`. Otherwise, a single scalar tensor is returned.
+  """
+  if ((isinstance(weights, float) and weights != 0.0) or
+      (context.executing_eagerly() and weights._rank() == 0  # pylint: disable=protected-access
+       and not math_ops.equal(weights, 0.0))):
+    return _num_elements(losses)
+  with ops.name_scope(None, "num_present", (losses, weights)) as scope:
+    weights = math_ops.cast(weights, dtype=dtypes.float32)
+    present = array_ops.where(
+        math_ops.equal(weights, 0.0),
+        array_ops.zeros_like(weights),
+        array_ops.ones_like(weights))
+    present = weights_broadcast_ops.broadcast_weights(present, losses)
+    if per_batch:
+      return math_ops.reduce_sum(
+          present,
+          axis=math_ops.range(1, array_ops.rank(present)),
+          keepdims=True,
+          name=scope)
+    return math_ops.reduce_sum(present, name=scope)
+
+
+def _num_elements(losses):
+  """Computes the number of elements in `losses` tensor."""
+  with ops.name_scope(None, "num_elements", values=[losses]) as scope:
+    return math_ops.cast(array_ops.size(losses, name=scope), dtype=losses.dtype)
+
+
+@tf_export(v1=["losses.compute_weighted_loss"])
+@dispatch.add_dispatch_support
+def compute_weighted_loss(
+    losses, weights=1.0, scope=None, loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Computes the weighted loss.
+
+  Args:
+    losses: `Tensor` of shape `[batch_size, d1, ... dN]`.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `losses`, and must be broadcastable to `losses` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    scope: the scope for the operations performed in computing the loss.
+    loss_collection: the loss will be added to these collections.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss `Tensor` of the same type as `losses`. If `reduction` is
+    `NONE`, this has the same shape as `losses`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If `weights` is `None` or the shape is not compatible with
+      `losses`, or if the number of dimensions (rank) of either `losses` or
+      `weights` is missing.
+
+  Note:
+    When calculating the gradient of a weighted loss contributions from
+    both `losses` and `weights` are considered. If your `weights` depend
+    on some model parameters but you do not want this to affect the loss
+    gradient, you need to apply `tf.stop_gradient` to `weights` before
+    passing them to `compute_weighted_loss`.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  Reduction.validate(reduction)
+  with ops.name_scope(scope, "weighted_loss", (losses, weights)):
+    # Save the `reduction` argument for loss normalization when distributing
+    # to multiple replicas. Used only for estimator + v1 optimizer flow.
+    ops.get_default_graph()._last_loss_reduction = reduction  # pylint: disable=protected-access
+
+    def compute_loss(losses, weights, loss_collection, reduction):
+      losses = ops.convert_to_tensor(losses)
+      input_dtype = losses.dtype
+      losses = math_ops.cast(losses, dtype=dtypes.float32)
+      weights = math_ops.cast(weights, dtype=dtypes.float32)
+      weighted_losses = math_ops.multiply(losses, weights)
+      if reduction == Reduction.NONE:
+        loss = weighted_losses
+      else:
+        loss = math_ops.reduce_sum(weighted_losses)
+        if reduction == Reduction.MEAN:
+          loss = _safe_mean(
+              loss, math_ops.reduce_sum(array_ops.ones_like(losses) * weights))
+        elif (reduction == Reduction.SUM_BY_NONZERO_WEIGHTS or
+              reduction == Reduction.SUM_OVER_NONZERO_WEIGHTS):
+          loss = _safe_mean(loss, _num_present(losses, weights))
+        elif reduction == Reduction.SUM_OVER_BATCH_SIZE:
+          loss = _safe_mean(loss, _num_elements(losses))
+
+      # Convert the result back to the input type.
+      loss = math_ops.cast(loss, input_dtype)
+      util.add_loss(loss, loss_collection)
+      return loss
+
+    # Skip the assert_broadcastable in XLA context because asserts are not
+    # supported so it only causes unnecessary ops. Also skip it because it uses
+    # a DenseToDenseSetOperation op that is incompatible with XLA when
+    # the shape(s) are dynamic.
+    if control_flow_ops.get_enclosing_xla_context() is not None:
+      return compute_loss(losses, weights, loss_collection, reduction)
+    else:
+      with ops.control_dependencies(
+          (weights_broadcast_ops.assert_broadcastable(weights, losses),)):
+        return compute_loss(losses, weights, loss_collection, reduction)
+
+
+@tf_export(v1=["losses.absolute_difference"])
+@dispatch.add_dispatch_support
+def absolute_difference(
+    labels, predictions, weights=1.0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Adds an Absolute Difference loss to the training procedure.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided, then
+  the loss is simply scaled by the given value. If `weights` is a `Tensor` of
+  shape `[batch_size]`, then the total loss for each sample of the batch is
+  rescaled by the corresponding element in the `weights` vector. If the shape of
+  `weights` matches the shape of `predictions`, then the loss of each
+  measurable element of `predictions` is scaled by the corresponding value of
+  `weights`.
+
+  Args:
+    labels: The ground truth output tensor, same dimensions as 'predictions'.
+    predictions: The predicted outputs.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which this loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss float `Tensor`. If `reduction` is `NONE`, this has the same
+    shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shape of `predictions` doesn't match that of
+      `labels` or if the shape of `weights` is invalid or if `labels`
+      or `predictions` is None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if predictions is None:
+    raise ValueError("Argument `predictions` must not be None.")
+  with ops.name_scope(scope, "absolute_difference",
+                      (predictions, labels, weights)) as scope:
+    predictions = math_ops.cast(predictions, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
+    losses = math_ops.abs(math_ops.subtract(predictions, labels))
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+@tf_export(v1=["losses.cosine_distance"])
+@dispatch.add_dispatch_support
+@deprecated_args(None, "dim is deprecated, use axis instead", "dim")
+def cosine_distance(
+    labels, predictions, axis=None, weights=1.0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS,
+    dim=None):
+  """Adds a cosine-distance loss to the training procedure.
+
+  Note that the function assumes that `predictions` and `labels` are already
+  unit-normalized.
+
+  Args:
+    labels: `Tensor` whose shape matches 'predictions'
+    predictions: An arbitrary matrix.
+    axis: The dimension along which the cosine distance is computed.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which this loss will be added.
+    reduction: Type of reduction to apply to loss.
+    dim: The old (deprecated) name for `axis`.
+
+  Returns:
+    Weighted loss float `Tensor`. If `reduction` is `NONE`, this has the same
+    shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If `predictions` shape doesn't match `labels` shape, or
+      `axis`, `labels`, `predictions` or `weights` is `None`.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  axis = deprecated_argument_lookup("axis", axis, "dim", dim)
+  if axis is None:
+    raise ValueError("You must specify argument `axis`.")
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if predictions is None:
+    raise ValueError("Argument `predictions` must not be None.")
+  with ops.name_scope(scope, "cosine_distance_loss",
+                      (predictions, labels, weights)) as scope:
+    predictions = math_ops.cast(predictions, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
+
+    radial_diffs = math_ops.multiply(predictions, labels)
+    losses = 1 - math_ops.reduce_sum(radial_diffs, axis=(axis,), keepdims=True)
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+@tf_export(v1=["losses.hinge_loss"])
+@dispatch.add_dispatch_support
+def hinge_loss(labels, logits, weights=1.0, scope=None,
+               loss_collection=ops.GraphKeys.LOSSES,
+               reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Adds a hinge loss to the training procedure.
+
+  Args:
+    labels: The ground truth output tensor. Its shape should match the shape of
+      logits. The values of the tensor are expected to be 0.0 or 1.0. Internally
+      the {0,1} labels are converted to {-1,1} when calculating the hinge loss.
+    logits: The logits, a float tensor. Note that logits are assumed to be
+      unbounded and 0-centered. A value > 0 (resp. < 0) is considered a positive
+      (resp. negative) binary prediction.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss float `Tensor`. If `reduction` is `NONE`, this has the same
+    shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shapes of `logits` and `labels` don't match or
+      if `labels` or `logits` is None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if logits is None:
+    raise ValueError("Argument `logits` must not be None.")
+  with ops.name_scope(scope, "hinge_loss", (logits, labels, weights)) as scope:
+    logits = math_ops.cast(logits, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+    logits.get_shape().assert_is_compatible_with(labels.get_shape())
+    # We first need to convert binary labels to -1/1 labels (as floats).
+    all_ones = array_ops.ones_like(labels)
+    labels = math_ops.subtract(2 * labels, all_ones)
+    losses = nn_ops.relu(
+        math_ops.subtract(all_ones, math_ops.multiply(labels, logits)))
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+@tf_export(v1=["losses.huber_loss"])
+@dispatch.add_dispatch_support
+def huber_loss(labels, predictions, weights=1.0, delta=1.0, scope=None,
+               loss_collection=ops.GraphKeys.LOSSES,
+               reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Adds a [Huber Loss](https://en.wikipedia.org/wiki/Huber_loss) term to the training procedure.
+
+  For each value x in `error=labels-predictions`, the following is calculated:
+
+  ```
+    0.5 * x^2                  if |x| <= d
+    0.5 * d^2 + d * (|x| - d)  if |x| > d
+  ```
+
+  where d is `delta`.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided, then
+  the loss is simply scaled by the given value. If `weights` is a tensor of size
+  `[batch_size]`, then the total loss for each sample of the batch is rescaled
+  by the corresponding element in the `weights` vector. If the shape of
+  `weights` matches the shape of `predictions`, then the loss of each
+  measurable element of `predictions` is scaled by the corresponding value of
+  `weights`.
+
+  Args:
+    labels: The ground truth output tensor, same dimensions as 'predictions'.
+    predictions: The predicted outputs.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    delta: `float`, the point where the huber loss function changes from a
+      quadratic to linear.
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss float `Tensor`. If `reduction` is `NONE`, this has the same
+    shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shape of `predictions` doesn't match that of `labels` or
+      if the shape of `weights` is invalid.  Also if `labels` or
+     `predictions` is None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if predictions is None:
+    raise ValueError("Argument `predictions` must not be None.")
+  with ops.name_scope(scope, "huber_loss",
+                      (predictions, labels, weights)) as scope:
+    predictions = math_ops.cast(predictions, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
+    error = math_ops.subtract(predictions, labels)
+    abs_error = math_ops.abs(error)
+    quadratic = math_ops.minimum(abs_error, delta)
+    # The following expression is the same in value as
+    # tf.maximum(abs_error - delta, 0), but importantly the gradient for the
+    # expression when abs_error == delta is 0 (for tf.maximum it would be 1).
+    # This is necessary to avoid doubling the gradient, since there is already a
+    # nonzero contribution to the gradient from the quadratic term.
+    linear = math_ops.subtract(abs_error, quadratic)
+    losses = math_ops.add(
+        math_ops.multiply(
+            ops.convert_to_tensor(0.5, dtype=quadratic.dtype),
+            math_ops.multiply(quadratic, quadratic)),
+        math_ops.multiply(delta, linear))
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+@tf_export(v1=["losses.log_loss"])
+@dispatch.add_dispatch_support
+def log_loss(labels, predictions, weights=1.0, epsilon=1e-7, scope=None,
+             loss_collection=ops.GraphKeys.LOSSES,
+             reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Adds a Log Loss term to the training procedure.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided, then
+  the loss is simply scaled by the given value. If `weights` is a tensor of size
+  `[batch_size]`, then the total loss for each sample of the batch is rescaled
+  by the corresponding element in the `weights` vector. If the shape of
+  `weights` matches the shape of `predictions`, then the loss of each
+  measurable element of `predictions` is scaled by the corresponding value of
+  `weights`.
+
+  Args:
+    labels: The ground truth output tensor, same dimensions as 'predictions'.
+    predictions: The predicted outputs.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    epsilon: A small increment to add to avoid taking a log of zero.
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss float `Tensor`. If `reduction` is `NONE`, this has the same
+    shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shape of `predictions` doesn't match that of `labels` or
+      if the shape of `weights` is invalid.  Also if `labels` or `predictions`
+      is None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if predictions is None:
+    raise ValueError("Argument `predictions` must not be None.")
+  with ops.name_scope(scope, "log_loss",
+                      (predictions, labels, weights)) as scope:
+    predictions = math_ops.cast(predictions, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
+    losses = -math_ops.multiply(
+        labels,
+        math_ops.log(predictions + epsilon)) - math_ops.multiply(
+            (1 - labels), math_ops.log(1 - predictions + epsilon))
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+# TODO(b/37208492): Add reduction arg.
+@tf_export(v1=["losses.mean_pairwise_squared_error"])
+@dispatch.add_dispatch_support
+def mean_pairwise_squared_error(
+    labels, predictions, weights=1.0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES):
+  """Adds a pairwise-errors-squared loss to the training procedure.
+
+  Unlike `mean_squared_error`, which is a measure of the differences between
+  corresponding elements of `predictions` and `labels`,
+  `mean_pairwise_squared_error` is a measure of the differences between pairs of
+  corresponding elements of `predictions` and `labels`.
+
+  For example, if `labels`=[a, b, c] and `predictions`=[x, y, z], there are
+  three pairs of differences are summed to compute the loss:
+    loss = [ ((a-b) - (x-y)).^2 + ((a-c) - (x-z)).^2 + ((b-c) - (y-z)).^2 ] / 3
+
+  Note that since the inputs are of shape `[batch_size, d0, ... dN]`, the
+  corresponding pairs are computed within each batch sample but not across
+  samples within a batch. For example, if `predictions` represents a batch of
+  16 grayscale images of dimension [batch_size, 100, 200], then the set of pairs
+  is drawn from each image, but not across images.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided, then
+  the loss is simply scaled by the given value. If `weights` is a tensor of size
+  `[batch_size]`, then the total loss for each sample of the batch is rescaled
+  by the corresponding element in the `weights` vector.
+
+  Args:
+    labels: The ground truth output tensor, whose shape must match the shape of
+      `predictions`.
+    predictions: The predicted outputs, a tensor of size
+      `[batch_size, d0, .. dN]` where N+1 is the total number of dimensions in
+      `predictions`.
+    weights: Coefficients for the loss a scalar, a tensor of shape
+      `[batch_size]` or a tensor whose shape matches `predictions`.
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+
+  Returns:
+    A scalar `Tensor` that returns the weighted loss.
+
+  Raises:
+    ValueError: If the shape of `predictions` doesn't match that of `labels` or
+      if the shape of `weights` is invalid.  Also if `labels` or `predictions`
+      is None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if predictions is None:
+    raise ValueError("Argument `predictions` must not be None.")
+  with ops.name_scope(scope, "mean_pairwise_squared_error",
+                      (predictions, labels, weights)) as scope:
+    weights = math_ops.cast(weights, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+
+    def compute_loss(labels, predictions, weights, loss_collection):
+      predictions = math_ops.cast(predictions, dtype=dtypes.float32)
+      predictions.get_shape().assert_is_compatible_with(labels.get_shape())
+
+      diffs = math_ops.subtract(predictions, labels)
+
+      axis = math_ops.range(1, array_ops.rank(diffs))
+
+      sum_squares_diff_per_batch = math_ops.reduce_sum(
+          math_ops.square(diffs), axis=axis, keepdims=True)
+      num_present_per_batch = _num_present(diffs, weights, per_batch=True)
+
+      term1 = 2.0 * math_ops.div_no_nan(
+          sum_squares_diff_per_batch,
+          math_ops.maximum(num_present_per_batch - 1, 0),
+          name="value")
+
+      sum_diff = math_ops.reduce_sum(diffs, axis=axis, keepdims=True)
+      term2 = 2.0 * math_ops.div_no_nan(
+          math_ops.square(sum_diff),
+          math_ops.maximum(
+              math_ops.multiply(num_present_per_batch,
+                                num_present_per_batch - 1), 0),
+          name="value")
+
+      weighted_losses = math_ops.multiply(term1 - term2, weights)
+      loss = math_ops.reduce_sum(weighted_losses)
+
+      mean_loss = array_ops.where(
+          math_ops.reduce_sum(num_present_per_batch) > 0,
+          loss,
+          array_ops.zeros_like(loss),
+          name="value")
+      util.add_loss(mean_loss, loss_collection)
+      return mean_loss
+
+    # Skip the assert_broadcastable in XLA context because asserts are not
+    # supported so it only causes unnecessary ops. Also skip it because it uses
+    # a DenseToDenseSetOperation op that is incompatible with XLA when
+    # the shape(s) are dynamic.
+    if control_flow_ops.get_enclosing_xla_context() is not None:
+      return compute_loss(labels, predictions, weights, loss_collection)
+    else:
+      with ops.control_dependencies(
+          (weights_broadcast_ops.assert_broadcastable(weights, labels),)):
+        return compute_loss(labels, predictions, weights, loss_collection)
+
+
+@tf_export(v1=["losses.mean_squared_error"])
+@dispatch.add_dispatch_support
+def mean_squared_error(
+    labels, predictions, weights=1.0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Adds a Sum-of-Squares loss to the training procedure.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided, then
+  the loss is simply scaled by the given value. If `weights` is a tensor of size
+  `[batch_size]`, then the total loss for each sample of the batch is rescaled
+  by the corresponding element in the `weights` vector. If the shape of
+  `weights` matches the shape of `predictions`, then the loss of each
+  measurable element of `predictions` is scaled by the corresponding value of
+  `weights`.
+
+  Args:
+    labels: The ground truth output tensor, same dimensions as 'predictions'.
+    predictions: The predicted outputs.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+      `labels`, and must be broadcastable to `labels` (i.e., all dimensions must
+      be either `1`, or the same as the corresponding `losses` dimension).
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss float `Tensor`. If `reduction` is `NONE`, this has the same
+    shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shape of `predictions` doesn't match that of `labels` or
+      if the shape of `weights` is invalid.  Also if `labels` or `predictions`
+      is None.
+
+  @compatibility(TF2)
+
+  `tf.compat.v1.losses.mean_squared_error` is mostly compatible with eager
+  execution and `tf.function`. But, the `loss_collection` argument is
+  ignored when executing eagerly and no loss will be written to the loss
+  collections. You will need to either hold on to the return value manually
+  or rely on `tf.keras.Model` loss tracking.
+
+
+  To switch to native TF2 style, instantiate the
+   `tf.keras.losses.MeanSquaredError` class and call the object instead.
+
+
+  #### Structural Mapping to Native TF2
+
+  Before:
+
+  ```python
+  loss = tf.compat.v1.losses.mean_squared_error(
+    labels=labels,
+    predictions=predictions,
+    weights=weights,
+    reduction=reduction)
+  ```
+
+  After:
+
+  ```python
+  loss_fn = tf.keras.losses.MeanSquaredError(
+    reduction=reduction)
+  loss = loss_fn(
+    y_true=labels,
+    y_pred=predictions,
+    sample_weight=weights)
+  ```
+
+  #### How to Map Arguments
+
+  | TF1 Arg Name          | TF2 Arg Name     | Note                       |
+  | :-------------------- | :--------------- | :------------------------- |
+  | `labels`              | `y_true`         | In `__call__()` method     |
+  | `predictions`         | `y_pred`         | In `__call__()` method     |
+  | `weights`             | `sample_weight`  | In `__call__()` method.    |
+  : : : The shape requirements for `sample_weight` is different from      :
+  : : : `weights`. Please check the [argument definition][api_docs] for   :
+  : : : details.                                                          :
+  | `scope`               | Not supported    | -                          |
+  | `loss_collection`     | Not supported    | Losses should be tracked   |
+  : : : explicitly or with Keras APIs, for example, [add_loss][add_loss], :
+  : : : instead of via collections                                        :
+  | `reduction`           | `reduction`      | In constructor. Value of   |
+  : : : `tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE`,              :
+  : : : `tf.compat.v1.losses.Reduction.SUM`,                              :
+  : : : `tf.compat.v1.losses.Reduction.NONE` in                           :
+  : : : `tf.compat.v1.losses.softmax_cross_entropy` correspond to         :
+  : : : `tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE`,                  :
+  : : : `tf.keras.losses.Reduction.SUM`,                                  :
+  : : : `tf.keras.losses.Reduction.NONE`, respectively. If you            :
+  : : : used other value for `reduction`, including the default value     :
+  : : :  `tf.compat.v1.losses.Reduction.SUM_BY_NONZERO_WEIGHTS`, there is :
+  : : : no directly corresponding value. Please modify the loss           :
+  : : : implementation manually.                                          :
+
+  [add_loss]:https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#add_loss
+  [api_docs]:https://www.tensorflow.org/api_docs/python/tf/keras/losses/MeanSquaredError#__call__
+
+
+  #### Before & After Usage Example
+
+  Before:
+
+  >>> y_true = [1, 2, 3]
+  >>> y_pred = [1, 3, 5]
+  >>> weights = [0, 1, 0.25]
+  >>> # samples with zero-weight are excluded from calculation when `reduction`
+  >>> # argument is set to default value `Reduction.SUM_BY_NONZERO_WEIGHTS`
+  >>> tf.compat.v1.losses.mean_squared_error(
+  ...    labels=y_true,
+  ...    predictions=y_pred,
+  ...    weights=weights).numpy()
+  1.0
+
+  >>> tf.compat.v1.losses.mean_squared_error(
+  ...    labels=y_true,
+  ...    predictions=y_pred,
+  ...    weights=weights,
+  ...    reduction=tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE).numpy()
+  0.66667
+
+  After:
+
+  >>> y_true = [[1.0], [2.0], [3.0]]
+  >>> y_pred = [[1.0], [3.0], [5.0]]
+  >>> weights = [1, 1, 0.25]
+  >>> mse = tf.keras.losses.MeanSquaredError(
+  ...    reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE)
+  >>> mse(y_true=y_true, y_pred=y_pred, sample_weight=weights).numpy()
+  0.66667
+
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if predictions is None:
+    raise ValueError("Argument `predictions` must not be None.")
+  with ops.name_scope(scope, "mean_squared_error",
+                      (predictions, labels, weights)) as scope:
+    predictions = math_ops.cast(predictions, dtype=dtypes.float32)
+    labels = math_ops.cast(labels, dtype=dtypes.float32)
+    predictions.get_shape().assert_is_compatible_with(labels.get_shape())
+    losses = math_ops.squared_difference(predictions, labels)
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+@tf_export(v1=["losses.sigmoid_cross_entropy"])
+@dispatch.add_dispatch_support
+def sigmoid_cross_entropy(
+    multi_class_labels, logits, weights=1.0, label_smoothing=0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Creates a cross-entropy loss using tf.nn.sigmoid_cross_entropy_with_logits.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided,
+  then the loss is simply scaled by the given value. If `weights` is a
+  tensor of shape `[batch_size]`, then the loss weights apply to each
+  corresponding sample.
+
+  If `label_smoothing` is nonzero, smooth the labels towards 1/2:
+
+      new_multiclass_labels = multiclass_labels * (1 - label_smoothing)
+                              + 0.5 * label_smoothing
+
+  Args:
+    multi_class_labels: `[batch_size, num_classes]` target integer labels in
+      `{0, 1}`.
+    logits: Float `[batch_size, num_classes]` logits outputs of the network.
+    weights: Optional `Tensor` whose rank is either 0, or the same rank as
+    `multi_class_labels`, and must be broadcastable to `multi_class_labels`
+    (i.e., all dimensions must be either `1`, or the same as the
+    corresponding `losses` dimension).
+    label_smoothing: If greater than `0` then smooth the labels.
+    scope: The scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss `Tensor` of the same type as `logits`. If `reduction` is
+    `NONE`, this has the same shape as `logits`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shape of `logits` doesn't match that of
+      `multi_class_labels` or if the shape of `weights` is invalid, or if
+      `weights` is None.  Also if `multi_class_labels` or `logits` is None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if multi_class_labels is None:
+    raise ValueError("Argument `multi_class_labels` must not be None.")
+  if logits is None:
+    raise ValueError("Argument `logits` must not be None.")
+  with ops.name_scope(scope, "sigmoid_cross_entropy_loss",
+                      (logits, multi_class_labels, weights)) as scope:
+    logits = ops.convert_to_tensor(logits)
+    multi_class_labels = math_ops.cast(multi_class_labels, logits.dtype)
+    logits.get_shape().assert_is_compatible_with(multi_class_labels.get_shape())
+
+    if label_smoothing > 0:
+      multi_class_labels = (multi_class_labels * (1 - label_smoothing) +
+                            0.5 * label_smoothing)
+
+    losses = nn.sigmoid_cross_entropy_with_logits(labels=multi_class_labels,
+                                                  logits=logits,
+                                                  name="xentropy")
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+@tf_export(v1=["losses.softmax_cross_entropy"])
+@dispatch.add_dispatch_support
+def softmax_cross_entropy(
+    onehot_labels, logits, weights=1.0, label_smoothing=0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  r"""Creates a cross-entropy loss using tf.nn.softmax_cross_entropy_with_logits_v2.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided,
+  then the loss is simply scaled by the given value. If `weights` is a
+  tensor of shape `[batch_size]`, then the loss weights apply to each
+  corresponding sample.
+
+  If `label_smoothing` is nonzero, smooth the labels towards 1/num_classes:
+      new_onehot_labels = onehot_labels * (1 - label_smoothing)
+                          + label_smoothing / num_classes
+
+  Note that `onehot_labels` and `logits` must have the same shape,
+  e.g. `[batch_size, num_classes]`. The shape of `weights` must be
+  broadcastable to loss, whose shape is decided by the shape of `logits`.
+  In case the shape of `logits` is `[batch_size, num_classes]`, loss is
+  a `Tensor` of shape `[batch_size]`.
+
+  Args:
+    onehot_labels: One-hot-encoded labels.
+    logits: Logits outputs of the network.
+    weights: Optional `Tensor` that is broadcastable to loss.
+    label_smoothing: If greater than 0 then smooth the labels.
+    scope: the scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss `Tensor` of the same type as `logits`. If `reduction` is
+    `NONE`, this has shape `[batch_size]`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shape of `logits` doesn't match that of `onehot_labels`
+      or if the shape of `weights` is invalid or if `weights` is None.  Also if
+      `onehot_labels` or `logits` is None.
+
+  @compatibility(TF2)
+
+  `tf.compat.v1.losses.softmax_cross_entropy` is mostly compatible with eager
+  execution and `tf.function`. But, the `loss_collection` argument is
+  ignored when executing eagerly and no loss will be written to the loss
+  collections. You will need to either hold on to the return value manually
+  or rely on `tf.keras.Model` loss tracking.
+
+
+  To switch to native TF2 style, instantiate the
+   `tf.keras.losses.CategoricalCrossentropy` class with `from_logits` set
+  as `True` and call the object instead.
+
+
+  #### Structural Mapping to Native TF2
+
+  Before:
+
+  ```python
+  loss = tf.compat.v1.losses.softmax_cross_entropy(
+    onehot_labels=onehot_labels,
+    logits=logits,
+    weights=weights,
+    label_smoothing=smoothing)
+  ```
+
+  After:
+
+  ```python
+  loss_fn = tf.keras.losses.CategoricalCrossentropy(
+    from_logits=True,
+    label_smoothing=smoothing)
+  loss = loss_fn(
+    y_true=onehot_labels,
+    y_pred=logits,
+    sample_weight=weights)
+  ```
+
+  #### How to Map Arguments
+
+  | TF1 Arg Name          | TF2 Arg Name     | Note                       |
+  | :-------------------- | :--------------- | :------------------------- |
+  |  -                    | `from_logits`    | Set `from_logits` as True  |
+  :                       :                  : to have identical behavior :
+  | `onehot_labels`       | `y_true`         | In `__call__()` method     |
+  | `logits`              | `y_pred`         | In `__call__()` method     |
+  | `weights`             | `sample_weight`  | In `__call__()` method     |
+  | `label_smoothing`     | `label_smoothing`| In constructor             |
+  | `scope`               | Not supported    | -                          |
+  | `loss_collection`     | Not supported    | Losses should be tracked   |
+  :                       :                  : explicitly or with Keras   :
+  :                       :                  : APIs, for example,         :
+  :                       :                  : [add_loss][add_loss],      :
+  :                       :                  : instead of via collections :
+  | `reduction`           | `reduction`      | In constructor. Value of   |
+  : : : `tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE`,              :
+  : : : `tf.compat.v1.losses.Reduction.SUM`,                              :
+  : : : `tf.compat.v1.losses.Reduction.NONE` in                           :
+  : : : `tf.compat.v1.losses.softmax_cross_entropy` correspond to         :
+  : : : `tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE`,                  :
+  : : : `tf.keras.losses.Reduction.SUM`,                                  :
+  : : : `tf.keras.losses.Reduction.NONE`, respectively. If you            :
+  : : : used other value for `reduction`, including the default value     :
+  : : :  `tf.compat.v1.losses.Reduction.SUM_BY_NONZERO_WEIGHTS`, there is :
+  : : : no directly corresponding value. Please modify the loss           :
+  : : : implementation manually.                                          :
+
+  [add_loss]:https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#add_loss
+
+
+  #### Before & After Usage Example
+
+  Before:
+
+  >>> y_true = [[0, 1, 0], [0, 0, 1]]
+  >>> y_pred = [[0.05, 0.95, 0], [0.1, 0.8, 0.1]]
+  >>> weights = [0.3, 0.7]
+  >>> smoothing = 0.2
+  >>> tf.compat.v1.losses.softmax_cross_entropy(y_true, y_pred, weights=weights,
+  ...   label_smoothing=smoothing).numpy()
+  0.57618
+
+  After:
+
+  >>> cce = tf.keras.losses.CategoricalCrossentropy(from_logits=True,
+  ...   label_smoothing=smoothing)
+  >>> cce(y_true, y_pred, sample_weight=weights).numpy()
+  0.57618
+
+  @end_compatibility
+  """
+  if onehot_labels is None:
+    raise ValueError("Argument `onehot_labels` must not be None.")
+  if logits is None:
+    raise ValueError("Argument `logits` must not be None.")
+  with ops.name_scope(scope, "softmax_cross_entropy_loss",
+                      (logits, onehot_labels, weights)) as scope:
+    logits = ops.convert_to_tensor(logits)
+    onehot_labels = math_ops.cast(onehot_labels, logits.dtype)
+    logits.get_shape().assert_is_compatible_with(onehot_labels.get_shape())
+
+    if label_smoothing > 0:
+      num_classes = math_ops.cast(
+          array_ops.shape(onehot_labels)[-1], logits.dtype)
+      smooth_positives = 1.0 - label_smoothing
+      smooth_negatives = label_smoothing / num_classes
+      onehot_labels = onehot_labels * smooth_positives + smooth_negatives
+
+    onehot_labels = array_ops.stop_gradient(
+        onehot_labels, name="labels_stop_gradient")
+    losses = nn.softmax_cross_entropy_with_logits_v2(
+        labels=onehot_labels, logits=logits, name="xentropy")
+
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)
+
+
+# TODO(ptucker): Merge this with similar method in metrics_impl.
+def _remove_squeezable_dimensions(
+    labels, predictions, weights=None, expected_rank_diff=0):
+  """Internal version of _remove_squeezable_dimensions which handles weights.
+
+  Squeezes `predictions` and `labels` if their ranks differ from expected by
+  exactly 1.
+  Squeezes `weights` if its rank is 1 more than the new rank of `predictions`
+
+  This will use static shape if available. Otherwise, it will add graph
+  operations, which could result in a performance hit.
+
+  Args:
+    labels: Label values, a `Tensor` whose dimensions match `predictions`.
+    predictions: Predicted values, a `Tensor` of arbitrary dimensions.
+    weights: Optional weight `Tensor`. It will be squeezed if it's not scalar,
+      and its rank is 1 more than the new rank of `labels`.
+    expected_rank_diff: Expected result of `rank(predictions) - rank(labels)`.
+
+  Returns:
+    Tuple of `predictions`, `labels` and `weights`, possibly with the last
+    dimension squeezed.
+  """
+  labels, predictions = confusion_matrix.remove_squeezable_dimensions(
+      labels, predictions, expected_rank_diff=expected_rank_diff)
+
+  if weights is not None:
+    weights = ops.convert_to_tensor(weights)
+    labels_rank = labels.get_shape().ndims
+    weights_shape = weights.get_shape()
+    weights_rank = weights_shape.ndims
+
+    if (labels_rank is not None) and (weights_rank is not None):
+      # Use static rank.
+      rank_diff = weights_rank - labels_rank
+      if rank_diff == 1:
+        weights = array_ops.squeeze(weights, [-1])
+      return labels, predictions, weights
+
+    # Use dynamic rank.
+    rank_diff = array_ops.rank(weights) - array_ops.rank(labels)
+    if (weights_rank is None) or (
+        weights_rank > 0 and weights_shape.dims[-1].is_compatible_with(1)):
+      weights = cond.cond(
+          math_ops.equal(1, rank_diff),
+          lambda: array_ops.squeeze(weights, [-1]),
+          lambda: weights)
+
+  return labels, predictions, weights
+
+
+@tf_export(v1=["losses.sparse_softmax_cross_entropy"])
+@dispatch.add_dispatch_support
+def sparse_softmax_cross_entropy(
+    labels, logits, weights=1.0, scope=None,
+    loss_collection=ops.GraphKeys.LOSSES,
+    reduction=Reduction.SUM_BY_NONZERO_WEIGHTS):
+  """Cross-entropy loss using `tf.nn.sparse_softmax_cross_entropy_with_logits`.
+
+  `weights` acts as a coefficient for the loss. If a scalar is provided,
+  then the loss is simply scaled by the given value. If `weights` is a
+  tensor of shape `[batch_size]`, then the loss weights apply to each
+  corresponding sample.
+
+  Args:
+    labels: `Tensor` of shape `[d_0, d_1, ..., d_{r-1}]` (where `r` is rank of
+      `labels` and result) and dtype `int32` or `int64`. Each entry in `labels`
+      must be an index in `[0, num_classes)`. Other values will raise an
+      exception when this op is run on CPU, and return `NaN` for corresponding
+      loss and gradient rows on GPU.
+    logits: Unscaled log probabilities of shape
+      `[d_0, d_1, ..., d_{r-1}, num_classes]` and dtype `float16`, `float32` or
+      `float64`.
+    weights: Coefficients for the loss. This must be scalar or broadcastable to
+      `labels` (i.e. same rank and each dimension is either 1 or the same).
+    scope: the scope for the operations performed in computing the loss.
+    loss_collection: collection to which the loss will be added.
+    reduction: Type of reduction to apply to loss.
+
+  Returns:
+    Weighted loss `Tensor` of the same type as `logits`. If `reduction` is
+    `NONE`, this has the same shape as `labels`; otherwise, it is scalar.
+
+  Raises:
+    ValueError: If the shapes of `logits`, `labels`, and `weights` are
+      incompatible, or if any of them are None.
+
+  @compatibility(eager)
+  The `loss_collection` argument is ignored when executing eagerly. Consider
+  holding on to the return value or collecting losses via a `tf.keras.Model`.
+  @end_compatibility
+  """
+  if labels is None:
+    raise ValueError("Argument `labels` must not be None.")
+  if logits is None:
+    raise ValueError("Argument `logits` must not be None.")
+  with ops.name_scope(scope, "sparse_softmax_cross_entropy_loss",
+                      (logits, labels, weights)) as scope:
+    # As documented above in Args, labels contain class IDs and logits contains
+    # 1 probability per class ID, so we expect rank(logits) - rank(labels) == 1;
+    # therefore, expected_rank_diff=1.
+    labels, logits, weights = _remove_squeezable_dimensions(
+        labels, logits, weights, expected_rank_diff=1)
+    losses = nn.sparse_softmax_cross_entropy_with_logits(labels=labels,
+                                                         logits=logits,
+                                                         name="xentropy")
+    return compute_weighted_loss(
+        losses, weights, scope, loss_collection, reduction=reduction)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/losses/util.py

@@ -0,0 +1,263 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for manipulating the loss collections."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import confusion_matrix
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util.tf_export import tf_export
+
+
+def squeeze_or_expand_dimensions(y_pred, y_true=None, sample_weight=None):
+  """Squeeze or expand last dimension if needed.
+
+  1. Squeezes last dim of `y_pred` or `y_true` if their rank differs by 1
+  (using `confusion_matrix.remove_squeezable_dimensions`).
+  2. Squeezes or expands last dim of `sample_weight` if its rank differs by 1
+  from the new rank of `y_pred`.
+  If `sample_weight` is scalar, it is kept scalar.
+
+  This will use static shape if available. Otherwise, it will add graph
+  operations, which could result in a performance hit.
+
+  Args:
+    y_pred: Predicted values, a `Tensor` of arbitrary dimensions.
+    y_true: Optional label `Tensor` whose dimensions match `y_pred`.
+    sample_weight: Optional weight scalar or `Tensor` whose dimensions match
+      `y_pred`.
+
+  Returns:
+    Tuple of `y_pred`, `y_true` and `sample_weight`. Each of them possibly has
+    the last dimension squeezed,
+    `sample_weight` could be extended by one dimension.
+    If `sample_weight` is None, (y_pred, y_true) is returned.
+  """
+  y_pred_shape = y_pred.shape
+  y_pred_rank = y_pred_shape.ndims
+  if y_true is not None:
+
+    # If sparse matrix is provided as `y_true`, the last dimension in `y_pred`
+    # may be > 1. Eg: y_true = [0, 1, 2] (shape=(3,)),
+    # y_pred = [[.9, .05, .05], [.5, .89, .6], [.05, .01, .94]] (shape=(3, 3))
+    # In this case, we should not try to remove squeezable dimension.
+    y_true_shape = y_true.shape
+    y_true_rank = y_true_shape.ndims
+    if (y_true_rank is not None) and (y_pred_rank is not None):
+      # Use static rank for `y_true` and `y_pred`.
+      if (y_pred_rank - y_true_rank != 1) or y_pred_shape[-1] == 1:
+        y_true, y_pred = confusion_matrix.remove_squeezable_dimensions(
+            y_true, y_pred)
+    else:
+      # Use dynamic rank.
+      rank_diff = array_ops.rank(y_pred) - array_ops.rank(y_true)
+      squeeze_dims = lambda: confusion_matrix.remove_squeezable_dimensions(  # pylint: disable=g-long-lambda
+          y_true, y_pred)
+      is_last_dim_1 = math_ops.equal(1, array_ops.shape(y_pred)[-1])
+      maybe_squeeze_dims = lambda: cond.cond(  # pylint: disable=g-long-lambda
+          is_last_dim_1, squeeze_dims, lambda: (y_true, y_pred))
+      y_true, y_pred = cond.cond(
+          math_ops.equal(1, rank_diff), maybe_squeeze_dims, squeeze_dims)
+
+  if sample_weight is None:
+    return y_pred, y_true
+
+  weights_shape = sample_weight.shape
+  weights_rank = weights_shape.ndims
+  if weights_rank == 0:  # If weights is scalar, do nothing.
+    return y_pred, y_true, sample_weight
+
+  if (y_pred_rank is not None) and (weights_rank is not None):
+    # Use static rank.
+    if weights_rank - y_pred_rank == 1:
+      sample_weight = array_ops.squeeze(sample_weight, [-1])
+    elif y_pred_rank - weights_rank == 1:
+      sample_weight = array_ops.expand_dims(sample_weight, [-1])
+    return y_pred, y_true, sample_weight
+
+  # Use dynamic rank.
+  weights_rank_tensor = array_ops.rank(sample_weight)
+  rank_diff = weights_rank_tensor - array_ops.rank(y_pred)
+  maybe_squeeze_weights = lambda: array_ops.squeeze(sample_weight, [-1])
+
+  def _maybe_expand_weights():
+    expand_weights = lambda: array_ops.expand_dims(sample_weight, [-1])
+    return cond.cond(
+        math_ops.equal(rank_diff, -1), expand_weights, lambda: sample_weight)
+
+  def _maybe_adjust_weights():
+    return cond.cond(
+        math_ops.equal(rank_diff, 1), maybe_squeeze_weights,
+        _maybe_expand_weights)
+
+  # squeeze or expand last dim of `sample_weight` if its rank differs by 1
+  # from the new rank of `y_pred`.
+  sample_weight = cond.cond(
+      math_ops.equal(weights_rank_tensor, 0), lambda: sample_weight,
+      _maybe_adjust_weights)
+  return y_pred, y_true, sample_weight
+
+
+def scale_losses_by_sample_weight(losses, sample_weight):
+  """Scales loss values by the given sample weights.
+
+  `sample_weight` dimensions are updated to match with the dimension of `losses`
+  if possible by using squeeze/expand/broadcast.
+
+  Args:
+    losses: Loss tensor.
+    sample_weight: Sample weights tensor.
+
+  Returns:
+    `losses` scaled by `sample_weight` with dtype float32.
+  """
+  # TODO(psv): Handle the casting here in a better way, eg. if losses is float64
+  # we do not want to lose precision.
+  losses = math_ops.cast(losses, dtypes.float32)
+  sample_weight = math_ops.cast(sample_weight, dtypes.float32)
+
+  # Update dimensions of `sample_weight` to match with `losses` if possible.
+  losses, _, sample_weight = squeeze_or_expand_dimensions(
+      losses, None, sample_weight)
+  return math_ops.multiply(losses, sample_weight)
+
+
+@tf_contextlib.contextmanager
+def check_per_example_loss_rank(per_example_loss):
+  """Context manager that checks that the rank of per_example_loss is at least 1.
+
+  Args:
+    per_example_loss: Per example loss tensor.
+
+  Yields:
+    A context manager.
+  """
+  loss_rank = per_example_loss.shape.rank
+  if loss_rank is not None:
+    # Handle static rank.
+    if loss_rank == 0:
+      raise ValueError(
+          "Invalid value passed for `per_example_loss`. Expected a tensor with "
+          f"at least rank 1. Received per_example_loss={per_example_loss} with "
+          f"rank {loss_rank}")
+    yield
+  else:
+    # Handle dynamic rank.
+    with ops.control_dependencies([
+        check_ops.assert_greater_equal(
+            array_ops.rank(per_example_loss),
+            math_ops.cast(1, dtype=dtypes.int32),
+            message="Invalid value passed for `per_example_loss`. Expected a "
+            "tensor with at least rank 1.")
+    ]):
+      yield
+
+
+@tf_export(v1=["losses.add_loss"])
+def add_loss(loss, loss_collection=ops.GraphKeys.LOSSES):
+  """Adds a externally defined loss to the collection of losses.
+
+  Args:
+    loss: A loss `Tensor`.
+    loss_collection: Optional collection to add the loss to.
+  """
+  # Since we have no way of figuring out when a training iteration starts or
+  # ends, holding on to a loss when executing eagerly is indistinguishable from
+  # leaking memory. We instead leave the collection empty.
+  if loss_collection and not context.executing_eagerly():
+    ops.add_to_collection(loss_collection, loss)
+
+
+@tf_export(v1=["losses.get_losses"])
+def get_losses(scope=None, loss_collection=ops.GraphKeys.LOSSES):
+  """Gets the list of losses from the loss_collection.
+
+  Args:
+    scope: An optional scope name for filtering the losses to return.
+    loss_collection: Optional losses collection.
+
+  Returns:
+    a list of loss tensors.
+  """
+  return ops.get_collection(loss_collection, scope)
+
+
+@tf_export(v1=["losses.get_regularization_losses"])
+def get_regularization_losses(scope=None):
+  """Gets the list of regularization losses.
+
+  Args:
+    scope: An optional scope name for filtering the losses to return.
+
+  Returns:
+    A list of regularization losses as Tensors.
+  """
+  return ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES, scope)
+
+
+@tf_export(v1=["losses.get_regularization_loss"])
+def get_regularization_loss(scope=None, name="total_regularization_loss"):
+  """Gets the total regularization loss.
+
+  Args:
+    scope: An optional scope name for filtering the losses to return.
+    name: The name of the returned tensor.
+
+  Returns:
+    A scalar regularization loss.
+  """
+  losses = get_regularization_losses(scope)
+  if losses:
+    return math_ops.add_n(losses, name=name)
+  else:
+    return constant_op.constant(0.0)
+
+
+@tf_export(v1=["losses.get_total_loss"])
+def get_total_loss(add_regularization_losses=True,
+                   name="total_loss",
+                   scope=None):
+  """Returns a tensor whose value represents the total loss.
+
+  In particular, this adds any losses you have added with `tf.add_loss()` to
+  any regularization losses that have been added by regularization parameters
+  on layers constructors e.g. `tf.layers`. Be very sure to use this if you
+  are constructing a loss_op manually. Otherwise regularization arguments
+  on `tf.layers` methods will not function.
+
+  Args:
+    add_regularization_losses: A boolean indicating whether or not to use the
+      regularization losses in the sum.
+    name: The name of the returned tensor.
+    scope: An optional scope name for filtering the losses to return. Note that
+      this filters the losses added with `tf.add_loss()` as well as the
+      regularization losses to that scope.
+
+  Returns:
+    A `Tensor` whose value represents the total loss.
+
+  Raises:
+    ValueError: if `losses` is not iterable.
+  """
+  losses = get_losses(scope=scope)
+  if add_regularization_losses:
+    losses += get_regularization_losses(scope=scope)
+  return math_ops.add_n(losses, name=name)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/__init__.py

@@ -0,0 +1,171 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""# tf.experimental.numpy: NumPy API on TensorFlow.
+
+This module provides a subset of NumPy API, built on top of TensorFlow
+operations. APIs are based on and have been tested with NumPy 1.16 version.
+
+The set of supported APIs may be expanded over time. Also future releases may
+change the baseline version of NumPy API being supported. A list of some
+systematic differences with NumPy is listed later in the "Differences with
+NumPy" section.
+
+## Getting Started
+
+Please also see [TensorFlow NumPy Guide](
+https://www.tensorflow.org/guide/tf_numpy).
+
+In the code snippets below, we will assume that `tf.experimental.numpy` is
+imported as `tnp` and NumPy is imported as `np`
+
+```python
+print(tnp.ones([2,1]) + np.ones([1, 2]))
+```
+
+## Types
+
+The module provides an `ndarray` class which wraps an immutable `tf.Tensor`.
+Additional functions are provided which accept array-like objects. Here
+array-like objects include `ndarrays` as defined by this module, as well as
+`tf.Tensor`, in addition to types accepted by NumPy.
+
+A subset of NumPy dtypes are supported. Type promotion* follows NumPy
+semantics.
+
+**Note**: A new type promotion that offers a lot of advantages over the old
+type promotion is now available. Learn more about enabling the new
+type promotion
+[here](https://www.tensorflow.org/guide/tf_numpy_type_promotion).
+
+```python
+print(tnp.ones([1, 2], dtype=tnp.int16) + tnp.ones([2, 1], dtype=tnp.uint8))
+```
+
+## Array Interface
+
+The `ndarray` class implements the `__array__` interface. This should allow
+these objects to be passed into contexts that expect a NumPy or array-like
+object (e.g. matplotlib).
+
+```python
+np.sum(tnp.ones([1, 2]) + np.ones([2, 1]))
+```
+
+
+## TF Interoperability
+
+The TF-NumPy API calls can be interleaved with TensorFlow calls
+without incurring Tensor data copies. This is true even if the `ndarray` or
+`tf.Tensor` is placed on a non-CPU device.
+
+In general, the expected behavior should be on par with that of code involving
+`tf.Tensor` and running stateless TensorFlow functions on them.
+
+```python
+tnp.sum(tnp.ones([1, 2]) + tf.ones([2, 1]))
+```
+
+Note that the `__array_priority__` is currently chosen to be lower than
+`tf.Tensor`. Hence the `+` operator above returns a `tf.Tensor`.
+
+Additional examples of interoperability include:
+
+*  using `with tf.GradientTape()` scope to compute gradients through the
+  TF-NumPy API calls.
+*  using `tf.distribution.Strategy` scope for distributed execution
+*  using `tf.vectorized_map()` for speeding up code using auto-vectorization
+
+
+
+## Device Support
+
+Given that `ndarray` and functions wrap TensorFlow constructs, the code will
+have GPU and TPU support on par with TensorFlow. Device placement can be
+controlled by using `with tf.device` scopes. Note that these devices could
+be local or remote.
+
+```python
+with tf.device("GPU:0"):
+  x = tnp.ones([1, 2])
+print(tf.convert_to_tensor(x).device)
+```
+
+## Graph and Eager Modes
+
+Eager mode execution should typically match NumPy semantics of executing
+op-by-op. However the same code can be executed in graph mode, by putting it
+inside a `tf.function`. The function body can contain NumPy code, and the inputs
+can be `ndarray` as well.
+
+```python
+@tf.function
+def f(x, y):
+  return tnp.sum(x + y)
+
+f(tnp.ones([1, 2]), tf.ones([2, 1]))
+```
+Python control flow based on `ndarray` values will be translated by
+[autograph](https://www.tensorflow.org/code/tensorflow/python/autograph/g3doc/reference/index.md)
+into `tf.cond` and `tf.while_loop` constructs. The code can be XLA compiled
+for further optimizations.
+
+However, note that graph mode execution can change behavior of certain
+operations since symbolic execution may not have information that is computed
+during runtime. Some differences are:
+
+*   Shapes can be incomplete or unknown in graph mode. This means that
+    `ndarray.shape`, `ndarray.size` and `ndarray.ndim` can return `ndarray`
+    objects instead of returning integer (or tuple of integer) values.
+*   `__len__`, `__iter__` and `__index__` properties of `ndarray`
+    may similarly not be supported in graph mode. Code using these
+    may need to change to explicit shape operations or control flow
+    constructs.
+*   Also note the [autograph limitations](
+https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md).
+
+
+## Mutation and Variables
+
+`ndarrays` currently wrap immutable `tf.Tensor`. Hence mutation
+operations like slice assigns are not supported. This may change in the future.
+Note however that one can directly construct a `tf.Variable` and use that with
+the TF-NumPy APIs.
+
+```python
+tf_var = tf.Variable(2.0)
+tf_var.assign_add(tnp.square(tf_var))
+```
+
+## Differences with NumPy
+
+Here is a non-exhaustive list of differences:
+
+*   Not all dtypes are currently supported. e.g. `np.float96`, `np.float128`.
+    `np.object_`, `np.str_`, `np.recarray` types are not supported.
+*   `ndarray` storage is in C order only. Fortran order, views, `stride_tricks`
+    are not supported.
+*   Only a subset of functions and modules are supported. This set will be
+    expanded over time. For supported functions, some arguments or argument
+    values may not be supported. These differences are generally provided in the
+    function comments. Full `ufunc` support is also not provided.
+*   Buffer mutation is currently not supported. `ndarrays` wrap immutable
+    tensors. This means that output buffer arguments (e.g. `out` in ufuncs) are
+    not supported.
+*   NumPy C API is not supported. NumPy's Cython and Swig integration are not
+    supported.
+
+API docstring: tensorflow.experimental.numpy
+"""
+# TODO(wangpeng): Append `tf_export`ed symbols to the comments above.

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_array_ops.py

@@ -0,0 +1,2116 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Common array methods."""
+# pylint: disable=g-direct-tensorflow-import
+
+import builtins
+import enum
+import functools
+import math
+import numbers
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import manip_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sort_ops
+from tensorflow.python.ops.numpy_ops import np_arrays
+from tensorflow.python.ops.numpy_ops import np_dtypes
+from tensorflow.python.ops.numpy_ops import np_utils
+from tensorflow.python.types import core as core_tf_types
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_export
+
+
+newaxis = np.newaxis
+tf_export.tf_export('experimental.numpy.newaxis', v1=[]).export_constant(
+    __name__, 'newaxis'
+)
+
+
+@tf_export.tf_export('experimental.numpy.empty', v1=[])
+@np_utils.np_doc('empty')
+def empty(shape, dtype=float):  # pylint: disable=redefined-outer-name
+  return zeros(shape, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.empty_like', v1=[])
+@np_utils.np_doc('empty_like')
+def empty_like(a, dtype=None):
+  return zeros_like(a, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.zeros', v1=[])
+@np_utils.np_doc('zeros')
+def zeros(shape, dtype=float):  # pylint: disable=redefined-outer-name
+  dtype = (
+      np_utils.result_type(dtype) if dtype else np_dtypes.default_float_type()
+  )
+  return array_ops.zeros(shape, dtype=dtype)
+
+
+@tf_export.tf_export('experimental.numpy.zeros_like', v1=[])
+@np_utils.np_doc('zeros_like')
+def zeros_like(a, dtype=None):  # pylint: disable=missing-docstring
+  dtype = np_utils.result_type_unary(a, dtype)
+
+  dtype = dtypes.as_dtype(dtype)  # Work around b/149877262
+  return array_ops.zeros_like(a, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.ones', v1=[])
+@np_utils.np_doc('ones')
+def ones(shape, dtype=float):  # pylint: disable=redefined-outer-name
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  return array_ops.ones(shape, dtype=dtype)
+
+
+@tf_export.tf_export('experimental.numpy.ones_like', v1=[])
+@np_utils.np_doc('ones_like')
+def ones_like(a, dtype=None):
+  dtype = np_utils.result_type_unary(a, dtype)
+  return array_ops.ones_like(a, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.eye', v1=[])
+@np_utils.np_doc('eye')
+def eye(N, M=None, k=0, dtype=float):  # pylint: disable=invalid-name,missing-docstring
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  if not M:
+    M = N
+  # Making sure N, M and k are `int`
+  N = int(N)
+  M = int(M)
+  k = int(k)
+  if k >= M or -k >= N:
+    # tf.linalg.diag will raise an error in this case
+    return zeros([N, M], dtype=dtype)
+  if k == 0:
+    return linalg_ops.eye(N, M, dtype=dtype)
+  # We need the precise length, otherwise tf.linalg.diag will raise an error
+  diag_len = builtins.min(N, M)
+  if k > 0:
+    if N >= M:
+      diag_len -= k
+    elif N + k > M:
+      diag_len = M - k
+  elif k <= 0:
+    if M >= N:
+      diag_len += k
+    elif M - k > N:
+      diag_len = N + k
+  diagonal_ = array_ops.ones([diag_len], dtype=dtype)
+  return array_ops.matrix_diag(diagonal=diagonal_, num_rows=N, num_cols=M, k=k)
+
+
+@tf_export.tf_export('experimental.numpy.identity', v1=[])
+@np_utils.np_doc('identity')
+def identity(n, dtype=float):
+  return eye(N=n, M=n, dtype=dtype)
+
+
+@tf_export.tf_export('experimental.numpy.full', v1=[])
+@np_utils.np_doc('full')
+def full(shape, fill_value, dtype=None):  # pylint: disable=redefined-outer-name
+  if not isinstance(shape, np_arrays.ndarray):
+    shape = asarray(np_arrays.convert_to_tensor(shape, dtype_hint=np.int32))
+  shape = atleast_1d(shape)
+  fill_value = asarray(fill_value, dtype=dtype)
+  return array_ops.broadcast_to(fill_value, shape)
+
+
+# Using doc only here since np full_like signature doesn't seem to have the
+# shape argument (even though it exists in the documentation online).
+@tf_export.tf_export('experimental.numpy.full_like', v1=[])
+@np_utils.np_doc_only('full_like')
+def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):  # pylint: disable=missing-docstring,redefined-outer-name
+  """order, subok and shape arguments mustn't be changed."""
+  if order != 'K':
+    raise ValueError('Non-standard orders are not supported.')
+  if not subok:
+    raise ValueError('subok being False is not supported.')
+  if shape:
+    raise ValueError('Overriding the shape is not supported.')
+
+  a = asarray(a)
+  dtype = dtype or np_utils.result_type(a)
+  fill_value = asarray(fill_value, dtype=dtype)
+  return array_ops.broadcast_to(fill_value, array_ops.shape(a))
+
+
+def _array_internal(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-outer-name
+  """Main implementation of np.array()."""
+  result_t = val
+
+  if not isinstance(result_t, tensor_lib.Tensor):
+    dtype = np_utils.result_type_unary(result_t, dtype)
+    # We can't call `convert_to_tensor(result_t, dtype=dtype)` here because
+    # convert_to_tensor doesn't allow incompatible arguments such as (5.5, int)
+    # while np.array allows them. We need to convert-then-cast.
+
+    # EagerTensor conversion complains about "mixed types" when converting
+    # tensors with no dtype information. This is because it infers types based
+    # on one selected item in the list. So e.g. when converting [2., 2j]
+    # to a tensor, it will select float32 as the inferred type and not be able
+    # to convert the list to a float 32 tensor.
+    # Since we have some information about the final dtype we care about, we
+    # supply that information so that convert_to_tensor will do best-effort
+    # conversion to that dtype first.
+    result_t = np_arrays.convert_to_tensor(result_t, dtype_hint=dtype)
+    result_t = math_ops.cast(result_t, dtype=dtype)
+  elif dtype:
+    result_t = math_ops.cast(result_t, dtype)
+
+  if copy:
+    result_t = array_ops.identity(result_t)
+
+  max_ndmin = 32
+  if ndmin > max_ndmin:
+    raise ValueError(
+        f'ndmin bigger than allowable number of dimensions: {max_ndmin}.'
+    )
+
+  if ndmin == 0:
+    return result_t
+
+  ndims = array_ops.rank(result_t)
+
+  def true_fn():
+    old_shape = array_ops.shape(result_t)
+    new_shape = array_ops.concat(
+        [array_ops.ones(ndmin - ndims, dtypes.int32), old_shape], axis=0
+    )
+    return array_ops.reshape(result_t, new_shape)
+
+  result_t = np_utils.cond(
+      np_utils.greater(ndmin, ndims), true_fn, lambda: result_t
+  )
+  return result_t
+
+
+# TODO(wangpeng): investigate whether we can make `copy` default to False.
+# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
+@tf_export.tf_export('experimental.numpy.array', v1=[])
+@np_utils.np_doc_only('array')
+def array(val, dtype=None, copy=True, ndmin=0):  # pylint: disable=redefined-outer-name
+  """Since Tensors are immutable, a copy is made only if val is placed on a
+
+  different device than the current one. Even if `copy` is False, a new Tensor
+  may need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
+  is an ndarray or a Tensor.
+  """  # pylint:disable=g-docstring-missing-newline
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  return _array_internal(val, dtype, copy, ndmin)
+
+
+# pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-return-or-yield,g-doc-args
+
+
+@tf_export.tf_export('experimental.numpy.asarray', v1=[])
+@np_utils.np_doc('asarray')
+def asarray(a, dtype=None):
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  if isinstance(a, np_arrays.ndarray) and (
+      not dtype or dtype == a.dtype.as_numpy_dtype
+  ):
+    return a
+  return array(a, dtype, copy=False)
+
+
+@tf_export.tf_export('experimental.numpy.asanyarray', v1=[])
+@np_utils.np_doc('asanyarray')
+def asanyarray(a, dtype=None):
+  return asarray(a, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.ascontiguousarray', v1=[])
+@np_utils.np_doc('ascontiguousarray')
+def ascontiguousarray(a, dtype=None):
+  return array(a, dtype, ndmin=1)
+
+
+# Numerical ranges.
+@tf_export.tf_export('experimental.numpy.arange', v1=[])
+@np_utils.np_doc('arange')
+def arange(start, stop=None, step=1, dtype=None):
+  """Returns `step`-separated values in the range [start, stop).
+
+  Args:
+    start: Start of the interval. Included in the range.
+    stop: End of the interval. If not specified, `start` is treated as 0 and
+      `start` value is used as `stop`. If specified, it is not included in the
+      range if `step` is integer. When `step` is floating point, it may or may
+      not be included.
+    step: The difference between 2 consecutive values in the output range. It is
+      recommended to use `linspace` instead of using non-integer values for
+      `step`.
+    dtype: Optional. Type of the resulting ndarray. Could be a python type, a
+      NumPy type or a TensorFlow `DType`. If not provided, the largest type of
+      `start`, `stop`, `step` is used.
+
+  Raises:
+    ValueError: If step is zero.
+  """
+  if not step:
+    raise ValueError('step must be non-zero.')
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  else:
+    if stop is None:
+      dtype = np_utils.result_type(start, step)
+    else:
+      dtype = np_utils.result_type(start, step, stop)
+  if step > 0 and (
+      (stop is not None and start > stop) or (stop is None and start < 0)
+  ):
+    return array([], dtype=dtype)
+  if step < 0 and (
+      (stop is not None and start < stop) or (stop is None and start > 0)
+  ):
+    return array([], dtype=dtype)
+  # TODO(srbs): There are some bugs when start or stop is float type and dtype
+  # is integer type.
+  return math_ops.cast(
+      math_ops.range(start, limit=stop, delta=step), dtype=dtype
+  )
+
+
+# Building matrices.
+@tf_export.tf_export('experimental.numpy.diag', v1=[])
+@np_utils.np_doc('diag')
+def diag(v, k=0):  # pylint: disable=missing-docstring
+  """Raises an error if input is not 1- or 2-d."""
+  v = asarray(v)
+  v_rank = array_ops.rank(v)
+
+  v.shape.with_rank_at_most(2)
+
+  # TODO(nareshmodi): Consider a np_utils.Assert version that will fail during
+  # tracing time if the shape is known.
+  control_flow_assert.Assert(
+      np_utils.logical_or(math_ops.equal(v_rank, 1), math_ops.equal(v_rank, 2)),
+      [v_rank],
+  )
+
+  def _diag(v, k):
+    return np_utils.cond(
+        math_ops.equal(array_ops.size(v), 0),
+        lambda: array_ops.zeros([abs(k), abs(k)], dtype=v.dtype),
+        lambda: array_ops.matrix_diag(v, k=k),
+    )
+
+  def _diag_part(v, k):
+    v_shape = array_ops.shape(v)
+    v, k = np_utils.cond(
+        np_utils.logical_or(
+            np_utils.less_equal(k, -1 * np_utils.getitem(v_shape, 0)),
+            np_utils.greater_equal(k, np_utils.getitem(v_shape, 1)),
+        ),
+        lambda: (array_ops.zeros([0, 0], dtype=v.dtype), 0),
+        lambda: (v, k),
+    )
+    result = array_ops.matrix_diag_part(v, k=k)
+    return result
+
+  result = np_utils.cond(
+      math_ops.equal(v_rank, 1), lambda: _diag(v, k), lambda: _diag_part(v, k)
+  )
+  return result
+
+
+@tf_export.tf_export('experimental.numpy.diagonal', v1=[])
+@np_utils.np_doc('diagonal')
+def diagonal(a, offset=0, axis1=0, axis2=1):  # pylint: disable=missing-docstring
+  a = asarray(a)
+
+  maybe_rank = a.shape.rank
+  if (
+      maybe_rank is not None
+      and offset == 0
+      and (axis1 == maybe_rank - 2 or axis1 == -2)
+      and (axis2 == maybe_rank - 1 or axis2 == -1)
+  ):
+    return array_ops.matrix_diag_part(a)
+
+  a = moveaxis(a, (axis1, axis2), (-2, -1))
+
+  a_shape = array_ops.shape(a)
+
+  def _zeros():  # pylint: disable=missing-docstring
+    return (
+        array_ops.zeros(
+            array_ops.concat([a_shape[:-1], [0]], 0), dtype=a.dtype
+        ),
+        0,
+    )
+
+  # All zeros since diag_part doesn't handle all possible k (aka offset).
+  # Written this way since cond will run shape inference on both branches,
+  # and diag_part shape inference will fail when offset is out of bounds.
+  a, offset = np_utils.cond(
+      np_utils.logical_or(
+          np_utils.less_equal(offset, -1 * np_utils.getitem(a_shape, -2)),
+          np_utils.greater_equal(offset, np_utils.getitem(a_shape, -1)),
+      ),
+      _zeros,
+      lambda: (a, offset),
+  )
+
+  a = array_ops.matrix_diag_part(a, k=offset)
+  return a
+
+
+@tf_export.tf_export('experimental.numpy.diagflat', v1=[])
+@np_utils.np_doc('diagflat')
+def diagflat(v, k=0):
+  v = asarray(v)
+  return diag(array_ops.reshape(v, [-1]), k)
+
+
+def _promote_dtype(*arrays):
+  dtype = np_utils.result_type(*arrays)
+
+  def _fast_asarray(a):
+    if isinstance(a, np_arrays.ndarray) and dtype == a.dtype.as_numpy_dtype:
+      return a
+    return _array_internal(a, dtype=dtype, copy=False)
+
+  return [_fast_asarray(a) for a in arrays]
+
+
+def _promote_dtype_binary(t1, t2):
+  dtype = np_utils._result_type_binary(t1, t2)  # pylint: disable=protected-access
+  if not (
+      isinstance(t1, np_arrays.ndarray) and dtype == t1.dtype.as_numpy_dtype
+  ):
+    t1 = _array_internal(t1, dtype=dtype, copy=False)
+  if not (
+      isinstance(t2, np_arrays.ndarray) and dtype == t2.dtype.as_numpy_dtype
+  ):
+    t2 = _array_internal(t2, dtype=dtype, copy=False)
+  return t1, t2
+
+
+@tf_export.tf_export('experimental.numpy.all', v1=[])
+@np_utils.np_doc('all')
+def all(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
+  a = asarray(a, dtype=bool)
+  return math_ops.reduce_all(input_tensor=a, axis=axis, keepdims=keepdims)
+
+
+@tf_export.tf_export('experimental.numpy.any', v1=[])
+@np_utils.np_doc('any')
+def any(a, axis=None, keepdims=None):  # pylint: disable=redefined-builtin
+  a = asarray(a, dtype=bool)
+  return math_ops.reduce_any(input_tensor=a, axis=axis, keepdims=keepdims)
+
+
+@tf_export.tf_export('experimental.numpy.compress', v1=[])
+@np_utils.np_doc('compress')
+def compress(condition, a, axis=None):  # pylint: disable=redefined-outer-name,missing-function-docstring
+  condition = asarray(condition, dtype=bool)
+  a = asarray(a)
+
+  if condition.ndim != 1:
+    raise ValueError('condition must be a 1-d array.')
+  # `np.compress` treats scalars as 1-d arrays.
+  if a.ndim == 0:
+    a = ravel(a)
+
+  if axis is None:
+    a = ravel(a)
+    axis = 0
+
+  if axis < 0:
+    axis += a.ndim
+
+  assert axis >= 0 and axis < a.ndim
+
+  # `tf.boolean_mask` requires the first dimensions of array and condition to
+  # match. `np.compress` pads condition with False when it is shorter.
+  condition_t = condition
+  a_t = a
+  if condition.shape[0] < a.shape[axis]:
+    padding = array_ops.fill([a.shape[axis] - condition.shape[0]], False)
+    condition_t = array_ops.concat([condition_t, padding], axis=0)
+  return array_ops.boolean_mask(tensor=a_t, mask=condition_t, axis=axis)
+
+
+@tf_export.tf_export('experimental.numpy.copy', v1=[])
+@np_utils.np_doc('copy')
+def copy(a):
+  return array(a, copy=True)
+
+
+def _maybe_promote_to_int(a):
+  if dtypes.as_dtype(a.dtype).is_integer:
+    # If a is an integer type and its precision is less than that of `int`,
+    # the output type will be `int`.
+    a_numpy_dtype = a.dtype.as_numpy_dtype
+    output_type = np.promote_types(a_numpy_dtype, int)
+    if output_type != a_numpy_dtype:
+      a = asarray(a, dtype=output_type)
+
+  return a
+
+
+@tf_export.tf_export('experimental.numpy.cumprod', v1=[])
+@np_utils.np_doc('cumprod')
+def cumprod(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
+  a = asarray(a, dtype=dtype)
+
+  if dtype is None:
+    a = _maybe_promote_to_int(a)
+
+  # If axis is None, the input is flattened.
+  if axis is None:
+    a = ravel(a)
+    axis = 0
+  elif axis < 0:
+    axis += array_ops.rank(a)
+  return math_ops.cumprod(a, axis)
+
+
+@tf_export.tf_export('experimental.numpy.cumsum', v1=[])
+@np_utils.np_doc('cumsum')
+def cumsum(a, axis=None, dtype=None):  # pylint: disable=missing-docstring
+  a = asarray(a, dtype=dtype)
+
+  if dtype is None:
+    a = _maybe_promote_to_int(a)
+
+  # If axis is None, the input is flattened.
+  if axis is None:
+    a = ravel(a)
+    axis = 0
+  elif axis < 0:
+    axis += array_ops.rank(a)
+  return math_ops.cumsum(a, axis)
+
+
+@tf_export.tf_export('experimental.numpy.imag', v1=[])
+@np_utils.np_doc('imag')
+def imag(val):
+  val = asarray(val)
+  # TODO(srbs): np.imag returns a scalar if `val` is a scalar, whereas we always
+  # return an ndarray.
+  return math_ops.imag(val)
+
+
+_TO_INT_ = 0
+_TO_FLOAT = 1
+
+
+def _reduce(
+    tf_fn,
+    a,
+    axis=None,
+    dtype=None,
+    keepdims=None,
+    promote_int=_TO_INT_,
+    tf_bool_fn=None,
+    preserve_bool=False,
+):
+  """A general reduction function.
+
+  Args:
+    tf_fn: the TF reduction function.
+    a: the array to be reduced.
+    axis: (optional) the axis along which to do the reduction. If None, all
+      dimensions are reduced.
+    dtype: (optional) the dtype of the result.
+    keepdims: (optional) whether to keep the reduced dimension(s).
+    promote_int: how to promote integer and bool inputs. There are three
+      choices. (1) `_TO_INT_` always promotes them to np.int_ or np.uint; (2)
+      `_TO_FLOAT` always promotes them to a float type (determined by
+      dtypes.default_float_type); (3) None: don't promote.
+    tf_bool_fn: (optional) the TF reduction function for bool inputs. It will
+      only be used if `dtype` is explicitly set to `np.bool_` or if `a`'s dtype
+      is `np.bool_` and `preserve_bool` is True.
+    preserve_bool: a flag to control whether to use `tf_bool_fn` if `a`'s dtype
+      is `np.bool_` (some reductions such as np.sum convert bools to integers,
+      while others such as np.max preserve bools.
+
+  Returns:
+    An ndarray.
+  """
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  if keepdims is None:
+    keepdims = False
+  a = asarray(a, dtype=dtype)
+  if (
+      dtype == np.bool_ or preserve_bool and a.dtype == np.bool_
+  ) and tf_bool_fn is not None:
+    return tf_bool_fn(input_tensor=a, axis=axis, keepdims=keepdims)
+  if dtype is None:
+    dtype = a.dtype.as_numpy_dtype
+    if np.issubdtype(dtype, np.integer) or dtype == np.bool_:
+      if promote_int == _TO_INT_:
+        # If a is an integer/bool type and whose bit width is less than np.int_,
+        # numpy up-casts it to np.int_ based on the documentation at
+        # https://numpy.org/doc/1.18/reference/generated/numpy.sum.html
+        if dtype == np.bool_:
+          is_signed = True
+          width = 8  # We can use any number here that is less than 64
+        else:
+          is_signed = np.issubdtype(dtype, np.signedinteger)
+          width = np.iinfo(dtype).bits
+        # Numpy int_ and uint are defined as 'long' and 'unsigned long', so
+        # should have the same bit width.
+        if ops.is_auto_dtype_conversion_enabled():
+          # We default to 32 bits when using auto dtype conversion semantics.
+          if width < np.iinfo(np.int32).bits:
+            if is_signed:
+              dtype = np.int32
+            else:
+              dtype = np.uint32
+        else:
+          if width < np.iinfo(np.int_).bits:
+            if is_signed:
+              dtype = np.int_
+            else:
+              dtype = np.uint
+        a = math_ops.cast(a, dtype)
+      elif promote_int == _TO_FLOAT:
+        # Use a default float type.
+        a = math_ops.cast(a, np_utils.result_type(float))
+
+  if isinstance(axis, tensor_lib.Tensor) and axis.dtype not in (
+      dtypes.int32,
+      dtypes.int64,
+  ):
+    axis = math_ops.cast(axis, dtypes.int64)
+
+  return tf_fn(input_tensor=a, axis=axis, keepdims=keepdims)
+
+
+# TODO (DarrenZhang01): Add `axis` support to the `size` API.
+@tf_export.tf_export('experimental.numpy.size', v1=[])
+@np_utils.np_doc('size')
+def size(x, axis=None):  # pylint: disable=missing-docstring
+  if axis is not None:
+    raise NotImplementedError(
+        'axis argument is not supported in the current `np.size` implementation'
+    )
+  if isinstance(x, (int, float, np.int32, np.int64, np.float32, np.float64)):
+    return 1
+  x = asarray(x)
+  if x.shape.is_fully_defined():
+    return np.prod(x.shape.as_list(), dtype=int)
+  else:
+    return array_ops.size_v2(x)
+
+
+@tf_export.tf_export('experimental.numpy.sum', v1=[])
+@np_utils.np_doc('sum')
+def sum(a, axis=None, dtype=None, keepdims=None):  # pylint: disable=redefined-builtin
+  return _reduce(
+      math_ops.reduce_sum,
+      a,
+      axis=axis,
+      dtype=dtype,
+      keepdims=keepdims,
+      tf_bool_fn=math_ops.reduce_any,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.prod', v1=[])
+@np_utils.np_doc('prod')
+def prod(a, axis=None, dtype=None, keepdims=None):
+  return _reduce(
+      math_ops.reduce_prod,
+      a,
+      axis=axis,
+      dtype=dtype,
+      keepdims=keepdims,
+      tf_bool_fn=math_ops.reduce_all,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.mean', v1=[])
+@np_utils.np_doc('mean', unsupported_params=['out'])
+def mean(a, axis=None, dtype=None, out=None, keepdims=None):
+  if out is not None:
+    raise ValueError('Setting out is not supported.')
+  return _reduce(
+      math_ops.reduce_mean,
+      a,
+      axis=axis,
+      dtype=dtype,
+      keepdims=keepdims,
+      promote_int=_TO_FLOAT,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.amax', v1=[])
+@np_utils.np_doc('amax', unsupported_params=['out'])
+def amax(a, axis=None, out=None, keepdims=None):
+  if out is not None:
+    raise ValueError('Setting out is not supported.')
+  return _reduce(
+      math_ops.reduce_max,
+      a,
+      axis=axis,
+      dtype=None,
+      keepdims=keepdims,
+      promote_int=None,
+      tf_bool_fn=math_ops.reduce_any,
+      preserve_bool=True,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.amin', v1=[])
+@np_utils.np_doc('amin', unsupported_params=['out'])
+def amin(a, axis=None, out=None, keepdims=None):
+  if out is not None:
+    raise ValueError('Setting out is not supported.')
+  return _reduce(
+      math_ops.reduce_min,
+      a,
+      axis=axis,
+      dtype=None,
+      keepdims=keepdims,
+      promote_int=None,
+      tf_bool_fn=math_ops.reduce_all,
+      preserve_bool=True,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.var', v1=[])
+@np_utils.np_doc('var')
+def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: disable=missing-docstring
+  if dtype:
+    working_dtype = np_utils.result_type(a, dtype)
+  else:
+    working_dtype = None
+  if out is not None:
+    raise ValueError('Setting out is not supported.')
+  if ddof != 0:
+    # TF reduce_variance doesn't support ddof, so calculate it using raw ops.
+    def reduce_fn(input_tensor, axis, keepdims):
+      means = math_ops.reduce_mean(input_tensor, axis=axis, keepdims=True)
+      centered = input_tensor - means
+      if input_tensor.dtype in (dtypes.complex64, dtypes.complex128):
+        centered = math_ops.cast(
+            math_ops.real(centered * math_ops.conj(centered)),
+            input_tensor.dtype,
+        )
+      else:
+        centered = math_ops.square(centered)
+      squared_deviations = math_ops.reduce_sum(
+          centered, axis=axis, keepdims=keepdims
+      )
+
+      if axis is None:
+        n = array_ops.size(input_tensor)
+      else:
+        if axis < 0:
+          axis += array_ops.rank(input_tensor)
+        n = math_ops.reduce_prod(
+            array_ops.gather(array_ops.shape(input_tensor), axis)
+        )
+      n = math_ops.cast(n - ddof, input_tensor.dtype)
+
+      return math_ops.cast(math_ops.divide(squared_deviations, n), dtype)
+
+  else:
+    reduce_fn = math_ops.reduce_variance
+
+  result = _reduce(
+      reduce_fn,
+      a,
+      axis=axis,
+      dtype=working_dtype,
+      keepdims=keepdims,
+      promote_int=_TO_FLOAT,
+  )
+  if dtype:
+    result = math_ops.cast(result, dtype)
+  return result
+
+
+@tf_export.tf_export('experimental.numpy.std', v1=[])
+@np_utils.np_doc('std')
+def std(a, axis=None, keepdims=None):  # pylint: disable=missing-function-docstring
+  return _reduce(
+      math_ops.reduce_std,
+      a,
+      axis=axis,
+      dtype=None,
+      keepdims=keepdims,
+      promote_int=_TO_FLOAT,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.ravel', v1=[])
+@np_utils.np_doc('ravel')
+def ravel(a):  # pylint: disable=missing-docstring
+  a = asarray(a)
+  return array_ops.reshape(a, [-1])
+
+
+@tf_export.tf_export('experimental.numpy.real', v1=[])
+@np_utils.np_doc('real')
+def real(val):
+  val = asarray(val)
+  # TODO(srbs): np.real returns a scalar if val is a scalar, whereas we always
+  # return an ndarray.
+  return math_ops.real(val)
+
+
+@tf_export.tf_export('experimental.numpy.repeat', v1=[])
+@np_utils.np_doc('repeat')
+def repeat(a, repeats, axis=None):  # pylint: disable=missing-docstring
+  a = asarray(a)
+  original_shape = a._shape_as_list()  # pylint: disable=protected-access
+  # Best effort recovery of the shape.
+  known_shape = original_shape is not None and None not in original_shape
+  if known_shape:
+    if not original_shape:
+      original_shape = (repeats,)
+    else:
+      repeats_np = np.ravel(np.array(repeats))
+      if repeats_np.size == 1:
+        repeats_np = repeats_np.item()
+        if axis is None:
+          original_shape = (repeats_np * np.prod(original_shape),)
+        else:
+          original_shape[axis] = repeats_np * original_shape[axis]
+      else:
+        if axis is None:
+          original_shape = (repeats_np.sum(),)
+        else:
+          original_shape[axis] = repeats_np.sum()
+
+  repeats = asarray(repeats)
+  result = array_ops.repeat(a, repeats, axis)
+  if known_shape:
+    result.set_shape(original_shape)
+
+  return result
+
+
+@tf_export.tf_export('experimental.numpy.around', v1=[])
+@np_utils.np_doc('around')
+def around(a, decimals=0):  # pylint: disable=missing-docstring
+  a = asarray(a)
+  dtype = a.dtype.as_numpy_dtype
+  factor = math.pow(10, decimals)
+  if np.issubdtype(dtype, np.inexact):
+    factor = math_ops.cast(factor, dtype)
+  else:
+    # Use float as the working dtype when a.dtype is exact (e.g. integer),
+    # because `decimals` can be negative.
+    float_dtype = np_utils.result_type(float)
+    a = a.astype(float_dtype)
+    factor = math_ops.cast(factor, float_dtype)
+  a = math_ops.multiply(a, factor)
+  a = math_ops.round(a)
+  a = math_ops.divide(a, factor)
+  return a.astype(dtype)
+
+
+setattr(np_arrays.ndarray, '__round__', around)
+
+
+@tf_export.tf_export('experimental.numpy.reshape', v1=[])
+@np_utils.np_doc('reshape')
+def reshape(a, newshape, order='C'):
+  """order argument can only b 'C' or 'F'."""
+  if order not in {'C', 'F'}:
+    raise ValueError('Unsupported order argument {}'.format(order))
+
+  a = asarray(a)
+  if isinstance(newshape, int):
+    newshape = [newshape]
+
+  if order == 'F':
+    r = array_ops.transpose(
+        array_ops.reshape(array_ops.transpose(a), newshape[::-1])
+    )
+  else:
+    r = array_ops.reshape(a, newshape)
+
+  return r
+
+
+def _reshape_method_wrapper(a, *newshape, **kwargs):
+  order = kwargs.pop('order', 'C')
+  if kwargs:
+    raise ValueError('Unsupported arguments: {}'.format(kwargs.keys()))
+
+  if len(newshape) == 1 and not isinstance(newshape[0], int):
+    newshape = newshape[0]
+
+  return reshape(a, newshape, order=order)
+
+
+@tf_export.tf_export('experimental.numpy.expand_dims', v1=[])
+@np_utils.np_doc('expand_dims')
+def expand_dims(a, axis):
+  a = asarray(a)
+  return array_ops.expand_dims(a, axis=axis)
+
+
+@tf_export.tf_export('experimental.numpy.squeeze', v1=[])
+@np_utils.np_doc('squeeze')
+def squeeze(a, axis=None):
+  a = asarray(a)
+  return array_ops.squeeze(a, axis)
+
+
+@tf_export.tf_export('experimental.numpy.flatten', v1=[])
+@np_utils.np_doc('flatten', link=np_utils.NoLink())
+def flatten(a, order='C'):
+  a = asarray(a)
+  if order == 'C' or order == 'A' or order == 'K':
+    # Row major.
+    return array_ops.reshape(a, [-1])
+  elif order == 'F':
+    # Column major
+    return array_ops.reshape(array_ops.transpose(a), [-1])
+  else:
+    raise ValueError(
+        'order can only be C, A, K (all row major) or F (column major).'
+    )
+
+
+@tf_export.tf_export('experimental.numpy.transpose', v1=[])
+@np_utils.np_doc('transpose')
+def transpose(a, axes=None):
+  a = asarray(a)
+  if axes is not None:
+    axes = asarray(axes)
+  return array_ops.transpose(a=a, perm=axes)
+
+
+@tf_export.tf_export('experimental.numpy.swapaxes', v1=[])
+@np_utils.np_doc('swapaxes')
+def swapaxes(a, axis1, axis2):  # pylint: disable=missing-docstring
+  a = asarray(a)
+
+  def adjust_axes(axes, rank):
+    def f(x):
+      if isinstance(x, int):
+        if x < 0:
+          x = x + rank
+      else:
+        x = array_ops.where_v2(x < 0, np_utils.add(x, a_rank), x)
+      return x
+
+    return nest.map_structure(f, axes)
+
+  if (
+      a.shape.rank is not None
+      and isinstance(axis1, int)
+      and isinstance(axis2, int)
+  ):
+    # This branch makes sure `perm` is statically known, to avoid a
+    # not-compile-time-constant XLA error.
+    a_rank = a.shape.rank
+    axis1, axis2 = adjust_axes((axis1, axis2), a_rank)
+    perm = list(range(a_rank))
+    perm[axis1] = axis2
+    perm[axis2] = axis1
+  else:
+    a_rank = array_ops.rank(a)
+    axis1, axis2 = adjust_axes((axis1, axis2), a_rank)
+    perm = math_ops.range(a_rank)
+    perm = array_ops.tensor_scatter_update(
+        perm, [[axis1], [axis2]], [axis2, axis1]
+    )
+  a = array_ops.transpose(a, perm)
+  return a
+
+
+@tf_export.tf_export('experimental.numpy.moveaxis', v1=[])
+@np_utils.np_doc('moveaxis')
+def moveaxis(a, source, destination):  # pylint: disable=missing-docstring
+  """Raises ValueError if source, destination not in (-ndim(a), ndim(a))."""
+  if not source and not destination:
+    return a
+
+  a = asarray(a)
+
+  if isinstance(source, int):
+    source = (source,)
+  if isinstance(destination, int):
+    destination = (destination,)
+  if len(source) != len(destination):
+    raise ValueError('The lengths of source and destination must equal')
+
+  a_rank = np_utils._maybe_static(array_ops.rank(a))  # pylint: disable=protected-access
+
+  def _correct_axis(axis, rank):
+    if axis < 0:
+      return axis + rank
+    return axis
+
+  source = tuple(_correct_axis(axis, a_rank) for axis in source)
+  destination = tuple(_correct_axis(axis, a_rank) for axis in destination)
+
+  if a.shape.rank is not None:
+    perm = [i for i in range(a_rank) if i not in source]
+    for dest, src in sorted(zip(destination, source)):
+      assert dest <= len(perm)
+      perm.insert(dest, src)
+  else:
+    r = math_ops.range(a_rank)
+
+    def _remove_indices(a, b):
+      """Remove indices (`b`) from `a`."""
+      items = array_ops_stack.unstack(
+          sort_ops.sort(array_ops_stack.stack(b)), num=len(b)
+      )
+
+      i = 0
+      result = []
+
+      for item in items:
+        result.append(a[i:item])
+        i = item + 1
+
+      result.append(a[i:])
+
+      return array_ops.concat(result, 0)
+
+    minus_sources = _remove_indices(r, source)
+    minus_dest = _remove_indices(r, destination)
+
+    perm = array_ops.scatter_nd(
+        array_ops.expand_dims(minus_dest, 1), minus_sources, [a_rank]
+    )
+    perm = array_ops.tensor_scatter_update(
+        perm, array_ops.expand_dims(destination, 1), source
+    )
+  a = array_ops.transpose(a, perm)
+
+  return a
+
+
+@tf_export.tf_export('experimental.numpy.pad', v1=[])
+@np_utils.np_doc('pad')
+def pad(array, pad_width, mode, **kwargs):  # pylint: disable=redefined-outer-name
+  """Only supports modes 'constant', 'reflect' and 'symmetric' currently."""
+  constant_values = kwargs.get('constant_values', 0)
+  if not (mode == 'constant' or mode == 'reflect' or mode == 'symmetric'):
+    raise ValueError('Unsupported padding mode: ' + mode)
+  mode = mode.upper()
+  array = asarray(array)
+  pad_width = asarray(pad_width, dtype=dtypes.int32)
+  return array_ops.pad(
+      tensor=array,
+      paddings=pad_width,
+      mode=mode,
+      constant_values=constant_values,
+  )
+
+
+@tf_export.tf_export('experimental.numpy.take', v1=[])
+@np_utils.np_doc('take')
+def take(a, indices, axis=None, out=None, mode='clip'):
+  """out argument is not supported, and default mode is clip."""
+  if out is not None:
+    raise ValueError('out argument is not supported in take.')
+
+  if mode not in {'raise', 'clip', 'wrap'}:
+    raise ValueError("Invalid mode '{}' for take".format(mode))
+
+  a = asarray(a)
+  indices = asarray(indices)
+
+  if axis is None:
+    a = array_ops.reshape(a, [-1])
+    axis = 0
+
+  axis_size = array_ops.shape(a, out_type=indices.dtype)[axis]
+  if mode == 'clip':
+    indices = clip_ops.clip_by_value(indices, 0, axis_size - 1)
+  elif mode == 'wrap':
+    indices = math_ops.floormod(indices, axis_size)
+  else:
+    raise ValueError("The 'raise' mode to take is not supported.")
+
+  return array_ops.gather(a, indices, axis=axis)
+
+
+@tf_export.tf_export('experimental.numpy.where', v1=[])
+@np_utils.np_doc_only('where')
+def where(condition, x=None, y=None):
+  """Raises ValueError if exactly one of x or y is not None."""
+  condition = asarray(condition, dtype=np.bool_)
+  if x is None and y is None:
+    return nonzero(condition)
+  elif x is not None and y is not None:
+    x, y = _promote_dtype(x, y)
+    return array_ops.where_v2(condition, x, y)
+  raise ValueError('Both x and y must be ndarrays, or both must be None.')
+
+
+@tf_export.tf_export('experimental.numpy.select', v1=[])
+@np_utils.np_doc('select')
+def select(condlist, choicelist, default=0):  # pylint: disable=missing-docstring
+  if len(condlist) != len(choicelist):
+    msg = 'condlist must have length equal to choicelist ({} vs {})'
+    raise ValueError(msg.format(len(condlist), len(choicelist)))
+  if not condlist:
+    raise ValueError('condlist must be non-empty')
+  choices = _promote_dtype(default, *choicelist)
+  choicelist = choices[1:]
+  output = choices[0]
+  # The traversal is in reverse order so we can return the first value in
+  # choicelist where condlist is True.
+  for cond, choice in zip(condlist[::-1], choicelist[::-1]):
+    output = where(cond, choice, output)
+  return output
+
+
+@tf_export.tf_export('experimental.numpy.shape', v1=[])
+@np_utils.np_doc(
+    'shape',
+    link=np_utils.Link(
+        'https://numpy.org/doc/1.18/reference/generated/numpy.shape.html'
+    ),
+)
+def shape(a):
+  a = asarray(a)
+  return a.shape
+
+
+@tf_export.tf_export('experimental.numpy.ndim', v1=[])
+@np_utils.np_doc('ndim', link=np_utils.NoLink())
+def ndim(a):
+  a = asarray(a)
+  return a.ndim
+
+
+@tf_export.tf_export('experimental.numpy.isscalar', v1=[])
+@np_utils.np_doc('isscalar')
+def isscalar(num):
+  return ndim(num) == 0
+
+
+def _boundaries_to_sizes(a, boundaries, axis):
+  """Converting boundaries of splits to sizes of splits.
+
+  Args:
+    a: the array to be split.
+    boundaries: the boundaries, as in np.split.
+    axis: the axis along which to split.
+
+  Returns:
+    A list of sizes of the splits, as in tf.split.
+  """
+  if axis >= len(a.shape):
+    raise ValueError('axis %s is out of bound for shape %s' % (axis, a.shape))
+  total_size = a.shape[axis]
+  sizes = []
+  sizes_sum = 0
+  prev = 0
+  for i, b in enumerate(boundaries):
+    size = b - prev
+    if size < 0:
+      raise ValueError(
+          'The %s-th boundary %s is smaller than the previous boundary %s'
+          % (i, b, prev)
+      )
+    size = builtins.min(size, builtins.max(0, total_size - sizes_sum))
+    sizes.append(size)
+    sizes_sum += size
+    prev = b
+  sizes.append(builtins.max(0, total_size - sizes_sum))
+  return sizes
+
+
+@tf_export.tf_export('experimental.numpy.split', v1=[])
+@np_utils.np_doc('split')
+def split(ary, indices_or_sections, axis=0):
+  ary = asarray(ary)
+  if not isinstance(indices_or_sections, int):
+    indices_or_sections = _boundaries_to_sizes(ary, indices_or_sections, axis)
+  return array_ops.split(ary, indices_or_sections, axis=axis)
+
+
+def _split_on_axis(np_fun_name, axis):  # pylint: disable=missing-function-docstring
+  @np_utils.np_doc(np_fun_name)
+  def f(ary, indices_or_sections):
+    # for 1-D array, hsplit becomes vsplit
+    new_axis = np_utils.cond(
+        math_ops.equal(axis, 1),
+        lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+            math_ops.equal(array_ops.rank(ary), 1), lambda: 0, lambda: axis
+        ),
+        lambda: axis,
+    )
+    if isinstance(indices_or_sections, int):
+      ary_shape = ary.shape[new_axis]
+      if ary_shape is not None and ary_shape % indices_or_sections:
+        raise ValueError('array split does not result in an equal division')
+    return split(ary, indices_or_sections, axis=new_axis)
+
+  return f
+
+
+vsplit = tf_export.tf_export('experimental.numpy.vsplit', v1=[])(
+    _split_on_axis('vsplit', axis=0)
+)
+hsplit = tf_export.tf_export('experimental.numpy.hsplit', v1=[])(
+    _split_on_axis('hsplit', axis=1)
+)
+dsplit = tf_export.tf_export('experimental.numpy.dsplit', v1=[])(
+    _split_on_axis('dsplit', axis=2)
+)
+
+
+@tf_export.tf_export('experimental.numpy.broadcast_to', v1=[])
+@np_utils.np_doc('broadcast_to')
+def broadcast_to(array, shape):  # pylint: disable=redefined-outer-name
+  return full(shape, array)
+
+
+@tf_export.tf_export('experimental.numpy.stack', v1=[])
+@np_utils.np_doc('stack')
+def stack(arrays, axis=0):  # pylint: disable=missing-function-docstring
+  if isinstance(arrays, (np_arrays.ndarray, tensor_lib.Tensor)):
+    arrays = asarray(arrays)
+    if axis == 0:
+      return arrays
+    else:
+      return swapaxes(arrays, 0, axis)
+  arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
+  unwrapped_arrays = [
+      a if isinstance(a, np_arrays.ndarray) else a for a in arrays
+  ]
+  return asarray(array_ops_stack.stack(unwrapped_arrays, axis))
+
+
+@tf_export.tf_export('experimental.numpy.hstack', v1=[])
+@np_utils.np_doc('hstack')
+def hstack(tup):
+  arrays = [atleast_1d(a) for a in tup]
+  arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
+  unwrapped_arrays = [
+      a if isinstance(a, np_arrays.ndarray) else a for a in arrays
+  ]
+  rank = array_ops.rank(unwrapped_arrays[0])
+  return np_utils.cond(
+      math_ops.equal(rank, 1),
+      lambda: array_ops.concat(unwrapped_arrays, axis=0),
+      lambda: array_ops.concat(unwrapped_arrays, axis=1),
+  )
+
+
+@tf_export.tf_export('experimental.numpy.vstack', v1=[])
+@np_utils.np_doc('vstack')
+def vstack(tup):
+  arrays = [atleast_2d(a) for a in tup]
+  arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
+  unwrapped_arrays = [
+      a if isinstance(a, np_arrays.ndarray) else a for a in arrays
+  ]
+  return array_ops.concat(unwrapped_arrays, axis=0)
+
+
+@tf_export.tf_export('experimental.numpy.dstack', v1=[])
+@np_utils.np_doc('dstack')
+def dstack(tup):
+  arrays = [atleast_3d(a) for a in tup]
+  arrays = _promote_dtype(*arrays)  # pylint: disable=protected-access
+  unwrapped_arrays = [
+      a if isinstance(a, np_arrays.ndarray) else a for a in arrays
+  ]
+  return array_ops.concat(unwrapped_arrays, axis=2)
+
+
+def _pad_left_to(n, old_shape):
+  old_shape = asarray(old_shape, dtype=np.int32)
+  new_shape = array_ops.pad(
+      old_shape,
+      [[math_ops.maximum(n - array_ops.size(old_shape), 0), 0]],
+      constant_values=1,
+  )
+  return asarray(new_shape)
+
+
+def _atleast_nd(n, new_shape, *arys):
+  """Reshape arrays to be at least `n`-dimensional.
+
+  Args:
+    n: The minimal rank.
+    new_shape: a function that takes `n` and the old shape and returns the
+      desired new shape.
+    *arys: ndarray(s) to be reshaped.
+
+  Returns:
+    The reshaped array(s).
+  """
+
+  def f(x):
+    # pylint: disable=g-long-lambda
+    x = asarray(x)
+    return asarray(
+        np_utils.cond(
+            np_utils.greater(n, array_ops.rank(x)),
+            lambda: reshape(x, new_shape(n, array_ops.shape(x))),
+            lambda: x,
+        )
+    )
+
+  arys = list(map(f, arys))
+  if len(arys) == 1:
+    return arys[0]
+  else:
+    return arys
+
+
+@tf_export.tf_export('experimental.numpy.atleast_1d', v1=[])
+@np_utils.np_doc('atleast_1d')
+def atleast_1d(*arys):
+  return _atleast_nd(1, _pad_left_to, *arys)
+
+
+@tf_export.tf_export('experimental.numpy.atleast_2d', v1=[])
+@np_utils.np_doc('atleast_2d')
+def atleast_2d(*arys):
+  return _atleast_nd(2, _pad_left_to, *arys)
+
+
+@tf_export.tf_export('experimental.numpy.atleast_3d', v1=[])
+@np_utils.np_doc('atleast_3d')
+def atleast_3d(*arys):  # pylint: disable=missing-docstring
+  def new_shape(_, old_shape):
+    # pylint: disable=g-long-lambda
+    ndim_ = array_ops.size(old_shape)
+    return np_utils.cond(
+        math_ops.equal(ndim_, 0),
+        lambda: constant_op.constant([1, 1, 1], dtype=dtypes.int32),
+        lambda: np_utils.cond(
+            math_ops.equal(ndim_, 1),
+            lambda: array_ops.pad(old_shape, [[1, 1]], constant_values=1),
+            lambda: array_ops.pad(old_shape, [[0, 1]], constant_values=1),
+        ),
+    )
+
+  return _atleast_nd(3, new_shape, *arys)
+
+
+@tf_export.tf_export('experimental.numpy.nonzero', v1=[])
+@np_utils.np_doc('nonzero')
+def nonzero(a):
+  a = atleast_1d(a)
+  if a.shape.rank is None:
+    raise ValueError(
+        "The rank of `a` is unknown, so we can't decide how many "
+        'arrays to return.'
+    )
+  return array_ops_stack.unstack(
+      array_ops.where_v2(math_ops.cast(a, dtypes.bool)), a.shape.rank, axis=1
+  )
+
+
+@tf_export.tf_export('experimental.numpy.diag_indices', v1=[])
+@np_utils.np_doc('diag_indices')
+def diag_indices(n, ndim=2):  # pylint: disable=missing-docstring,redefined-outer-name
+  if n < 0:
+    raise ValueError(
+        'n argument to diag_indices must be nonnegative, got {}'.format(n)
+    )
+  if ndim < 0:
+    raise ValueError(
+        'ndim argument to diag_indices must be nonnegative, got {}'.format(ndim)
+    )
+
+  return (math_ops.range(n),) * ndim
+
+
+@tf_export.tf_export('experimental.numpy.tri', v1=[])
+@np_utils.np_doc('tri')
+def tri(N, M=None, k=0, dtype=None):  # pylint: disable=invalid-name,missing-docstring
+  M = M if M is not None else N
+  if dtype is not None:
+    dtype = np_utils.result_type(dtype)
+  else:
+    # Use a default float type.
+    dtype = np_utils.result_type(float)
+
+  if k < 0:
+    lower = -k - 1
+    if lower > N:
+      r = array_ops.zeros([N, M], dtype)
+    else:
+      # Keep as tf bool, since we create an upper triangular matrix and invert
+      # it.
+      o = array_ops.ones([N, M], dtype=dtypes.bool)
+      r = math_ops.cast(
+          math_ops.logical_not(array_ops.matrix_band_part(o, lower, -1)), dtype
+      )
+  else:
+    o = array_ops.ones([N, M], dtype)
+    if k > M:
+      r = o
+    else:
+      r = array_ops.matrix_band_part(o, -1, k)
+  return r
+
+
+@tf_export.tf_export('experimental.numpy.tril', v1=[])
+@np_utils.np_doc('tril')
+def tril(m, k=0):  # pylint: disable=missing-docstring
+  m = asarray(m)
+  if m.shape.ndims is None:
+    raise ValueError('Argument to tril should have known rank')
+  m_shape = m.shape.as_list()
+
+  if len(m_shape) < 2:
+    raise ValueError('Argument to tril must have rank at least 2')
+
+  if m_shape[-1] is None or m_shape[-2] is None:
+    raise ValueError(
+        'Currently, the last two dimensions of the input array '
+        'need to be known.'
+    )
+
+  z = constant_op.constant(0, m.dtype)
+
+  mask = tri(*m_shape[-2:], k=k, dtype=bool)
+  return array_ops.where_v2(
+      array_ops.broadcast_to(mask, array_ops.shape(m)), m, z
+  )
+
+
+@tf_export.tf_export('experimental.numpy.triu', v1=[])
+@np_utils.np_doc('triu')
+def triu(m, k=0):  # pylint: disable=missing-docstring
+  m = asarray(m)
+  if m.shape.ndims is None:
+    raise ValueError('Argument to triu should have known rank')
+  m_shape = m.shape.as_list()
+
+  if len(m_shape) < 2:
+    raise ValueError('Argument to triu must have rank at least 2')
+
+  if m_shape[-1] is None or m_shape[-2] is None:
+    raise ValueError(
+        'Currently, the last two dimensions of the input array '
+        'need to be known.'
+    )
+
+  z = constant_op.constant(0, m.dtype)
+
+  mask = tri(*m_shape[-2:], k=k - 1, dtype=bool)
+  return array_ops.where_v2(
+      array_ops.broadcast_to(mask, array_ops.shape(m)), z, m
+  )
+
+
+@tf_export.tf_export('experimental.numpy.flip', v1=[])
+@np_utils.np_doc('flip')
+def flip(m, axis=None):  # pylint: disable=missing-docstring
+  m = asarray(m)
+
+  if axis is None:
+    return array_ops.reverse(m, math_ops.range(array_ops.rank(m)))
+
+  axis = np_utils._canonicalize_axis(axis, array_ops.rank(m))  # pylint: disable=protected-access
+
+  return array_ops.reverse(m, [axis])
+
+
+@tf_export.tf_export('experimental.numpy.flipud', v1=[])
+@np_utils.np_doc('flipud')
+def flipud(m):  # pylint: disable=missing-docstring
+  return flip(m, 0)
+
+
+@tf_export.tf_export('experimental.numpy.fliplr', v1=[])
+@np_utils.np_doc('fliplr')
+def fliplr(m):  # pylint: disable=missing-docstring
+  return flip(m, 1)
+
+
+@tf_export.tf_export('experimental.numpy.roll', v1=[])
+@np_utils.np_doc('roll')
+def roll(a, shift, axis=None):  # pylint: disable=missing-docstring
+  a = asarray(a)
+
+  if axis is not None:
+    return manip_ops.roll(a, shift, axis)
+
+  # If axis is None, the roll happens as a 1-d tensor.
+  original_shape = array_ops.shape(a)
+  a = manip_ops.roll(array_ops.reshape(a, [-1]), shift, 0)
+  return array_ops.reshape(a, original_shape)
+
+
+@tf_export.tf_export('experimental.numpy.rot90', v1=[])
+@np_utils.np_doc('rot90')
+def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
+  m_rank = array_ops.rank(m)
+  ax1, ax2 = np_utils._canonicalize_axes(axes, m_rank)  # pylint: disable=protected-access
+
+  k = k % 4
+  if k == 0:
+    return m
+  elif k == 2:
+    return flip(flip(m, ax1), ax2)
+  else:
+    perm = math_ops.range(m_rank)
+    perm = array_ops.tensor_scatter_update(perm, [[ax1], [ax2]], [ax2, ax1])
+
+    if k == 1:
+      return transpose(flip(m, ax2), perm)
+    else:
+      return flip(transpose(m, perm), ax2)
+
+
+@tf_export.tf_export('experimental.numpy.vander', v1=[])
+@np_utils.np_doc('vander')
+def vander(x, N=None, increasing=False):  # pylint: disable=missing-docstring,invalid-name
+  x = asarray(x)
+
+  x_shape = array_ops.shape(x)
+  if N is None:
+    N = x_shape[0]
+
+  N_temp = np_utils.get_static_value(N)  # pylint: disable=invalid-name
+  if N_temp is not None:
+    N = N_temp
+    if N < 0:
+      raise ValueError('N must be nonnegative')
+  else:
+    control_flow_assert.Assert(N >= 0, [N])
+
+  rank = array_ops.rank(x)
+  rank_temp = np_utils.get_static_value(rank)
+  if rank_temp is not None:
+    rank = rank_temp
+    if rank != 1:
+      raise ValueError('x must be a one-dimensional array')
+  else:
+    control_flow_assert.Assert(math_ops.equal(rank, 1), [rank])
+
+  if increasing:
+    start = 0
+    limit = N
+    delta = 1
+  else:
+    start = N - 1
+    limit = -1
+    delta = -1
+
+  x = array_ops.expand_dims(x, -1)
+  return math_ops.pow(
+      x, math_ops.cast(math_ops.range(start, limit, delta), dtype=x.dtype)
+  )
+
+
+@tf_export.tf_export('experimental.numpy.ix_', v1=[])
+@np_utils.np_doc('ix_')
+def ix_(*args):  # pylint: disable=missing-docstring
+  n = len(args)
+  output = []
+  for i, a in enumerate(args):
+    a = asarray(a)
+    a_rank = array_ops.rank(a)
+    a_rank_temp = np_utils.get_static_value(a_rank)
+    if a_rank_temp is not None:
+      a_rank = a_rank_temp
+      if a_rank != 1:
+        raise ValueError(
+            'Arguments must be 1-d, got arg {} of rank {}'.format(i, a_rank)
+        )
+    else:
+      control_flow_assert.Assert(math_ops.equal(a_rank, 1), [a_rank])
+
+    new_shape = [1] * n
+    new_shape[i] = -1
+    dtype = a.dtype
+    if dtype == dtypes.bool:
+      output.append(array_ops.reshape(nonzero(a)[0], new_shape))
+    elif dtype.is_integer:
+      output.append(array_ops.reshape(a, new_shape))
+    else:
+      raise ValueError(
+          'Only integer and bool dtypes are supported, got {}'.format(dtype)
+      )
+
+  return output
+
+
+@tf_export.tf_export('experimental.numpy.broadcast_arrays', v1=[])
+@np_utils.np_doc('broadcast_arrays')
+def broadcast_arrays(*args, **kwargs):  # pylint: disable=missing-docstring
+  subok = kwargs.pop('subok', False)
+  if subok:
+    raise ValueError('subok=True is not supported.')
+  if kwargs:
+    raise ValueError('Received unsupported arguments {}'.format(kwargs.keys()))
+
+  args = [asarray(arg) for arg in args]
+  return np_utils.tf_broadcast(*args)
+
+
+@tf_export.tf_export('experimental.numpy.sign', v1=[])
+@np_utils.np_doc_only('sign')
+def sign(x, out=None, where=None, **kwargs):  # pylint: disable=missing-docstring,redefined-outer-name
+  if out:
+    raise ValueError('tf.numpy doesnt support setting out.')
+  if where:
+    raise ValueError('tf.numpy doesnt support setting where.')
+  if kwargs:
+    raise ValueError('tf.numpy doesnt support setting {}'.format(kwargs.keys()))
+
+  x = asarray(x)
+
+  # Numpy 2.x and later uses the same definition of sign.
+  if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
+    return math_ops.sign(x)
+
+  dtype = x.dtype.as_numpy_dtype
+  if np.issubdtype(dtype, np.complexfloating):
+    result = math_ops.cast(math_ops.sign(math_ops.real(x)), dtype)
+  else:
+    result = math_ops.sign(x)
+
+  return result
+
+
+# Note that np.take_along_axis may not be present in some supported versions of
+# numpy.
+@tf_export.tf_export('experimental.numpy.take_along_axis', v1=[])
+@np_utils.np_doc('take_along_axis')
+def take_along_axis(arr, indices, axis):  # pylint: disable=missing-docstring
+  arr = asarray(arr)
+  indices = asarray(indices)
+
+  if axis is None:
+    return take_along_axis(arr.ravel(), indices, 0)
+
+  rank = array_ops.rank(arr)
+  axis = axis + rank if axis < 0 else axis
+
+  # Broadcast shapes to match, ensure that the axis of interest is not
+  # broadcast.
+  arr_shape_original = array_ops.shape(arr, out_type=indices.dtype)
+  indices_shape_original = array_ops.shape(indices, out_type=indices.dtype)
+  arr_shape = array_ops.tensor_scatter_update(arr_shape_original, [[axis]], [1])
+  indices_shape = array_ops.tensor_scatter_update(
+      indices_shape_original, [[axis]], [1]
+  )
+  broadcasted_shape = array_ops.broadcast_dynamic_shape(
+      arr_shape, indices_shape
+  )
+  arr_shape = array_ops.tensor_scatter_update(
+      broadcasted_shape, [[axis]], [arr_shape_original[axis]]
+  )
+  indices_shape = array_ops.tensor_scatter_update(
+      broadcasted_shape, [[axis]], [indices_shape_original[axis]]
+  )
+  arr = array_ops.broadcast_to(arr, arr_shape)
+  indices = array_ops.broadcast_to(indices, indices_shape)
+
+  # Save indices shape so we can restore it later.
+  possible_result_shape = indices.shape
+
+  # Correct indices since gather doesn't correctly handle negative indices.
+  indices = array_ops.where_v2(indices < 0, indices + arr_shape[axis], indices)
+
+  swapaxes_ = lambda t: swapaxes(t, axis, -1)
+
+  dont_move_axis_to_end = math_ops.equal(axis, np_utils.subtract(rank, 1))
+  arr = np_utils.cond(
+      dont_move_axis_to_end, lambda: arr, lambda: swapaxes_(arr)
+  )
+  indices = np_utils.cond(
+      dont_move_axis_to_end, lambda: indices, lambda: swapaxes_(indices)
+  )
+
+  arr_shape = array_ops.shape(arr)
+  arr = array_ops.reshape(arr, [-1, arr_shape[-1]])
+
+  indices_shape = array_ops.shape(indices)
+  indices = array_ops.reshape(indices, [-1, indices_shape[-1]])
+
+  result = array_ops.gather(arr, indices, batch_dims=1)
+  result = array_ops.reshape(result, indices_shape)
+  result = np_utils.cond(
+      dont_move_axis_to_end, lambda: result, lambda: swapaxes_(result)
+  )
+  result.set_shape(possible_result_shape)
+
+  return result
+
+
+# pylint: disable=redefined-builtin,undefined-variable
+@tf_export.tf_export('experimental.numpy.max', v1=[])
+@np_utils.np_doc('max', link=np_utils.AliasOf('amax'))
+def max(a, axis=None, keepdims=None):
+  return amax(a, axis=axis, keepdims=keepdims)
+
+
+@tf_export.tf_export('experimental.numpy.min', v1=[])
+@np_utils.np_doc('min', link=np_utils.AliasOf('amin'))
+def min(a, axis=None, keepdims=None):
+  return amin(a, axis=axis, keepdims=keepdims)
+
+
+@tf_export.tf_export('experimental.numpy.round', v1=[])
+@np_utils.np_doc('round', link=np_utils.AliasOf('around'))
+def round(a, decimals=0):
+  return around(a, decimals=decimals)
+
+
+# pylint: enable=redefined-builtin,undefined-variable
+
+
+_SLICE_ERROR = (
+    'only integers, slices (`:`), ellipsis (`...`), '
+    'numpy.newaxis (`None`) and integer or boolean arrays are valid indices'
+)
+
+
+def _as_index(idx, need_scalar=True):
+  """Helper function to parse idx as an index.
+
+  Args:
+    idx: index
+    need_scalar: If idx needs to be a scalar value.
+
+  Returns:
+    A pair, (indx, bool). First one is the parsed index and can be a tensor,
+    or scalar integer / Dimension. Second one is True if rank is known to be 0.
+
+  Raises:
+    IndexError: For incorrect indices.
+  """
+  if isinstance(idx, (numbers.Integral, tensor_shape.Dimension)):
+    return idx, True
+  data = asarray(idx)
+  if data.dtype == dtypes.bool:
+    if data.shape.ndims != 1:
+      # TODO(agarwal): handle higher rank boolean masks.
+      raise NotImplementedError('Need rank 1 for bool index %s' % idx)
+    data = array_ops.where_v2(data)
+    data = array_ops.reshape(data, [-1])
+  if need_scalar and data.shape.rank not in (None, 0):
+    raise IndexError(_SLICE_ERROR + ', got {!r}'.format(idx))
+  np_dtype = data.dtype.as_numpy_dtype
+  if not np.issubdtype(np_dtype, np.integer):
+    raise IndexError(_SLICE_ERROR + ', got {!r}'.format(idx))
+  if data.dtype not in (dtypes.int64, dtypes.int32):
+    # TF slicing can only handle int32/int64. So we need to cast.
+    promoted_dtype = np.promote_types(np.int32, np_dtype)
+    if promoted_dtype == np.int32:
+      data = math_ops.cast(data, dtypes.int32)
+    elif promoted_dtype == np.int64:
+      data = math_ops.cast(data, dtypes.int64)
+    else:
+      raise IndexError(_SLICE_ERROR + ', got {!r}'.format(idx))
+  return data, data.shape.rank == 0
+
+
+class _UpdateMethod(enum.Enum):
+  UPDATE = 0
+  ADD = 1
+  MIN = 2
+  MAX = 3
+
+
+def _slice_helper(tensor, slice_spec, update_method=None, updates=None):
+  """Helper function for __getitem__ and _with_index_update_helper.
+
+  This function collects the indices in `slice_spec` into two buckets, which we
+  can call "idx1" and "idx2" here. idx1 is intended for `strided_slice`, idx2
+  `gather`.  They also correspond to "basic indices" and "advanced indices" in
+  numpy.  This function supports both reading and writing at the indices. The
+  reading path can be summarized as `gather(stride_slice(tensor, idx1),
+  idx2)`. The writing path can be summarized as `strided_slice_update(tensor,
+  idx1, scatter(strided_slice(tensor, idx1), idx2, updates))`.  (`gather` here
+  means `tf.gather` or `tf.gather_nd`; `scatter` here means
+  `tf.tensor_scatter_update`.)  The writing path is inefficient because it needs
+  to first read out a portion (probably much larger than `updates`) of `tensor`
+  using `strided_slice`, update it, and then write the portion back. An
+  alternative approach is to only use `scatter`, which amounts to using the
+  indexing mechanism of gather/scatter to implement
+  strided_slice/strided_slice_update. This is feasible for XLA Gather/Scatter
+  because they support spans (e.g. `2:5`) in indices (as begin/end pairs), but
+  not TF gather/scatter because they don't support spans (except those that
+  cover entire dimensions, i.e. `:`).  If we materialize spans into individual
+  indices, the size of the index tensor would explode.  (Note that XLA
+  Gather/Scatter have a similar problem for stride > 1 because they don't
+  support strides.  Indices such as `1:2:8` will need to be materialized into
+  individual indices such as [1, 3, 5, 7].)
+
+  Args:
+    tensor: the tensor to be read from or write into.
+    slice_spec: the indices.
+    update_method: (optional) a member of `_UpdateMethod`, indicating how to
+      update the values (replacement, add, etc.). `None` indicates just reading.
+    updates: (optional) the new values to write into `tensor`. It must have the
+      same dtype as `tensor`.
+
+  Returns:
+    The result of reading (if `update_method` is `None`) or the updated `tensor`
+    after writing.
+  """
+  begin, end, strides = [], [], []
+  new_axis_mask, shrink_axis_mask = 0, 0
+  begin_mask, end_mask = 0, 0
+  ellipsis_mask = 0
+  advanced_indices = []
+  shrink_indices = []
+  for index, s in enumerate(slice_spec):
+    if isinstance(s, slice):
+      if s.start is not None:
+        begin.append(_as_index(s.start)[0])
+      else:
+        begin.append(0)
+        begin_mask |= 1 << index
+      if s.stop is not None:
+        end.append(_as_index(s.stop)[0])
+      else:
+        end.append(0)
+        end_mask |= 1 << index
+      if s.step is not None:
+        strides.append(_as_index(s.step)[0])
+      else:
+        strides.append(1)
+    elif s is Ellipsis:
+      begin.append(0)
+      end.append(0)
+      strides.append(1)
+      ellipsis_mask |= 1 << index
+    elif s is array_ops.newaxis:
+      begin.append(0)
+      end.append(0)
+      strides.append(1)
+      new_axis_mask |= 1 << index
+    else:
+      s, is_scalar = _as_index(s, False)
+      if is_scalar:
+        begin.append(s)
+        end.append(s + 1)
+        strides.append(1)
+        shrink_axis_mask |= 1 << index
+        shrink_indices.append(index)
+      else:
+        begin.append(0)
+        end.append(0)
+        strides.append(1)
+        begin_mask |= 1 << index
+        end_mask |= 1 << index
+        advanced_indices.append((index, s, ellipsis_mask != 0))
+
+  # stack possibly involves no tensors, so we must use op_scope correct graph.
+  with ops.name_scope(
+      None,
+      'strided_slice',
+      [tensor] + begin + end + strides,
+      skip_on_eager=False,
+  ) as name:
+    if begin:
+      packed_begin, packed_end, packed_strides = (
+          array_ops_stack.stack(begin),
+          array_ops_stack.stack(end),
+          array_ops_stack.stack(strides),
+      )
+      if (
+          packed_begin.dtype == dtypes.int64
+          or packed_end.dtype == dtypes.int64
+          or packed_strides.dtype == dtypes.int64
+      ):
+        if packed_begin.dtype != dtypes.int64:
+          packed_begin = math_ops.cast(packed_begin, dtypes.int64)
+        if packed_end.dtype != dtypes.int64:
+          packed_end = math_ops.cast(packed_end, dtypes.int64)
+        if packed_strides.dtype != dtypes.int64:
+          packed_strides = math_ops.cast(packed_strides, dtypes.int64)
+    else:
+      var_empty = constant_op.constant([], dtype=dtypes.int32)
+      packed_begin = packed_end = packed_strides = var_empty
+    if update_method == _UpdateMethod.UPDATE and not advanced_indices:
+      return array_ops.tensor_strided_slice_update(
+          tensor,
+          packed_begin,
+          packed_end,
+          packed_strides,
+          updates,
+          begin_mask=begin_mask,
+          end_mask=end_mask,
+          shrink_axis_mask=shrink_axis_mask,
+          new_axis_mask=new_axis_mask,
+          ellipsis_mask=ellipsis_mask,
+          name=name,
+      )
+    else:
+      # TODO(b/164251540): Find a better way to support update that does not
+      #   involve one read + two writes.
+      if updates is not None:
+        original_tensor = tensor
+      # TODO(agarwal): set_shape on tensor to set rank.
+      tensor = array_ops.strided_slice(
+          tensor,
+          packed_begin,
+          packed_end,
+          packed_strides,
+          begin_mask=begin_mask,
+          end_mask=end_mask,
+          shrink_axis_mask=shrink_axis_mask,
+          new_axis_mask=new_axis_mask,
+          ellipsis_mask=ellipsis_mask,
+          name=name,
+      )
+    if not advanced_indices:
+      if update_method is None:
+        return tensor
+      assert update_method != _UpdateMethod.UPDATE
+      # TF lacks TensorStridedSliceAdd and alike, so we need to do
+      # read+add+update.
+      if update_method == _UpdateMethod.ADD:
+        update_op = math_ops.add
+      elif update_method == _UpdateMethod.MIN:
+        update_op = math_ops.minimum
+      elif update_method == _UpdateMethod.MAX:
+        update_op = math_ops.maximum
+      return array_ops.tensor_strided_slice_update(
+          original_tensor,
+          packed_begin,
+          packed_end,
+          packed_strides,
+          update_op(tensor, updates),
+          begin_mask=begin_mask,
+          end_mask=end_mask,
+          shrink_axis_mask=shrink_axis_mask,
+          new_axis_mask=new_axis_mask,
+          ellipsis_mask=ellipsis_mask,
+          name=name + '_2',
+      )
+    advanced_indices_map = {}
+    for index, data, had_ellipsis in advanced_indices:
+      if had_ellipsis:
+        num_shrink = len([x for x in shrink_indices if x > index])
+        dim = index - len(slice_spec) + num_shrink
+      else:
+        num_shrink = len([x for x in shrink_indices if x < index])
+        dim = index - num_shrink
+      advanced_indices_map[dim] = data
+    dims = sorted(advanced_indices_map.keys())
+    dims_contiguous = True
+    if len(dims) > 1:
+      if dims[0] < 0 and dims[-1] >= 0:  # not all same sign
+        dims_contiguous = False
+      else:
+        for i in range(len(dims) - 1):
+          if dims[i] + 1 != dims[i + 1]:
+            dims_contiguous = False
+            break
+    indices = [advanced_indices_map[x] for x in dims]
+    indices = _promote_dtype(*indices)
+    indices = np_utils.tf_broadcast(*indices)
+    stacked_indices = array_ops_stack.stack(indices, axis=-1)
+    # Skip the contiguous-dims optimization for update because there is no
+    # tf.*scatter* op that supports the `axis` argument.
+    if not dims_contiguous or updates is not None:
+      if range(len(dims)) != dims:
+        tensor = moveaxis(tensor, dims, range(len(dims)))
+      tensor_shape_prefix = array_ops.shape(
+          tensor, out_type=stacked_indices.dtype
+      )[: len(dims)]
+      stacked_indices = array_ops.where_v2(
+          stacked_indices < 0,
+          stacked_indices + tensor_shape_prefix,
+          stacked_indices,
+      )
+      if updates is None:
+        return array_ops.gather_nd(tensor, stacked_indices)
+      else:
+        # We only need to move-axis `updates` in the contiguous case becausce
+        # only in this case the result dimensions of advanced indexing are in
+        # the middle of `updates`. In the non-contiguous case, those dimensions
+        # are always at the front.
+        if dims_contiguous:
+          # TODO(wangpeng): Support unknown rank (e.g. by partially flattening
+          #   `updates`)
+          if stacked_indices.shape.rank is None:
+            raise NotImplementedError(
+                'Rank of the advanced indices must currently be known'
+            )
+          batch_size = stacked_indices.shape.rank - 1
+          batch_start = dims[0]
+          if batch_start < 0:
+            batch_start += len(dims) - batch_size
+
+          def range_(start, length):
+            return range(start, start + length)
+
+          updates = moveaxis(
+              updates, range_(batch_start, batch_size), range(batch_size)
+          )
+        if update_method == _UpdateMethod.UPDATE:
+          update_op = array_ops.tensor_scatter_update
+        elif update_method == _UpdateMethod.ADD:
+          update_op = array_ops.tensor_scatter_add
+        elif update_method == _UpdateMethod.MIN:
+          update_op = array_ops.tensor_scatter_min
+        elif update_method == _UpdateMethod.MAX:
+          update_op = array_ops.tensor_scatter_max
+        tensor = update_op(tensor, stacked_indices, updates)
+        if range(len(dims)) != dims:
+          tensor = moveaxis(tensor, range(len(dims)), dims)
+        return array_ops.tensor_strided_slice_update(
+            original_tensor,
+            packed_begin,
+            packed_end,
+            packed_strides,
+            tensor,
+            begin_mask=begin_mask,
+            end_mask=end_mask,
+            shrink_axis_mask=shrink_axis_mask,
+            new_axis_mask=new_axis_mask,
+            ellipsis_mask=ellipsis_mask,
+            name=name + '_2',
+        )
+    # Note that gather_nd does not support gathering from inside the array.
+    # To avoid shuffling data back and forth, we transform the indices and
+    # do a gather instead.
+    rank = np_utils._maybe_static(array_ops.rank(tensor))  # pylint: disable=protected-access
+    dims = [(x + rank if x < 0 else x) for x in dims]
+    shape_tensor = array_ops.shape(tensor)
+    dim_sizes = array_ops.gather(shape_tensor, dims)
+    if len(dims) == 1:
+      stacked_indices = indices[0]
+    stacked_indices = math_ops.cast(stacked_indices, dtypes.int32)
+    stacked_indices = array_ops.where_v2(
+        stacked_indices < 0, stacked_indices + dim_sizes, stacked_indices
+    )
+    axis = dims[0]
+    if len(dims) > 1:
+      index_scaling = math_ops.cumprod(dim_sizes, reverse=True, exclusive=True)
+
+      def _tensordot(a, b):
+        # TODO(b/168657656): This function should be replaced by
+        # tensordot(axis=1) once MatMul has int32 XLA kernel.
+        b = array_ops.broadcast_to(b, array_ops.shape(a))
+        return math_ops.reduce_sum(a * b, axis=-1)
+
+      stacked_indices = _tensordot(stacked_indices, index_scaling)
+      flat_shape = array_ops.concat(
+          [shape_tensor[:axis], [-1], shape_tensor[axis + len(dims) :]], axis=0
+      )
+      tensor = array_ops.reshape(tensor, flat_shape)
+
+    return array_ops.gather(tensor, stacked_indices, axis=axis)
+
+
+def _as_spec_tuple(slice_spec):
+  """Convert slice_spec to tuple."""
+  if isinstance(slice_spec, (list, tuple)) and not isinstance(
+      slice_spec, np.ndarray
+  ):
+    is_index = True
+    for s in slice_spec:
+      if s is None or s is Ellipsis or isinstance(s, (list, tuple, slice)):
+        is_index = False
+        break
+      elif isinstance(s, (np_arrays.ndarray, np.ndarray)) and s.ndim != 0:
+        is_index = False
+        break
+    if not is_index:
+      return tuple(slice_spec)
+  return (slice_spec,)
+
+
+def _getitem(self, slice_spec):
+  """Implementation of ndarray.__getitem__."""
+  if (
+      isinstance(slice_spec, bool)
+      or (
+          isinstance(slice_spec, core_tf_types.Tensor)
+          and slice_spec.dtype == dtypes.bool
+      )
+      or (
+          isinstance(slice_spec, (np.ndarray, np_arrays.ndarray))
+          and slice_spec.dtype == np.bool_
+      )
+  ):
+    return array_ops.boolean_mask(tensor=self, mask=slice_spec)
+
+  if not isinstance(slice_spec, tuple):
+    slice_spec = _as_spec_tuple(slice_spec)
+
+  result_t = _slice_helper(self, slice_spec)
+  return result_t
+
+
+def _with_index_update_helper(update_method, a, slice_spec, updates):
+  """Implementation of ndarray._with_index_*."""
+  if (
+      isinstance(slice_spec, bool)
+      or (
+          isinstance(slice_spec, core_tf_types.Tensor)
+          and slice_spec.dtype == dtypes.bool
+      )
+      or (
+          isinstance(slice_spec, (np.ndarray, np_arrays.ndarray))
+          and slice_spec.dtype == np.bool_
+      )
+  ):
+    slice_spec = nonzero(slice_spec)
+
+  if not isinstance(slice_spec, tuple):
+    slice_spec = _as_spec_tuple(slice_spec)
+
+  a_dtype = a.dtype
+  a, updates = _promote_dtype_binary(a, updates)
+  result_t = _slice_helper(a, slice_spec, update_method, updates)
+  return result_t.astype(a_dtype)
+
+
+setattr(np_arrays.ndarray, '_numpy_style_getitem', _getitem)
+setattr(
+    np_arrays.ndarray,
+    '_with_index_update',
+    functools.partial(_with_index_update_helper, _UpdateMethod.UPDATE),
+)
+setattr(
+    np_arrays.ndarray,
+    '_with_index_add',
+    functools.partial(_with_index_update_helper, _UpdateMethod.ADD),
+)
+setattr(
+    np_arrays.ndarray,
+    '_with_index_min',
+    functools.partial(_with_index_update_helper, _UpdateMethod.MIN),
+)
+setattr(
+    np_arrays.ndarray,
+    '_with_index_max',
+    functools.partial(_with_index_update_helper, _UpdateMethod.MAX),
+)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_arrays.py

@@ -0,0 +1,50 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""ndarray class."""
+
+# pylint: disable=g-direct-tensorflow-import
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.ops.numpy_ops import np_dtypes
+
+
+def convert_to_tensor(value, dtype=None, dtype_hint=None):
+  """Wrapper over `tf.convert_to_tensor`.
+
+  Args:
+    value: value to convert
+    dtype: (optional) the type we would like it to be converted to.
+    dtype_hint: (optional) soft preference for the type we would like it to be
+      converted to. `tf.convert_to_tensor` will attempt to convert value to this
+      type first, but will not fail if conversion is not possible falling back
+      to inferring the type instead.
+
+  Returns:
+    Value converted to tf.Tensor.
+  """
+  # A safer version of `tf.convert_to_tensor` to work around b/149876037.
+  # TODO(wangpeng): Remove this function once the bug is fixed.
+  if (dtype is None and isinstance(value, int) and
+      value >= 2**63):
+    dtype = dtypes.uint64
+  elif dtype is None and dtype_hint is None and isinstance(value, float):
+    dtype = np_dtypes.default_float_type()
+  return tensor_conversion.convert_to_tensor_v2_with_dispatch(
+      value, dtype=dtype, dtype_hint=dtype_hint)
+
+
+ndarray = tensor.Tensor

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_config.py

@@ -0,0 +1,58 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Config functions for TF NumPy."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import weak_tensor_ops  # pylint: disable=unused-import
+from tensorflow.python.ops.numpy_ops import np_dtypes
+from tensorflow.python.ops.numpy_ops import np_math_ops
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.util import tf_export
+
+
+@tf_export.tf_export(
+    "experimental.numpy.experimental_enable_numpy_behavior", v1=[]
+)
+def enable_numpy_behavior(prefer_float32=False, dtype_conversion_mode="legacy"):
+  """Enable NumPy behavior on Tensors.
+
+  Enabling NumPy behavior has three effects:
+  * It adds to `tf.Tensor` some common NumPy methods such as `T`,
+    `reshape` and `ravel`.
+  * It changes dtype promotion in `tf.Tensor` operators to be
+    compatible with NumPy. For example,
+    `tf.ones([], tf.int32) + tf.ones([], tf.float32)` used to throw a
+    "dtype incompatible" error, but after this it will return a
+    float64 tensor (obeying NumPy's promotion rules).
+  * It enhances `tf.Tensor`'s indexing capability to be on par with
+    [NumPy's](https://numpy.org/doc/stable/reference/arrays.indexing.html).
+
+  Args:
+    prefer_float32: Controls whether dtype inference will use float32 for Python
+      floats, or float64 (the default and the NumPy-compatible behavior).
+    dtype_conversion_mode: a string that specifies promotion mode. This string
+      corresponds to a PromoMode Enum and can be 'off', 'legacy', 'safe', or
+      'all'. 'safe' or 'all' mode enables the auto dtype conversion semantics.
+  """
+  if dtype_conversion_mode == "safe" or dtype_conversion_mode == "all":
+    tf_logging.warning(
+        "UserWarning: enabling the new type promotion must happen at the"
+        " beginning of the program. Please ensure no TF APIs have been used"
+        " yet."
+    )
+  ops.set_dtype_conversion_mode(dtype_conversion_mode)
+  ops.enable_numpy_style_slicing()
+  np_math_ops.enable_numpy_methods_on_tensor()
+  np_dtypes.set_prefer_float32(prefer_float32)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_dtypes.py

@@ -0,0 +1,216 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Dtypes and dtype utilities."""
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.util import tf_export
+
+
+# We use numpy's dtypes instead of TF's, because the user expects to use them
+# with numpy facilities such as `np.dtype(np.int64)` and
+# `if x.dtype.type is np.int64`.
+bool_ = np.bool_
+tf_export.tf_export('experimental.numpy.bool_', v1=[]).export_constant(
+    __name__, 'bool_'
+)
+complex128 = np.complex128
+tf_export.tf_export('experimental.numpy.complex128', v1=[]).export_constant(
+    __name__, 'complex128'
+)
+complex64 = np.complex64
+tf_export.tf_export('experimental.numpy.complex64', v1=[]).export_constant(
+    __name__, 'complex64'
+)
+float16 = np.float16
+tf_export.tf_export('experimental.numpy.float16', v1=[]).export_constant(
+    __name__, 'float16'
+)
+float32 = np.float32
+tf_export.tf_export('experimental.numpy.float32', v1=[]).export_constant(
+    __name__, 'float32'
+)
+float64 = np.float64
+tf_export.tf_export('experimental.numpy.float64', v1=[]).export_constant(
+    __name__, 'float64'
+)
+inexact = np.inexact
+tf_export.tf_export('experimental.numpy.inexact', v1=[]).export_constant(
+    __name__, 'inexact'
+)
+int_ = np.int_
+tf_export.tf_export('experimental.numpy.int_', v1=[]).export_constant(
+    __name__, 'int_'
+)
+int16 = np.int16
+tf_export.tf_export('experimental.numpy.int16', v1=[]).export_constant(
+    __name__, 'int16'
+)
+int32 = np.int32
+tf_export.tf_export('experimental.numpy.int32', v1=[]).export_constant(
+    __name__, 'int32'
+)
+int64 = np.int64
+tf_export.tf_export('experimental.numpy.int64', v1=[]).export_constant(
+    __name__, 'int64'
+)
+int8 = np.int8
+tf_export.tf_export('experimental.numpy.int8', v1=[]).export_constant(
+    __name__, 'int8'
+)
+object_ = np.object_
+tf_export.tf_export('experimental.numpy.object_', v1=[]).export_constant(
+    __name__, 'object_'
+)
+# np.string_ is aliased to np.bytes_ and depercated in numpy 2.0.
+string_ = np.bytes_
+tf_export.tf_export('experimental.numpy.string_', v1=[]).export_constant(
+    __name__, 'string_'
+)
+uint16 = np.uint16
+tf_export.tf_export('experimental.numpy.uint16', v1=[]).export_constant(
+    __name__, 'uint16'
+)
+uint32 = np.uint32
+tf_export.tf_export('experimental.numpy.uint32', v1=[]).export_constant(
+    __name__, 'uint32'
+)
+uint64 = np.uint64
+tf_export.tf_export('experimental.numpy.uint64', v1=[]).export_constant(
+    __name__, 'uint64'
+)
+uint8 = np.uint8
+tf_export.tf_export('experimental.numpy.uint8', v1=[]).export_constant(
+    __name__, 'uint8'
+)
+# np.unicode_ is aliased to np.str_ and depercated in numpy 2.0.
+unicode_ = np.str_
+tf_export.tf_export('experimental.numpy.unicode_', v1=[]).export_constant(
+    __name__, 'unicode_'
+)
+if int(np.__version__.split('.')[0]) < 2:
+  complex_ = np.complex_
+  float_ = np.float_
+else:
+  # Aliases np.complex_ and np.float_ have been removed in Numpy 2.0. Use
+  # np.complex128 and np.float64 instead.
+  complex_ = np.complex128
+  float_ = np.float64
+tf_export.tf_export('experimental.numpy.complex_', v1=[]).export_constant(
+    __name__, 'complex_'
+)
+tf_export.tf_export('experimental.numpy.float_', v1=[]).export_constant(
+    __name__, 'float_'
+)
+
+
+iinfo = np.iinfo
+tf_export.tf_export('experimental.numpy.iinfo', v1=[]).export_constant(
+    __name__, 'iinfo'
+)
+
+
+issubdtype = tf_export.tf_export('experimental.numpy.issubdtype', v1=[])(
+    np.issubdtype
+)
+
+
+_to_float32 = {
+    np.dtype('float64'): np.dtype('float32'),
+    np.dtype('complex128'): np.dtype('complex64'),
+}
+
+
+_cached_np_dtypes = {}
+
+
+# Difference between is_prefer_float32 and is_allow_float64: is_prefer_float32
+# only decides which dtype to use for Python floats; is_allow_float64 decides
+# whether float64 dtypes can ever appear in programs. The latter is more
+# restrictive than the former.
+_prefer_float32 = False
+
+
+# TODO(b/178862061): Consider removing this knob
+_allow_float64 = True
+
+
+def is_prefer_float32():
+  return _prefer_float32
+
+
+def set_prefer_float32(b):
+  global _prefer_float32
+  _prefer_float32 = b
+
+
+def is_allow_float64():
+  return _allow_float64
+
+
+def set_allow_float64(b):
+  global _allow_float64
+  _allow_float64 = b
+
+
+def canonicalize_dtype(dtype):
+  if not _allow_float64:
+    try:
+      return _to_float32[dtype]
+    except KeyError:
+      pass
+  return dtype
+
+
+def _result_type(*arrays_and_dtypes):
+  """Returns the resulting type given a set of arrays."""
+
+  def preprocess_float(x):
+    if is_prefer_float32():
+      if isinstance(x, float):
+        return np.float32(x)
+      elif isinstance(x, complex):
+        return np.complex64(x)
+    return x
+
+  arrays_and_dtypes = [preprocess_float(x) for x in arrays_and_dtypes]
+  dtype = np.result_type(*arrays_and_dtypes)
+  return dtypes.as_dtype(canonicalize_dtype(dtype))
+
+
+def _get_cached_dtype(dtype):
+  """Returns an np.dtype for the TensorFlow DType."""
+  global _cached_np_dtypes
+  try:
+    return _cached_np_dtypes[dtype]
+  except KeyError:
+    pass
+  cached_dtype = np.dtype(dtype.as_numpy_dtype)
+  _cached_np_dtypes[dtype] = cached_dtype
+  return cached_dtype
+
+
+def default_float_type():
+  """Gets the default float type.
+
+  Returns:
+    If `is_prefer_float32()` is false and `is_allow_float64()` is true, returns
+    float64; otherwise returns float32.
+  """
+  if not is_prefer_float32() and is_allow_float64():
+    return float64
+  else:
+    return float32

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_math_ops.py

@@ -0,0 +1,1652 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Mathematical operations."""
+# pylint: disable=g-direct-tensorflow-import
+
+import numbers
+import sys
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import bitwise_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import sort_ops
+from tensorflow.python.ops import special_math_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.ops.numpy_ops import np_array_ops
+from tensorflow.python.ops.numpy_ops import np_arrays
+from tensorflow.python.ops.numpy_ops import np_dtypes
+from tensorflow.python.ops.numpy_ops import np_utils
+from tensorflow.python.util import tf_export
+
+
+pi = np.pi
+tf_export.tf_export('experimental.numpy.pi', v1=[]).export_constant(
+    __name__, 'pi'
+)
+e = np.e
+tf_export.tf_export('experimental.numpy.e', v1=[]).export_constant(
+    __name__, 'e'
+)
+inf = np.inf
+tf_export.tf_export('experimental.numpy.inf', v1=[]).export_constant(
+    __name__, 'inf'
+)
+
+
+@tf_export.tf_export('experimental.numpy.dot', v1=[])
+@np_utils.np_doc_only('dot')
+def dot(a, b):  # pylint: disable=missing-docstring
+  def f(a, b):  # pylint: disable=missing-docstring
+    return np_utils.cond(
+        np_utils.logical_or(
+            math_ops.equal(array_ops.rank(a), 0),
+            math_ops.equal(array_ops.rank(b), 0),
+        ),
+        lambda: a * b,
+        lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+            math_ops.equal(array_ops.rank(b), 1),
+            lambda: math_ops.tensordot(a, b, axes=[[-1], [-1]]),
+            lambda: math_ops.tensordot(a, b, axes=[[-1], [-2]]),
+        ),
+    )
+
+  return _bin_op(f, a, b)
+
+
+# TODO(wangpeng): Make element-wise ops `ufunc`s
+def _bin_op(tf_fun, a, b, promote=True):
+  if promote:
+    a, b = np_array_ops._promote_dtype_binary(a, b)  # pylint: disable=protected-access
+  else:
+    a = np_array_ops.array(a)
+    b = np_array_ops.array(b)
+  return tf_fun(a, b)
+
+
+@tf_export.tf_export('experimental.numpy.add', v1=[])
+@np_utils.np_doc('add')
+def add(x1, x2):
+  def add_or_or(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      return math_ops.logical_or(x1, x2)
+    return math_ops.add(x1, x2)
+
+  return _bin_op(add_or_or, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.subtract', v1=[])
+@np_utils.np_doc('subtract')
+def subtract(x1, x2):
+  return _bin_op(math_ops.subtract, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.multiply', v1=[])
+@np_utils.np_doc('multiply')
+def multiply(x1, x2):
+  def mul_or_and(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      return math_ops.logical_and(x1, x2)
+    return math_ops.multiply(x1, x2)
+
+  return _bin_op(mul_or_and, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.true_divide', v1=[])
+@np_utils.np_doc('true_divide')
+def true_divide(x1, x2):  # pylint: disable=missing-function-docstring
+  def _avoid_float64(x1, x2):
+    if x1.dtype == x2.dtype and x1.dtype in (dtypes.int32, dtypes.int64):
+      x1 = math_ops.cast(x1, dtype=dtypes.float32)
+      x2 = math_ops.cast(x2, dtype=dtypes.float32)
+    return x1, x2
+
+  def f(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      float_ = np_utils.result_type(float)
+      x1 = math_ops.cast(x1, float_)
+      x2 = math_ops.cast(x2, float_)
+    if not np_dtypes.is_allow_float64():
+      # math_ops.truediv in Python3 produces float64 when both inputs are int32
+      # or int64. We want to avoid that when is_allow_float64() is False.
+      x1, x2 = _avoid_float64(x1, x2)
+    return math_ops.truediv(x1, x2)
+
+  return _bin_op(f, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.divide', v1=[])
+@np_utils.np_doc('divide')
+def divide(x1, x2):  # pylint: disable=missing-function-docstring
+  return true_divide(x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.floor_divide', v1=[])
+@np_utils.np_doc('floor_divide')
+def floor_divide(x1, x2):  # pylint: disable=missing-function-docstring
+  def f(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      x1 = math_ops.cast(x1, dtypes.int8)
+      x2 = math_ops.cast(x2, dtypes.int8)
+    return math_ops.floordiv(x1, x2)
+
+  return _bin_op(f, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.mod', v1=[])
+@np_utils.np_doc('mod')
+def mod(x1, x2):  # pylint: disable=missing-function-docstring
+  def f(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      x1 = math_ops.cast(x1, dtypes.int8)
+      x2 = math_ops.cast(x2, dtypes.int8)
+    return math_ops.mod(x1, x2)
+
+  return _bin_op(f, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.remainder', v1=[])
+@np_utils.np_doc('remainder')
+def remainder(x1, x2):  # pylint: disable=missing-function-docstring
+  return mod(x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.divmod', v1=[])
+@np_utils.np_doc('divmod')
+def divmod(x1, x2):  # pylint: disable=redefined-builtin
+  return floor_divide(x1, x2), mod(x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.maximum', v1=[])
+@np_utils.np_doc('maximum')
+def maximum(x1, x2):  # pylint: disable=missing-function-docstring
+  # Fast path for when maximum is used as relu.
+  if (
+      isinstance(x2, numbers.Real)
+      and not isinstance(x2, bool)
+      and x2 == 0
+      and isinstance(x1, np_arrays.ndarray)
+      and x1.dtype != dtypes.bool
+  ):
+    return nn_ops.relu(np_array_ops.asarray(x1))
+
+  def max_or_or(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      return math_ops.logical_or(x1, x2)
+    return math_ops.maximum(x1, x2)
+
+  return _bin_op(max_or_or, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.minimum', v1=[])
+@np_utils.np_doc('minimum')
+def minimum(x1, x2):
+  def min_or_and(x1, x2):
+    if x1.dtype == dtypes.bool:
+      assert x2.dtype == dtypes.bool
+      return math_ops.logical_and(x1, x2)
+    return math_ops.minimum(x1, x2)
+
+  return _bin_op(min_or_and, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.clip', v1=[])
+@np_utils.np_doc('clip')
+def clip(a, a_min, a_max):  # pylint: disable=missing-docstring
+  if a_min is None and a_max is None:
+    raise ValueError('Not more than one of `a_min` and `a_max` may be `None`.')
+  if a_min is None:
+    return minimum(a, a_max)
+  elif a_max is None:
+    return maximum(a, a_min)
+  else:
+    a, a_min, a_max = np_array_ops._promote_dtype(a, a_min, a_max)  # pylint: disable=protected-access
+    return clip_ops.clip_by_value(*np_utils.tf_broadcast(a, a_min, a_max))
+
+
+@tf_export.tf_export('experimental.numpy.matmul', v1=[])
+@np_utils.np_doc('matmul')
+def matmul(x1, x2):  # pylint: disable=missing-docstring
+  def f(x1, x2):
+    try:
+      if x1._rank() == 2 and x2._rank() == 2:  # pylint: disable=protected-access
+        # Fast path for known ranks.
+        return gen_math_ops.mat_mul(x1, x2)
+      return np_utils.cond(
+          math_ops.equal(np_utils.tf_rank(x2), 1),
+          lambda: math_ops.tensordot(x1, x2, axes=1),
+          lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+              math_ops.equal(np_utils.tf_rank(x1), 1),
+              lambda: math_ops.tensordot(  # pylint: disable=g-long-lambda
+                  x1, x2, axes=[[0], [-2]]
+              ),
+              lambda: math_ops.matmul(x1, x2),
+          ),
+      )
+    except errors.InvalidArgumentError as err:
+      raise ValueError(str(err)).with_traceback(sys.exc_info()[2])
+
+  return _bin_op(f, x1, x2)
+
+
+# Exported so it can be called from Tensor.__matmul__. NumPy's matmul handles
+# batched matmul as well, so simply including promotion in TF's current
+# __matmul__ implementation was not sufficient.
+setattr(np_arrays.ndarray, '_matmul', matmul)
+
+
+@tf_export.tf_export('experimental.numpy.tensordot', v1=[])
+@np_utils.np_doc('tensordot')
+def tensordot(a, b, axes=2):
+  return _bin_op(lambda a, b: math_ops.tensordot(a, b, axes=axes), a, b)
+
+
+@tf_export.tf_export('experimental.numpy.inner', v1=[])
+@np_utils.np_doc_only('inner')
+def inner(a, b):  # pylint: disable=missing-function-docstring
+  def f(a, b):
+    return np_utils.cond(
+        np_utils.logical_or(
+            math_ops.equal(array_ops.rank(a), 0),
+            math_ops.equal(array_ops.rank(b), 0),
+        ),
+        lambda: a * b,
+        lambda: math_ops.tensordot(a, b, axes=[[-1], [-1]]),
+    )
+
+  return _bin_op(f, a, b)
+
+
+@tf_export.tf_export('experimental.numpy.cross', v1=[])
+@np_utils.np_doc('cross')
+def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):  # pylint: disable=missing-docstring
+  def f(a, b):  # pylint: disable=missing-docstring
+    # We can't assign to captured variable `axisa`, so make a new variable
+    if axis is None:
+      axis_a = axisa
+      axis_b = axisb
+      axis_c = axisc
+    else:
+      axis_a = axis
+      axis_b = axis
+      axis_c = axis
+    if axis_a < 0:
+      axis_a = np_utils.add(axis_a, array_ops.rank(a))
+    if axis_b < 0:
+      axis_b = np_utils.add(axis_b, array_ops.rank(b))
+
+    def maybe_move_axis_to_last(a, axis):
+      def move_axis_to_last(a, axis):
+        return array_ops.transpose(
+            a,
+            array_ops.concat(
+                [
+                    math_ops.range(axis),
+                    math_ops.range(axis + 1, array_ops.rank(a)),
+                    [axis],
+                ],
+                axis=0,
+            ),
+        )
+
+      return np_utils.cond(
+          axis == np_utils.subtract(array_ops.rank(a), 1),
+          lambda: a,
+          lambda: move_axis_to_last(a, axis),
+      )
+
+    a = maybe_move_axis_to_last(a, axis_a)
+    b = maybe_move_axis_to_last(b, axis_b)
+    a_dim = np_utils.getitem(array_ops.shape(a), -1)
+    b_dim = np_utils.getitem(array_ops.shape(b), -1)
+
+    def maybe_pad_0(a, size_of_last_dim):
+      def pad_0(a):
+        return array_ops.pad(
+            a,
+            array_ops.concat(
+                [
+                    array_ops.zeros([array_ops.rank(a) - 1, 2], dtypes.int32),
+                    constant_op.constant([[0, 1]], dtypes.int32),
+                ],
+                axis=0,
+            ),
+        )
+
+      return np_utils.cond(
+          math_ops.equal(size_of_last_dim, 2), lambda: pad_0(a), lambda: a
+      )
+
+    a = maybe_pad_0(a, a_dim)
+    b = maybe_pad_0(b, b_dim)
+    c = math_ops.cross(*np_utils.tf_broadcast(a, b))
+    if axis_c < 0:
+      axis_c = np_utils.add(axis_c, array_ops.rank(c))
+
+    def move_last_to_axis(a, axis):
+      r = array_ops.rank(a)
+      return array_ops.transpose(
+          a,
+          array_ops.concat(
+              [math_ops.range(axis), [r - 1], math_ops.range(axis, r - 1)],
+              axis=0,
+          ),
+      )
+
+    c = np_utils.cond(
+        (a_dim == 2) & (b_dim == 2),
+        lambda: c[..., 2],
+        lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+            axis_c == np_utils.subtract(array_ops.rank(c), 1),
+            lambda: c,
+            lambda: move_last_to_axis(c, axis_c),
+        ),
+    )
+    return c
+
+  return _bin_op(f, a, b)
+
+
+@tf_export.tf_export('experimental.numpy.vdot', v1=[])
+@np_utils.np_doc_only('vdot')
+def vdot(a, b):  # pylint: disable=missing-docstring
+  a, b = np_array_ops._promote_dtype(a, b)  # pylint: disable=protected-access
+  a = np_array_ops.reshape(a, [-1])
+  b = np_array_ops.reshape(b, [-1])
+  if a.dtype == np_dtypes.complex128 or a.dtype == np_dtypes.complex64:
+    a = conj(a)
+  return dot(a, b)
+
+
+@tf_export.tf_export('experimental.numpy.power', v1=[])
+@np_utils.np_doc('power')
+def power(x1, x2):
+  return _bin_op(math_ops.pow, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.float_power', v1=[])
+@np_utils.np_doc('float_power')
+def float_power(x1, x2):
+  return power(x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.arctan2', v1=[])
+@np_utils.np_doc('arctan2')
+def arctan2(x1, x2):
+  return _bin_op(math_ops.atan2, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.nextafter', v1=[])
+@np_utils.np_doc('nextafter')
+def nextafter(x1, x2):
+  return _bin_op(math_ops.nextafter, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.heaviside', v1=[])
+@np_utils.np_doc('heaviside')
+def heaviside(x1, x2):  # pylint: disable=missing-function-docstring
+  def f(x1, x2):
+    return array_ops.where_v2(
+        x1 < 0,
+        constant_op.constant(0, dtype=x2.dtype),
+        array_ops.where_v2(x1 > 0, constant_op.constant(1, dtype=x2.dtype), x2),
+    )
+
+  y = _bin_op(f, x1, x2)
+  if not np.issubdtype(y.dtype.as_numpy_dtype, np.inexact):
+    y = y.astype(np_utils.result_type(float))
+  return y
+
+
+@tf_export.tf_export('experimental.numpy.hypot', v1=[])
+@np_utils.np_doc('hypot')
+def hypot(x1, x2):
+  return sqrt(square(x1) + square(x2))
+
+
+@tf_export.tf_export('experimental.numpy.kron', v1=[])
+@np_utils.np_doc('kron')
+def kron(a, b):  # pylint: disable=missing-function-docstring
+  # pylint: disable=protected-access,g-complex-comprehension
+  a, b = np_array_ops._promote_dtype(a, b)
+  t_a = np_utils.cond(
+      a.shape.rank < b.shape.rank,
+      lambda: np_array_ops.reshape(  # pylint: disable=g-long-lambda
+          a, np_array_ops._pad_left_to(b.shape.rank, a.shape)
+      ),
+      lambda: a,
+  )
+  t_b = np_utils.cond(
+      b.shape.rank < a.shape.rank,
+      lambda: np_array_ops.reshape(  # pylint: disable=g-long-lambda
+          b, np_array_ops._pad_left_to(a.shape.rank, b.shape)
+      ),
+      lambda: b,
+  )
+
+  def _make_shape(shape, prepend):
+    ones = array_ops.ones_like(shape)
+    if prepend:
+      shapes = [ones, shape]
+    else:
+      shapes = [shape, ones]
+    return array_ops.reshape(array_ops_stack.stack(shapes, axis=1), [-1])
+
+  a_shape = array_ops.shape(t_a)
+  b_shape = array_ops.shape(t_b)
+  a_reshaped = np_array_ops.reshape(t_a, _make_shape(a_shape, False))
+  b_reshaped = np_array_ops.reshape(t_b, _make_shape(b_shape, True))
+  out_shape = a_shape * b_shape
+  return np_array_ops.reshape(a_reshaped * b_reshaped, out_shape)
+
+
+@tf_export.tf_export('experimental.numpy.outer', v1=[])
+@np_utils.np_doc('outer')
+def outer(a, b):
+  def f(a, b):
+    return array_ops.reshape(a, [-1, 1]) * array_ops.reshape(b, [-1])
+
+  return _bin_op(f, a, b)
+
+
+# This can also be implemented via tf.reduce_logsumexp
+@tf_export.tf_export('experimental.numpy.logaddexp', v1=[])
+@np_utils.np_doc('logaddexp')
+def logaddexp(x1, x2):
+  amax = maximum(x1, x2)
+  delta = x1 - x2
+  return np_array_ops.where(
+      isnan(delta),
+      x1 + x2,  # NaNs or infinities of the same sign.
+      amax + log1p(exp(-abs(delta))),
+  )
+
+
+@tf_export.tf_export('experimental.numpy.logaddexp2', v1=[])
+@np_utils.np_doc('logaddexp2')
+def logaddexp2(x1, x2):
+  amax = maximum(x1, x2)
+  delta = x1 - x2
+  return np_array_ops.where(
+      isnan(delta),
+      x1 + x2,  # NaNs or infinities of the same sign.
+      amax + log1p(exp2(-abs(delta))) / np.log(2),
+  )
+
+
+@tf_export.tf_export('experimental.numpy.polyval', v1=[])
+@np_utils.np_doc('polyval')
+def polyval(p, x):  # pylint: disable=missing-function-docstring
+  def f(p, x):
+    if p.shape.rank == 0:
+      p = array_ops.reshape(p, [1])
+    p = array_ops_stack.unstack(p)
+    # TODO(wangpeng): Make tf version take a tensor for p instead of a list.
+    y = math_ops.polyval(p, x)
+    # If the polynomial is 0-order, numpy requires the result to be broadcast to
+    # `x`'s shape.
+    if len(p) == 1:
+      y = array_ops.broadcast_to(y, x.shape)
+    return y
+
+  return _bin_op(f, p, x)
+
+
+@tf_export.tf_export('experimental.numpy.isclose', v1=[])
+@np_utils.np_doc('isclose')
+def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):  # pylint: disable=missing-docstring
+  def f(a, b):  # pylint: disable=missing-docstring
+    dtype = a.dtype
+    if np.issubdtype(dtype.as_numpy_dtype, np.inexact):
+      rtol_ = ops.convert_to_tensor(rtol, dtype.real_dtype)
+      atol_ = ops.convert_to_tensor(atol, dtype.real_dtype)
+      result = math_ops.abs(a - b) <= atol_ + rtol_ * math_ops.abs(b)
+      if equal_nan:
+        result = result | (math_ops.is_nan(a) & math_ops.is_nan(b))
+      return result
+    else:
+      return a == b
+
+  return _bin_op(f, a, b)
+
+
+@tf_export.tf_export('experimental.numpy.allclose', v1=[])
+@np_utils.np_doc('allclose')
+def allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
+  return np_array_ops.all(
+      isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
+  )
+
+
+def _tf_gcd(x1, x2):  # pylint: disable=missing-function-docstring
+  def _gcd_cond_fn(_, x2):
+    return math_ops.reduce_any(x2 != 0)
+
+  def _gcd_body_fn(x1, x2):
+    # math_ops.mod will raise an error when any element of x2 is 0. To avoid
+    # that, we change those zeros to ones. Their values don't matter because
+    # they won't be used.
+    x2_safe = array_ops.where_v2(x2 != 0, x2, constant_op.constant(1, x2.dtype))
+    x1, x2 = (
+        array_ops.where_v2(x2 != 0, x2, x1),
+        array_ops.where_v2(
+            x2 != 0,
+            math_ops.mod(x1, x2_safe),
+            constant_op.constant(0, x2.dtype),
+        ),
+    )
+    return (
+        array_ops.where_v2(x1 < x2, x2, x1),
+        array_ops.where_v2(x1 < x2, x1, x2),
+    )
+
+  if not np.issubdtype(
+      x1.dtype.as_numpy_dtype, np.integer
+  ) or not np.issubdtype(x2.dtype.as_numpy_dtype, np.integer):
+    raise ValueError('Arguments to gcd must be integers.')
+  shape = array_ops.broadcast_dynamic_shape(
+      array_ops.shape(x1), array_ops.shape(x2)
+  )
+  x1 = array_ops.broadcast_to(x1, shape)
+  x2 = array_ops.broadcast_to(x2, shape)
+  value, _ = while_loop.while_loop(
+      _gcd_cond_fn, _gcd_body_fn, (math_ops.abs(x1), math_ops.abs(x2))
+  )
+  return value
+
+
+# Note that np.gcd may not be present in some supported versions of numpy.
+@tf_export.tf_export('experimental.numpy.gcd', v1=[])
+@np_utils.np_doc('gcd')
+def gcd(x1, x2):
+  return _bin_op(_tf_gcd, x1, x2)
+
+
+# Note that np.lcm may not be present in some supported versions of numpy.
+@tf_export.tf_export('experimental.numpy.lcm', v1=[])
+@np_utils.np_doc('lcm')
+def lcm(x1, x2):  # pylint: disable=missing-function-docstring
+  def f(x1, x2):
+    d = _tf_gcd(x1, x2)
+    # Same as the `x2_safe` trick above
+    d_safe = array_ops.where_v2(
+        math_ops.equal(d, 0), constant_op.constant(1, d.dtype), d
+    )
+    x1 = math_ops.abs(x1)
+    x2 = math_ops.abs(x2)
+    return array_ops.where_v2(
+        math_ops.equal(d, 0),
+        constant_op.constant(0, d.dtype),
+        x1 * (x2 // d_safe),
+    )
+
+  return _bin_op(f, x1, x2)
+
+
+def _bitwise_binary_op(tf_fn, x1, x2):  # pylint: disable=missing-function-docstring
+  def f(x1, x2):
+    is_bool = x1.dtype == dtypes.bool
+    if is_bool:
+      assert x2.dtype == dtypes.bool
+      x1 = math_ops.cast(x1, dtypes.int8)
+      x2 = math_ops.cast(x2, dtypes.int8)
+    r = tf_fn(x1, x2)
+    if is_bool:
+      r = math_ops.cast(r, dtypes.bool)
+    return r
+
+  return _bin_op(f, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.bitwise_and', v1=[])
+@np_utils.np_doc('bitwise_and')
+def bitwise_and(x1, x2):
+  return _bitwise_binary_op(bitwise_ops.bitwise_and, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.bitwise_or', v1=[])
+@np_utils.np_doc('bitwise_or')
+def bitwise_or(x1, x2):
+  return _bitwise_binary_op(bitwise_ops.bitwise_or, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.bitwise_xor', v1=[])
+@np_utils.np_doc('bitwise_xor')
+def bitwise_xor(x1, x2):
+  return _bitwise_binary_op(bitwise_ops.bitwise_xor, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.bitwise_not', v1=[])
+@np_utils.np_doc('bitwise_not', link=np_utils.AliasOf('invert'))
+def bitwise_not(x):
+  def f(x):
+    if x.dtype == dtypes.bool:
+      return math_ops.logical_not(x)
+    return bitwise_ops.invert(x)
+
+  return _scalar(f, x)
+
+
+def _scalar(tf_fn, x, promote_to_float=False):
+  """Computes the tf_fn(x) for each element in `x`.
+
+  Args:
+    tf_fn: function that takes a single Tensor argument.
+    x: array_like. Could be an ndarray, a Tensor or any object that can be
+      converted to a Tensor using `ops.convert_to_tensor`.
+    promote_to_float: whether to cast the argument to a float dtype if it is not
+      already.
+
+  Returns:
+    An ndarray with the same shape as `x`. The default output dtype is
+    determined by `np_utils.result_type(float)`, unless x is an ndarray with a
+    floating point type, in which case the output type is same as x.dtype.
+  """
+  x = np_array_ops.asarray(x)
+  if promote_to_float and not np.issubdtype(x.dtype.as_numpy_dtype, np.inexact):
+    x = x.astype(np_utils.result_type(float))
+  return tf_fn(x)
+
+
+@tf_export.tf_export('experimental.numpy.log', v1=[])
+@np_utils.np_doc('log')
+def log(x):
+  return _scalar(math_ops.log, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.exp', v1=[])
+@np_utils.np_doc('exp')
+def exp(x):
+  return _scalar(math_ops.exp, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.sqrt', v1=[])
+@np_utils.np_doc('sqrt')
+def sqrt(x):
+  return _scalar(math_ops.sqrt, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.abs', v1=[])
+@np_utils.np_doc('abs', link=np_utils.AliasOf('absolute'))
+def abs(x):  # pylint: disable=redefined-builtin
+  return _scalar(math_ops.abs, x)
+
+
+@tf_export.tf_export('experimental.numpy.absolute', v1=[])
+@np_utils.np_doc('absolute')
+def absolute(x):
+  return abs(x)
+
+
+@tf_export.tf_export('experimental.numpy.fabs', v1=[])
+@np_utils.np_doc('fabs')
+def fabs(x):
+  return abs(x)
+
+
+@tf_export.tf_export('experimental.numpy.ceil', v1=[])
+@np_utils.np_doc('ceil')
+def ceil(x):
+  return _scalar(math_ops.ceil, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.floor', v1=[])
+@np_utils.np_doc('floor')
+def floor(x):
+  return _scalar(math_ops.floor, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.conj', v1=[])
+@np_utils.np_doc('conj')
+def conj(x):
+  return _scalar(math_ops.conj, x)
+
+
+@tf_export.tf_export('experimental.numpy.negative', v1=[])
+@np_utils.np_doc('negative')
+def negative(x):
+  return _scalar(math_ops.negative, x)
+
+
+@tf_export.tf_export('experimental.numpy.reciprocal', v1=[])
+@np_utils.np_doc('reciprocal')
+def reciprocal(x):
+  return _scalar(math_ops.reciprocal, x)
+
+
+@tf_export.tf_export('experimental.numpy.signbit', v1=[])
+@np_utils.np_doc('signbit')
+def signbit(x):
+  def f(x):
+    if x.dtype == dtypes.bool:
+      return array_ops.fill(array_ops.shape(x), False)
+    return x < 0
+
+  return _scalar(f, x)
+
+
+@tf_export.tf_export('experimental.numpy.sin', v1=[])
+@np_utils.np_doc('sin')
+def sin(x):
+  return _scalar(math_ops.sin, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.cos', v1=[])
+@np_utils.np_doc('cos')
+def cos(x):
+  return _scalar(math_ops.cos, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.tan', v1=[])
+@np_utils.np_doc('tan')
+def tan(x):
+  return _scalar(math_ops.tan, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.sinh', v1=[])
+@np_utils.np_doc('sinh')
+def sinh(x):
+  return _scalar(math_ops.sinh, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.cosh', v1=[])
+@np_utils.np_doc('cosh')
+def cosh(x):
+  return _scalar(math_ops.cosh, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.tanh', v1=[])
+@np_utils.np_doc('tanh')
+def tanh(x):
+  return _scalar(math_ops.tanh, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.arcsin', v1=[])
+@np_utils.np_doc('arcsin')
+def arcsin(x):
+  return _scalar(math_ops.asin, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.arccos', v1=[])
+@np_utils.np_doc('arccos')
+def arccos(x):
+  return _scalar(math_ops.acos, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.arctan', v1=[])
+@np_utils.np_doc('arctan')
+def arctan(x):
+  return _scalar(math_ops.atan, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.arcsinh', v1=[])
+@np_utils.np_doc('arcsinh')
+def arcsinh(x):
+  return _scalar(math_ops.asinh, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.arccosh', v1=[])
+@np_utils.np_doc('arccosh')
+def arccosh(x):
+  return _scalar(math_ops.acosh, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.arctanh', v1=[])
+@np_utils.np_doc('arctanh')
+def arctanh(x):
+  return _scalar(math_ops.atanh, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.deg2rad', v1=[])
+@np_utils.np_doc('deg2rad')
+def deg2rad(x):
+  def f(x):
+    return x * (np.pi / 180.0)
+
+  return _scalar(f, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.rad2deg', v1=[])
+@np_utils.np_doc('rad2deg')
+def rad2deg(x):
+  return x * (180.0 / np.pi)
+
+
+_tf_float_types = [
+    dtypes.bfloat16,
+    dtypes.float16,
+    dtypes.float32,
+    dtypes.float64,
+]
+
+
+@tf_export.tf_export('experimental.numpy.angle', v1=[])
+@np_utils.np_doc('angle')
+def angle(z, deg=False):  # pylint: disable=missing-function-docstring
+  def f(x):
+    if x.dtype in _tf_float_types:
+      # Workaround for b/147515503
+      return array_ops.where_v2(x < 0, np.pi, 0)
+    else:
+      return math_ops.angle(x)
+
+  y = _scalar(f, z, True)
+  if deg:
+    y = rad2deg(y)
+  return y
+
+
+@tf_export.tf_export('experimental.numpy.cbrt', v1=[])
+@np_utils.np_doc('cbrt')
+def cbrt(x):
+  def f(x):
+    # __pow__ can't handle negative base, so we use `abs` here.
+    rt = math_ops.abs(x) ** (1.0 / 3)
+    return array_ops.where_v2(x < 0, -rt, rt)
+
+  return _scalar(f, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.conjugate', v1=[])
+@np_utils.np_doc('conjugate', link=np_utils.AliasOf('conj'))
+def conjugate(x):
+  return _scalar(math_ops.conj, x)
+
+
+@tf_export.tf_export('experimental.numpy.exp2', v1=[])
+@np_utils.np_doc('exp2')
+def exp2(x):
+  def f(x):
+    return 2**x
+
+  return _scalar(f, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.expm1', v1=[])
+@np_utils.np_doc('expm1')
+def expm1(x):
+  return _scalar(math_ops.expm1, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.fix', v1=[])
+@np_utils.np_doc('fix')
+def fix(x):
+  def f(x):
+    return array_ops.where_v2(x < 0, math_ops.ceil(x), math_ops.floor(x))
+
+  return _scalar(f, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.iscomplex', v1=[])
+@np_utils.np_doc('iscomplex')
+def iscomplex(x):
+  return np_array_ops.imag(x) != 0
+
+
+@tf_export.tf_export('experimental.numpy.isreal', v1=[])
+@np_utils.np_doc('isreal')
+def isreal(x):
+  return np_array_ops.imag(x) == 0
+
+
+@tf_export.tf_export('experimental.numpy.iscomplexobj', v1=[])
+@np_utils.np_doc('iscomplexobj')
+def iscomplexobj(x):
+  x = np_array_ops.array(x)
+  return np.issubdtype(x.dtype.as_numpy_dtype, np.complexfloating)
+
+
+@tf_export.tf_export('experimental.numpy.isrealobj', v1=[])
+@np_utils.np_doc('isrealobj')
+def isrealobj(x):
+  return not iscomplexobj(x)
+
+
+@tf_export.tf_export('experimental.numpy.isnan', v1=[])
+@np_utils.np_doc('isnan')
+def isnan(x):
+  return _scalar(math_ops.is_nan, x, True)
+
+
+def _make_nan_reduction(np_fun_name, reduction, init_val):
+  """Helper to generate nan* functions."""
+
+  @np_utils.np_doc(np_fun_name)
+  def nan_reduction(a, axis=None, dtype=None, keepdims=False):
+    a = np_array_ops.array(a)
+    v = np_array_ops.array(init_val, dtype=a.dtype)
+    return reduction(
+        np_array_ops.where(isnan(a), v, a),
+        axis=axis,
+        dtype=dtype,
+        keepdims=keepdims,
+    )
+
+  return nan_reduction
+
+
+nansum = tf_export.tf_export('experimental.numpy.nansum', v1=[])(
+    _make_nan_reduction('nansum', np_array_ops.sum, 0)
+)
+nanprod = tf_export.tf_export('experimental.numpy.nanprod', v1=[])(
+    _make_nan_reduction('nanprod', np_array_ops.prod, 1)
+)
+
+
+@tf_export.tf_export('experimental.numpy.nanmean', v1=[])
+@np_utils.np_doc('nanmean')
+def nanmean(a, axis=None, dtype=None, keepdims=None):  # pylint: disable=missing-docstring
+  a = np_array_ops.array(a)
+  if np.issubdtype(a.dtype.as_numpy_dtype, np.bool_) or np.issubdtype(
+      a.dtype.as_numpy_dtype, np.integer
+  ):
+    return np_array_ops.mean(a, axis=axis, dtype=dtype, keepdims=keepdims)
+  nan_mask = logical_not(isnan(a))
+  if dtype is None:
+    dtype = a.dtype.as_numpy_dtype
+  normalizer = np_array_ops.sum(
+      nan_mask, axis=axis, dtype=dtype, keepdims=keepdims
+  )
+  return nansum(a, axis=axis, dtype=dtype, keepdims=keepdims) / normalizer
+
+
+@tf_export.tf_export('experimental.numpy.isfinite', v1=[])
+@np_utils.np_doc('isfinite')
+def isfinite(x):
+  return _scalar(math_ops.is_finite, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.isinf', v1=[])
+@np_utils.np_doc('isinf')
+def isinf(x):
+  if x.dtype.is_floating:
+    return _scalar(math_ops.is_inf, x, True)
+  return False
+
+
+@tf_export.tf_export('experimental.numpy.isneginf', v1=[])
+@np_utils.np_doc('isneginf')
+def isneginf(x):
+  if x.dtype.is_floating:
+    return x == np_array_ops.full_like(x, -np.inf)
+  return False
+
+
+@tf_export.tf_export('experimental.numpy.isposinf', v1=[])
+@np_utils.np_doc('isposinf')
+def isposinf(x):
+  if x.dtype.is_floating:
+    return x == np_array_ops.full_like(x, np.inf)
+  return False
+
+
+@tf_export.tf_export('experimental.numpy.log2', v1=[])
+@np_utils.np_doc('log2')
+def log2(x):
+  return log(x) / np.log(2)
+
+
+@tf_export.tf_export('experimental.numpy.log10', v1=[])
+@np_utils.np_doc('log10')
+def log10(x):
+  return log(x) / np.log(10)
+
+
+@tf_export.tf_export('experimental.numpy.log1p', v1=[])
+@np_utils.np_doc('log1p')
+def log1p(x):
+  return _scalar(math_ops.log1p, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.positive', v1=[])
+@np_utils.np_doc('positive')
+def positive(x):
+  return _scalar(lambda x: x, x)
+
+
+@tf_export.tf_export('experimental.numpy.sinc', v1=[])
+@np_utils.np_doc('sinc')
+def sinc(x):
+  def f(x):
+    pi_x = x * np.pi
+    return array_ops.where_v2(
+        x == 0, array_ops.ones_like(x), math_ops.sin(pi_x) / pi_x
+    )
+
+  return _scalar(f, x, True)
+
+
+@tf_export.tf_export('experimental.numpy.square', v1=[])
+@np_utils.np_doc('square')
+def square(x):
+  return _scalar(math_ops.square, x)
+
+
+@tf_export.tf_export('experimental.numpy.diff', v1=[])
+@np_utils.np_doc('diff')
+def diff(a, n=1, axis=-1):  # pylint: disable=missing-function-docstring
+  def f(a):
+    # TODO(agarwal): transpose and reshape to N, H, 1 and do a 1D convolution
+    # TODO(agarwal): avoid depending on static rank.
+    nd = a.shape.rank
+    if nd is None:
+      raise ValueError(
+          'Function `diff` currently requires a known rank for input `a`. '
+          f'Received: a={a} (unknown rank)'
+      )
+    if (axis + nd if axis < 0 else axis) >= nd:
+      raise ValueError(
+          f'Argument `axis` (received axis={axis}) is out of bounds '
+          f'for input {a} of rank {nd}.'
+      )
+    if n < 0:
+      raise ValueError(
+          f'Argument `order` must be a non-negative integer. Received: axis={n}'
+      )
+    slice1 = [slice(None)] * nd
+    slice2 = [slice(None)] * nd
+    slice1[axis] = slice(1, None)
+    slice2[axis] = slice(None, -1)
+    slice1 = tuple(slice1)
+    slice2 = tuple(slice2)
+    op = math_ops.not_equal if a.dtype == dtypes.bool else math_ops.subtract
+    for _ in range(n):
+      a = op(a[slice1], a[slice2])
+    return a
+
+  return _scalar(f, a)
+
+
+def _wrap(f, reverse=False):
+  """Wraps binary ops so they can be added as operator overloads on ndarray."""
+
+  def _f(a, b):
+    if reverse:
+      a, b = b, a
+
+    if (
+        getattr(b, '__array_priority__', 0)
+        > np_arrays.ndarray.__array_priority__
+    ):
+      return NotImplemented
+
+    return f(a, b)
+
+  return _f
+
+
+def _comparison(tf_fun, x1, x2, cast_bool_to_int=False):
+  """Helper function for comparision."""
+  dtype = np_utils.result_type(x1, x2)
+  # Cast x1 and x2 to the result_type if needed.
+  x1 = np_array_ops.array(x1, dtype=dtype)
+  x2 = np_array_ops.array(x2, dtype=dtype)
+  if cast_bool_to_int and x1.dtype == dtypes.bool:
+    x1 = math_ops.cast(x1, dtypes.int32)
+    x2 = math_ops.cast(x2, dtypes.int32)
+  return tf_fun(x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.equal', v1=[])
+@np_utils.np_doc('equal')
+def equal(x1, x2):
+  return _comparison(math_ops.equal, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.not_equal', v1=[])
+@np_utils.np_doc('not_equal')
+def not_equal(x1, x2):
+  return _comparison(math_ops.not_equal, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.greater', v1=[])
+@np_utils.np_doc('greater')
+def greater(x1, x2):
+  return _comparison(math_ops.greater, x1, x2, True)
+
+
+@tf_export.tf_export('experimental.numpy.greater_equal', v1=[])
+@np_utils.np_doc('greater_equal')
+def greater_equal(x1, x2):
+  return _comparison(math_ops.greater_equal, x1, x2, True)
+
+
+@tf_export.tf_export('experimental.numpy.less', v1=[])
+@np_utils.np_doc('less')
+def less(x1, x2):
+  return _comparison(math_ops.less, x1, x2, True)
+
+
+@tf_export.tf_export('experimental.numpy.less_equal', v1=[])
+@np_utils.np_doc('less_equal')
+def less_equal(x1, x2):
+  return _comparison(math_ops.less_equal, x1, x2, True)
+
+
+@tf_export.tf_export('experimental.numpy.array_equal', v1=[])
+@np_utils.np_doc('array_equal')
+def array_equal(a1, a2):  # pylint: disable=missing-function-docstring
+  def f(x1, x2):
+    return np_utils.cond(
+        math_ops.equal(array_ops.rank(x1), array_ops.rank(x2)),
+        lambda: np_utils.cond(  # pylint: disable=g-long-lambda
+            np_utils.reduce_all(
+                math_ops.equal(array_ops.shape(x1), array_ops.shape(x2))
+            ),
+            lambda: math_ops.reduce_all(math_ops.equal(x1, x2)),
+            lambda: constant_op.constant(False),
+        ),
+        lambda: constant_op.constant(False),
+    )
+
+  return _comparison(f, a1, a2)
+
+
+def _logical_binary_op(tf_fun, x1, x2):
+  x1 = np_array_ops.array(x1, dtype=np.bool_)
+  x2 = np_array_ops.array(x2, dtype=np.bool_)
+  return tf_fun(x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.logical_and', v1=[])
+@np_utils.np_doc('logical_and')
+def logical_and(x1, x2):
+  return _logical_binary_op(math_ops.logical_and, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.logical_or', v1=[])
+@np_utils.np_doc('logical_or')
+def logical_or(x1, x2):
+  return _logical_binary_op(math_ops.logical_or, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.logical_xor', v1=[])
+@np_utils.np_doc('logical_xor')
+def logical_xor(x1, x2):
+  return _logical_binary_op(math_ops.logical_xor, x1, x2)
+
+
+@tf_export.tf_export('experimental.numpy.logical_not', v1=[])
+@np_utils.np_doc('logical_not')
+def logical_not(x):
+  x = np_array_ops.array(x, dtype=np.bool_)
+  return math_ops.logical_not(x)
+
+
+@tf_export.tf_export('experimental.numpy.linspace', v1=[])
+@np_utils.np_doc('linspace')
+def linspace(  # pylint: disable=missing-docstring
+    start, stop, num=50, endpoint=True, retstep=False, dtype=float, axis=0
+):
+  if dtype:
+    # In numpy 2.x, the result type of np.linspace is based off of `start` and
+    # `end`. We mimic the behavior.
+    if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
+      dtype = np_utils.result_type([start * 1.0, stop * 1.0])
+    else:
+      dtype = np_utils.result_type(dtype)
+  start = np_array_ops.array(start, dtype=dtype)
+  stop = np_array_ops.array(stop, dtype=dtype)
+  if num < 0:
+    raise ValueError(
+        'Argument `num` (number of samples) must be a non-negative integer. '
+        f'Received: num={num}'
+    )
+  step = ops.convert_to_tensor(np.nan)
+  if endpoint:
+    result = math_ops.linspace(start, stop, num, axis=axis)
+    if num > 1:
+      step = (stop - start) / (num - 1)
+  else:
+    # math_ops.linspace does not support endpoint=False so we manually handle it
+    # here.
+    if num > 0:
+      step = (stop - start) / num
+    if num > 1:
+      new_stop = math_ops.cast(stop, step.dtype) - step
+      start = math_ops.cast(start, new_stop.dtype)
+      result = math_ops.linspace(start, new_stop, num, axis=axis)
+    else:
+      result = math_ops.linspace(start, stop, num, axis=axis)
+  if dtype:
+    if dtype.is_integer:
+      # Since numpy 1.20, linspace's rounding is towards -inf instead of 0
+      result = math_ops.floor(result)
+    result = math_ops.cast(result, dtype)
+  if retstep:
+    return (result, step)
+  else:
+    return result
+
+
+@tf_export.tf_export('experimental.numpy.logspace', v1=[])
+@np_utils.np_doc('logspace')
+def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
+  # In numpy 2.x, the result type of np.logspace is based off of `start` and
+  # `end`. We mimic the behavior.
+  if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
+    dtype = np_utils.result_type([start * 1.0, stop * 1.0])
+  else:
+    dtype = np_utils.result_type(start, stop, dtype)
+  result = linspace(
+      start, stop, num=num, endpoint=endpoint, dtype=dtype, axis=axis
+  )
+  result = math_ops.pow(math_ops.cast(base, result.dtype), result)
+  if dtype:
+    result = math_ops.cast(result, dtype)
+  return result
+
+
+@tf_export.tf_export('experimental.numpy.geomspace', v1=[])
+@np_utils.np_doc('geomspace')
+def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):  # pylint: disable=missing-docstring
+  dtype = (
+      dtypes.as_dtype(dtype)  # pylint: disable=g-long-ternary
+      if dtype
+      else np_utils.result_type(
+          start, stop, float(num), np_array_ops.zeros((), dtype)
+      )
+  )
+  computation_dtype = np.promote_types(dtype.as_numpy_dtype, np.float32)
+  start = np_array_ops.asarray(start, dtype=computation_dtype)
+  stop = np_array_ops.asarray(stop, dtype=computation_dtype)
+  # follow the numpy geomspace convention for negative and complex endpoints
+  start_sign = 1 - np_array_ops.sign(np_array_ops.real(start))
+  stop_sign = 1 - np_array_ops.sign(np_array_ops.real(stop))
+  signflip = 1 - start_sign * stop_sign // 2
+  res = signflip * logspace(
+      log10(signflip * start),
+      log10(signflip * stop),
+      num,
+      endpoint=endpoint,
+      base=10.0,
+      dtype=computation_dtype,
+      axis=0,
+  )
+  if axis != 0:
+    res = np_array_ops.moveaxis(res, 0, axis)
+  return math_ops.cast(res, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.ptp', v1=[])
+@np_utils.np_doc('ptp')
+def ptp(a, axis=None, keepdims=None):
+  return np_array_ops.amax(a, axis=axis, keepdims=keepdims) - np_array_ops.amin(
+      a, axis=axis, keepdims=keepdims
+  )
+
+
+@tf_export.tf_export('experimental.numpy.concatenate', v1=[])
+@np_utils.np_doc_only('concatenate')
+def concatenate(arys, axis=0):  # pylint: disable=missing-function-docstring
+  if not isinstance(arys, (list, tuple)):
+    arys = [arys]
+  if not arys:
+    raise ValueError(
+        'Need at least one array to concatenate. Received empty '
+        f'input: arys={arys}'
+    )
+  dtype = np_utils.result_type(*arys)
+  arys = [np_array_ops.array(array, dtype=dtype) for array in arys]
+  return array_ops.concat(arys, axis)
+
+
+@tf_export.tf_export('experimental.numpy.tile', v1=[])
+@np_utils.np_doc_only('tile')
+def tile(a, reps):  # pylint: disable=missing-function-docstring
+  a = np_array_ops.array(a)
+  reps = array_ops.reshape(np_array_ops.array(reps, dtype=dtypes.int32), [-1])
+
+  a_rank = array_ops.rank(a)
+  reps_size = array_ops.size(reps)
+  reps = array_ops.pad(
+      reps, [[math_ops.maximum(a_rank - reps_size, 0), 0]], constant_values=1
+  )
+  a_shape = array_ops.pad(
+      array_ops.shape(a),
+      [[math_ops.maximum(reps_size - a_rank, 0), 0]],
+      constant_values=1,
+  )
+  a = array_ops.reshape(a, a_shape)
+
+  return array_ops.tile(a, reps)
+
+
+@tf_export.tf_export('experimental.numpy.count_nonzero', v1=[])
+@np_utils.np_doc('count_nonzero')
+def count_nonzero(a, axis=None):
+  return math_ops.count_nonzero(np_array_ops.array(a), axis)
+
+
+@tf_export.tf_export('experimental.numpy.argsort', v1=[])
+@np_utils.np_doc('argsort')
+def argsort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missing-docstring
+  # TODO(nareshmodi): make string tensors also work.
+  if kind not in ('quicksort', 'stable'):
+    raise ValueError(
+        'Invalid value for argument `kind`. '
+        'Only kind="quicksort" and kind="stable" are supported. '
+        f'Received: kind={kind}'
+    )
+  if order is not None:
+    raise ValueError('The `order` argument is not supported. Pass order=None')
+  stable = kind == 'stable'
+
+  a = np_array_ops.array(a)
+
+  def _argsort(a, axis, stable):
+    if axis is None:
+      a = array_ops.reshape(a, [-1])
+      axis = 0
+
+    return sort_ops.argsort(a, axis, stable=stable)
+
+  tf_ans = np_utils.cond(
+      math_ops.equal(array_ops.rank(a), 0),
+      lambda: constant_op.constant([0]),
+      lambda: _argsort(a, axis, stable),
+  )
+
+  if ops.is_auto_dtype_conversion_enabled():
+    return np_array_ops.array(tf_ans, dtype=int)
+  else:
+    return np_array_ops.array(tf_ans, dtype=np.intp)
+
+
+@tf_export.tf_export('experimental.numpy.sort', v1=[])
+@np_utils.np_doc('sort')
+def sort(a, axis=-1, kind='quicksort', order=None):  # pylint: disable=missing-docstring
+  if kind != 'quicksort':
+    raise ValueError(
+        'Invalid value for argument `kind`. '
+        'Only kind="quicksort" is supported. '
+        f'Received: kind={kind}'
+    )
+  if order is not None:
+    raise ValueError('The `order` argument is not supported. Pass order=None')
+
+  a = np_array_ops.array(a)
+
+  if axis is None:
+    return sort_ops.sort(array_ops.reshape(a, [-1]), 0)
+  else:
+    return sort_ops.sort(a, axis)
+
+
+def _argminmax(fn, a, axis=None):
+  a = np_array_ops.array(a)
+  if axis is None:
+    # When axis is None numpy flattens the array.
+    a_t = array_ops.reshape(a, [-1])
+  else:
+    a_t = np_array_ops.atleast_1d(a)
+  return fn(input=a_t, axis=axis)
+
+
+@tf_export.tf_export('experimental.numpy.argmax', v1=[])
+@np_utils.np_doc('argmax')
+def argmax(a, axis=None):
+  return _argminmax(math_ops.argmax, a, axis)
+
+
+@tf_export.tf_export('experimental.numpy.argmin', v1=[])
+@np_utils.np_doc('argmin')
+def argmin(a, axis=None):
+  return _argminmax(math_ops.argmin, a, axis)
+
+
+@tf_export.tf_export('experimental.numpy.append', v1=[])
+@np_utils.np_doc('append')
+def append(arr, values, axis=None):
+  if axis is None:
+    return concatenate([np_array_ops.ravel(arr), np_array_ops.ravel(values)], 0)
+  else:
+    return concatenate([arr, values], axis=axis)
+
+
+@tf_export.tf_export('experimental.numpy.average', v1=[])
+@np_utils.np_doc('average')
+def average(a, axis=None, weights=None, returned=False):  # pylint: disable=missing-docstring
+  if axis is not None and not isinstance(axis, int):
+    # TODO(wangpeng): Support tuple of ints as `axis`
+    raise ValueError(
+        'Argument `axis` must be an integer. '
+        f'Received axis={axis} (of type {type(axis)})'
+    )
+  a = np_array_ops.array(a)
+  default_float_type = np_utils.result_type(float)
+  if weights is None:  # Treat all weights as 1
+    if not np.issubdtype(a.dtype.as_numpy_dtype, np.inexact):
+      a = a.astype(np_utils.result_type(a.dtype, default_float_type))
+    avg = math_ops.reduce_mean(a, axis=axis)
+    if returned:
+      if axis is None:
+        weights_sum = array_ops.size(a)
+      else:
+        weights_sum = array_ops.shape(a)[axis]
+      weights_sum = math_ops.cast(weights_sum, a.dtype)
+  else:
+    if np.issubdtype(a.dtype.as_numpy_dtype, np.inexact):
+      out_dtype = np_utils.result_type(a.dtype, weights)
+    else:
+      out_dtype = np_utils.result_type(a.dtype, weights, default_float_type)
+    a = np_array_ops.array(a, out_dtype)
+    weights = np_array_ops.array(weights, out_dtype)
+
+    def rank_equal_case():
+      control_flow_assert.Assert(
+          math_ops.reduce_all(array_ops.shape(a) == array_ops.shape(weights)),
+          [array_ops.shape(a), array_ops.shape(weights)],
+      )
+      weights_sum = math_ops.reduce_sum(weights, axis=axis)
+      avg = math_ops.reduce_sum(a * weights, axis=axis) / weights_sum
+      return avg, weights_sum
+
+    if axis is None:
+      avg, weights_sum = rank_equal_case()
+    else:
+
+      def rank_not_equal_case():
+        control_flow_assert.Assert(
+            array_ops.rank(weights) == 1, [array_ops.rank(weights)]
+        )
+        weights_sum = math_ops.reduce_sum(weights)
+        axes = ops.convert_to_tensor([[axis], [0]])
+        avg = math_ops.tensordot(a, weights, axes) / weights_sum
+        return avg, weights_sum
+
+      # We condition on rank rather than shape equality, because if we do the
+      # latter, when the shapes are partially unknown but the ranks are known
+      # and different, np_utils.cond will run shape checking on the true branch,
+      # which will raise a shape-checking error.
+      avg, weights_sum = np_utils.cond(
+          math_ops.equal(array_ops.rank(a), array_ops.rank(weights)),
+          rank_equal_case,
+          rank_not_equal_case,
+      )
+
+  avg = np_array_ops.array(avg)
+  if returned:
+    weights_sum = np_array_ops.broadcast_to(weights_sum, array_ops.shape(avg))
+    return avg, weights_sum
+  return avg
+
+
+@tf_export.tf_export('experimental.numpy.trace', v1=[])
+@np_utils.np_doc('trace')
+def trace(a, offset=0, axis1=0, axis2=1, dtype=None):  # pylint: disable=missing-docstring
+  if dtype:
+    dtype = np_utils.result_type(dtype)
+  a = np_array_ops.asarray(a, dtype)
+
+  if offset == 0:
+    a_shape = a.shape
+    if a_shape.rank is not None:
+      rank = len(a_shape)
+      if (axis1 == -2 or axis1 == rank - 2) and (
+          axis2 == -1 or axis2 == rank - 1
+      ):
+        return math_ops.trace(a)
+
+  a = np_array_ops.diagonal(a, offset, axis1, axis2)
+  return np_array_ops.sum(a, -1, dtype)
+
+
+@tf_export.tf_export('experimental.numpy.meshgrid', v1=[])
+@np_utils.np_doc('meshgrid')
+def meshgrid(*xi, **kwargs):
+  """This currently requires copy=True and sparse=False."""
+  sparse = kwargs.get('sparse', False)
+  if sparse:
+    raise ValueError(
+        'Function `meshgrid` does not support returning sparse arrays yet. '
+        f'Received: sparse={sparse}'
+    )
+
+  copy = kwargs.get('copy', True)
+  if not copy:
+    raise ValueError(
+        f'Function `meshgrid` only supports copy=True. Received: copy={copy}'
+    )
+
+  indexing = kwargs.get('indexing', 'xy')
+
+  xi = [np_array_ops.asarray(arg) for arg in xi]
+  kwargs = {'indexing': indexing}
+
+  outputs = array_ops.meshgrid(*xi, **kwargs)
+
+  return outputs
+
+
+# Uses np_doc_only here because np.einsum (in 1.16) doesn't have argument
+# `subscripts`, even though the doc says it has.
+@tf_export.tf_export('experimental.numpy.einsum', v1=[])
+@np_utils.np_doc_only('einsum')
+def einsum(subscripts, *operands, **kwargs):  # pylint: disable=missing-docstring
+  casting = kwargs.get('casting', 'safe')
+  optimize = kwargs.get('optimize', False)
+  if casting == 'safe':
+    operands = np_array_ops._promote_dtype(*operands)  # pylint: disable=protected-access
+  elif casting == 'no':
+    operands = [np_array_ops.asarray(x) for x in operands]
+  else:
+    raise ValueError(
+        'Invalid value for argument `casting`. '
+        f'Expected casting="safe" or casting="no". Received: casting={casting}'
+    )
+  if not optimize:
+    # TF doesn't have a "no optimization" option.
+    # TODO(wangpeng): Print a warning that np and tf use different
+    #   optimizations.
+    tf_optimize = 'greedy'
+  elif optimize == True:  # pylint: disable=singleton-comparison,g-explicit-bool-comparison
+    tf_optimize = 'greedy'
+  elif optimize == 'greedy':
+    tf_optimize = 'greedy'
+  elif optimize == 'optimal':
+    tf_optimize = 'optimal'
+  else:
+    raise ValueError(
+        'Invalid value for argument `optimize`. '
+        'Expected one of {True, "greedy", "optimal"}. '
+        f'Received: optimize={optimize}'
+    )
+
+  res = special_math_ops.einsum(subscripts, *operands, optimize=tf_optimize)
+  return res
+
+
+def _tensor_t(self):
+  """Returns a Tensor which is the transpose of this Tensor."""
+  return self.transpose()
+
+
+def _tensor_ndim(self):
+  """Returns the rank of the Tensor."""
+  return self.shape.ndims
+
+
+def _tensor_pos(self):
+  """Returns self, for unary operator `+`."""
+  return self
+
+
+def _tensor_size(self):
+  """Returns the number of elements in this Tensor, if fully known."""
+  if not self.shape.is_fully_defined():
+    return None
+  return np.prod(self.shape.as_list())
+
+
+def _tensor_tolist(self):
+  if ops.is_symbolic_tensor(self):
+    raise ValueError('Symbolic Tensors do not support the tolist API.')
+
+  return self._numpy().tolist()  # pylint: disable=protected-access
+
+
+def _enable_numpy_methods(tensor_class):
+  """A helper method for adding additional NumPy methods."""
+  t = property(_tensor_t)
+  setattr(tensor_class, 'T', t)
+
+  ndim = property(_tensor_ndim)
+  setattr(tensor_class, 'ndim', ndim)
+
+  size = property(_tensor_size)
+  setattr(tensor_class, 'size', size)
+
+  setattr(tensor_class, '__pos__', _tensor_pos)
+  setattr(tensor_class, 'tolist', _tensor_tolist)
+
+  # TODO(b/178540516): Make a custom `setattr` that changes the method's
+  #   docstring to the TF one.
+  setattr(tensor_class, 'transpose', np_array_ops.transpose)
+  setattr(tensor_class, 'flatten', np_array_ops.flatten)
+  setattr(tensor_class, 'reshape', np_array_ops._reshape_method_wrapper)  # pylint: disable=protected-access
+  setattr(tensor_class, 'ravel', np_array_ops.ravel)
+  setattr(tensor_class, 'clip', clip)
+  setattr(tensor_class, 'astype', math_ops.cast)
+  setattr(tensor_class, '__round__', np_array_ops.around)
+  setattr(tensor_class, 'max', np_array_ops.amax)
+  setattr(tensor_class, 'mean', np_array_ops.mean)
+  setattr(tensor_class, 'min', np_array_ops.amin)
+
+  # TODO(wangpeng): Remove `data` when all uses of it are removed
+  data = property(lambda self: self)
+  setattr(tensor_class, 'data', data)
+
+
+def enable_numpy_methods_on_tensor():
+  """Adds additional NumPy methods on tf.Tensor class."""
+  _enable_numpy_methods(tensor.Tensor)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_random.py

@@ -0,0 +1,137 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Random functions."""
+
+# pylint: disable=g-direct-tensorflow-import
+
+import numpy as onp
+
+from tensorflow.python.framework import random_seed
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.numpy_ops import np_array_ops
+from tensorflow.python.ops.numpy_ops import np_dtypes
+from tensorflow.python.ops.numpy_ops import np_utils
+from tensorflow.python.util import tf_export
+
+# TODO(agarwal): deprecate this.
+DEFAULT_RANDN_DTYPE = onp.float32
+
+
+@tf_export.tf_export('experimental.numpy.random.seed', v1=[])
+@np_utils.np_doc('random.seed')
+def seed(s):
+  """Sets the seed for the random number generator.
+
+  Uses `tf.set_random_seed`.
+
+  Args:
+    s: an integer.
+  """
+  try:
+    s = int(s)
+  except TypeError:
+    # TODO(wangpeng): support this?
+    raise ValueError(
+        f'Argument `s` got an invalid value {s}. Only integers are supported.'
+    )
+  random_seed.set_seed(s)
+
+
+@tf_export.tf_export('experimental.numpy.random.randn', v1=[])
+@np_utils.np_doc('random.randn')
+def randn(*args):
+  """Returns samples from a normal distribution.
+
+  Uses `tf.random_normal`.
+
+  Args:
+    *args: The shape of the output array.
+
+  Returns:
+    An ndarray with shape `args` and dtype `float64`.
+  """
+  return standard_normal(size=args)
+
+
+@tf_export.tf_export('experimental.numpy.random.standard_normal', v1=[])
+@np_utils.np_doc('random.standard_normal')
+def standard_normal(size=None):
+  # TODO(wangpeng): Use new stateful RNG
+  if size is None:
+    size = ()
+  elif np_utils.isscalar(size):
+    size = (size,)
+  dtype = np_utils.result_type(float)
+  return random_ops.random_normal(size, dtype=dtype)
+
+
+@tf_export.tf_export('experimental.numpy.random.uniform', v1=[])
+@np_utils.np_doc('random.uniform')
+def uniform(low=0.0, high=1.0, size=None):
+  dtype = np_utils.result_type(float)
+  low = np_array_ops.asarray(low, dtype=dtype)
+  high = np_array_ops.asarray(high, dtype=dtype)
+  if size is None:
+    size = array_ops.broadcast_dynamic_shape(low.shape, high.shape)
+  return random_ops.random_uniform(
+      shape=size, minval=low, maxval=high, dtype=dtype
+  )
+
+
+@tf_export.tf_export('experimental.numpy.random.poisson', v1=[])
+@np_utils.np_doc('random.poisson')
+def poisson(lam=1.0, size=None):
+  if size is None:
+    size = ()
+  elif np_utils.isscalar(size):
+    size = (size,)
+  return random_ops.random_poisson(shape=size, lam=lam, dtype=np_dtypes.int_)
+
+
+@tf_export.tf_export('experimental.numpy.random.random', v1=[])
+@np_utils.np_doc('random.random')
+def random(size=None):
+  return uniform(0.0, 1.0, size)
+
+
+@tf_export.tf_export('experimental.numpy.random.rand', v1=[])
+@np_utils.np_doc('random.rand')
+def rand(*size):
+  return uniform(0.0, 1.0, size)
+
+
+@tf_export.tf_export('experimental.numpy.random.randint', v1=[])
+@np_utils.np_doc('random.randint')
+def randint(low, high=None, size=None, dtype=onp.int64):  # pylint: disable=missing-function-docstring
+  low = int(low)
+  if high is None:
+    high = low
+    low = 0
+  if size is None:
+    size = ()
+  elif isinstance(size, int):
+    size = (size,)
+  dtype_orig = dtype
+  dtype = np_utils.result_type(dtype)
+  accepted_dtypes = (onp.int32, onp.int64)
+  if dtype not in accepted_dtypes:
+    raise ValueError(
+        f'Argument `dtype` got an invalid value {dtype_orig}. Only those '
+        f'convertible to {accepted_dtypes} are supported.'
+    )
+  return random_ops.random_uniform(
+      shape=size, minval=low, maxval=high, dtype=dtype
+  )

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/numpy_ops/np_utils.py

@@ -0,0 +1,715 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions for internal use."""
+# pylint: disable=g-direct-tensorflow-import
+
+import inspect
+import numbers
+import os
+import re
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import flexible_dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond as tf_cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.numpy_ops import np_arrays
+from tensorflow.python.ops.numpy_ops import np_dtypes
+from tensorflow.python.types import core
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_export
+
+
+def _canonicalize_axis(axis, rank):
+  return _canonicalize_axes([axis], rank)[0]
+
+
+def _canonicalize_axes(axes, rank):
+  rank = _maybe_static(rank)
+
+  if isinstance(rank, core.Tensor):
+    canonicalizer = lambda axis: cond(  # pylint: disable=g-long-lambda
+        axis < 0, lambda: axis + rank, lambda: axis
+    )
+  else:
+    canonicalizer = lambda axis: axis + rank if axis < 0 else axis
+
+  return [canonicalizer(axis) for axis in axes]
+
+
+def _supports_signature():
+  return hasattr(inspect, 'signature')
+
+
+def _to_tf_type(dtype):
+  """Converts a native python or numpy type to TF DType.
+
+  Args:
+    dtype: Could be a python type, a numpy type or a TF DType.
+
+  Returns:
+    A tensorflow `DType`.
+  """
+  return dtypes.as_dtype(dtype)
+
+
+def _to_numpy_type(dtype):
+  """Converts a native python or TF DType to numpy type.
+
+  Args:
+    dtype: Could be a python type, a numpy type or a TF DType.
+
+  Returns:
+    A NumPy `dtype`.
+  """
+  if isinstance(dtype, dtypes.DType):
+    return dtype.as_numpy_dtype
+  return np.dtype(dtype)
+
+
+def isscalar(val):
+  """Returns whether `val` is a scalar value or scalar Tensor."""
+  if isinstance(val, np_arrays.ndarray):
+    val = val.data
+  if isinstance(val, core.Tensor):
+    ndims = val.shape.ndims
+    if ndims is not None:
+      return ndims == 0
+    else:
+      return math_ops.equal(array_ops.rank(val), 0)
+  else:
+    return np.isscalar(val)
+
+
+def _has_docstring(f):
+  return (
+      f and hasattr(f, '__doc__') and isinstance(f.__doc__, str) and f.__doc__
+  )
+
+
+def _add_blank_line(s):
+  if s.endswith('\n'):
+    return s + '\n'
+  else:
+    return s + '\n\n'
+
+
+def _np_signature(f):
+  """An enhanced inspect.signature that can handle numpy.ufunc."""
+  # TODO(wangpeng): consider migrating away from inspect.signature.
+  # inspect.signature is supported in Python 3.3.
+  if not hasattr(inspect, 'signature'):
+    return None
+  if f is None:
+    return None
+  if not isinstance(f, np.ufunc):
+    try:
+      return inspect.signature(f)
+    except ValueError:
+      return None
+
+  def names_from_num(prefix, n):
+    if n <= 0:
+      return []
+    elif n == 1:
+      return [prefix]
+    else:
+      return [prefix + str(i + 1) for i in range(n)]
+
+  input_names = names_from_num('x', f.nin)
+  output_names = names_from_num('out', f.nout)
+  keyword_only_params = [
+      ('where', True),
+      ('casting', 'same_kind'),
+      ('order', 'K'),
+      ('dtype', None),
+      ('subok', True),
+      ('signature', None),
+      ('extobj', None),
+  ]
+  params = []
+  params += [
+      inspect.Parameter(name, inspect.Parameter.POSITIONAL_ONLY)
+      for name in input_names
+  ]
+  if f.nout > 1:
+    params += [
+        inspect.Parameter(name, inspect.Parameter.POSITIONAL_ONLY, default=None)
+        for name in output_names
+    ]
+  params += [
+      inspect.Parameter(
+          'out',
+          inspect.Parameter.POSITIONAL_OR_KEYWORD,
+          default=None if f.nout == 1 else (None,) * f.nout,
+      )
+  ]
+  params += [
+      inspect.Parameter(name, inspect.Parameter.KEYWORD_ONLY, default=default)
+      for name, default in keyword_only_params
+  ]
+  return inspect.Signature(params)
+
+
+# Python 2 doesn't allow keyword-only argument. Python prior to 3.8 doesn't
+# allow positional-only argument. So we conflate positional-only, keyword-only
+# and positional-or-keyword arguments here.
+def _is_compatible_param_kind(a, b):
+  def relax(k):
+    if k in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.KEYWORD_ONLY):
+      return inspect.Parameter.POSITIONAL_OR_KEYWORD
+    return k
+
+  return relax(a) == relax(b)
+
+
+def _prepare_np_fun_name_and_fun(np_fun_name, np_fun):
+  """Mutually propagates information between `np_fun_name` and `np_fun`.
+
+  If one is None and the other is not, we'll try to make the former not None in
+  a best effort.
+
+  Args:
+    np_fun_name: name for the np_fun symbol. At least one of np_fun or
+      np_fun_name shoud be set.
+    np_fun: the numpy function whose docstring will be used.
+
+  Returns:
+    Processed `np_fun_name` and `np_fun`.
+  """
+  if np_fun_name is not None:
+    assert isinstance(np_fun_name, str)
+  if np_fun is not None:
+    assert not isinstance(np_fun, str)
+  if np_fun is None:
+    assert np_fun_name is not None
+    try:
+      np_fun = getattr(np, str(np_fun_name))
+    except AttributeError:
+      np_fun = None
+  if np_fun_name is None:
+    assert np_fun is not None
+    np_fun_name = np_fun.__name__
+  return np_fun_name, np_fun
+
+
+def _np_doc_helper(
+    f, np_f, np_fun_name=None, unsupported_params=None, link=None
+):
+  """Helper to get docs."""
+  assert np_f or np_fun_name
+  if not np_fun_name:
+    np_fun_name = np_f.__name__
+  doc = "TensorFlow variant of NumPy's `%s`.\n\n" % np_fun_name
+  if unsupported_params:
+    doc += (
+        'Unsupported arguments: '
+        + ', '.join('`' + name + '`' for name in unsupported_params)
+        + '.\n\n'
+    )
+  if _has_docstring(f):
+    doc += f.__doc__
+    doc = _add_blank_line(doc)
+  # TODO(wangpeng): Re-enable the following and choose inlined vs. link to numpy
+  #   doc according to some global switch.
+  doc = _add_np_doc(doc, np_fun_name, np_f, link=link)
+  return doc
+
+
+_np_doc_form = os.getenv('TF_NP_DOC_FORM', 'stable')
+
+
+def get_np_doc_form():
+  """Gets the form of the original numpy docstrings.
+
+  Returns:
+    See `set_np_doc_form` for the list of valid values.
+  """
+  return _np_doc_form
+
+
+def set_np_doc_form(value):
+  r"""Selects the form of the original numpy docstrings.
+
+  This function sets a global variable that controls how a tf-numpy symbol's
+  docstring should refer to the original numpy docstring. If `value` is
+  `'inlined'`, the numpy docstring will be verbatim copied into the tf-numpy
+  docstring. Otherwise, a link to the original numpy docstring will be
+  added. Which numpy version the link points to depends on `value`:
+  * `'stable'`: the current stable version;
+  * `'dev'`: the current development version;
+  * pattern `\d+(\.\d+(\.\d+)?)?`: `value` will be treated as a version number,
+    e.g. '1.16'.
+
+  Args:
+    value: the value to set the global variable to.
+  """
+  global _np_doc_form
+  _np_doc_form = value
+
+
+class Link:
+
+  def __init__(self, v):
+    self.value = v
+
+
+class AliasOf:
+
+  def __init__(self, v):
+    self.value = v
+
+
+class NoLink:
+  pass
+
+
+def generate_link(flag, np_fun_name):
+  """Generates link from numpy function name.
+
+  Args:
+    flag: the flag to control link form. See `set_np_doc_form`.
+    np_fun_name: the numpy function name.
+
+  Returns:
+    A string.
+  """
+  # Only adds link in this case
+  if flag == 'dev':
+    template = 'https://numpy.org/devdocs/reference/generated/numpy.%s.html'
+  elif flag == 'stable':
+    template = 'https://numpy.org/doc/stable/reference/generated/numpy.%s.html'
+  elif re.match(r'\d+(\.\d+(\.\d+)?)?$', flag):
+    # `flag` is the version number
+    template = f'https://numpy.org/doc/{flag}/reference/generated/numpy.%s.html'
+  else:
+    return None
+  return template % np_fun_name
+
+
+_is_check_link = os.getenv('TF_NP_CHECK_LINK', 'False') in ('True', 'true', '1')
+
+
+def is_check_link():
+  return _is_check_link
+
+
+def set_check_link(value):
+  global _is_check_link
+  _is_check_link = value
+
+
+def _add_np_doc(doc, np_fun_name, np_f, link):
+  """Appends the numpy docstring to `doc`, according to `set_np_doc_form`.
+
+  See `set_np_doc_form` for how it controls the form of the numpy docstring.
+
+  Args:
+    doc: the docstring to be appended to.
+    np_fun_name: the name of the numpy function.
+    np_f: (optional) the numpy function.
+    link: (optional) which link to use. See `np_doc` for details.
+
+  Returns:
+    `doc` with numpy docstring appended.
+  """
+  flag = get_np_doc_form()
+  if flag == 'inlined':
+    if _has_docstring(np_f):
+      doc += 'Documentation for `numpy.%s`:\n\n' % np_fun_name
+      # TODO(wangpeng): It looks like code snippets in numpy doc don't work
+      # correctly with doctest. Fix that and remove the reformatting of the np_f
+      # comment.
+      doc += np_f.__doc__.replace('>>>', '>')
+  elif isinstance(flag, str):
+    if link is None:
+      url = generate_link(flag, np_fun_name)
+    elif isinstance(link, AliasOf):
+      url = generate_link(flag, link.value)
+    elif isinstance(link, Link):
+      url = link.value
+    else:
+      url = None
+    if url is not None:
+      if is_check_link():
+        # Imports locally because some builds may not have `requests`
+        import requests  # pylint: disable=g-import-not-at-top
+
+        r = requests.head(url)
+        if r.status_code != 200:
+          raise ValueError(
+              f'Check link failed at [{url}] with status code {r.status_code}. '
+              f'Argument `np_fun_name` is {np_fun_name}.'
+          )
+      doc += 'See the NumPy documentation for [`numpy.%s`](%s).' % (
+          np_fun_name,
+          url,
+      )
+  return doc
+
+
+_is_sig_mismatch_an_error = os.getenv(
+    'TF_NP_SIG_MISMATCH_IS_ERROR', 'False'
+) in ('True', 'true', '1')
+
+
+def is_sig_mismatch_an_error():
+  return _is_sig_mismatch_an_error
+
+
+def set_is_sig_mismatch_an_error(value):
+  global _is_sig_mismatch_an_error
+  _is_sig_mismatch_an_error = value
+
+
+def np_doc(np_fun_name, np_fun=None, unsupported_params=None, link=None):
+  """Attachs numpy docstring to a function.
+
+  Args:
+    np_fun_name: name for the np_fun symbol. At least one of np_fun or
+      np_fun_name shoud be set.
+    np_fun: (optional) the numpy function whose docstring will be used.
+    unsupported_params: (optional) the list of parameters not supported by
+      tf.numpy.
+    link: (optional) which link to use. If `None`, a default link generated from
+      `np_fun_name` will be used. If an instance of `AliasOf`, `link.value` will
+      be used in place of `np_fun_name` for the link generation. If an instance
+      of `Link`, `link.value` will be used as the whole link. If an instance of
+      `NoLink`, no link will be added.
+
+  Returns:
+    A function decorator that attaches the docstring from `np_fun` to the
+    decorated function.
+  """
+  np_fun_name_orig, np_fun_orig = np_fun_name, np_fun
+  np_fun_name, np_fun = _prepare_np_fun_name_and_fun(np_fun_name, np_fun)
+  np_sig = _np_signature(np_fun)
+  if unsupported_params is None:
+    unsupported_params = []
+
+  def decorator(f):
+    """The decorator."""
+    if hasattr(inspect, 'signature') and np_sig is not None:
+      try:
+        sig = inspect.signature(f)
+      except ValueError:
+        sig = None
+      if sig is not None:
+        for name, param in sig.parameters.items():
+          np_param = np_sig.parameters.get(name)
+          if np_param is None:
+            if is_sig_mismatch_an_error():
+              raise TypeError(
+                  f"Cannot find parameter {name} in the numpy function's "
+                  'signature (which has these parameters: '
+                  f'{list(np_sig.parameters.keys())}). Argument `np_fun_name` '
+                  f'is {np_fun_name_orig}. Argument `np_fun` is {np_fun_orig}.'
+              )
+            else:
+              continue
+          if is_sig_mismatch_an_error() and not _is_compatible_param_kind(
+              param.kind, np_param.kind
+          ):
+            raise TypeError(
+                f'Parameter {name} is of kind {param.kind} while in numpy it '
+                f'is of kind {np_param.kind}. Argument `np_fun_name` is '
+                f'{np_fun_name_orig}. Argument `np_fun` is {np_fun_orig}.'
+            )
+          has_default = param.default != inspect.Parameter.empty
+          np_has_default = np_param.default != inspect.Parameter.empty
+          if is_sig_mismatch_an_error() and has_default != np_has_default:
+            raise TypeError(
+                'Parameter {} should{} have a default value. Argument '
+                '`np_fun_name` is {}. Argument `np_fun` is {}.'.format(
+                    name,
+                    '' if np_has_default else ' not',
+                    np_fun_name_orig,
+                    np_fun_orig,
+                )
+            )
+        for name in np_sig.parameters:
+          if name not in sig.parameters:
+            unsupported_params.append(name)
+    f.__doc__ = _np_doc_helper(
+        f,
+        np_fun,
+        np_fun_name=np_fun_name,
+        unsupported_params=unsupported_params,
+        link=link,
+    )
+    return f
+
+  return decorator
+
+
+def np_doc_only(np_fun_name, np_fun=None):
+  """Attachs numpy docstring to a function.
+
+  This differs from np_doc in that it doesn't check for a match in signature.
+
+  Args:
+    np_fun_name: name for the np_fun symbol. At least one of np_fun or
+      np_fun_name shoud be set.
+    np_fun: (optional) the numpy function whose docstring will be used.
+
+  Returns:
+    A function decorator that attaches the docstring from `np_fun` to the
+    decorated function.
+  """
+  np_fun_name, np_fun = _prepare_np_fun_name_and_fun(np_fun_name, np_fun)
+
+  def decorator(f):
+    f.__doc__ = _np_doc_helper(f, np_fun, np_fun_name=np_fun_name)
+    return f
+
+  return decorator
+
+
+# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-docstring-missing-newline,g-doc-return-or-yield,g-doc-args
+@tf_export.tf_export('experimental.numpy.finfo', v1=[])
+@np_doc('finfo')
+def finfo(dtype):
+  """Note that currently it just forwards to the numpy namesake, while
+
+  tensorflow and numpy dtypes may have different properties.
+  """
+  return np.finfo(_to_numpy_type(dtype))
+
+
+# pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-docstring-missing-newline,g-doc-return-or-yield,g-doc-args
+
+
+def _maybe_get_dtype(x):
+  """Returns a numpy type if available from x. Skips if x is numpy.ndarray."""
+  # Don't put np.ndarray in this list, because np.result_type looks at the
+  # value (not just dtype) of np.ndarray to decide the result type.
+  if isinstance(x, numbers.Real):
+    return x
+  if isinstance(x, indexed_slices.IndexedSlices) or tensor_util.is_tf_type(x):
+    return _to_numpy_type(x.dtype)
+  if isinstance(x, dtypes.DType):
+    return x.as_numpy_dtype
+  if isinstance(x, (list, tuple)):
+    raise ValueError(
+        'Cannot find dtype for type inference from argument `x` of a sequence '
+        f'type {type(x)}. For sequences, please call this function on each '
+        'element individually.'
+    )
+  return x
+
+
+@tf_export.tf_export('experimental.numpy.result_type', v1=[])
+# Can't use np_doc because np.result_type is a builtin function.
+@np_doc_only('result_type')
+def result_type(*arrays_and_dtypes):  # pylint: disable=missing-function-docstring
+  if ops.is_auto_dtype_conversion_enabled():
+    # Use auto dtype conversion semantics for type inference.
+    dtype, _ = flexible_dtypes.result_type(*arrays_and_dtypes)
+    return dtype
+  arrays_and_dtypes = [
+      _maybe_get_dtype(x) for x in nest.flatten(arrays_and_dtypes)
+  ]
+  if not arrays_and_dtypes:
+    # If arrays_and_dtypes is an empty list, let numpy decide what the dtype is.
+    arrays_and_dtypes = [np.asarray([])]
+  return np_dtypes._result_type(*arrays_and_dtypes)  # pylint: disable=protected-access
+
+
+def result_type_unary(a, dtype):  # pylint: disable=missing-function-docstring
+  """Find the result type from a single input and a dtype."""
+  if dtype:
+    # We need to let np_utils.result_type decide the dtype, not tf.zeros_like
+    return result_type(dtype)
+
+  # np_utils.result_type treats string inputs as dtype strings, not as strings.
+  # but for unary we want to treat it as a string input.
+  if isinstance(a, str):
+    return np.str_
+  elif isinstance(a, bytes):
+    return np.bytes_
+
+  # TF and numpy has different interpretations of Python types such as
+  # `float`, so we let `np_utils.result_type` decide.
+  return result_type(a)
+
+
+def _result_type_binary(t1, t2):  # pylint: disable=missing-function-docstring
+  """A specialization of result_type for 2 arguments for performance reasons."""
+  try:
+    return np_dtypes._result_type(  # pylint: disable=protected-access
+        _maybe_get_dtype(t1),
+        _maybe_get_dtype(t2),
+    )
+  except ValueError:
+    return result_type(t1, t2)
+
+
+@tf_export.tf_export('experimental.numpy.promote_types', v1=[])
+@np_doc('promote_types')
+def promote_types(type1, type2):  # pylint: disable=missing-function-docstring
+  type1 = _to_numpy_type(type1)
+  type2 = _to_numpy_type(type2)
+  return np_dtypes.canonicalize_dtype(np.promote_types(type1, type2))
+
+
+def tf_broadcast(*args):
+  """Broadcast tensors.
+
+  Args:
+    *args: a list of tensors whose shapes are broadcastable against each other.
+
+  Returns:
+    Tensors broadcasted to the common shape.
+  """
+  if len(args) <= 1:
+    return args
+  sh = array_ops.shape(args[0])
+  for arg in args[1:]:
+    sh = array_ops.broadcast_dynamic_shape(sh, array_ops.shape(arg))
+  return [array_ops.broadcast_to(arg, sh) for arg in args]
+
+
+# TODO(wangpeng): Move the following functions to a separate file and check for
+#   float dtypes in each of them.
+
+
+def get_static_value(x):
+  """A version of tf.get_static_value that returns None on float dtypes.
+
+  It returns None on float dtypes in order to avoid breaking gradients.
+
+  Args:
+    x: a tensor.
+
+  Returns:
+    Same as `tf.get_static_value`, except that it returns None when `x` has a
+    float dtype.
+  """
+  if isinstance(x, core.Tensor) and (x.dtype.is_floating or x.dtype.is_complex):
+    return None
+  return tensor_util.constant_value(x)
+
+
+def _maybe_static(x):
+  value = get_static_value(x)
+  if value is None:
+    return x
+  else:
+    return value
+
+
+# All the following functions exist becaues get_static_value can't handle
+# their TF counterparts.
+
+
+def cond(pred, true_fn, false_fn):
+  """A version of tf.cond that tries to evaluate the condition."""
+  v = get_static_value(pred)
+  if v is None:
+    return tf_cond.cond(pred, true_fn, false_fn)
+  if v:
+    return true_fn()
+  else:
+    return false_fn()
+
+
+def add(a, b):
+  """A version of tf.add that eagerly evaluates if possible."""
+  return _maybe_static(a) + _maybe_static(b)
+
+
+def subtract(a, b):
+  """A version of tf.subtract that eagerly evaluates if possible."""
+  return _maybe_static(a) - _maybe_static(b)
+
+
+def greater(a, b):
+  """A version of tf.greater that eagerly evaluates if possible."""
+  return _maybe_static(a) > _maybe_static(b)
+
+
+def greater_equal(a, b):
+  """A version of tf.greater_equal that eagerly evaluates if possible."""
+  return _maybe_static(a) >= _maybe_static(b)
+
+
+def less_equal(a, b):
+  """A version of tf.less_equal that eagerly evaluates if possible."""
+  return _maybe_static(a) <= _maybe_static(b)
+
+
+def logical_and(a, b):
+  """A version of tf.logical_and that eagerly evaluates if possible."""
+  a_value = get_static_value(a)
+  if a_value is not None:
+    if np.isscalar(a_value):
+      if a_value:
+        return _maybe_static(b)
+      else:
+        return a_value
+    else:
+      return a_value & _maybe_static(b)
+  else:
+    return a & _maybe_static(b)
+
+
+def logical_or(a, b):
+  """A version of tf.logical_or that eagerly evaluates if possible."""
+  a_value = get_static_value(a)
+  if a_value is not None:
+    if np.isscalar(a_value):
+      if a_value:
+        return a_value
+      else:
+        return _maybe_static(b)
+    else:
+      return a_value | _maybe_static(b)
+  else:
+    return a | _maybe_static(b)
+
+
+def getitem(a, slice_spec):
+  """A version of __getitem__ that eagerly evaluates if possible."""
+  return _maybe_static(a)[slice_spec]
+
+
+def reduce_all(input_tensor, axis=None, keepdims=False):
+  """A version of tf.reduce_all that eagerly evaluates if possible."""
+  v = get_static_value(input_tensor)
+  if v is None:
+    return math_ops.reduce_all(input_tensor, axis=axis, keepdims=keepdims)
+  else:
+    return v.all(axis=axis, keepdims=keepdims)
+
+
+def reduce_any(input_tensor, axis=None, keepdims=False):
+  """A version of tf.reduce_any that eagerly evaluates if possible."""
+  v = get_static_value(input_tensor)
+  if v is None:
+    return math_ops.reduce_any(input_tensor, axis=axis, keepdims=keepdims)
+  else:
+    return v.any(axis=axis, keepdims=keepdims)
+
+
+def tf_rank(t):
+  r = t.shape.rank
+  if r is not None:
+    return r
+  return array_ops.rank(t)

SwarmUI/dlbackend/ComfyUI/venv/lib/python3.10/site-packages/tensorflow/python/ops/ragged/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Ragged Tensors.
+
+This package defines ops for manipulating ragged tensors (`tf.RaggedTensor`),
+which are tensors with non-uniform shapes.  In particular, each `RaggedTensor`
+has one or more *ragged dimensions*, which are dimensions whose slices may have
+different lengths.  For example, the inner (column) dimension of
+`rt=[[3, 1, 4, 1], [], [5, 9, 2], [6], []]` is ragged, since the column slices
+(`rt[0, :]`, ..., `rt[4, :]`) have different lengths.  For a more detailed
+description of ragged tensors, see the `tf.RaggedTensor` class documentation
+and the [Ragged Tensor Guide](/guide/ragged_tensor).
+
+API docstring: tensorflow.ragged
+"""
+from tensorflow.python.ops.ragged import ragged_tensor