wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/ops/numpy_ops/BUILD
+287
View File
@@ -0,0 +1,287 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
"//third_party/py/tensorflow_numerics:__subpackages__",
],
licenses = ["notice"],
)
py_library(
name = "numpy",
srcs = ["__init__.py"],
strict_deps = True,
deps = [
":np_array_ops",
":np_arrays",
":np_config",
":np_dtypes",
":np_math_ops",
":np_random",
":np_utils",
"//tensorflow/python/ops:array_ops",
],
)
py_library(
name = "np_utils",
srcs = ["np_utils.py"],
strict_deps = True,
tags = [
"ignore_for_dep=third_party.py.requests",
],
deps = [
":np_arrays",
":np_dtypes",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:flexible_dtypes",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/types:core",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "np_array_ops",
srcs = ["np_array_ops.py"],
strict_deps = True,
deps = [
":np_arrays",
":np_dtypes",
":np_utils",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:clip_ops",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sort_ops",
"//tensorflow/python/types:core",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "np_config",
srcs = ["np_config.py"],
strict_deps = True,
deps = [
":np_dtypes",
":np_math_ops",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:weak_tensor_ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "np_dtypes",
srcs = ["np_dtypes.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "np_random",
srcs = ["np_random.py"],
strict_deps = True,
deps = [
":np_array_ops",
":np_dtypes",
":np_utils",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "np_math_ops",
srcs = ["np_math_ops.py"],
strict_deps = True,
deps = [
":np_array_ops",
":np_arrays",
":np_dtypes",
":np_utils",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:bitwise_ops",
"//tensorflow/python/ops:clip_ops",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:sort_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_library(
name = "np_arrays",
srcs = ["np_arrays.py"],
strict_deps = True,
deps = [
":np_dtypes",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_conversion",
],
)
cuda_py_strict_test(
name = "np_dtypes_test",
srcs = ["np_dtypes_test.py"],
deps = [
":np_dtypes",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "np_arrays_test",
srcs = ["np_arrays_test.py"],
deps = [
":np_arrays",
":np_math_ops",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "np_array_ops_test",
srcs = ["np_array_ops_test.py"],
tags = [
"no_windows", # TODO(b/215381493)
],
deps = [
":np_array_ops",
":np_arrays",
":np_math_ops",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "np_logic_test",
srcs = ["np_logic_test.py"],
deps = [
":np_array_ops",
":np_arrays",
":np_math_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "np_math_ops_test",
srcs = ["np_math_ops_test.py"],
deps = [
":np_array_ops",
":np_arrays",
":np_math_ops",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "np_random_test",
srcs = ["np_random_test.py"],
deps = [
":np_array_ops",
":np_dtypes",
":np_math_ops",
":np_random",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "np_utils_test",
srcs = ["np_utils_test.py"],
deps = [
":np_utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "np_interop_test",
srcs = ["np_interop_test.py"],
deps = [
":np_math_ops",
"//tensorflow:tensorflow_py",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//third_party/py/numpy",
],
)
tensorflow/python/ops/numpy_ops/README.md
+5
View File
@@ -0,0 +1,5 @@
This module implements `tf.experimental.numpy` APIs, which provide NumPy APIs
implemented on top of TensorFlow.
Please see [TensorFlow NumPy API
Documentation](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy).
tensorflow/python/ops/numpy_ops/__init__.py
+171
View File
@@ -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.
tensorflow/python/ops/numpy_ops/g3doc/TensorFlow_NumPy_Keras_and_Distribution_Strategy.ipynb
+318
View File
@@ -0,0 +1,318 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "TWLmmaQpX-i1"
},
"source": [
"# TensorFlow NumPy: Keras and Distribution Strategy"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "fmGBjt1arUk7"
},
"source": [
"## Overview\n",
"\n",
"TensorFlow Numpy provides an implementation of a subset of NumPy API on top of TensorFlow backend. Please see [TF NumPy API documentation](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) and \n",
" [TensorFlow NumPy Guide](https://www.tensorflow.org/guide/tf_numpy).\n",
"\n",
"This document shows how TensorFlow NumPy interoperates with TensorFlow's high level APIs like DistributionStrategky and Keras."
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "eAf_CAIerkPZ"
},
"source": [
"## Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "OG0u3eVdSOAk"
},
"outputs": [],
"source": [
"!pip install --quiet --upgrade tf-nightly"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "YjQUVUd3X325"
},
"outputs": [],
"source": [
"import tensorflow as tf\n",
"import tensorflow.experimental.numpy as tnp\n",
"\n",
"# Creates 3 logical GPU devices for demonstrating distribution.\n",
"gpu_device = tf.config.list_physical_devices(\"GPU\")[0]\n",
"tf.config.set_logical_device_configuration(\n",
" gpu_device, [tf.config.LogicalDeviceConfiguration(128)] * 3)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "UTZPYMaPr_oU"
},
"source": [
"## TF NumPy and Keras\n",
"\n",
"TF NumPy can be used to create custom Keras layers. These layers interoperate with and behave like regular Keras layers. Here are some things to note to understand how these layers work.\n",
"\n",
"- Existing Keras layers can be invoked with ND Array inputs, in addition to other input types like `tf.Tensor`, `np.ndarray`, python literals, etc. All these types will be internally convert to a `tf.Tensor` before the layer's `call` method is invoked\n",
"- Existing Keras layers will continue to output `tf.Tensor` values. Custom layers could output ND Array or `tf.Tensor`. \n",
"- Custom and existing Keras layers should be freely composable.\n",
"\n",
"Checkout the examples below that demonstrate the above.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "gsZLC4eEsm8P"
},
"source": [
"### ND Array inputs\n",
"\n",
"Create and call an existing Keras layers with ND Array inputs. Note that the layer outputs a `tf.Tensor`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "CTiylo_UrxW7"
},
"outputs": [],
"source": [
"dense_layer = tf.keras.layers.Dense(5)\n",
"inputs = tnp.random.randn(2, 3).astype(tnp.float32)\n",
"outputs = dense_layer(inputs)\n",
"print(\"Shape:\", outputs.shape)\n",
"print(\"Class:\", outputs.__class__)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "vltJnASzXJNq"
},
"source": [
"### Custom Keras Layer\n",
"\n",
"Create a new Keras layer as below using TensorFlow NumPy methods. Note that the layer's call method receives a `tf.tensor` value as input. It can convert to `ndarray` using `tnp.asarray`. However this conversion may not be needed since TF NumPy APIs can handle `tf.Tensor` inputs."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "0i7lOWJwsVMy"
},
"outputs": [],
"source": [
"class ProjectionLayer(tf.keras.layers.Layer):\n",
" \"\"\"Linear projection layer using TF NumPy.\"\"\"\n",
"\n",
" def __init__(self, units):\n",
" super(ProjectionLayer, self).__init__()\n",
" self._units = units\n",
"\n",
" def build(self, input_shape):\n",
" stddev = tnp.sqrt(self._units).astype(tnp.float32)\n",
" initial_value = tnp.random.randn(input_shape[1], self._units).astype(\n",
" tnp.float32) / stddev\n",
" # Note that TF NumPy can interoperate with tf.Variable.\n",
" self.w = tf.Variable(initial_value, trainable=True)\n",
"\n",
" def call(self, inputs):\n",
" return tnp.matmul(inputs, self.w)\n",
"\n",
"# Call with ndarray inputs\n",
"layer = ProjectionLayer(2)\n",
"tnp_inputs = tnp.random.randn(2, 4).astype(tnp.float32)\n",
"print(\"output:\", layer(tnp_inputs))\n",
"\n",
"# Call with tf.Tensor inputs\n",
"tf_inputs = tf.random.uniform([2, 4])\n",
"print(\"\\noutput: \", layer(tf_inputs))"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "UExEbq1EENLB"
},
"source": [
"### Composing layers\n",
"\n",
"Next create a Keras model by composing the `ProjectionLayer` defined above with a `Dense` layer."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "qbTkqFgDDXaw"
},
"outputs": [],
"source": [
"batch_size = 3\n",
"units = 5\n",
"model = tf.keras.Sequential([tf.keras.layers.Dense(units),\n",
" ProjectionLayer(2)])\n",
"\n",
"print(\"Calling with ND Array inputs\")\n",
"tnp_inputs = tnp.random.randn(batch_size, units).astype(tnp.float32)\n",
"output = model.call(tnp_inputs)\n",
"print(\"Output shape %s.\\nOutput class: %s\\n\" % (output.shape, output.__class__))\n",
"\n",
"print(\"Calling with tensor inputs\")\n",
"tf_inputs = tf.convert_to_tensor(tnp_inputs)\n",
"output = model.call(tf_inputs)\n",
"print(\"Output shape %s.\\nOutput class: %s\" % (output.shape, output.__class__))\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "QeooMJZdYbXq"
},
"source": [
"## Distributed Strategy: tf.distribution\n",
"\n",
"[TensorFlow NumPy Guide](https://colab.sandbox.google.com/drive/15AshdHLS_xTMohWDleTiAgyPdRt6JQJJ#scrollTo=s2enCDi_FvCR) shows how `tf.device` API can be used to place individual operations on specific devices. Note that this works for remote devices as well.\n",
"\n",
"\n",
"TensorFlow also has higher level distribution APIs that make it easy to replicate computation across devices. \n",
"Here we will show how to place TensorFlow NumPy code in a Distribution Strategy context to easily perform replicated computation.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "tOTNvkTxZ-ok"
},
"outputs": [],
"source": [
"# Initialize the strategy\n",
"gpus = tf.config.list_logical_devices(\"GPU\")\n",
"print(\"Using following GPUs\", gpus)\n",
"\n",
"strategy = tf.distribute.MirroredStrategy(gpus)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "Zlmeo8i7Euq0"
},
"source": [
"### Simple replication example\n",
"\n",
"First try running a simple NumPy function in `strategy` context."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "u3ZLh3_ZB8mk"
},
"outputs": [],
"source": [
"@tf.function\n",
"def replica_fn():\n",
" replica_id = tf.distribute.get_replica_context().replica_id_in_sync_group\n",
" print(\"Running on device %s\" % replica_id.device)\n",
" return tnp.asarray(replica_id) * 5\n",
"\n",
"print(strategy.run(replica_fn).values)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "UyyZBpLyE9LG"
},
"source": [
"### Replicated model execution\n",
"\n",
"Next run the model defined earlier under `strategy` scope."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "6VeBFzTCCbZk"
},
"outputs": [],
"source": [
"# Test running the model in a distributed setting.\n",
"model = tf.keras.Sequential([tf.keras.layers.Dense(units), ProjectionLayer(2)])\n",
"\n",
"@tf.function\n",
"def model_replica_fn():\n",
" inputs = tnp.random.randn(batch_size, units).astype(tnp.float32)\n",
" return model.call(inputs)\n",
"\n",
"print(\"Outputs:\\n\", strategy.run(model_replica_fn).values)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "TensorFlow NumPy: Keras and Distribution Strategy",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tensorflow/python/ops/numpy_ops/g3doc/TensorFlow_NumPy_Text_Generation.ipynb
+1088
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/g3doc/TensorFlow_Numpy_Distributed_Image_Classification.ipynb
+552
View File
@@ -0,0 +1,552 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "KQALG9h23b0R"
},
"source": [
"##### Copyright 2020 The TensorFlow Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "both",
"id": "U34SJW0W3dg_"
},
"outputs": [],
"source": [
"#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VIX1XZHJ3gFo"
},
"source": [
"# TensorFlow NumPy: Distributed Image Classification Tutorial"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f7NApJ7R3ndN"
},
"source": [
"## Overview\n",
"\n",
"TensorFlow implements a subset of the [NumPy API](https://numpy.org/doc/1.16), available as `tf.experimental.numpy`. This allows running NumPy code, accelerated by TensorFlow together with access to all of TensorFlow's APIs. Please see [TensorFlow NumPy Guide](https://www.tensorflow.org/guide/tf_numpy) to get started.\n",
"\n",
"Here you will learn how to build a deep model for an image classification task by using TensorFlow Numpy APIs. For using higher level `tf.keras` APIs, see the following [tutorial](https://www.tensorflow.org/tutorials/quickstart/beginner)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IYDdfih63rSG"
},
"source": [
"## Setup\n",
"\n",
"tf.experimental.numpy will be available in the stable branch starting from TensorFlow 2.4. For now, it is available in nightly."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "3IlLM-YlTMv5"
},
"outputs": [],
"source": [
"!pip install --quiet --upgrade tf-nightly\n",
"!pip install --quiet --upgrade tensorflow-datasets"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "U13hRXHKTcsE"
},
"outputs": [],
"source": [
"import collections\n",
"import functools\n",
"import matplotlib.pyplot as plt\n",
"import os\n",
"import tempfile\n",
"import tensorflow as tf\n",
"import tensorflow.experimental.numpy as tnp\n",
"import tensorflow_datasets as tfds\n",
"\n",
"gpus = tf.config.list_physical_devices('GPU')\n",
"if gpus:\n",
" tf.config.set_logical_device_configuration(gpus[0], [\n",
" tf.config.LogicalDeviceConfiguration(memory_limit=128),\n",
" tf.config.LogicalDeviceConfiguration(memory_limit=128)])\n",
" devices = tf.config.list_logical_devices('GPU')\n",
"else:\n",
" cpus = tf.config.list_physical_devices('CPU')\n",
" tf.config.set_logical_device_configuration(cpus[0], [\n",
" tf.config.LogicalDeviceConfiguration(),\n",
" tf.config.LogicalDeviceConfiguration()])\n",
" devices = tf.config.list_logical_devices('CPU')\n",
"\n",
"print(\"Using following virtual devices\", devices)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AxNuZSqZKcdM"
},
"source": [
"## Mnist dataset\n",
"\n",
"Mnist contains 28 * 28 images of digits from 0 to 9. The task is to classify the images as these 10 possible classes.\n",
"\n",
"Below, load the dataset and examine a few samples."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yKf9Tm5OjwGK"
},
"outputs": [],
"source": [
"NUM_CLASSES = 10\n",
"BATCH_SIZE = 64\n",
"INPUT_SIZE = 28 * 28\n",
"\n",
"def process_data(data_dict):\n",
" images = tnp.asarray(data_dict['image']) / 255.0\n",
" images = images.reshape(-1, INPUT_SIZE).astype(tnp.float32)\n",
" labels = tnp.asarray(data_dict['label'])\n",
" labels = tnp.eye(NUM_CLASSES, dtype=tnp.float32)[labels]\n",
" return images, labels\n",
"\n",
"with tf.device(\"CPU:0\"):\n",
" train_dataset = tfds.load('mnist', split='train', shuffle_files=True, \n",
" batch_size=BATCH_SIZE).map(process_data)\n",
" test_dataset = tfds.load('mnist', split='test', shuffle_files=True, \n",
" batch_size=-1)\n",
" x_test, y_test = process_data(test_dataset)\n",
"\n",
" # Plots some examples.\n",
" images, labels = next(iter(train_dataset.take(1)))\n",
" _, axes = plt.subplots(1, 8, figsize=(12, 96))\n",
" for i, ax in enumerate(axes):\n",
" ax.imshow(images[i].reshape(28, 28), cmap='gray')\n",
" ax.axis(\"off\")\n",
" ax.set_title(\"Label: %d\" % int(tnp.argmax(labels[i])))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZDJQp4i00qaJ"
},
"source": [
"## Define layers and model\n",
"\n",
"Here, you will implement a multi-layer perceptron model that trains on the MNIST data. First, define a `Dense` class which applies a linear transform followed by a \"relu\" non-linearity."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "44yzAmBFreyg"
},
"outputs": [],
"source": [
"class Dense(tf.Module):\n",
"\n",
" def __init__(self, units, use_relu=True):\n",
" self.wt = None\n",
" self.bias = None\n",
" self._use_relu = use_relu\n",
" self._built = False\n",
" self._units = units\n",
"\n",
" def __call__(self, inputs):\n",
" if not self._built:\n",
" self._build(inputs.shape)\n",
" x = tnp.add(tnp.matmul(inputs, self.wt), self.bias)\n",
" if self._use_relu:\n",
" return tnp.maximum(x, 0.)\n",
" else:\n",
" return x\n",
"\n",
" @property\n",
" def params(self):\n",
" assert self._built\n",
" return [self.wt, self.bias]\n",
"\n",
" def _build(self, input_shape):\n",
" size = input_shape[1]\n",
" stddev = 1 / tnp.sqrt(size)\n",
" # Note that model parameters are `tf.Variable` since they requires\n",
" # mutation, which is currently unsupported by TensorFlow NumPy.\n",
" # Also note interoperation with TensorFlow APIs below.\n",
" self.wt = tf.Variable(\n",
" tf.random.truncated_normal(\n",
" [size, self._units], stddev=stddev, dtype=tf.float32))\n",
" self.bias = tf.Variable(tf.zeros([self._units], dtype=tf.float32))\n",
" self._built = True"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wfKpg3adUCy9"
},
"source": [
"Next, create a `Model` object that applies two non-linear `Dense` transforms,\n",
"followed by a linear transform."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "NdrdxKB7SenC"
},
"outputs": [],
"source": [
"class Model(tf.Module):\n",
" \"\"\"A three layer neural network.\"\"\"\n",
"\n",
" def __init__(self):\n",
" self.layer1 = Dense(128)\n",
" self.layer2 = Dense(32)\n",
" self.layer3 = Dense(NUM_CLASSES, use_relu=False)\n",
"\n",
" def __call__(self, inputs):\n",
" x = self.layer1(inputs)\n",
" x = self.layer2(x)\n",
" return self.layer3(x)\n",
"\n",
" @property\n",
" def params(self):\n",
" return self.layer1.params + self.layer2.params + self.layer3.params"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Hoxh5Z7E_9Pv"
},
"source": [
"## Training and evaluation\n",
"\n",
"Checkout the following methods for performing training and evaluation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hOxqjE7rZPdr"
},
"outputs": [],
"source": [
"def forward(model, inputs, labels):\n",
" \"\"\"Computes prediction and loss.\"\"\"\n",
" logits = model(inputs)\n",
" # TensorFlow's loss function has numerically stable implementation of forward\n",
" # pass and gradients. So we prefer that here.\n",
" loss = tf.nn.softmax_cross_entropy_with_logits(labels, logits)\n",
" mean_loss = tnp.mean(loss)\n",
" return logits, mean_loss\n",
"\n",
"def compute_gradients(model, inputs, labels):\n",
" \"\"\"Computes gradients of loss based on `labels` and prediction on `inputs`.\"\"\"\n",
" with tf.GradientTape() as tape:\n",
" tape.watch(inputs)\n",
" _, loss = forward(model, inputs, labels)\n",
" gradients = tape.gradient(loss, model.params)\n",
" return gradients\n",
"\n",
"def compute_sgd_updates(gradients, learning_rate):\n",
" \"\"\"Computes parameter updates based on SGD update rule.\"\"\"\n",
" return [-learning_rate * grad for grad in gradients]\n",
"\n",
"def apply_updates(model, updates):\n",
" \"\"\"Applies `update` to `model.params`.\"\"\"\n",
" for param, update in zip(model.params, updates):\n",
" param.assign_add(update)\n",
"\n",
"def evaluate(model, images, labels):\n",
" \"\"\"Evaluates accuracy for `model`'s predictions.\"\"\"\n",
" prediction = model(images)\n",
" predicted_class = tnp.argmax(prediction, axis=-1)\n",
" actual_class = tnp.argmax(labels, axis=-1)\n",
" return float(tnp.mean(predicted_class == actual_class))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8t70b5d6XCs7"
},
"source": [
"### Single GPU training"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "HrhS_M6kALeP"
},
"outputs": [],
"source": [
"NUM_EPOCHS = 10\n",
"\n",
"@tf.function\n",
"def train_step(model, input, labels, learning_rate):\n",
" gradients = compute_gradients(model, input, labels)\n",
" updates = compute_sgd_updates(gradients, learning_rate)\n",
" apply_updates(model, updates)\n",
"\n",
"# Creates and build a model.\n",
"model = Model()\n",
"\n",
"accuracies = []\n",
"for _ in range(NUM_EPOCHS):\n",
" for inputs, labels in train_dataset:\n",
" train_step(model, inputs, labels, learning_rate=0.1)\n",
" accuracies.append(evaluate(model, x_test, y_test))\n",
"\n",
"def plot_accuracies(accuracies):\n",
" plt.plot(accuracies)\n",
" plt.xlabel(\"epoch\")\n",
" plt.ylabel(\"accuracy\")\n",
" plt.title(\"Eval accuracy vs epoch\")\n",
"\n",
"plot_accuracies(accuracies)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Dw7RwQmKcYK9"
},
"source": [
"#### Saving the models to disk"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rmk2xLQLcXkl"
},
"outputs": [],
"source": [
"# A temporary directory to save our models into.\n",
"dir = tempfile.TemporaryDirectory()\n",
"\n",
"# We take our model, and create a wrapper for it.\n",
"class SaveableModel(Model):\n",
" @tf.function\n",
" def __call__(self, inputs):\n",
" return super().__call__(inputs)\n",
"\n",
"saveable_model = SaveableModel()\n",
"\n",
"# This saves a concrete function that we care about.\n",
"outputs = saveable_model(x_test)\n",
"\n",
"# This saves the model to disk.\n",
"tf.saved_model.save(saveable_model, dir.name)\n",
"\n",
"loaded = tf.saved_model.load(dir.name)\n",
"outputs_loaded = loaded(x_test)\n",
"\n",
"# Ensure that the loaded model preserves the weights\n",
"# of the saved model.\n",
"assert tnp.allclose(outputs, outputs_loaded)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ak_hCOkGXXfl"
},
"source": [
"### Multi GPU runs\n",
"\n",
"Next, run mirrored training on multiple GPUs. Note that the GPUs used here are virtual and map to the same physical GPU.\n",
"\n",
"First, define a few utilities to run replicated computation and reductions."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ujbeT5p6Xm7k"
},
"source": [
"#### Distribution primitives\n",
"\n",
"Checkout primitives below for function replication and distributed reduction."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MZ6hivj-ZIRo"
},
"outputs": [],
"source": [
"import threading\n",
"import queue\n",
"\n",
"# Note that this code currently relies on dispatching operations from python\n",
"# threads.\n",
"class ReplicatedFunction(object):\n",
" \"\"\"Creates a callable that will run `fn` on each device in `devices`.\"\"\"\n",
"\n",
" def __init__(self, fn, devices, **kw_args):\n",
" self._shutdown = False\n",
"\n",
" def _replica_fn(device, input_queue, output_queue):\n",
" while not self._shutdown:\n",
" inputs = input_queue.get()\n",
" with tf.device(device):\n",
" output_queue.put(fn(*inputs, **kw_args))\n",
"\n",
" self.threads = []\n",
" self.input_queues = [queue.Queue() for _ in devices]\n",
" self.output_queues = [queue.Queue() for _ in devices]\n",
" for i, device in enumerate(devices):\n",
" thread = threading.Thread(\n",
" target=_replica_fn,\n",
" args=(device, self.input_queues[i], self.output_queues[i]))\n",
" thread.start()\n",
" self.threads.append(thread)\n",
"\n",
" def __call__(self, *inputs):\n",
" all_inputs = zip(*inputs)\n",
" for input_queue, replica_input, in zip(self.input_queues, all_inputs):\n",
" input_queue.put(replica_input)\n",
" return [q.get() for q in self.output_queues]\n",
"\n",
" def __del__(self):\n",
" self._shutdown = True\n",
" for t in self.threads:\n",
" t.join(3)\n",
" self.threads = None\n",
"\n",
"def collective_mean(inputs, num_devices):\n",
" \"\"\"Performs collective mean reduction on inputs.\"\"\"\n",
" outputs = []\n",
" for instance_key, inp in enumerate(inputs):\n",
" outputs.append(tnp.asarray(\n",
" tf.raw_ops.CollectiveReduce(\n",
" input=inp, group_size=num_devices, group_key=0,\n",
" instance_key=instance_key, merge_op='Add', final_op='Div',\n",
" subdiv_offsets=[])))\n",
" return outputs"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1ZiN1rpJYHLu"
},
"source": [
"#### Distributed training "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "A6ZHYmLapunm"
},
"outputs": [],
"source": [
"# This is similar to `train_step` except for an extra collective reduction of\n",
"# gradients\n",
"@tf.function\n",
"def replica_step(model, inputs, labels,\n",
" learning_rate=None, num_devices=None):\n",
" gradients = compute_gradients(model, inputs, labels)\n",
" # Note that each replica performs a reduction to compute mean of gradients.\n",
" reduced_gradients = collective_mean(gradients, num_devices)\n",
" updates = compute_sgd_updates(reduced_gradients, learning_rate)\n",
" apply_updates(model, updates)\n",
"\n",
"models = [Model() for _ in devices]\n",
"\n",
"# The code below builds all the model objects and copies model parameters from\n",
"# the first model to all the replicas.\n",
"def init_model(model):\n",
" model(tnp.zeros((1, INPUT_SIZE), dtype=tnp.float32))\n",
" if model != models[0]:\n",
" # Copy the first models weights into the other models.\n",
" for p1, p2 in zip(model.params, models[0].params):\n",
" p1.assign(p2)\n",
"\n",
"with tf.device(devices[0]):\n",
" init_model(models[0])\n",
"# Replicate and run the parameter initialization.\n",
"ReplicatedFunction(init_model, devices[1:])(models[1:])\n",
"\n",
"# Replicate the training step\n",
"replicated_step = ReplicatedFunction(\n",
" replica_step, devices, learning_rate=0.1, num_devices=len(devices))\n",
"\n",
"accuracies = []\n",
"print(\"Running distributed training on devices: %s\" % devices)\n",
"for _ in range(NUM_EPOCHS):\n",
" for inputs, labels in train_dataset:\n",
" replicated_step(models,\n",
" tnp.split(inputs, len(devices)),\n",
" tnp.split(labels, len(devices)))\n",
" accuracies.append(evaluate(models[0], x_test, y_test))\n",
"\n",
"plot_accuracies(accuracies)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"KQALG9h23b0R",
"f7NApJ7R3ndN"
],
"name": "TensorFlow Numpy: Distributed Image Classification",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tensorflow/python/ops/numpy_ops/integration_test/BUILD
+25
View File
@@ -0,0 +1,25 @@
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
licenses(["notice"])
py_test(
name = "public_symbol_test",
srcs = ["public_symbol_test.py"],
strict_deps = True,
deps = [
"//tensorflow:tensorflow_py",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "np_config_test",
srcs = ["np_config_test.py"],
deps = [
"//tensorflow:tensorflow_py_no_contrib",
"//tensorflow/python/ops/numpy_ops:np_config",
],
)
tensorflow/python/ops/numpy_ops/integration_test/benchmarks/BUILD
+38
View File
@@ -0,0 +1,38 @@
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
py_binary(
name = "micro_benchmarks",
srcs = ["micro_benchmarks.py"],
strict_deps = True,
deps = [
":numpy_mlp",
":tf_numpy_mlp",
"//tensorflow:tensorflow_py_no_contrib",
"//tensorflow/python:extra_py_tests_deps",
"//third_party/py/numpy",
"@absl_py//absl/flags",
"@absl_py//absl/logging",
],
)
py_library(
name = "numpy_mlp",
srcs = ["numpy_mlp.py"],
strict_deps = True,
deps = [
"//third_party/py/numpy",
],
)
py_library(
name = "tf_numpy_mlp",
srcs = ["tf_numpy_mlp.py"],
strict_deps = True,
deps = ["//tensorflow:tensorflow_py_no_contrib"],
)
tensorflow/python/ops/numpy_ops/integration_test/benchmarks/micro_benchmarks.py
+163
View File
@@ -0,0 +1,163 @@
# 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.
# ==============================================================================
r"""Micro benchmark.
bazel run -c opt --config=cuda \
//third_party/tensorflow/python/ops/numpy_ops/integration_test/benchmarks:micro_benchmarks -- \
--number=100 --repeat=100 \
--benchmark_filter=.
"""
import gc
import time
from absl import flags
from absl import logging
import numpy as np # pylint: disable=unused-import
import tensorflow.compat.v2 as tf
from tensorflow.python.ops.numpy_ops.integration_test.benchmarks import numpy_mlp
from tensorflow.python.ops.numpy_ops.integration_test.benchmarks import tf_numpy_mlp
FLAGS = flags.FLAGS
# Used instead of "import tensorflow(dot)experimental.numpy as tfnp" due to
# copybara issues.
tfnp = tf.experimental.numpy
flags.DEFINE_integer('repeat', 100, '#Measurements per benchmark.')
flags.DEFINE_integer('number', 100, '#Runs per a measure.')
class MicroBenchmarks(tf.test.Benchmark):
"""Main micro benchmark class."""
def _benchmark_and_report(
self,
name,
fn,
repeat=None,
number=None):
"""Run fn repeat * number times, report time, and return fastest time."""
# Can't make these default above since the flags may not have been parsed
# at module import time.
repeat = repeat or int(FLAGS.repeat)
number = number or int(FLAGS.number)
# Warmup
fn()
times = []
for _ in range(repeat):
gc.disable()
start = time.time()
for _ in range(number):
fn()
times.append(time.time() - start)
gc.enable()
gc.collect()
# Regular benchmark to report numbers.
fastest_time_us = min(times) * 1e6 / number
total_time = sum(times)
self.report_benchmark(name=name,
wall_time=total_time,
extras={'fastest_time_us': fastest_time_us})
return fastest_time_us
def benchmark_tf_np_mlp_inference_batch_1_cpu(self):
with tf.device('/CPU:0'):
model = tf_numpy_mlp.MLP()
x = tfnp.ones(shape=(1, 10)).astype(np.float32)
self._benchmark_and_report(self._get_name(), lambda: model.inference(x))
def benchmark_tf_np_tf_function_mlp_inference_batch_1_cpu(self):
with tf.device('/CPU:0'):
model = tf_numpy_mlp.MLP()
x = tfnp.ones(shape=(1, 10)).astype(np.float32)
self._benchmark_and_report(
self._get_name(), tf.function(lambda: model.inference(x)))
def benchmark_numpy_mlp_inference_batch_1_cpu(self):
model = numpy_mlp.MLP()
x = np.random.uniform(size=(1, 10)).astype(np.float32, copy=False)
self._benchmark_and_report(self._get_name(), lambda: model.inference(x))
def _benchmark_np_and_tf_np(self, name, op, args, repeat=None): # pylint: disable=redefined-builtin
fn = getattr(np, op)
assert fn is not None
np_time = self._benchmark_and_report(
'{}_numpy'.format(name), lambda: fn(*args), repeat=repeat)
fn = getattr(tfnp, op)
assert fn is not None
with tf.device('CPU:0'):
tf_time = self._benchmark_and_report(
'{}_tfnp_cpu'.format(name), lambda: fn(*args), repeat=repeat)
return np_time, tf_time
def _print_times(self, op, sizes, times):
# For easy reporting.
print('For np.{}:'.format(op))
print('{:<15} {:>11} {:>11}'.format('Size', 'NP time', 'TF NP Time'))
for size, (np_time, tf_time) in zip(sizes, times):
print('{:<15} {:>10.5}us {:>10.5}us'.format(
str(size), np_time, tf_time))
print()
def _benchmark_np_and_tf_np_unary(self, op):
sizes = [(100,), (10000,), (1000000,)]
repeats = [FLAGS.repeat] * 2 + [10]
times = []
for size, repeat in zip(sizes, repeats):
x = np.random.uniform(size=size).astype(np.float32, copy=False)
name = '{}_{}'.format(self._get_name(), size)
times.append(self._benchmark_np_and_tf_np(name, op, (x,), repeat))
self._print_times(op, sizes, times)
def benchmark_count_nonzero(self):
self._benchmark_np_and_tf_np_unary('count_nonzero')
def benchmark_log(self):
self._benchmark_np_and_tf_np_unary('log')
def benchmark_exp(self):
self._benchmark_np_and_tf_np_unary('exp')
def benchmark_tanh(self):
self._benchmark_np_and_tf_np_unary('tanh')
def benchmark_matmul(self):
sizes = [(2, 2), (10, 10), (100, 100), (200, 200), (1000, 1000)]
# Override repeat flag since this can be very slow.
repeats = [FLAGS.repeat] * 3 + [50, 10]
times = []
for size, repeat in zip(sizes, repeats):
x = np.random.uniform(size=size).astype(np.float32, copy=False)
name = '{}_{}'.format(self._get_name(), size)
times.append(
self._benchmark_np_and_tf_np(name, 'matmul', (x, x), repeat=repeat))
self._print_times('matmul', sizes, times)
if __name__ == '__main__':
logging.set_verbosity(logging.WARNING)
tf.enable_v2_behavior()
tf.test.main()
tensorflow/python/ops/numpy_ops/integration_test/benchmarks/numpy_mlp.py
+47
View File
@@ -0,0 +1,47 @@
# 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.
# ==============================================================================
"""Builds the MLP network."""
import numpy as np
NUM_CLASSES = 3
INPUT_SIZE = 10
HIDDEN_UNITS = 10
class MLP:
"""MLP model.
T = Relu(Add(MatMul(A, B), C))
R = Relu(Add(MatMul(T, D), E))
"""
def __init__(self, num_classes=NUM_CLASSES, input_size=INPUT_SIZE,
hidden_units=HIDDEN_UNITS):
self.w1 = np.random.uniform(size=[input_size, hidden_units]).astype(
np.float32, copy=False)
self.w2 = np.random.uniform(size=[hidden_units, num_classes]).astype(
np.float32, copy=False)
self.b1 = np.random.uniform(size=[1, hidden_units]).astype(
np.float32, copy=False)
self.b2 = np.random.uniform(size=[1, num_classes]).astype(
np.float32, copy=False)
def inference(self, inputs):
return self._forward(inputs, self.w1, self.w2, self.b1, self.b2)
def _forward(self, x, w1, w2, b1, b2):
x = np.maximum(np.matmul(x, w1) + b1, 0.)
x = np.maximum(np.matmul(x, w2) + b2, 0.)
return x
tensorflow/python/ops/numpy_ops/integration_test/benchmarks/tf_numpy_mlp.py
+49
View File
@@ -0,0 +1,49 @@
# 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.
# ==============================================================================
"""Builds the MLP network."""
import tensorflow.compat.v2 as tf
np = tf.experimental.numpy
NUM_CLASSES = 3
INPUT_SIZE = 10
HIDDEN_UNITS = 10
class MLP:
"""MLP model.
T = Relu(Add(MatMul(A, B), C))
R = Relu(Add(MatMul(T, D), E))
"""
def __init__(self, num_classes=NUM_CLASSES, input_size=INPUT_SIZE,
hidden_units=HIDDEN_UNITS):
self.w1 = np.random.uniform(size=[input_size, hidden_units]).astype(
np.float32)
self.w2 = np.random.uniform(size=[hidden_units, num_classes]).astype(
np.float32)
self.b1 = np.random.uniform(size=[1, hidden_units]).astype(
np.float32)
self.b2 = np.random.uniform(size=[1, num_classes]).astype(
np.float32)
def inference(self, inputs):
return self._forward(inputs, self.w1, self.w2, self.b1, self.b2)
def _forward(self, x, w1, w2, b1, b2):
x = np.maximum(np.matmul(x, w1) + b1, 0.)
x = np.maximum(np.matmul(x, w2) + b2, 0.)
return x
tensorflow/python/ops/numpy_ops/integration_test/np_config_test.py
+40
View File
@@ -0,0 +1,40 @@
# 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.
# ==============================================================================
"""Tests that an error is raised when numpy functions are called."""
import tensorflow.compat.v2 as tf
from tensorflow.python.ops.numpy_ops import np_config
class ConfigTest(tf.test.TestCase):
def testMethods(self):
a = tf.constant(1.)
for name in {'T', 'astype', 'ravel', 'transpose', 'reshape', 'clip', 'size',
'tolist'}:
with self.assertRaisesRegex(AttributeError, 'enable_numpy_behavior'):
getattr(a, name)
np_config.enable_numpy_behavior()
for name in {'T', 'astype', 'ravel', 'transpose', 'reshape', 'clip', 'size',
'tolist'}:
_ = getattr(a, name)
if __name__ == '__main__':
tf.compat.v1.enable_eager_execution()
tf.test.main()
tensorflow/python/ops/numpy_ops/integration_test/public_symbol_test.py
+34
View File
@@ -0,0 +1,34 @@
# 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.
# ==============================================================================
"""Tests using module `tf.experimental.numpy` via an alias."""
import numpy as onp
import tensorflow as tf
np = tf.experimental.numpy
class PublicSymbolTest(tf.test.TestCase):
def testSimple(self):
a = 0.1
b = 0.2
self.assertAllClose(onp.add(a, b), np.add(a, b))
if __name__ == "__main__":
tf.compat.v1.enable_eager_execution()
tf.test.main()
tensorflow/python/ops/numpy_ops/np_array_ops.py
+2116
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+1329
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/np_arrays.py
+50
View File
@@ -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
tensorflow/python/ops/numpy_ops/np_arrays_test.py
+215
View File
@@ -0,0 +1,215 @@
# 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.
# ==============================================================================
"""Tests for ndarray."""
import numpy as np
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops.numpy_ops import np_arrays
# Required for operator overloads
from tensorflow.python.ops.numpy_ops import np_math_ops # pylint: disable=unused-import
from tensorflow.python.platform import test
from tensorflow.python.util import nest
class ArrayTest(test.TestCase):
def testDtype(self):
a = array_ops.zeros(shape=[1, 2], dtype=dtypes.int64)
self.assertIs(a.dtype.as_numpy_dtype, np.int64)
np_dt = a.dtype.as_numpy_dtype
self.assertAllEqual(0, np_dt(0))
def testAstype(self):
a = ops.convert_to_tensor(value=1.1, dtype=dtypes.float32).astype(np.int32)
self.assertIs(a.dtype.as_numpy_dtype, np.int32)
self.assertAllEqual(1, a)
a = ops.convert_to_tensor(value=[0.0, 1.1], dtype=dtypes.float32).astype(
np.bool_)
self.assertIs(a.dtype.as_numpy_dtype, np.bool_)
self.assertAllEqual([False, True], a)
def testNeg(self):
a = ops.convert_to_tensor(value=[1.0, 2.0])
self.assertAllEqual([-1.0, -2.0], -a) # pylint: disable=invalid-unary-operand-type
def _testBinOp(self, a, b, out, f, types=None):
a = ops.convert_to_tensor(value=a, dtype=np.int32)
b = ops.convert_to_tensor(value=b, dtype=np.int32)
if not isinstance(out, np_arrays.ndarray):
out = ops.convert_to_tensor(value=out, dtype=np.int32)
if types is None:
types = [[np.int32, np.int32, np.int32], [np.int64, np.int32, np.int64],
[np.int32, np.int64, np.int64],
[np.float32, np.int32, np.float64],
[np.int32, np.float32, np.float64],
[np.float32, np.float32, np.float32],
[np.float64, np.float32, np.float64],
[np.float32, np.float64, np.float64]]
for a_type, b_type, out_type in types:
o = f(a.astype(a_type), b.astype(b_type))
self.assertIs(o.dtype.as_numpy_dtype, out_type)
out = out.astype(out_type)
if np.issubdtype(out_type, np.inexact):
self.assertAllClose(out, o)
else:
self.assertAllEqual(out, o)
def testAdd(self):
self._testBinOp([1, 2], [3, 4], [4, 6], lambda a, b: a.__add__(b))
def testRadd(self):
self._testBinOp([1, 2], [3, 4], [4, 6], lambda a, b: b.__radd__(a))
def testSub(self):
self._testBinOp([1, 2], [3, 5], [-2, -3], lambda a, b: a.__sub__(b))
def testRsub(self):
self._testBinOp([1, 2], [3, 5], [-2, -3], lambda a, b: b.__rsub__(a))
def testMul(self):
self._testBinOp([1, 2], [3, 4], [3, 8], lambda a, b: a.__mul__(b))
def testRmul(self):
self._testBinOp([1, 2], [3, 4], [3, 8], lambda a, b: b.__rmul__(a))
def testPow(self):
self._testBinOp([4, 5], [3, 2], [64, 25], lambda a, b: a.__pow__(b))
def testRpow(self):
self._testBinOp([4, 5], [3, 2], [64, 25], lambda a, b: b.__rpow__(a))
_truediv_types = [[np.int32, np.int32, np.float64],
[np.int64, np.int32, np.float64],
[np.int32, np.int64, np.float64],
[np.float32, np.int32, np.float64],
[np.int32, np.float32, np.float64],
[np.float32, np.float32, np.float32],
[np.float64, np.float32, np.float64],
[np.float32, np.float64, np.float64]]
def testTruediv(self):
self._testBinOp([3, 5], [2, 4],
ops.convert_to_tensor(value=[1.5, 1.25]),
lambda a, b: a.__truediv__(b),
types=self._truediv_types)
def testRtruediv(self):
self._testBinOp([3, 5], [2, 4],
ops.convert_to_tensor(value=[1.5, 1.25]),
lambda a, b: b.__rtruediv__(a),
types=self._truediv_types)
def _testCmp(self, a, b, out, f):
a = ops.convert_to_tensor(value=a, dtype=np.int32)
b = ops.convert_to_tensor(value=b, dtype=np.int32)
types = [[np.int32, np.int32], [np.int64, np.int32], [np.int32, np.int64],
[np.float32, np.int32], [np.int32, np.float32],
[np.float32, np.float32], [np.float64, np.float32],
[np.float32, np.float64]]
for a_type, b_type in types:
o = f(a.astype(a_type), b.astype(b_type))
self.assertAllEqual(out, o)
def testLt(self):
self._testCmp([1, 2, 3], [3, 2, 1], [True, False, False],
lambda a, b: a.__lt__(b))
def testLe(self):
self._testCmp([1, 2, 3], [3, 2, 1], [True, True, False],
lambda a, b: a.__le__(b))
def testGt(self):
self._testCmp([1, 2, 3], [3, 2, 1], [False, False, True],
lambda a, b: a.__gt__(b))
def testGe(self):
self._testCmp([1, 2, 3], [3, 2, 1], [False, True, True],
lambda a, b: a.__ge__(b))
def testEq(self):
self._testCmp([1, 2, 3], [3, 2, 1], [False, True, False],
lambda a, b: a.__eq__(b))
def testNe(self):
self._testCmp([1, 2, 3], [3, 2, 1], [True, False, True],
lambda a, b: a.__ne__(b))
def testInt(self):
v = 10
u = int(ops.convert_to_tensor(value=v))
self.assertIsInstance(u, int)
self.assertAllEqual(v, u)
def testFloat(self):
v = 21.32
u = float(ops.convert_to_tensor(value=v))
self.assertIsInstance(u, float)
self.assertAllClose(v, u)
def testBool(self):
b = bool(ops.convert_to_tensor(value=10))
self.assertIsInstance(b, bool)
self.assertTrue(b)
self.assertFalse(bool(ops.convert_to_tensor(value=0)))
self.assertTrue(bool(ops.convert_to_tensor(value=0.1)))
self.assertFalse(bool(ops.convert_to_tensor(value=0.0)))
def testHash(self):
a = ops.convert_to_tensor(value=10)
def eager():
hash(a)
def graph():
@def_function.function
def f(x):
hash(x)
f(a)
for f in [eager, graph]:
with self.assertRaisesRegex(
TypeError,
r'Tensor is unhashable. Instead, use tensor.ref\(\) as the key.'):
f()
def testFromToCompositeTensor(self):
tensors = [ops.convert_to_tensor(0.1), ops.convert_to_tensor(0.2)]
flattened = nest.flatten(tensors, expand_composites=True)
# Each ndarray contains only one tensor, so the flattened output should be
# just 2 tensors in a list.
self.assertLen(flattened, 2)
self.assertIsInstance(flattened[0], tensor.Tensor)
self.assertIsInstance(flattened[1], tensor.Tensor)
repacked = nest.pack_sequence_as(tensors, flattened, expand_composites=True)
self.assertLen(repacked, 2)
self.assertIsInstance(repacked[0], np_arrays.ndarray)
self.assertIsInstance(repacked[1], np_arrays.ndarray)
self.assertAllClose(tensors, repacked)
if __name__ == '__main__':
# TODO(wangpeng): Test in graph mode as well. Also test in V2 (the requirement
# for setting _USE_EQUALITY points to V2 behavior not being on).
ops.enable_eager_execution()
tensor.Tensor._USE_EQUALITY = True
ops.set_dtype_conversion_mode('legacy')
np_math_ops.enable_numpy_methods_on_tensor()
test.main()
tensorflow/python/ops/numpy_ops/np_config.py
+58
View File
@@ -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)
tensorflow/python/ops/numpy_ops/np_dtypes.py
+216
View File
@@ -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
tensorflow/python/ops/numpy_ops/np_dtypes_test.py
+59
View File
@@ -0,0 +1,59 @@
# 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.
# ==============================================================================
"""Tests for tf-numpy dtype utilities."""
from absl.testing import parameterized
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops.numpy_ops import np_dtypes
from tensorflow.python.platform import test
class DTypeTest(test.TestCase, parameterized.TestCase):
@parameterized.parameters([False, True])
def testAllowF64False(self, prefer_f32):
np_dtypes.set_allow_float64(False)
np_dtypes.set_prefer_float32(prefer_f32)
self.assertEqual(dtypes.float32, np_dtypes.default_float_type())
self.assertEqual(dtypes.float32,
np_dtypes._result_type(np.zeros([], np.float64), 1.1))
def testAllowF64TruePreferF32False(self):
np_dtypes.set_allow_float64(True)
np_dtypes.set_prefer_float32(False)
self.assertEqual(dtypes.float64, np_dtypes.default_float_type())
self.assertEqual(dtypes.float64, np_dtypes._result_type(1.1))
self.assertEqual(dtypes.complex128, np_dtypes._result_type(1.j))
def testAllowF64TruePreferF32True(self):
np_dtypes.set_allow_float64(True)
np_dtypes.set_prefer_float32(True)
self.assertEqual(dtypes.float32, np_dtypes.default_float_type())
self.assertEqual(dtypes.float32, np_dtypes._result_type(1.1))
self.assertEqual(dtypes.float64,
np_dtypes._result_type(np.zeros([], np.float64), 1.1))
self.assertEqual(dtypes.complex64, np_dtypes._result_type(1.1j))
self.assertEqual(dtypes.complex128,
np_dtypes._result_type(np.zeros([], np.complex128), 1.1j))
self.assertEqual(dtypes.complex64,
np_dtypes._result_type(np.zeros([], np.float32), 1.1j))
if __name__ == '__main__':
ops.enable_eager_execution()
test.main()
tensorflow/python/ops/numpy_ops/np_interop_test.py
+461
View File
@@ -0,0 +1,461 @@
# 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.
# ==============================================================================
"""Tests for interop between TF ops, numpy_ops, and numpy methods."""
import numpy as onp
import tensorflow.compat.v2 as tf
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops.numpy_ops import np_math_ops
# Used instead of "import tensorflow(dot)experimental.numpy as np" due to
# copybara issues.
np = tf.experimental.numpy
# Tests for code snippet put in README.md
class ReadmeTest(tf.test.TestCase):
def testBroadcastAdd(self):
x_np = np.ones([2, 1]) + np.ones([1, 2])
x_onp = onp.ones([2, 1]) + onp.ones([1, 2])
self.assertAllClose(x_onp, x_np)
def testTypePromotion(self):
x_np = np.ones([1, 2], dtype=np.int16) + np.ones([2, 1], dtype=np.uint8)
x_onp = np.ones([1, 2], dtype=np.int16) + np.ones([2, 1], dtype=np.uint8)
self.assertEqual(x_onp.dtype, x_np.dtype)
self.assertAllClose(x_onp, x_np)
def testTFInterop(self):
x_np = np.sum(np.ones([1, 2]) + tf.ones([2, 1]))
x_onp = onp.sum(onp.ones([1, 2]) + onp.ones([2, 1]))
self.assertAllClose(x_onp, x_np)
def testOnpInterop(self):
x_np = onp.sum(np.ones([1, 2]) + onp.ones([2, 1]))
x_onp = onp.sum(onp.ones([1, 2]) + onp.ones([2, 1]))
self.assertAllClose(x_onp, x_np)
def testDevice(self):
if tf.test.is_gpu_available():
with tf.device('GPU:0'):
x = np.ones([1, 2])
self.assertIn('GPU', tf.convert_to_tensor(x).device)
with tf.device('CPU:0'):
x = np.ones([1, 2])
self.assertIn('CPU', tf.convert_to_tensor(x).device)
def testFunction(self):
@tf.function
def f(x, y):
return np.sum(x + y)
x_np = f(np.ones([1, 2]), tf.ones([2, 1]))
x_onp = onp.sum(onp.ones([1, 2]) + onp.ones([2, 1]))
self.assertAllClose(x_onp, x_np)
class InteropTest(tf.test.TestCase):
def setUp(self):
super(InteropTest, self).setUp()
physical_devices = tf.config.list_physical_devices('CPU')
configs = tf.config.get_logical_device_configuration(physical_devices[0])
if configs is None:
logical_devices = [
tf.config.LogicalDeviceConfiguration() for _ in range(3)
]
tf.config.set_logical_device_configuration(physical_devices[0],
logical_devices)
def testGradientTapeInterop(self):
with tf.GradientTape() as t:
x = np.asarray(3.0)
y = np.asarray(2.0)
t.watch([x, y])
xx = 2 * x
yy = 3 * y
dx, dy = t.gradient([xx, yy], [x, y])
self.assertIsInstance(dx, np.ndarray)
self.assertIsInstance(dy, np.ndarray)
self.assertAllClose(dx, 2.0)
self.assertAllClose(dy, 3.0)
def testGradientTapeNoneGradients(self):
y = np.asarray(2.0)
with tf.GradientTape() as t:
x = np.asarray(3.0)
t.watch([x])
z = 2 * x
dz = t.gradient(z, y)
self.assertIsNone(dz)
def testCondInterop(self):
x = np.asarray(3.0)
def fn(x):
x_plus_1 = tf.cond(x > 0, lambda: x + 1, lambda: x + 2)
x_plus_2 = tf.cond(x < 0, lambda: x + 1, lambda: x + 2)
return x_plus_1, x_plus_2
raw_x_plus_1, raw_x_plus_2 = fn(x)
fn_x_plus_1, fn_x_plus_2 = tf.function(fn)(x)
self.assertAllClose(raw_x_plus_1, x + 1)
self.assertAllClose(raw_x_plus_2, x + 2)
self.assertAllClose(fn_x_plus_1, x + 1)
self.assertAllClose(fn_x_plus_2, x + 2)
def testWhileInterop(self):
def fn():
x = np.asarray(0)
c = lambda x: x < 10000
b = lambda x: [x + 1]
return tf.while_loop(c, b, [x], parallel_iterations=20)
self.assertEqual(10000, fn()[0])
self.assertEqual(10000, tf.function(fn)()[0])
def testTensorTFNPArrayInterop(self):
arr = np.asarray(0.)
t = tf.constant(10.)
arr_plus_t = arr + t
t_plus_arr = t + arr
self.assertIsInstance(arr_plus_t, tf.Tensor)
self.assertIsInstance(t_plus_arr, tf.Tensor)
self.assertEqual(10., arr_plus_t.numpy())
self.assertEqual(10., t_plus_arr.numpy())
def testTensorTFNPOp(self):
t = tf.constant(10.)
sq = np.square(t)
self.assertIsInstance(sq, np.ndarray)
self.assertEqual(100., sq)
def testTFNPArrayTFOpInterop(self):
arr = np.asarray(10.)
# TODO(nareshmodi): Test more ops.
sq = tf.square(arr)
self.assertIsInstance(sq, tf.Tensor)
self.assertEqual(100., sq.numpy())
def testTFNPArrayNPOpInterop(self):
arr = np.asarray([10.])
# TODO(nareshmodi): Test more ops.
sq = onp.square(arr)
self.assertIsInstance(sq, onp.ndarray)
self.assertEqual(100., sq[0])
# TODO(b/171313773): why doesn't tensor have __array_module__
def testArrayModule(self):
self.skipTest("Tensor doesn't have __array_module__")
arr = np.asarray([10])
module = arr.__array_module__((tf.Tensor,))
self.assertIs(module, tf.experimental.numpy)
class Dummy:
pass
module = arr.__array_module__((tf.Tensor, Dummy))
self.assertIs(module, NotImplemented)
# TODO(nareshmodi): Fails since the autopacking code doesn't use
# nest.flatten.
# def testAutopacking(self):
# arr1 = np.asarray(1.)
# arr2 = np.asarray(2.)
# arr3 = np.asarray(3.)
# t = ops.convert_to_tensor_v2([arr1, arr2, arr3])
# self.assertEqual(t.numpy(), [1., 2., 3.])
def testDistStratInterop(self):
strategy = tf.distribute.MirroredStrategy(
devices=['CPU:0', 'CPU:1', 'CPU:2'])
multiplier = np.asarray(5.)
@tf.function
def run():
ctx = tf.distribute.get_replica_context()
val = np.asarray(ctx.replica_id_in_sync_group)
return val * multiplier
distributed_values = strategy.run(run)
reduced = strategy.reduce(
tf.distribute.ReduceOp.SUM, distributed_values, axis=None)
values = strategy.experimental_local_results(distributed_values)
# Note that this should match the number of virtual CPUs.
self.assertLen(values, 3)
self.assertIsInstance(values[0], np.ndarray)
self.assertIsInstance(values[1], np.ndarray)
self.assertIsInstance(values[2], np.ndarray)
self.assertAllClose(values[0], 0)
self.assertAllClose(values[1], 5)
self.assertAllClose(values[2], 10)
# "strategy.reduce" doesn't rewrap in ndarray.
# self.assertIsInstance(reduced, np.ndarray)
self.assertAllClose(reduced, 15)
@test_util.disable_tfrt('b/180469928')
def testPyFuncInterop(self):
def py_func_fn(a, b):
return a + b
@tf.function
def fn(a, b):
result = tf.py_function(py_func_fn, [a, b], a.dtype)
return np.asarray(result)
a = np.asarray(1.)
b = np.asarray(2.)
result = fn(a, b)
self.assertIsInstance(result, np.ndarray)
self.assertAllClose(result, 3.)
def testDatasetInterop(self):
values = [1, 2, 3, 4, 5, 6]
values_as_array = np.asarray(values)
# Tensor dataset
dataset = tf.data.Dataset.from_tensors(values_as_array)
for value, value_from_dataset in zip([values_as_array], dataset):
self.assertIsInstance(value_from_dataset, np.ndarray)
self.assertAllEqual(value_from_dataset, value)
# Tensor slice dataset
dataset = tf.data.Dataset.from_tensor_slices(values_as_array)
for value, value_from_dataset in zip(values, dataset):
self.assertIsInstance(value_from_dataset, np.ndarray)
self.assertAllEqual(value_from_dataset, value)
# # TODO(nareshmodi): as_numpy_iterator() doesn't work.
# items = list(dataset.as_numpy_iterator())
# Map over a dataset.
dataset = dataset.map(lambda x: np.add(x, 1))
for value, value_from_dataset in zip(values, dataset):
self.assertIsInstance(value_from_dataset, np.ndarray)
self.assertAllEqual(value_from_dataset, value + 1)
# Batch a dataset.
dataset = tf.data.Dataset.from_tensor_slices(values_as_array).batch(2)
for value, value_from_dataset in zip([[1, 2], [3, 4], [5, 6]], dataset):
self.assertIsInstance(value_from_dataset, np.ndarray)
self.assertAllEqual(value_from_dataset, value)
def testKerasInterop(self):
# Return an ndarray from the model.
inputs = tf.keras.layers.Input(shape=(10,))
output_layer = tf.keras.layers.Lambda(np.square)(inputs)
model = tf.keras.Model([inputs], output_layer)
values = onp.arange(10, dtype=onp.float32).reshape((1, 10))
values_as_array = np.asarray(values)
result = model(values)
self.assertIsInstance(result, np.ndarray)
self.assertAllClose(result, onp.square(values))
result = model(values_as_array)
self.assertIsInstance(result, np.ndarray)
self.assertAllClose(result, onp.square(values))
def testKerasInteropSequential(self):
class ProjectionLayer(tf.keras.layers.Layer):
"""Linear projection layer using TF NumPy."""
def __init__(self, units):
super(ProjectionLayer, self).__init__()
self._units = units
def build(self, input_shape):
stddev = np.sqrt(self._units).astype(np.float32)
initial_value = np.random.randn(input_shape[1], self._units).astype(
np.float32) / stddev
# Note that TF NumPy can interoperate with tf.Variable.
self.w = tf.Variable(initial_value, trainable=True)
def call(self, inputs):
return np.matmul(inputs, self.w)
model = tf.keras.Sequential(
[tf.keras.layers.Dense(100), ProjectionLayer(2)])
output = model.call(np.random.randn(10, 100).astype(np.float32))
self.assertIsInstance(output, np.ndarray)
dense_layer = tf.keras.layers.Dense(100)
output = dense_layer(np.random.randn(10, 100).astype(np.float32))
def testPForInterop(self):
def outer_product(a):
return np.tensordot(a, a, 0)
batch_size = 100
a = np.ones((batch_size, 32, 32))
c = tf.vectorized_map(outer_product, a)
self.assertIsInstance(c, np.ndarray)
self.assertEqual(c.shape, (batch_size, 32, 32, 32, 32))
c = tf.vectorized_map(lambda x: x.T, a)
self.assertIsInstance(c, np.ndarray)
self.assertEqual(c.shape, (batch_size, 32, 32))
def testJacobian(self):
with tf.GradientTape() as g:
x = np.asarray([1., 2.])
y = np.asarray([3., 4.])
g.watch(x)
g.watch(y)
z = x * x * y
jacobian = g.jacobian(z, [x, y])
answer = [tf.linalg.diag(2 * x * y), tf.linalg.diag(x * x)]
self.assertIsInstance(jacobian[0], np.ndarray)
self.assertIsInstance(jacobian[1], np.ndarray)
self.assertAllClose(jacobian, answer)
def testBatchJacobian(self):
with tf.GradientTape() as g:
x = np.asarray([[1., 2.], [3., 4.]])
y = np.asarray([[3., 4.], [5., 6.]])
g.watch(x)
g.watch(y)
z = x * x * y
batch_jacobian = g.batch_jacobian(z, x)
answer = tf.stack(
[tf.linalg.diag(2 * x[0] * y[0]),
tf.linalg.diag(2 * x[1] * y[1])])
self.assertIsInstance(batch_jacobian, np.ndarray)
self.assertAllClose(batch_jacobian, answer)
def testForwardprop(self):
x = np.asarray([1., 2.])
xt = np.asarray([3., 4.])
with tf.autodiff.ForwardAccumulator(x, xt) as acc:
y = x * 2.
yt = acc.jvp(y)
self.assertIsInstance(yt, np.ndarray)
self.assertAllClose([6., 8.], yt)
z = np.asarray([1.])
self.assertIsNone(acc.jvp(z))
def testMapFn(self):
x = np.asarray([1., 2.])
mapped_x = tf.map_fn(lambda x: (x[0]+1, x[1]+1), (x, x))
self.assertIsInstance(mapped_x[0], np.ndarray)
self.assertIsInstance(mapped_x[1], np.ndarray)
self.assertAllClose(mapped_x[0], [2., 3.])
self.assertAllClose(mapped_x[1], [2., 3.])
class FunctionTest(InteropTest):
def testFunctionInterop(self):
x = np.asarray(3.0)
y = np.asarray(2.0)
add = lambda x, y: x + y
add_fn = tf.function(add)
raw_result = add(x, y)
fn_result = add_fn(x, y)
self.assertIsInstance(raw_result, np.ndarray)
self.assertIsInstance(fn_result, np.ndarray)
self.assertAllClose(raw_result, fn_result)
def testLen(self):
# len can be fixed by autograph.
# TODO(wangpeng): this test can just be removed
@tf.function(autograph=False)
def f(x):
# Note that shape of input to len is data dependent.
return len(np.where(x)[0])
t = np.asarray([True, False, True])
with self.assertRaises(TypeError):
f(t)
def testIter(self):
@tf.function
def f(x):
y, z = x
return y, z
with self.assertRaises(TypeError):
f(np.asarray([3, 4]))
def testIndex(self):
@tf.function
def f(x):
return [0, 1][x]
with self.assertRaises(TypeError):
f(np.asarray([1]))
class VariableTest(InteropTest):
def test(self):
tf_var = tf.Variable(2.0)
value = np.square(tf_var)
self.assertIsInstance(value, np.ndarray)
self.assertAllClose(4.0, value)
with tf.control_dependencies([tf_var.assign_add(value)]):
tf_var_value = tf_var.read_value()
self.assertAllClose(6.0, tf_var_value)
if __name__ == '__main__':
ops.set_dtype_conversion_mode('legacy')
np_math_ops.enable_numpy_methods_on_tensor()
tf.compat.v1.enable_eager_execution()
tf.test.main()
tensorflow/python/ops/numpy_ops/np_logic_test.py
+104
View File
@@ -0,0 +1,104 @@
# 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.
# ==============================================================================
"""Tests for tf numpy logical methods."""
import numpy as np
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
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_math_ops
from tensorflow.python.platform import test
class LogicTest(test.TestCase):
def setUp(self):
super(LogicTest, self).setUp()
self.array_transforms = [
lambda x: x, # Identity,
ops.convert_to_tensor,
np.array,
lambda x: np.array(x, dtype=np.int32),
lambda x: np.array(x, dtype=np.int64),
lambda x: np.array(x, dtype=np.float32),
lambda x: np.array(x, dtype=np.float64),
np_array_ops.array,
lambda x: np_array_ops.array(x, dtype=dtypes.int32),
lambda x: np_array_ops.array(x, dtype=dtypes.int64),
lambda x: np_array_ops.array(x, dtype=dtypes.float32),
lambda x: np_array_ops.array(x, dtype=dtypes.float64),
]
def testEqual(self):
def run_test(x1, x2=None):
if x2 is None:
x2 = x1
for fn1 in self.array_transforms:
for fn2 in self.array_transforms:
arg1 = fn1(x1)
arg2 = fn2(x2)
self.match(
np_math_ops.equal(arg1, arg2),
np.equal(
make_numpy_compatible(arg1), make_numpy_compatible(arg2)))
run_test(1)
run_test(1, 2)
run_test([1, 2])
run_test([1, 2, 3], [2])
run_test([[1, 2], [3, 4]], [1, 2])
run_test([[1, 2], [1, 4]], [1, 2])
run_test([1, 2], [[1, 2], [1, 4]])
run_test([[1, 2], [3, 4]], [[1, 2], [3, 4]])
run_test([[1, 2], [3, 4]], [[1, 3], [3, 4]])
def match_shape(self, actual, expected, msg=None):
if msg:
msg = 'Shape match failed for: {}. Expected: {} Actual: {}'.format(
msg, expected.shape, actual.shape)
self.assertEqual(actual.shape, expected.shape, msg=msg)
def match_dtype(self, actual, expected, msg=None):
if msg:
msg = 'Dtype match failed for: {}. Expected: {} Actual: {}.'.format(
msg, expected.dtype, actual.dtype)
self.assertEqual(actual.dtype, expected.dtype, msg=msg)
def match(self, actual, expected, msg=None):
msg_ = 'Expected: {} Actual: {}'.format(expected, actual)
if msg:
msg = '{} {}'.format(msg_, msg)
else:
msg = msg_
self.assertIsInstance(actual, np_arrays.ndarray)
self.match_dtype(actual, expected, msg)
self.match_shape(actual, expected, msg)
if not actual.shape.rank:
self.assertEqual(actual.tolist(), expected.tolist())
else:
self.assertSequenceEqual(actual.tolist(), expected.tolist())
def make_numpy_compatible(s):
return s if not isinstance(s, np_arrays.ndarray) else s.numpy()
if __name__ == '__main__':
ops.enable_eager_execution()
np_math_ops.enable_numpy_methods_on_tensor()
test.main()
tensorflow/python/ops/numpy_ops/np_math_ops.py
+1658
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+399
View File
@@ -0,0 +1,399 @@
# 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.
# ==============================================================================
"""Tests for tf numpy mathematical methods."""
import itertools
from absl.testing import parameterized
import numpy as np
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
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_math_ops
from tensorflow.python.platform import test
class MathTest(test.TestCase, parameterized.TestCase):
def setUp(self):
super(MathTest, self).setUp()
self.array_transforms = [
lambda x: x, # Identity,
ops.convert_to_tensor,
np.array,
lambda x: np.array(x, dtype=np.float32),
lambda x: np.array(x, dtype=np.float64),
np_array_ops.array,
lambda x: np_array_ops.array(x, dtype=np.float32),
lambda x: np_array_ops.array(x, dtype=np.float64),
]
self.types = [np.int32, np.int64, np.float32, np.float64]
def _testBinaryOp(self,
math_fun,
np_fun,
name,
operands=None,
extra_operands=None,
check_promotion=True,
check_promotion_result_type=True):
def run_test(a, b):
for fn in self.array_transforms:
arg1 = fn(a)
arg2 = fn(b)
self.match(
math_fun(arg1, arg2),
np_fun(arg1, arg2),
msg='{}({}, {})'.format(name, arg1, arg2))
# Tests type promotion
for type_a in self.types:
for type_b in self.types:
if not check_promotion and type_a != type_b:
continue
arg1 = np_array_ops.array(a, dtype=type_a)
arg2 = np_array_ops.array(b, dtype=type_b)
self.match(
math_fun(arg1, arg2),
np_fun(arg1, arg2),
msg='{}({}, {})'.format(name, arg1, arg2),
check_dtype=check_promotion_result_type)
if operands is None:
operands = [(5, 2), (5, [2, 3]), (5, [[2, 3], [6, 7]]), ([1, 2, 3], 7),
([1, 2, 3], [5, 6, 7])]
for operand1, operand2 in operands:
run_test(operand1, operand2)
if extra_operands is not None:
for operand1, operand2 in extra_operands:
run_test(operand1, operand2)
def testDot(self):
extra_operands = [([1, 2], [[5, 6, 7], [8, 9, 10]]),
(np.arange(2 * 3 * 5).reshape([2, 3, 5]).tolist(),
np.arange(5 * 7 * 11).reshape([7, 5, 11]).tolist())]
return self._testBinaryOp(
np_math_ops.dot, np.dot, 'dot', extra_operands=extra_operands)
def testMinimum(self):
# The numpy version has strange result type when promotion happens,
# so set check_promotion_result_type to False.
return self._testBinaryOp(
np_math_ops.minimum,
np.minimum,
'minimum',
check_promotion_result_type=False)
def testMaximum(self):
# The numpy version has strange result type when promotion happens,
# so set check_promotion_result_type to False.
return self._testBinaryOp(
np_math_ops.maximum,
np.maximum,
'maximum',
check_promotion_result_type=False)
def testMatmul(self):
operands = [([[1, 2]], [[3, 4, 5], [6, 7, 8]])]
return self._testBinaryOp(
np_math_ops.matmul, np.matmul, 'matmul', operands=operands)
def testMatmulError(self):
with self.assertRaisesRegex(ValueError, r''):
np_math_ops.matmul(
np_array_ops.ones([], np.int32), np_array_ops.ones([2, 3], np.int32))
with self.assertRaisesRegex(ValueError, r''):
np_math_ops.matmul(
np_array_ops.ones([2, 3], np.int32), np_array_ops.ones([], np.int32))
def testVDot(self):
operands = [([[1, 2], [3, 4]], [[3, 4], [6, 7]]),
([[1, 2], [3, 4]], [3, 4, 6, 7])]
return self._testBinaryOp(
np_math_ops.vdot, np.vdot, 'vdot', operands=operands)
def testLcm(self):
a = np_array_ops.array(6, dtype=np.int8)
b = np_array_ops.array(22, dtype=np.int8)
res_tf = np_math_ops.lcm(a, b)
res_np = np.lcm(np.array(a), np.array(b))
self.assertEqual(res_tf, res_np)
def _testUnaryOp(self, math_fun, np_fun, name):
def run_test(a):
for fn in self.array_transforms:
arg1 = fn(a)
self.match(
math_fun(arg1), np_fun(arg1), msg='{}({})'.format(name, arg1))
run_test(5)
run_test([2, 3])
run_test([[2, -3], [-6, 7]])
def testLog(self):
self._testUnaryOp(np_math_ops.log, np.log, 'log')
def testExp(self):
self._testUnaryOp(np_math_ops.exp, np.exp, 'exp')
def testTanh(self):
self._testUnaryOp(np_math_ops.tanh, np.tanh, 'tanh')
def testSqrt(self):
self._testUnaryOp(np_math_ops.sqrt, np.sqrt, 'sqrt')
def match(self, actual, expected, msg='', check_dtype=True):
self.assertIsInstance(actual, np_arrays.ndarray)
if check_dtype:
self.assertEqual(
actual.dtype, expected.dtype,
'Dtype mismatch.\nActual: {}\nExpected: {}\n{}'.format(
actual.dtype.as_numpy_dtype, expected.dtype, msg))
self.assertEqual(
actual.shape, expected.shape,
'Shape mismatch.\nActual: {}\nExpected: {}\n{}'.format(
actual.shape, expected.shape, msg))
np.testing.assert_allclose(actual.tolist(), expected.tolist(), rtol=1e-6)
def testArgsort(self):
self._testUnaryOp(np_math_ops.argsort, np.argsort, 'argsort')
# Test stability
r = np.arange(100)
a = np.zeros(100)
np.testing.assert_equal(np_math_ops.argsort(a, kind='stable'), r)
def testArgsortRaisesErrorForComplexDtypes(self):
"""Test that argsort raises TypeError for complex64 and complex128."""
complex64_array = np.array([1 + 2j, 3 + 4j, 5 + 6j], dtype=np.complex64)
with self.assertRaisesRegex(
TypeError, 'argsort does not support complex64/complex128 dtypes'
):
np_math_ops.argsort(complex64_array)
complex128_array = np.array([1 + 2j, 3 + 4j, 5 + 6j], dtype=np.complex128)
with self.assertRaisesRegex(
TypeError, 'argsort does not support complex64/complex128 dtypes'
):
np_math_ops.argsort(complex128_array)
def testArgMaxArgMin(self):
data = [
0,
5,
[1],
[1, 2, 3],
[[1, 2, 3]],
[[4, 6], [7, 8]],
[[[4, 6], [9, 10]], [[7, 8], [12, 34]]],
]
for fn, d in itertools.product(self.array_transforms, data):
arr = fn(d)
self.match(np_math_ops.argmax(arr), np.argmax(arr))
self.match(np_math_ops.argmin(arr), np.argmin(arr))
if hasattr(arr, 'shape'):
ndims = len(arr.shape)
else:
ndims = np_array_ops.array(arr, copy=False).ndim
if ndims == 0:
# Numpy flattens the scalar ndarray and treats it as a 1-d array of
# size 1.
ndims = 1
for axis in range(-ndims, ndims):
self.match(
np_math_ops.argmax(arr, axis=axis), np.argmax(arr, axis=axis))
self.match(
np_math_ops.argmin(arr, axis=axis), np.argmin(arr, axis=axis))
@parameterized.parameters([False, True])
def testIsCloseEqualNan(self, equal_nan):
a = np.asarray([1, 1, np.nan, 1, np.nan], np.float32)
b = np.asarray([1, 2, 1, np.nan, np.nan], np.float32)
self.match(
np_math_ops.isclose(a, b, equal_nan=equal_nan),
np.isclose(a, b, equal_nan=equal_nan))
def testAverageWrongShape(self):
with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):
np_math_ops.average(np.ones([2, 3]), weights=np.ones([2, 4]))
with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):
np_math_ops.average(np.ones([2, 3]), axis=0, weights=np.ones([2, 4]))
with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):
np_math_ops.average(np.ones([2, 3]), axis=0, weights=np.ones([]))
with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):
np_math_ops.average(np.ones([2, 3]), axis=0, weights=np.ones([5]))
def testClip(self):
def run_test(arr, *args, **kwargs):
check_dtype = kwargs.pop('check_dtype', True)
for fn in self.array_transforms:
arr = fn(arr)
self.match(
np_math_ops.clip(arr, *args, **kwargs),
np.clip(arr, *args, **kwargs),
check_dtype=check_dtype)
# NumPy exhibits weird typing behavior when a/a_min/a_max are scalars v/s
# lists, e.g.,
#
# np.clip(np.array(0, dtype=np.int32), -5, 5).dtype == np.int64
# np.clip(np.array([0], dtype=np.int32), -5, 5).dtype == np.int32
# np.clip(np.array([0], dtype=np.int32), [-5], [5]).dtype == np.int64
#
# So we skip matching type. In tf-numpy the type of the output array is
# always the same as the input array.
run_test(0, -1, 5, check_dtype=False)
run_test(-1, -1, 5, check_dtype=False)
run_test(5, -1, 5, check_dtype=False)
run_test(-10, -1, 5, check_dtype=False)
run_test(10, -1, 5, check_dtype=False)
run_test(10, None, 5, check_dtype=False)
run_test(10, -1, None, check_dtype=False)
run_test([0, 20, -5, 4], -1, 5, check_dtype=False)
run_test([0, 20, -5, 4], None, 5, check_dtype=False)
run_test([0, 20, -5, 4], -1, None, check_dtype=False)
run_test([0.5, 20.2, -5.7, 4.4], -1.5, 5.1, check_dtype=False)
run_test([0, 20, -5, 4], [-5, 0, -5, 0], [0, 5, 0, 5], check_dtype=False)
run_test([[1, 2, 3], [4, 5, 6]], [2, 0, 2], 5, check_dtype=False)
run_test([[1, 2, 3], [4, 5, 6]], 0, [5, 3, 1], check_dtype=False)
def testPtp(self):
def run_test(arr, *args, **kwargs):
for fn in self.array_transforms:
arg = fn(arr)
self.match(
np_math_ops.ptp(arg, *args, **kwargs), np.ptp(arg, *args, **kwargs))
run_test([1, 2, 3])
run_test([1., 2., 3.])
run_test([[1, 2], [3, 4]], axis=1)
run_test([[1, 2], [3, 4]], axis=0)
run_test([[1, 2], [3, 4]], axis=-1)
run_test([[1, 2], [3, 4]], axis=-2)
def testLinSpace(self):
array_transforms = [
lambda x: x, # Identity,
ops.convert_to_tensor,
np.array,
lambda x: np.array(x, dtype=np.float32),
lambda x: np.array(x, dtype=np.float64),
np_array_ops.array,
lambda x: np_array_ops.array(x, dtype=np.float32),
lambda x: np_array_ops.array(x, dtype=np.float64)
]
def run_test(start, stop, **kwargs):
for fn1 in array_transforms:
for fn2 in array_transforms:
arg1 = fn1(start)
arg2 = fn2(stop)
self.match(
np_math_ops.linspace(arg1, arg2, **kwargs),
np.linspace(arg1, arg2, **kwargs),
msg='linspace({}, {})'.format(arg1, arg2))
run_test(0, 1)
run_test(0, 1, num=10)
run_test(0, 1, endpoint=False)
run_test(0, -1)
run_test(0, -1, num=10)
run_test(0, -1, endpoint=False)
def testLogSpace(self):
array_transforms = [
lambda x: x, # Identity,
ops.convert_to_tensor,
np.array,
lambda x: np.array(x, dtype=np.float32),
lambda x: np.array(x, dtype=np.float64),
np_array_ops.array,
lambda x: np_array_ops.array(x, dtype=np.float32),
lambda x: np_array_ops.array(x, dtype=np.float64)
]
def run_test(start, stop, **kwargs):
for fn1 in array_transforms:
for fn2 in array_transforms:
arg1 = fn1(start)
arg2 = fn2(stop)
self.match(
np_math_ops.logspace(arg1, arg2, **kwargs),
np.logspace(arg1, arg2, **kwargs),
msg='logspace({}, {})'.format(arg1, arg2))
run_test(0, 5)
run_test(0, 5, num=10)
run_test(0, 5, endpoint=False)
run_test(0, 5, base=2.0)
run_test(0, -5)
run_test(0, -5, num=10)
run_test(0, -5, endpoint=False)
run_test(0, -5, base=2.0)
def testGeomSpace(self):
def run_test(start, stop, **kwargs):
arg1 = start
arg2 = stop
self.match(
np_math_ops.geomspace(arg1, arg2, **kwargs),
np.geomspace(arg1, arg2, **kwargs),
msg='geomspace({}, {})'.format(arg1, arg2))
run_test(1, 1000, num=5)
run_test(1, 1000, num=5, endpoint=False)
run_test(-1, -1000, num=5)
run_test(-1, -1000, num=5, endpoint=False)
@parameterized.parameters([
'T', 'ndim', 'size', 'data', '__pos__', '__round__', 'tolist', 'flatten',
'transpose', 'reshape', 'ravel', 'clip', 'astype', 'max', 'mean', 'min'])
def testNumpyMethodsOnTensor(self, np_method):
a = ops.convert_to_tensor([1, 2])
self.assertTrue(hasattr(a, np_method))
def testFlatten(self):
a1 = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
a2 = ops.convert_to_tensor(a1)
self.assertAllEqual(a1.flatten('C'), a2.flatten('C'))
self.assertAllEqual(a1.flatten('F'), a2.flatten('F'))
self.assertAllEqual(a1.flatten('C'), a2.flatten('A'))
self.assertAllEqual(a1.flatten('C'), a2.flatten('K'))
with self.assertRaises(ValueError):
a2.flatten('invalid')
def testIsInf(self):
x1 = ops.convert_to_tensor(-2147483648)
x2 = ops.convert_to_tensor(2147483647)
self.assertFalse(np_math_ops.isinf(x1))
self.assertFalse(np_math_ops.isinf(x2))
self.assertFalse(np_math_ops.isposinf(x1))
self.assertFalse(np_math_ops.isposinf(x2))
self.assertFalse(np_math_ops.isneginf(x1))
self.assertFalse(np_math_ops.isneginf(x2))
if __name__ == '__main__':
tensor.enable_tensor_equality()
ops.enable_eager_execution()
ops.set_dtype_conversion_mode('legacy')
np_math_ops.enable_numpy_methods_on_tensor()
test.main()
tensorflow/python/ops/numpy_ops/np_random.py
+137
View File
@@ -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
)
tensorflow/python/ops/numpy_ops/np_random_test.py
+244
View File
@@ -0,0 +1,244 @@
# 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.
# ==============================================================================
"""Tests for tf numpy random number methods."""
# pylint: disable=g-direct-tensorflow-import
from absl.testing import parameterized
import numpy as onp
from tensorflow.python.framework import ops
# Needed for ndarray.reshape.
from tensorflow.python.ops.numpy_ops import np_array_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.ops.numpy_ops import np_random
from tensorflow.python.platform import test
class SeedTest(test.TestCase):
def test(self):
np_random.seed(1)
np_random.seed(np_dtypes.int32(1))
with self.assertRaises(ValueError):
np_random.seed((1, 3))
class RandomTestBase(test.TestCase, parameterized.TestCase):
def _test(self, *args, **kw_args):
onp_dtype = kw_args.pop('onp_dtype', None)
allow_float64 = kw_args.pop('allow_float64', True)
old_allow_float64 = np_dtypes.is_allow_float64()
np_dtypes.set_allow_float64(allow_float64)
old_func = getattr(self, 'onp_func', None)
# TODO(agarwal): Note that onp can return a scalar type while np returns
# ndarrays. Currently np does not support scalar types.
self.onp_func = lambda *args, **kwargs: onp.asarray( # pylint: disable=g-long-lambda
old_func(*args, **kwargs))
np_out = self.np_func(*args, **kw_args)
onp_out = onp.asarray(self.onp_func(*args, **kw_args))
if onp_dtype is not None:
onp_out = onp_out.astype(onp_dtype)
self.assertEqual(np_out.shape, onp_out.shape)
self.assertEqual(np_out.dtype, onp_out.dtype)
np_dtypes.set_allow_float64(old_allow_float64)
class RandNTest(RandomTestBase):
def setUp(self):
self.np_func = np_random.randn
self.onp_func = onp.random.randn
super(RandNTest, self).setUp()
@parameterized.parameters((), (2), (2, 3))
def test_float64(self, *dims):
self._test(*dims)
@parameterized.parameters((), (2), ((2,)), (2, 3))
def test_float32(self, *dims):
self._test(*dims, allow_float64=False, onp_dtype=np_dtypes.float32)
class StandardNormalTest(RandomTestBase):
def setUp(self):
self.np_func = np_random.standard_normal
self.onp_func = onp.random.standard_normal
super(StandardNormalTest, self).setUp()
@parameterized.parameters((None,), ((),), ((1,),), ((1, 2),))
def test(self, size):
self._test(size)
class UniformTest(RandomTestBase):
def setUp(self):
self.np_func = np_random.uniform
self.onp_func = onp.random.uniform
super(UniformTest, self).setUp()
@parameterized.parameters(
((), (), None),
(1, (), None),
((), 1, None),
(1, 1, None),
((1, 2), (2, 1), None),
((1, 2, 1), (2, 1, 1), (2, 2, 2)),
((), (), (2, 2, 2)),
)
def test_broadcast(self, low_shape, high_shape, size):
low = np_array_ops.zeros(low_shape).astype(np_dtypes.float64)
high = np_array_ops.ones(high_shape).astype(np_dtypes.float64)
self._test(low=low, high=high, size=size)
def test_float32(self):
self._test(0, 1, (1, 2), allow_float64=False, onp_dtype=np_dtypes.float32)
def test_dtype_cast(self):
self._test(np_dtypes.int8(0), np_dtypes.uint8(1), (1, 2))
class PoissonTest(RandomTestBase):
def setUp(self):
def np_wrapper(lam, size):
return np_random.poisson(lam, size).astype(np_dtypes.int64)
self.np_func = np_wrapper
self.onp_func = onp.random.poisson
super(PoissonTest, self).setUp()
@parameterized.parameters(
(1.0, None),
(1.0, 1),
(2.0, (3, 3)),
)
def test(self, lam, size):
self._test(lam, size, onp_dtype=np_dtypes.int64)
class RandomTest(RandomTestBase):
def setUp(self):
self.np_func = np_random.random
self.onp_func = onp.random.random
super(RandomTest, self).setUp()
@parameterized.parameters((None,), ((),), ((1,),), ((1, 2),))
def test(self, size):
self._test(size)
class RandTest(RandomTestBase):
def setUp(self):
self.np_func = np_random.rand
self.onp_func = onp.random.rand
super(RandTest, self).setUp()
@parameterized.parameters((), (1,), (1, 2))
def test(self, *size):
self._test(*size)
class RandIntTest(RandomTestBase):
def setUp(self):
self.np_func = np_random.randint
self.onp_func = onp.random.randint
super(RandIntTest, self).setUp()
@parameterized.parameters(
(0, 1, None, 'l'),
(0, 1, None, np_dtypes.int64),
(0, 1, 2, np_dtypes.int32),
(0, 1, (), np_dtypes.int32),
(0, 1, (2), np_dtypes.int64),
(0, 1, (2, 2), 'l'),
)
def test(self, low, high, size, dtype):
self._test(low, high, size=size, dtype=dtype)
class RandNDistriutionTest(test.TestCase):
def assertNotAllClose(self, a, b, **kwargs):
try:
self.assertAllClose(a, b, **kwargs)
except AssertionError:
return
raise AssertionError(
'The two values are close at all %d elements' % np_array_ops.size(a)
)
def testDistribution(self):
def run_test(*args):
num_samples = 1000
tol = 0.1 # High tolerance to keep the # of samples low else the test
# takes a long time to run.
np_random.seed(10)
outputs = [np_random.randn(*args) for _ in range(num_samples)]
# Test output shape.
for output in outputs:
self.assertEqual(output.shape, tuple(args))
default_dtype = (
np_dtypes.float64
if np_dtypes.is_allow_float64()
else np_dtypes.float32
)
self.assertEqual(output.dtype.as_numpy_dtype, default_dtype)
if np_array_ops.prod(args): # Don't bother with empty arrays.
outputs = [output.tolist() for output in outputs]
# Test that the properties of normal distribution are satisfied.
mean = np_array_ops.mean(outputs, axis=0)
stddev = np_array_ops.std(outputs, axis=0)
self.assertAllClose(mean, np_array_ops.zeros(args), atol=tol)
self.assertAllClose(stddev, np_array_ops.ones(args), atol=tol)
# Test that outputs are different with different seeds.
np_random.seed(20)
diff_seed_outputs = [
np_random.randn(*args).tolist() for _ in range(num_samples)
]
self.assertNotAllClose(outputs, diff_seed_outputs)
# Test that outputs are the same with the same seed.
np_random.seed(10)
same_seed_outputs = [
np_random.randn(*args).tolist() for _ in range(num_samples)
]
self.assertAllClose(outputs, same_seed_outputs)
run_test()
run_test(0)
run_test(1)
run_test(5)
run_test(2, 3)
run_test(0, 2, 3)
run_test(2, 0, 3)
run_test(2, 3, 0)
run_test(2, 3, 5)
if __name__ == '__main__':
ops.enable_eager_execution()
np_math_ops.enable_numpy_methods_on_tensor()
test.main()
tensorflow/python/ops/numpy_ops/np_utils.py
+717
View File
@@ -0,0 +1,717 @@
# 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 hasattr(x, 'dtype') and isinstance(x.dtype, np.dtype):
return x.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)
tensorflow/python/ops/numpy_ops/np_utils_test.py
+187
View File
@@ -0,0 +1,187 @@
# 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.
# ==============================================================================
"""Tests for utils.py."""
from absl.testing import parameterized
from tensorflow.python.framework import dtypes
from tensorflow.python.ops.numpy_ops import np_utils
from tensorflow.python.platform import test
class UtilsTest(test.TestCase, parameterized.TestCase):
def setUp(self):
super(UtilsTest, self).setUp()
self._old_np_doc_form = np_utils.get_np_doc_form()
self._old_is_sig_mismatch_an_error = np_utils.is_sig_mismatch_an_error()
def tearDown(self):
np_utils.set_np_doc_form(self._old_np_doc_form)
np_utils.set_is_sig_mismatch_an_error(self._old_is_sig_mismatch_an_error)
super(UtilsTest, self).tearDown()
# pylint: disable=unused-argument
def testNpDocInlined(self):
def np_fun(x, y, z):
"""np_fun docstring."""
return
np_utils.set_np_doc_form('inlined')
@np_utils.np_doc(None, np_fun=np_fun, unsupported_params=['x'])
def f(x, z):
"""f docstring."""
return
expected = """TensorFlow variant of NumPy's `np_fun`.
Unsupported arguments: `x`, `y`.
f docstring.
Documentation for `numpy.np_fun`:
np_fun docstring."""
self.assertEqual(expected, f.__doc__)
@parameterized.named_parameters([
(version, version, link) for version, link in # pylint: disable=g-complex-comprehension
[('dev',
'https://numpy.org/devdocs/reference/generated/numpy.np_fun.html'),
('stable',
'https://numpy.org/doc/stable/reference/generated/numpy.np_fun.html'),
('1.16',
'https://numpy.org/doc/1.16/reference/generated/numpy.np_fun.html')
]])
def testNpDocLink(self, version, link):
def np_fun(x, y, z):
"""np_fun docstring."""
return
np_utils.set_np_doc_form(version)
@np_utils.np_doc(None, np_fun=np_fun, unsupported_params=['x'])
def f(x, z):
"""f docstring."""
return
expected = """TensorFlow variant of NumPy's `np_fun`.
Unsupported arguments: `x`, `y`.
f docstring.
See the NumPy documentation for [`numpy.np_fun`](%s)."""
expected = expected % (link)
self.assertEqual(expected, f.__doc__)
@parameterized.parameters([None, 1, 'a', '1a', '1.1a', '1.1.1a'])
def testNpDocInvalid(self, invalid_flag):
def np_fun(x, y, z):
"""np_fun docstring."""
return
np_utils.set_np_doc_form(invalid_flag)
@np_utils.np_doc(None, np_fun=np_fun, unsupported_params=['x'])
def f(x, z):
"""f docstring."""
return
expected = """TensorFlow variant of NumPy's `np_fun`.
Unsupported arguments: `x`, `y`.
f docstring.
"""
self.assertEqual(expected, f.__doc__)
def testNpDocName(self):
np_utils.set_np_doc_form('inlined')
@np_utils.np_doc('foo')
def f():
"""f docstring."""
return
expected = """TensorFlow variant of NumPy's `foo`.
f docstring.
"""
self.assertEqual(expected, f.__doc__)
def testDtypeOfTensorLikeClass(self):
class TensorLike:
def __init__(self, dtype):
self._dtype = dtype
@property
def is_tensor_like(self):
return True
@property
def dtype(self):
return self._dtype
t = TensorLike(dtypes.float32)
self.assertEqual(np_utils._maybe_get_dtype(t), dtypes.float32)
# pylint: disable=unused-variable
def testSigMismatchIsError(self):
"""Tests that signature mismatch is an error (when configured so)."""
if not np_utils._supports_signature():
self.skipTest('inspect.signature not supported')
np_utils.set_is_sig_mismatch_an_error(True)
def np_fun(x, y=1, **kwargs):
return
with self.assertRaisesRegex(TypeError, 'Cannot find parameter'):
@np_utils.np_doc(None, np_fun=np_fun)
def f1(a):
return
with self.assertRaisesRegex(TypeError, 'is of kind'):
@np_utils.np_doc(None, np_fun=np_fun)
def f2(x, kwargs):
return
with self.assertRaisesRegex(
TypeError, 'Parameter y should have a default value'):
@np_utils.np_doc(None, np_fun=np_fun)
def f3(x, y):
return
def testSigMismatchIsNotError(self):
"""Tests that signature mismatch is not an error (when configured so)."""
np_utils.set_is_sig_mismatch_an_error(False)
def np_fun(x, y=1, **kwargs):
return
# The following functions all have signature mismatches, but they shouldn't
# throw errors when is_sig_mismatch_an_error() is False.
@np_utils.np_doc(None, np_fun=np_fun)
def f1(a):
return
def f2(x, kwargs):
return
@np_utils.np_doc(None, np_fun=np_fun)
def f3(x, y):
return
# pylint: enable=unused-variable
if __name__ == '__main__':
test.main()
tensorflow/python/ops/numpy_ops/tests/BUILD
+284
View File
@@ -0,0 +1,284 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
licenses(["notice"])
py_library(
name = "config",
srcs = ["config.py"],
strict_deps = True,
deps = [
"@absl_py//absl:app",
"@absl_py//absl/flags",
],
)
py_library(
name = "test_util",
srcs = ["test_util.py"],
strict_deps = True,
deps = [
":config",
":extensions",
"//tensorflow:tensorflow_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:gradient_checker_v2",
"//tensorflow/python/ops/numpy_ops:np_array_ops",
"//tensorflow/python/ops/numpy_ops:np_utils",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:numpy_compat",
"//third_party/py/numpy",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "np_wrapper",
srcs = ["np_wrapper.py"],
strict_deps = True,
visibility = [
"//visibility:public",
],
deps = [
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops/numpy_ops:np_array_ops",
"//tensorflow/python/ops/numpy_ops:np_arrays",
"//tensorflow/python/ops/numpy_ops:np_config",
"//tensorflow/python/ops/numpy_ops:np_dtypes",
"//tensorflow/python/ops/numpy_ops:np_math_ops",
"//tensorflow/python/ops/numpy_ops:np_random",
"//tensorflow/python/ops/numpy_ops:np_utils",
"//tensorflow/python/ops/numpy_ops:numpy",
"//third_party/py/numpy",
],
)
py_library(
name = "extensions",
srcs = ["extensions.py"],
strict_deps = True,
deps = [
":np_wrapper",
"//tensorflow:tensorflow_py",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager/polymorphic_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device_spec",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_conversion_registry",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:bitwise_ops_gen",
"//tensorflow/python/ops:clip_ops",
"//tensorflow/python/ops:collective_ops_gen",
"//tensorflow/python/ops:control_flow_assert",
"//tensorflow/python/ops:custom_gradient",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:stateless_random_ops",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/ops/parallel_for:control_flow_ops",
"//tensorflow/python/tpu:tpu_py",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
"@six_archive//:six",
],
)
# copybara:uncomment_begin(google-only)
# py_test(
# name = "extensions_test",
# srcs = ["extensions_test.py"],
# strict_deps = True,
# tags = [
# "gpu",
# "no_pip",
# "notap", # b/294137902
# "requires-gpu-nvidia",
# ],
# deps = [
# ":extensions",
# ":np_wrapper",
# "//learning/brain/research/jax:gpu_support",
# # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
# "//third_party/py/jax",
# "//third_party/py/numpy",
# "//tensorflow:tensorflow_py",
# "//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
# "//tensorflow/python/eager:backprop",
# "//tensorflow/python/eager:context",
# "//tensorflow/python/eager/polymorphic_function",
# "//tensorflow/python/framework:config",
# "//tensorflow/python/framework:constant_op",
# "//tensorflow/python/framework:dtypes",
# "//tensorflow/python/framework:ops",
# "//tensorflow/python/framework:tensor",
# "//tensorflow/python/framework:tensor_shape",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/ops:gradient_checker_v2",
# "//tensorflow/python/ops:math_ops",
# "//tensorflow/python/ops:nn_ops",
# "//tensorflow/python/ops:random_ops",
# "//tensorflow/python/ops:stateful_random_ops",
# "//tensorflow/python/ops/numpy_ops:numpy",
# "//tensorflow/python/platform:client_testlib",
# "//tensorflow/python/util:nest",
# "@absl_py//absl/flags",
# "@absl_py//absl/testing:parameterized",
# ],
# )
#
# py_test(
# name = "extensions_test_tpu",
# srcs = ["extensions_test.py"],
# args = [
# "--jax_allow_unused_tpus",
# "--requires_tpu",
# ],
# main = "extensions_test.py",
# strict_deps = True,
# tags = [
# "no_pip",
# "requires-tpu",
# ],
# deps = [
# ":extensions",
# ":np_wrapper",
# "//learning/brain/google/xla",
# # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
# "//third_party/py/jax",
# "//third_party/py/numpy",
# "//tensorflow:tensorflow_py",
# "//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
# "//tensorflow/python/eager:backprop",
# "//tensorflow/python/eager:context",
# "//tensorflow/python/eager/polymorphic_function",
# "//tensorflow/python/framework:config",
# "//tensorflow/python/framework:constant_op",
# "//tensorflow/python/framework:dtypes",
# "//tensorflow/python/framework:ops",
# "//tensorflow/python/framework:tensor",
# "//tensorflow/python/framework:tensor_shape",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/ops:gradient_checker_v2",
# "//tensorflow/python/ops:math_ops",
# "//tensorflow/python/ops:nn_ops",
# "//tensorflow/python/ops:random_ops",
# "//tensorflow/python/ops:stateful_random_ops",
# "//tensorflow/python/ops/numpy_ops:numpy",
# "//tensorflow/python/platform:client_testlib",
# "//tensorflow/python/util:nest",
# "@absl_py//absl/flags",
# "@absl_py//absl/testing:parameterized",
# ],
# )
# copybara:uncomment_end
py_test(
name = "np_test",
timeout = "long",
srcs = ["np_test.py"],
args = [
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 20,
strict_deps = True,
tags = [
"gpu",
"no_pip",
"requires-gpu-nvidia",
],
deps = [
":config",
":extensions",
":np_wrapper",
":test_util",
# copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
"//third_party/py/numpy",
"@six_archive//:six",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops/numpy_ops:np_config",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:numpy_compat",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
name = "np_indexing_test",
srcs = ["np_indexing_test.py"],
args = [
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 10,
strict_deps = True,
# TODO(b/164245103): Re-enable GPU once tf.tensor_strided_slice_update's segfault is fixed.
tags = [
"no_pip",
# "gpu",
# "requires-gpu-nvidia",
],
deps = [
":config",
":extensions",
":np_wrapper",
":test_util",
# copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
"//third_party/py/numpy",
"//tensorflow/python/framework:config",
"//tensorflow/python/ops/numpy_ops:numpy",
"//tensorflow/python/util:nest",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
name = "np_einsum_test",
srcs = ["np_einsum_test.py"],
args = [
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 20,
strict_deps = True,
tags = [
"gpu",
"no_pip",
"requires-gpu-nvidia",
],
deps = [
":config",
":np_wrapper",
":test_util",
# copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
"//third_party/py/numpy",
"//tensorflow/python/ops/numpy_ops:np_config",
"//tensorflow/python/ops/numpy_ops:numpy",
"@absl_py//absl/testing:absltest",
"@absl_py//absl/testing:parameterized",
],
)
tensorflow/python/ops/numpy_ops/tests/config.py
+141
View File
@@ -0,0 +1,141 @@
# Copyright 2023 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.
# ==============================================================================
"""Test configurations."""
import os
import sys
def bool_env(varname: str, default: bool) -> bool:
"""Read an environment variable and interpret it as a boolean.
True values are (case insensitive): 'y', 'yes', 't', 'true', 'on', and '1';
false values are 'n', 'no', 'f', 'false', 'off', and '0'.
Args:
varname: the name of the variable
default: the default boolean value
Raises: ValueError if the environment variable is anything else.
"""
val = os.getenv(varname, str(default))
val = val.lower()
if val in ('y', 'yes', 't', 'true', 'on', '1'):
return True
elif val in ('n', 'no', 'f', 'false', 'off', '0'):
return False
else:
raise ValueError(
'invalid truth value %r for environment %r' % (val, varname)
)
class Config(object):
def __init__(self):
self.values = {}
self.meta = {}
self.FLAGS = NameSpace(self.read)
self.use_absl = False
def update(self, name, val):
if self.use_absl:
setattr(self.absl_flags.FLAGS, name, val)
else:
self.check_exists(name)
if name not in self.values:
raise Exception("Unrecognized config option: {}".format(name))
self.values[name] = val
def read(self, name):
if self.use_absl:
return getattr(self.absl_flags.FLAGS, name)
else:
self.check_exists(name)
return self.values[name]
def add_option(self, name, default, opt_type, meta_args, meta_kwargs):
if name in self.values:
raise Exception("Config option {} already defined".format(name))
self.values[name] = default
self.meta[name] = (opt_type, meta_args, meta_kwargs)
def check_exists(self, name):
if name not in self.values:
raise Exception("Unrecognized config option: {}".format(name))
def DEFINE_bool(self, name, default, *args, **kwargs):
self.add_option(name, default, bool, args, kwargs)
def DEFINE_integer(self, name, default, *args, **kwargs):
self.add_option(name, default, int, args, kwargs)
def DEFINE_string(self, name, default, *args, **kwargs):
self.add_option(name, default, str, args, kwargs)
def DEFINE_enum(self, name, default, *args, **kwargs):
self.add_option(name, default, 'enum', args, kwargs)
def config_with_absl(self):
# Run this before calling `app.run(main)` etc
import absl.flags as absl_FLAGS
from absl import app, flags as absl_flags
self.use_absl = True
self.absl_flags = absl_flags
absl_defs = { bool: absl_flags.DEFINE_bool,
int: absl_flags.DEFINE_integer,
str: absl_flags.DEFINE_string,
'enum': absl_flags.DEFINE_enum }
for name, val in self.values.items():
flag_type, meta_args, meta_kwargs = self.meta[name]
absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)
app.call_after_init(lambda: self.complete_absl_config(absl_flags))
def complete_absl_config(self, absl_flags):
for name, _ in self.values.items():
self.update(name, getattr(absl_flags.FLAGS, name))
def parse_flags_with_absl(self):
global already_configured_with_absl
if not already_configured_with_absl:
import absl.flags
self.config_with_absl()
absl.flags.FLAGS(sys.argv, known_only=True)
self.complete_absl_config(absl.flags)
already_configured_with_absl = True
class NameSpace(object):
def __init__(self, getter):
self._getter = getter
def __getattr__(self, name):
return self._getter(name)
config = Config()
flags = config
FLAGS = flags.FLAGS
already_configured_with_absl = False
flags.DEFINE_bool(
'jax_enable_checks',
bool_env('JAX_ENABLE_CHECKS', False),
help='Turn on invariant checking (core.skip_checks = False)')
flags.DEFINE_bool('tf_numpy_additional_tests', True,
'Run tests added specifically for TF numpy')
tensorflow/python/ops/numpy_ops/tests/extensions.py
+2094
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/tests/np_einsum_test.py
+355
View File
@@ -0,0 +1,355 @@
# Copyright 2023 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.
# ==============================================================================
from collections import defaultdict # pylint: disable=g-importing-member
import itertools
from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from tensorflow.python.ops.numpy_ops.tests.config import config
import tensorflow.python.ops.numpy_ops.tests.np_wrapper as tnp
import tensorflow.python.ops.numpy_ops.tests.test_util as tntu
config.parse_flags_with_absl()
class EinsumTest(tntu.TestCase):
def _check(self, s, *ops):
a = np.einsum(s, *ops)
b = tnp.einsum(s, *ops)
self.assertAllClose(a, b, check_dtypes=True, atol=1e-4, rtol=1e-4)
def test_three_operands_1(self):
r = self.rng()
x = r.randn(3)
y = r.randn(4)
z = r.randn(5)
s = 'i,j,k->ijk'
self._check(s, x, y, z)
def test_three_operands_2(self):
r = self.rng()
x = r.randn(3)
y = r.randn(4)
z = r.randn(5)
s = 'i,j,k->ijk'
self._check(s, x, y, z)
def test_two_operands_1(self):
r = self.rng()
x = r.randn(3, 4)
y = r.randn(4)
s = 'ij,j->i'
self._check(s, x, y)
def test_two_operands_2(self):
r = self.rng()
x = r.randn(3, 4, 5)
y = r.randn(4)
s = 'ijk,j->i'
self._check(s, x, y)
def test_two_operands_3(self):
r = self.rng()
x = r.randn(3, 4, 3)
y = r.randn(3)
s = 'iji,i->j'
self._check(s, x, y)
def test_two_operands_4(self):
r = self.rng()
x = r.randn(3, 4)
y = r.randn(3, 4)
s = 'ij,ij->'
self._check(s, x, y)
def test_two_operands_5(self):
r = self.rng()
x = r.randn(10, 2, 3)
y = r.randn(3, 4)
s = 'nij,jk->nik'
self._check(s, x, y)
def test_two_operands_6(self):
# based on https://github.com/google/jax/issues/37#issuecomment-448572187
r = self.rng()
x = r.randn(2, 1)
y = r.randn(2, 3, 4)
s = 'sa,shb->shab'
self._check(s, x, y)
def test_one_operand_1(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->j'
self._check(s, x)
def test_one_operand_2(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->kij'
self._check(s, x)
def test_one_operand_3(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->ki'
self._check(s, x)
def test_one_operand_4(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = 'ijk->ki'
self._check(s, x)
def test_one_operand_5(self):
r = self.rng()
x = r.randn(2, 3, 4, 5)
s = '...ijk->...ki'
self._check(s, x)
def test_one_operand_6(self):
r = self.rng()
x = r.randn(3, 4, 5)
s = '...ijk->ki'
self._check(s, x)
def test_one_operand_7(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ii->'
self._check(s, x)
def test_one_operand_8(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ij->'
self._check(s, x)
def test_one_operand_9(self):
r = self.rng()
x = r.randn(3, 3, 3)
s = 'iii->'
self._check(s, x)
def test_one_operand_10(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ii->i'
self._check(s, x)
def test_one_operand_11(self):
r = self.rng()
x = r.randn(3, 3, 4)
s = 'iij->i'
self._check(s, x)
def test_one_operand_12(self):
r = self.rng()
x = r.randn(3, 3, 3)
s = 'iii->i'
self._check(s, x)
def test_one_operand_13(self):
r = self.rng()
x = r.randn(3, 3, 5, 4, 4)
s = 'iijkk->i'
self._check(s, x)
def test_one_operand_14(self):
r = self.rng()
x = r.randn(3, 3, 5, 4, 4)
s = 'iijkk->ik'
self._check(s, x)
def test_one_operand_15(self):
r = self.rng()
x = r.randn(3, 3, 5, 4, 4)
s = 'iijkl->il'
self._check(s, x)
def test_one_operand_16(self):
r = self.rng()
x = r.randn(3, 3)
s = 'ij->ij'
self._check(s, x)
def test_tf_unsupported_1(self):
# from https://www.tensorflow.org/api_docs/python/tf/einsum
r = self.rng()
x = r.randn(2, 3, 5, 1)
y = r.randn(3, 4, 5, 1)
s = 'ij...,jk...->ik...'
self._check(s, x, y)
def test_tf_unsupported_2(self):
# from https://www.tensorflow.org/api_docs/python/tf/einsum
r = self.rng()
x = r.randn(2, 3, 3)
y = r.randn(4)
s = 'ijj,k->ik'
self._check(s, x, y)
def test_tf_unsupported_3(self):
# from https://www.tensorflow.org/api_docs/python/tf/einsum
r = self.rng()
x = r.randn(2, 3)
y = r.randn(2, 3)
z = r.randn(3, 4)
s = 'ij,ij,jk->ik'
self._check(s, x, y, z)
# these tests are based on https://github.com/dask/dask/pull/3412/files
@parameterized.named_parameters(
{'testcase_name': '_{}_dtype={}'.format(einstr, dtype.__name__), # pylint: disable=g-complex-comprehension
'einstr': einstr, 'dtype': dtype}
for einstr in [
'abc,bad->abcd',
'abcdef,bcdfg->abcdeg',
'ea,fb,abcd,gc,hd->efgh',
'ab,b',
'aa',
'a,a->',
'a,a->a',
'a,a',
'a,b',
'a,b,c',
'a',
'ba,b',
'ba,b->',
'defab,fedbc->defac',
'ab...,bc...->ac...',
'a...a',
'abc...->cba...',
'...ab->...a',
'a...a->a...',
# Following 2 from # https://stackoverflow.com/a/19203475/1611416
'...abc,...abcd->...d',
'ab...,b->ab...',
# https://github.com/dask/dask/pull/3412#discussion_r182413444
'aa->a',
'ab,ab,c->c',
'aab,bc->ac',
'aab,bcc->ac',
'fdf,cdd,ccd,afe->ae',
'fff,fae,bef,def->abd',
]
# TODO(wangpeng): Add tnp.bool_ to dtype list
for dtype in [tnp.float32, tnp.int32, tnp.complex64])
def test_from_dask(self, einstr, dtype):
r = tntu.rand_default()
if '->' in einstr:
input_str, _ = einstr.split('->')
else:
input_str = einstr
input_names = input_str.split(',')
dims = itertools.cycle([2, 3, 4])
shapes = defaultdict(lambda: next(dims))
input_shapes = [tuple(shapes[c] for c in names.replace('...', '01'))
for names in input_names]
operands = [r(shape, dtype) for shape in input_shapes]
self._check(einstr, *operands)
def test_ordered_front_batch_dim_case(self):
x = np.ones((1, 8, 20, 4))
y = np.ones((1, 8, 20, 4))
s = 'ijkl,ijml->ijkm'
self._check(s, x, y)
# pylint: disable=invalid-name
def test_einsum_path(self):
# just check examples from np.einsum_path docstring
a = self.rng().rand(2, 2)
b = self.rng().rand(2, 5)
c = self.rng().rand(5, 2)
path_info = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy')
self.assertEqual(str(path_info[0]), "['einsum_path', (1, 2), (0, 1)]")
self.assertEqual(path_info[1].split('\n')[0],
' Complete contraction: ij,jk,kl->il')
# check this doesn't crash
I = self.rng().rand(10, 10, 10, 10)
C = self.rng().rand(10, 10)
np.einsum_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C, optimize='greedy')
@tntu.disable
def test_einsum_kpmurphy_example(self):
# code from an email with @murphyk
N = 2
C = 3
D = 4
K = 5
T = 6
r = self.rng()
S = r.randn(N, T, K)
W = r.randn(K, D)
V = r.randn(D, C)
L = np.zeros((N, C))
for n in range(N):
for c in range(C):
s = 0
for d in range(D):
for k in range(K):
for t in range(T):
s += S[n, t, k] * W[k, d] * V[d, c]
L[n, c] = s
path = tnp.einsum_path('ntk,kd,dc->nc', S, W, V, optimize='optimal')[0]
rtol = 1e-2 if tntu.device_under_test() == 'tpu' else None
self.assertAllClose(L, tnp.einsum('ntk,kd,dc->nc', S, W, V, optimize=path),
check_dtypes=False, rtol=rtol)
# pylint: enable=invalid-name
@tntu.disable
def test_contraction_broadcasting(self):
r = self.rng()
x = r.randn(3, 4, 5)
y = r.randn(3, 1, 6)
s = 'cij,cjk->cik'
self._check(s, x, y)
@tntu.disable
def test_batch_broadcasting(self):
r = self.rng()
x = r.randn(1, 4, 5)
y = r.randn(3, 5, 6)
s = 'cij,cjk->cik'
self._check(s, x, y)
@tntu.disable
def test_batch_and_contraction_broadcasting(self):
r = self.rng()
x = r.randn(1, 4, 5)
y = r.randn(3, 1, 6)
s = 'cij,cjk->cik'
self._check(s, x, y)
@tntu.disable
def test_broadcasting_issue_2189(self):
r = self.rng()
x = r.randn(2, 1, 3, 3)
y = r.randn(2, 4, 3)
s = '...ij,...j'
self._check(s, x, y)
if __name__ == '__main__':
absltest.main()
tensorflow/python/ops/numpy_ops/tests/np_indexing_test.py
+988
View File
@@ -0,0 +1,988 @@
# Copyright 2023 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.
# ==============================================================================
import collections
import enum
from functools import partial
import itertools
import unittest
from absl.testing import absltest
from absl.testing import parameterized
import numpy as onp
from tensorflow.python.framework import config as tf_config
from tensorflow.python.ops.numpy_ops.tests.config import config
import tensorflow.python.ops.numpy_ops.tests.extensions as nje
import tensorflow.python.ops.numpy_ops.tests.np_wrapper as tnp
import tensorflow.python.ops.numpy_ops.tests.test_util as jtu
from tensorflow.python.util import nest
config.parse_flags_with_absl()
# We disable the whitespace continuation check in this file because otherwise it
# makes the test name formatting unwieldy.
# pylint: disable=bad-continuation
# We also disable undefined-variable till we start enabling tests.
# pylint: disable=undefined-variable
def subvals(lst, replace):
lst = list(lst)
for i, v in replace:
lst[i] = v
return tuple(lst)
float_dtypes = [onp.float32, onp.float64]
int_dtypes = [onp.int32, onp.int64]
bool_types = [onp.bool_]
default_dtypes = float_dtypes + int_dtypes
all_dtypes = float_dtypes + int_dtypes + bool_types
IndexSpec = collections.namedtuple("IndexTest", ["shape", "indexer"])
suppress_deprecated_indexing_warnings = partial(
jtu.ignore_warning, category=FutureWarning,
message='Using a non-tuple sequence.*')
STATIC_INDEXING_TESTS = [
("OneIntIndex", [
IndexSpec(shape=(3,), indexer=1),
IndexSpec(shape=(3, 3), indexer=0),
IndexSpec(shape=(3, 4, 5), indexer=2),
IndexSpec(shape=(3,), indexer=-1),
IndexSpec(shape=(3,), indexer=-2),
]),
("TwoIntIndices", [
IndexSpec(shape=(3, 3), indexer=(2, 1)),
IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
IndexSpec(shape=(3, 4, 5), indexer=(-1, 2)),
]),
("ThreeIntIndices", [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
("OneSliceIndex", [
IndexSpec(shape=(10,), indexer=slice(1, 3)),
IndexSpec(shape=(10,), indexer=slice(1, -1)),
IndexSpec(shape=(10,), indexer=slice(None, -1)),
IndexSpec(shape=(10,), indexer=slice(None, None, None)),
IndexSpec(shape=(10, 8), indexer=slice(1, 3)),
IndexSpec(shape=(10, 8), indexer=slice(1, None)),
IndexSpec(shape=(10, 8), indexer=slice(None, 3)),
IndexSpec(shape=(10, 8), indexer=slice(-3, None)),
]),
("OneSliceIndexNegativeStride", [
IndexSpec(shape=(10,), indexer=slice(3, 1, -1)),
IndexSpec(shape=(10,), indexer=slice(1, 8, -1)), # empty result
IndexSpec(shape=(10,), indexer=slice(None, 1, -2)),
IndexSpec(shape=(10,), indexer=slice(None, None, -1)),
IndexSpec(shape=(10, 8), indexer=slice(3, 1, -1)),
IndexSpec(shape=(10, 8), indexer=slice(0, 8, -1)), # empty result
IndexSpec(shape=(10, 8), indexer=slice(None, None, -1)),
]),
("OneSliceIndexNonUnitStride", [
IndexSpec(shape=(10,), indexer=slice(0, 8, 2)),
IndexSpec(shape=(10,), indexer=slice(0, 8, 3)),
IndexSpec(shape=(10,), indexer=slice(1, 3, 2)),
IndexSpec(shape=(10,), indexer=slice(1, None, 2)),
IndexSpec(shape=(10,), indexer=slice(None, 1, -2)),
IndexSpec(shape=(10, 8), indexer=slice(1, 8, 3)),
IndexSpec(shape=(10, 8), indexer=slice(None, None, 2)),
IndexSpec(shape=(10, 8), indexer=slice(None, 1, -2)),
IndexSpec(shape=(10, 8), indexer=slice(None, None, -2)),
]),
("TwoSliceIndices", [
IndexSpec(shape=(10, 8), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(10, 8), indexer=(slice(1, None), slice(None, 2))),
IndexSpec(
shape=(10, 8), indexer=(slice(None, None, -1), slice(None, 2))),
IndexSpec(shape=(10, 8, 3), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(10, 8, 3), indexer=(slice(1, 3), slice(0, None))),
IndexSpec(shape=(10, 8, 3), indexer=(slice(1, None), slice(0, 2))),
]),
("OneColonIndex", [
IndexSpec(shape=(3,), indexer=slice(None)),
IndexSpec(shape=(3, 4), indexer=slice(None)),
]),
("MultipleColonIndices", [
IndexSpec(shape=(3, 4), indexer=(slice(None), slice(None))),
IndexSpec(shape=(3, 4, 5), indexer=(slice(None), slice(None))),
]),
("MixedSliceIndices", [
IndexSpec(shape=(10, 4), indexer=(slice(None), slice(0, 2))),
IndexSpec(shape=(10, 4), indexer=(1, slice(None))),
]),
("EllipsisIndex", [
IndexSpec(shape=(3,), indexer=Ellipsis),
IndexSpec(shape=(3, 4), indexer=Ellipsis),
IndexSpec(shape=(3, 4, 5), indexer=(0, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis, 2, 3)),
]),
("NoneIndex", [
IndexSpec(shape=(), indexer=None),
IndexSpec(shape=(), indexer=(None, None)),
IndexSpec(shape=(), indexer=(Ellipsis, None)),
IndexSpec(shape=(3,), indexer=None),
IndexSpec(shape=(3, 4), indexer=None),
IndexSpec(shape=(3, 4), indexer=(Ellipsis, None)),
IndexSpec(shape=(3, 4), indexer=(0, None, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(1, None, Ellipsis)),
]),
("EmptyIndex", [
IndexSpec(shape=(), indexer=()),
IndexSpec(shape=(3,), indexer=()),
IndexSpec(shape=(3, 4), indexer=()),
]),
]
STATIC_INDEXING_GRAD_TESTS = [
("OneIntIndex", [
IndexSpec(shape=(3,), indexer=1),
IndexSpec(shape=(3, 3), indexer=0),
IndexSpec(shape=(3, 4, 5), indexer=2),
IndexSpec(shape=(3,), indexer=-1),
IndexSpec(shape=(3,), indexer=-2),
]),
("TwoIntIndices", [
IndexSpec(shape=(3, 3), indexer=(2, 1)),
IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
IndexSpec(shape=(3, 4, 5), indexer=(-1, 2)),
]),
("ThreeIntIndices", [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
("OneSliceIndex", [
IndexSpec(shape=(5,), indexer=slice(1, 3)),
IndexSpec(shape=(5,), indexer=slice(1, -1)),
IndexSpec(shape=(5,), indexer=slice(None, -1)),
IndexSpec(shape=(5,), indexer=slice(None, None, None)),
IndexSpec(shape=(5, 4), indexer=slice(1, 3)),
IndexSpec(shape=(5, 4), indexer=slice(1, None)),
IndexSpec(shape=(5, 4), indexer=slice(None, 3)),
IndexSpec(shape=(5, 4), indexer=slice(-3, None)),
]),
("TwoSliceIndices", [
IndexSpec(shape=(5, 4), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(5, 4), indexer=(slice(1, None), slice(None, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, None))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, None), slice(0, 2))),
]),
("OneColonIndex", [
IndexSpec(shape=(3,), indexer=slice(None)),
IndexSpec(shape=(3, 4), indexer=slice(None)),
]),
("MultipleColonIndices", [
IndexSpec(shape=(3, 4), indexer=(slice(None), slice(None))),
IndexSpec(shape=(3, 4, 5), indexer=(slice(None), slice(None))),
]),
("MixedSliceIndices", [
IndexSpec(shape=(5, 4), indexer=(slice(None), slice(0, 2))),
IndexSpec(shape=(5, 4), indexer=(1, slice(None))),
]),
("EllipsisIndex", [
IndexSpec(shape=(3,), indexer=Ellipsis),
IndexSpec(shape=(3, 4), indexer=Ellipsis),
IndexSpec(shape=(3, 4, 5), indexer=(0, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis, 2, 3)),
]),
("NoneIndex", [
IndexSpec(shape=(), indexer=None),
IndexSpec(shape=(), indexer=(None, None)),
IndexSpec(shape=(), indexer=(Ellipsis, None)),
IndexSpec(shape=(3,), indexer=None),
IndexSpec(shape=(3, 4), indexer=None),
IndexSpec(shape=(3, 4), indexer=(Ellipsis, None)),
IndexSpec(shape=(3, 4), indexer=(0, None, Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(1, None, Ellipsis)),
]),
# TODO(mattjj): these fail for uninteresting dtype reasons
# ("EmptyIndex",
# [IndexSpec(shape=(), indexer=()),
# IndexSpec(shape=(3,), indexer=()),
# IndexSpec(shape=(3, 4), indexer=()),
# ]),
]
ADVANCED_INDEXING_TESTS = [
("One1DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 1])),
IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 0, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
IndexSpec(shape=(0,), indexer=onp.array([], dtype=onp.int32)),
]),
("One2DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([[0, 0]])),
IndexSpec(shape=(3, 3), indexer=onp.array([[1, 2, 1],
[0, 1, -1]])),
IndexSpec(shape=(3, 4, 5), indexer=onp.array([[0, 2, 0, 1],
[-1, -2, 1, 0]])),
]),
("Two1DIntArrayIndicesNoBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([0, 1]),
onp.array([1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, 0, 1]),
onp.array([-1, 0, -1, 2]))),
]),
("Two1DIntArrayIndicesWithBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([[0, 1]]),
onp.array([1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([[0, 2, 0, 1]]),
onp.array([-1, 0, -1, 2]))),
]),
("TupleOfListsOfPythonInts",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0, 3]])),
]),
("TupleOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=(0, onp.array([0, 1]))),
IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
onp.array([[2, 3, 0, 3]]))),
]),
("TupleOfListsOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1], onp.array([0]))),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
onp.array([[2, 3, 0, 3]]))),
]),
]
ADVANCED_INDEXING_TESTS_NO_REPEATS = [
("One1DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([0, 1])),
IndexSpec(shape=(3, 3), indexer=onp.array([1, 2, 0])),
IndexSpec(shape=(3, 4, 5), indexer=onp.array([0, 2, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-1, 1])),
IndexSpec(shape=(3,), indexer=onp.array([-2, -1])),
# Fails with a TF/XLA error.
# IndexSpec(shape=(0,), indexer=onp.array([], dtype=onp.int32)),
]),
("One2DIntArrayIndex",
[IndexSpec(shape=(3,), indexer=onp.array([[0, 1]])),
IndexSpec(shape=(6, 6), indexer=onp.array([[1, 2, 0],
[3, 4, -1]])),
]),
("Two1DIntArrayIndicesNoBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([0, 1]),
onp.array([1, 2]))),
IndexSpec(shape=(4, 5, 6), indexer=(onp.array([0, 2, 1, 3]),
onp.array([-1, 0, -2, 1]))),
]),
("Two1DIntArrayIndicesWithBroadcasting",
[IndexSpec(shape=(3, 3), indexer=(onp.array([[0, 1]]),
onp.array([1, 2]))),
IndexSpec(shape=(4, 5, 6), indexer=(onp.array([[0, 2, -1, 1]]),
onp.array([-1, 0, -2, 2]))),
]),
("TupleOfListsOfPythonInts",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0]])),
]),
("TupleOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=(0, onp.array([0, 1]))),
IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
onp.array([[2, 3, 0]]))),
]),
("TupleOfListsOfPythonIntsAndIntArrays",
[IndexSpec(shape=(3, 4, 5), indexer=([0, 1], onp.array([0]))),
IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
onp.array([[2, 3, 0]]))),
]),
]
MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS = [
("SlicesAndOneIntArrayIndex",
[IndexSpec(shape=(2, 3), indexer=(onp.array([0, 1]), slice(1, 2))),
IndexSpec(shape=(2, 3), indexer=(slice(0, 2),
onp.array([0, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([0, 2]),
slice(None))),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([[0, 2], [1, 3]]),
slice(None))),
]),
("SlicesAndTwoIntArrayIndices",
[IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([0, 2]),
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
Ellipsis,
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
onp.array([-1, 2]),
Ellipsis)),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
onp.array([-1, 2]),
slice(1, 3))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
slice(1, 3),
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, -2]),
slice(None, None, 2),
onp.array([-1, 2, 1]))),
]),
("NonesAndIntArrayIndices",
[IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
None,
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2]),
None,
None,
onp.array([-1, 2]))),
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([0, 2]),
None,
None,
onp.array([-1, 2]))),
]),
("IntArrayWithInt32Type",
[IndexSpec(shape=(3, 4), indexer=(Ellipsis, onp.array(1, dtype=onp.int32)))
]),
]
MIXED_ADVANCED_INDEXING_TESTS = MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS + [
("SlicesAndOneIntArrayIndex",
[
IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
onp.array([[0, 2], [1, 1]]),
slice(None))),
]),
("SlicesAndTwoIntArrayIndices",
[IndexSpec(shape=(3, 4, 5), indexer=(onp.array([0, 2, -2]),
slice(None, None, 2),
onp.array([-1, 2, -1]))),
IndexSpec(shape=(3, 4, 5), indexer=(onp.array([[0, 2], [2, 0]]),
Ellipsis,
onp.array([[1, 0], [1, 0]]))),
]),]
def dynamic_slice_reference(operand, start_indices, slice_sizes):
out = onp.zeros(slice_sizes, dtype=operand.dtype)
idx = tuple(slice(start, start+size)
for start, size in zip(start_indices, slice_sizes))
section = operand[idx]
out[tuple(slice(None, stop) for stop in section.shape)] = section
return out
def dynamic_update_slice_reference(operand, update, start_indices):
slices = tuple(map(
slice, start_indices, onp.add(start_indices, update.shape)))
updated_operand = onp.copy(operand)
updated_operand[slices] = update
return updated_operand
class IndexingTest(jtu.TestCase):
"""Tests for Numpy indexing translation rules."""
@parameterized.named_parameters(jtu.cases_from_list({
"testcase_name": "{}_inshape={}_indexer={}".format(
name, jtu.format_shape_dtype_string( shape, dtype), indexer),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory, "indexer": indexer
} for name, index_specs in STATIC_INDEXING_TESTS
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testStaticIndexing(self, shape, dtype, rng_factory, indexer):
# TODO(rohanj): Revisit passing in self.rng() to this to customize further.
# This would need updating lax_numpy_test as well.
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype)]
onp_fun = lambda x: x[indexer]
jnp_fun = lambda x: tnp.asarray(x)[indexer]
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True)
def _ReplaceSlicesWithTuples(self, idx):
"""Helper method to replace slices with tuples for dynamic indexing args."""
if isinstance(idx, slice):
triple = idx.start, idx.stop, idx.step
isnone = [i for i, elt in enumerate(triple) if elt is None]
zeros = itertools.repeat(0)
nones = itertools.repeat(None)
out = subvals(triple, zip(isnone, zeros))
return out, lambda out: slice(*subvals(out, zip(isnone, nones)))
elif isinstance(idx, (tuple, list)) and idx:
t = type(idx)
elts, packs = zip(*map(self._ReplaceSlicesWithTuples, idx))
return elts, lambda elts: t((pack(i) for pack, i in zip(packs, elts)))
else:
return idx, lambda x: x
@parameterized.named_parameters(
{"testcase_name": "{}_inshape={}_indexer={}"
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory, "indexer": indexer}
for name, index_specs in [
("OneSliceIndex",
[IndexSpec(shape=(5,), indexer=slice(1, 3)),
IndexSpec(shape=(5, 4), indexer=slice(1, 3))]),
("TwoSliceIndices",
[IndexSpec(shape=(5, 4), indexer=(slice(1, 3), slice(0, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, 2)))]),
("NonUnitStrides", [
IndexSpec(shape=(3,), indexer=slice(None, None, -1)),
IndexSpec(shape=(3, 3), indexer=slice(0, 3, -2)),
IndexSpec(shape=(3, 4, 5), indexer=slice(0, 4, 2))
]),
("OnlyStartOrStopDynamic", [
IndexSpec(shape=(5, 4), indexer=(slice(None, 3), slice(0, 2))),
IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, None)))
]),
]
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testDynamicIndexingWithSlices(self, shape, dtype, rng_factory, indexer):
rng = rng_factory()
unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)
def onp_fun(x, unpacked_indexer):
indexer = pack_indexer(unpacked_indexer)
return x[indexer]
jnp_fun = lambda x, idx: onp_fun(tnp.asarray(x), idx)
args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
# TODO(wangpeng): check_xla_forced_compile is turned off because some
# compile-time-constant requirements are violated. Investigate and turn it
# on.
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_eval_on_shapes=False,
check_incomplete_shape=True,
check_xla_forced_compile=False)
@parameterized.named_parameters(
{"testcase_name": "{}_inshape={}_indexer={}"
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory, "indexer": indexer}
for name, index_specs in [
("OneIntIndex",
[IndexSpec(shape=(3,), indexer=1),
IndexSpec(shape=(3, 3), indexer=0),
IndexSpec(shape=(3, 4, 5), indexer=2),
IndexSpec(shape=(3,), indexer=-1),
IndexSpec(shape=(3,), indexer=-2)]),
("TwoIntIndices",
[IndexSpec(shape=(3, 3), indexer=(2, 1)),
IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
IndexSpec(shape=(3, 4, 5), indexer=(-1, 2))]),
("ThreeIntIndices",
[IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
]
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testDynamicIndexingWithIntegers(self, shape, dtype, rng_factory, indexer):
# TODO(rohanj): Revisit passing in self.rng() to this to customize further.
# This would need updating lax_numpy_test as well.
rng = rng_factory()
unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)
def onp_fun(x, unpacked_indexer):
indexer = pack_indexer(unpacked_indexer)
return x[indexer]
jnp_fun = lambda x, idx: onp_fun(tnp.asarray(x), idx)
args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True)
@parameterized.named_parameters(
{"testcase_name": "_{}_inshape={}_indexer={}" # pylint: disable=g-complex-comprehension
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"name": name, "shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer}
for name, index_specs in ADVANCED_INDEXING_TESTS
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testAdvancedIntegerIndexing(self, name, shape, dtype, rng_factory,
indexer):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), indexer]
onp_fun = lambda x, idx: x[idx]
jnp_fun = lambda x, idx: onp_fun(tnp.asarray(x), idx)
self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
# TODO(wangpeng): check_xla_forced_compile is turned off for
# ListOfPythonIntsAndIntArrays because it throws "The number of output
# elements has to equal to number of input elements that are sliced when
# input indices are not constant". Investigate and turn it on.
check_xla = (name != "ListOfPythonIntsAndIntArrays")
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True,
check_xla_forced_compile=check_xla)
@parameterized.named_parameters(
{"testcase_name": "_{}_inshape={}_indexer={}" # pylint: disable=g-complex-comprehension
.format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
"name": name, "shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer}
for name, index_specs in MIXED_ADVANCED_INDEXING_TESTS
for shape, indexer in index_specs
for dtype in all_dtypes
for rng_factory in [jtu.rand_default])
def testMixedAdvancedIntegerIndexing(self, name, shape, dtype, rng_factory,
indexer):
rng = rng_factory()
indexer_with_dummies = [e if isinstance(e, onp.ndarray) else ()
for e in indexer]
substitutes = [(i, e) for i, e in enumerate(indexer)
if not isinstance(e, onp.ndarray)]
args_maker = lambda: [rng(shape, dtype), indexer_with_dummies]
def np_fun(x, indexer_with_dummies):
idx = type(indexer)(subvals(indexer_with_dummies, substitutes))
return x[idx]
jnp_fun = lambda x, idx: np_fun(tnp.asarray(x), idx)
self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
# TODO(wangpeng): check_xla_forced_compile is turned off for
# IntArrayWithInt32Type because it throws "The number of output elements has
# to equal to number of input elements that are sliced when input indices
# are not constant". Investigate and turn it on.
check_xla = (name != "IntArrayWithInt32Type")
self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True,
check_incomplete_shape=True,
check_xla_forced_compile=check_xla)
def testAdvancedIndexingManually(self):
x = onp.random.RandomState(0).randn(3, 4, 5)
index_array = onp.array([0, 2, -1, 0])
op = lambda x, index_array: x[..., index_array, :]
cop = nje.jit(op)
a1 = op(x, index_array)
a2 = cop(x, index_array)
self.assertAllClose(a1, a2, check_dtypes=True)
op = lambda x, index_array: x[..., index_array, :, index_array, None]
cop = nje.jit(op)
a1 = op(x, index_array)
a2 = cop(x, index_array)
self.assertAllClose(a1, a2, check_dtypes=True)
op = lambda x, index_array: x[index_array, ..., index_array[:, None], None]
cop = nje.jit(op)
a1 = op(x, index_array)
a2 = cop(x, index_array)
self.assertAllClose(a1, a2, check_dtypes=True)
# Note that we don't currently allow __iter__ in graph mode. So this test only
# iterates over eager tensor.
def testUnpacking(self):
def foo(x):
a, b, c = x
return a + b + c
a1 = foo(onp.arange(3))
a2 = foo(tnp.arange(3))
self.assertAllClose(a1, a2, check_dtypes=True)
def testBooleanIndexingArray1D(self):
idx = onp.array([True, True, False])
x = tnp.asarray(onp.arange(3))
ans = x[idx]
expected = onp.arange(3)[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingList1D(self):
idx = [True, True, False]
x = tnp.asarray(onp.arange(3))
ans = x[idx]
expected = onp.arange(3)[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingArray2DBroadcast(self):
idx = onp.array([True, True, False, True])
x = onp.arange(8).reshape(4, 2)
ans = tnp.asarray(x)[idx]
expected = x[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingList2DBroadcast(self):
idx = [True, True, False, True]
x = onp.arange(8).reshape(4, 2)
ans = tnp.asarray(x)[idx]
expected = x[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingArray2D(self):
idx = onp.array([[True, False],
[False, True],
[False, False],
[True, True]])
x = onp.arange(8).reshape(4, 2)
ans = tnp.asarray(x)[idx]
expected = x[idx]
self.assertAllClose(ans, expected, check_dtypes=False)
def testBooleanIndexingDynamicShape(self):
x = onp.zeros(3)
i = onp.array([True, True, False])
ans = x[i]
expected = tnp.asarray(x)[i]
self.assertAllClose(ans, expected, check_dtypes=True)
def testIssue187(self):
x = tnp.ones((5, 5))
x[[0, 2, 4], [0, 2, 4]] # doesn't crash
x = onp.arange(25).reshape((5, 5))
ans = nje.jit(lambda x: x[[0, 2, 4], [0, 2, 4]])(x)
expected = x[[0, 2, 4], [0, 2, 4]]
self.assertAllClose(ans, expected, check_dtypes=False)
# TODO(agarwal): Fix this use case.
@jtu.disable
def testIndexingEmptyDimension(self):
# Issue 2671: XLA error when indexing into dimension of size 0
x = tnp.ones((2, 0))
# The following work, even on axis 1 of size 0
_ = x[0, :] + x[0, None] + x[0, 1:] + x[0, 1:3:2]
with self.assertRaisesRegex(IndexError,
"index .* is out of bounds for axis .* with size 0"):
_ = onp.ones((2, 0))[0, 0] # The numpy error
with self.assertRaisesRegex(IndexError,
"index is out of bounds for axis .* with size 0"):
_ = x[0, 0] # JAX indexing
with self.assertRaisesRegex(IndexError,
"index is out of bounds for axis .* with size 0"):
nje.jit(lambda i: x[0, i])(0) # JAX indexing under jit
def testBooleanIndexingWithEmptyResult(self):
# based on a TensorFlow Probability test that started failing after #1623
x = tnp.array([-1])
mask = tnp.array([False])
ans = x[mask] # doesn't crash
expected = onp.array([-1])[onp.array([False])]
self.assertAllClose(ans, expected, check_dtypes=False)
def testFloatIndexingError(self):
error_regex = "only integers, slices.*are valid indices"
# Verify onp behavior
with self.assertRaisesRegex(IndexError, error_regex):
_ = onp.zeros((2, 2))[(0, 0.)]
# Test tnp
with self.assertRaisesRegex(IndexError, error_regex):
tnp.zeros(2)[0.] # pylint: disable=expression-not-assigned
with self.assertRaisesRegex(IndexError, error_regex):
tnp.zeros((2, 2))[(0, 0.)] # pylint: disable=expression-not-assigned
# Test with jit
with self.assertRaisesRegex(IndexError, error_regex):
nje.jit(lambda idx: tnp.zeros((2, 2))[idx])((0, 0.0))
def testIndexOutOfBounds(self): # https://github.com/google/jax/issues/2245
array = tnp.ones(5)
self.assertAllClose(array, array[:10], check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_size_indices={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, size_indices),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"size_indices": size_indices, "rng_factory": rng_factory}
for shape, start_indices, size_indices in [
[(3,), onp.array((1,)), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(5, 3), onp.array((1, 1)), (3, 1)],
[(7, 5, 3), onp.array((4, 1, 0)), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicSlice(self, shape, dtype, start_indices, size_indices,
rng_factory):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), onp.array(start_indices)]
op = lambda x, starts: nje.dynamic_slice(x, starts, size_indices)
self._CompileAndCheck(op, args_maker)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_size_indices={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, size_indices),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"size_indices": size_indices, "rng_factory": rng_factory}
for shape, start_indices, size_indices in [
[(3,), (1,), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(7, 5, 3), (4, 1, 0), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicSliceAgainstNumpy(self, shape, dtype, start_indices,
size_indices, rng_factory):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), onp.array(start_indices)]
op = lambda x, s: nje.dynamic_slice(x, s, size_indices)
numpy_op = lambda x, s: dynamic_slice_reference(x, s, size_indices)
self._CheckAgainstNumpy(numpy_op, op, args_maker)
def testDynamicSliceInDim(self):
rng = jtu.rand_default()
x = rng((6, 7), onp.int32)
self.assertAllClose(
nje.dynamic_slice_in_dim(x, 2, 3), x[2:5], check_dtypes=True
)
def _broadcastable_shapes(shape):
"""Returns all shapes that broadcast to `shape`."""
def f(rshape):
yield []
if rshape:
for s in f(rshape[1:]):
yield rshape[0:1] + s
if rshape[0] != 1:
for s in f(rshape[1:]):
yield [1] + s
for x in f(list(reversed(shape))):
yield list(reversed(x))
def _update_shape(shape, indexer):
return onp.zeros(shape)[indexer].shape
class UpdateOps(enum.Enum):
UPDATE = 0
ADD = 1
# MUL = 2
MIN = 3
MAX = 4
def np_fn(op, indexer, x, y): # pylint: disable=no-self-argument
x = x.copy()
x[indexer] = {
UpdateOps.UPDATE: lambda: y,
UpdateOps.ADD: lambda: x[indexer] + y,
# UpdateOps.MUL: lambda: x[indexer] * y,
UpdateOps.MIN: lambda: onp.minimum(x[indexer], y),
UpdateOps.MAX: lambda: onp.maximum(x[indexer], y),
}[op]()
return x
def tfnp_fn(op, indexer, x, y): # pylint: disable=no-self-argument
return {
UpdateOps.UPDATE: nje.index_update,
UpdateOps.ADD: nje.index_add,
# UpdateOps.MUL: nje.index_mul,
UpdateOps.MIN: nje.index_min,
UpdateOps.MAX: nje.index_max,
}[op](x, indexer, y)
# a test to workaround b/123559667
def has_non_trivial_stride(indexer):
def has(idx):
return isinstance(idx, slice) and idx.step not in (1, -1, None)
return any(has(idx) for idx in nest.flatten(indexer))
class IndexedUpdateTest(jtu.TestCase):
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in STATIC_INDEXING_TESTS
for shape, indexer in index_specs
for op in UpdateOps
for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testStaticIndexing(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
self._CheckAgainstNumpy(np_fn, tfnp_fn, args_maker)
# TODO(wangpeng): When indexer is slice(_, 8, -1), XLA throws error "Missing
# xla_context 0-th output from". Investigate.
check_xla = (not has_non_trivial_stride(indexer) and # b/123559667
not (isinstance(indexer, slice) and indexer.stop == 8 and
indexer.step == -1))
self._CompileAndCheck(tfnp_fn, args_maker, check_incomplete_shape=True,
check_experimental_compile=check_xla,
check_xla_forced_compile=check_xla)
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in ADVANCED_INDEXING_TESTS_NO_REPEATS
for shape, indexer in index_specs
for op in UpdateOps
for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
self._CheckAgainstNumpy(np_fn, tfnp_fn, args_maker)
self._CompileAndCheck(tfnp_fn, args_maker, check_incomplete_shape=True)
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS
for shape, indexer in index_specs
for op in UpdateOps
for dtype in (all_dtypes if op == UpdateOps.UPDATE else default_dtypes)
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in all_dtypes
for rng_factory in [jtu.rand_default]))
def testMixedAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
self._CheckAgainstNumpy(np_fn, tfnp_fn, args_maker)
check_xla = not has_non_trivial_stride(indexer) # b/123559667
self._CompileAndCheck(tfnp_fn, args_maker, check_incomplete_shape=True,
check_experimental_compile=check_xla,
check_xla_forced_compile=check_xla)
@parameterized.named_parameters(jtu.cases_from_list({ # pylint: disable=g-complex-comprehension
"testcase_name": "_{}_{}_{}_{}".format(
jtu.format_shape_dtype_string(shape, dtype), indexer,
jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
"shape": shape, "dtype": dtype, "rng_factory": rng_factory,
"indexer": indexer, "update_shape": update_shape,
"update_dtype": update_dtype, "op": op
} for name, index_specs in STATIC_INDEXING_TESTS
for shape, indexer in index_specs
for op in [UpdateOps.ADD, UpdateOps.UPDATE]
for dtype in float_dtypes
for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
for update_dtype in float_dtypes
for rng_factory in [jtu.rand_default]))
def testStaticIndexingGrads(self, shape, dtype, update_shape, update_dtype,
rng_factory, indexer, op):
rng = rng_factory()
tfnp_fn = lambda x, y: UpdateOps.tfnp_fn(op, indexer, x, y)
x = rng(shape, dtype)
y = rng(update_shape, update_dtype)
self.check_grads(tfnp_fn, (x, y), rtol=1e-3, atol=1e-3, delta=1.)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_update_shape={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, update_shape),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"update_shape": update_shape, "rng_factory": rng_factory}
for shape, start_indices, update_shape in [
[(3,), (1,), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(7, 5, 3), (4, 1, 0), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicUpdateSlice(self, shape, dtype, start_indices, update_shape,
rng_factory):
rng = rng_factory()
def args_maker():
return [rng(shape, dtype), rng(update_shape, dtype),
onp.array(start_indices)]
# update's shape must be fully known.
# TODO(wangpeng): Support turning off check_incomplete_shape for individual
# arguments.
self._CompileAndCheck(
nje.dynamic_update_slice, args_maker, check_incomplete_shape=False
)
@parameterized.named_parameters(jtu.cases_from_list(
{"testcase_name": "_shape={}_start_indices={}_update_shape={}".format( # pylint: disable=g-complex-comprehension
jtu.format_shape_dtype_string(shape, dtype),
start_indices, update_shape),
"shape": shape, "dtype": dtype, "start_indices": start_indices,
"update_shape": update_shape, "rng_factory": rng_factory}
for shape, start_indices, update_shape in [
[(3,), (1,), (1,)],
[(5, 3), (1, 1), (3, 1)],
[(5, 3), (1, -2), (3, 1)],
[(7, 5, 3), (4, 1, 0), (2, 0, 1)],
[(), (), ()],
]
for dtype in default_dtypes
for rng_factory in [jtu.rand_default]))
def testDynamicUpdateSliceAgainstNumpy(self, shape, dtype, start_indices,
update_shape, rng_factory):
rng = rng_factory()
def args_maker():
return [rng(shape, dtype), rng(update_shape, dtype),
onp.array(start_indices)]
self._CheckAgainstNumpy(
dynamic_update_slice_reference, nje.dynamic_update_slice, args_maker
)
def testDynamicUpdateSliceInDim(self):
rng = jtu.rand_default()
x = rng((6, 7), onp.int32)
y = rng((3, 7), onp.int32)
z = x.copy()
z[2:5] = y
self.assertAllClose(
nje.dynamic_update_slice_in_dim(x, y, 2, 0), z, check_dtypes=True
)
if __name__ == "__main__":
tf_config.set_soft_device_placement(False)
absltest.main()
tensorflow/python/ops/numpy_ops/tests/np_test.py
+3166
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/ops/numpy_ops/tests/np_wrapper.py
+43
View File
@@ -0,0 +1,43 @@
# Copyright 2023 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 NumPy API wrapper for the tests."""
# pylint: disable=wildcard-import
# pylint: disable=unused-import
# pylint: disable=g-importing-member
import numpy as onp
from tensorflow.python.compat import v2_compat
from tensorflow.python.framework.dtypes import bfloat16
from tensorflow.python.ops.numpy_ops import np_random as random
from tensorflow.python.ops.numpy_ops.np_array_ops import *
from tensorflow.python.ops.numpy_ops.np_arrays import ndarray
from tensorflow.python.ops.numpy_ops.np_config import enable_numpy_behavior
from tensorflow.python.ops.numpy_ops.np_dtypes import *
from tensorflow.python.ops.numpy_ops.np_dtypes import canonicalize_dtype
from tensorflow.python.ops.numpy_ops.np_dtypes import default_float_type
from tensorflow.python.ops.numpy_ops.np_dtypes import is_allow_float64
from tensorflow.python.ops.numpy_ops.np_dtypes import set_allow_float64
from tensorflow.python.ops.numpy_ops.np_math_ops import *
from tensorflow.python.ops.numpy_ops.np_utils import finfo
from tensorflow.python.ops.numpy_ops.np_utils import promote_types
from tensorflow.python.ops.numpy_ops.np_utils import result_type
random.DEFAULT_RANDN_DTYPE = onp.float32
# pylint: enable=unused-import
v2_compat.enable_v2_behavior()
# TODO(b/171429739): This should be moved to every individual file/test.
enable_numpy_behavior()
tensorflow/python/ops/numpy_ops/tests/test_util.py
+899
View File
@@ -0,0 +1,899 @@
# Copyright 2023 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.
# ==============================================================================
"""NumPy test utilities."""
from contextlib import contextmanager
import functools
from functools import partial
import re
import itertools as it
import os
from typing import Dict, Sequence, Union
import unittest
import warnings
import zlib
from absl.testing import absltest
from absl.testing import parameterized
import numpy as onp
import numpy.random as npr
from tensorflow.python.util import nest
from tensorflow.python.util import numpy_compat
from tensorflow.python.framework import tensor
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops.numpy_ops.tests.config import flags
import tensorflow.python.ops.numpy_ops.tests.extensions as nje
from tensorflow.python.ops.numpy_ops import np_utils
from tensorflow.python.ops.numpy_ops import np_array_ops
tree_map = nest.map_structure
tree_multimap = nest.map_structure
FLAGS = flags.FLAGS
# https://danielms.site/zet/2023/pythons-distutil-strtobool-replacement/
def strtobool(value: str) -> bool:
value = value.lower()
if value in ('y', 'yes', 'on', '1', 'true', 't'):
return True
return False
# TODO(wangpeng): Remove this flag after broken tests are fixed
flags.DEFINE_bool('enable_x64',
strtobool('False'),
'Enable 64-bit types to be used.')
flags.DEFINE_enum(
'test_dut', '',
enum_values=['', 'cpu', 'gpu', 'tpu'],
help=
'Describes the device under test in case special consideration is required.'
)
flags.DEFINE_integer(
'num_generated_cases',
10,
help='Number of generated cases to test')
EPS = 1e-4
# Default dtypes corresponding to Python scalars.
python_scalar_dtypes = {
bool: onp.dtype(onp.bool_),
int: onp.dtype(onp.int_),
float: onp.dtype(onp.float64),
complex: onp.dtype(onp.complex128),
}
def _dtype(x):
if isinstance(x, tensor.Tensor):
return x.dtype.as_numpy_dtype
return (getattr(x, 'dtype', None) or
onp.dtype(python_scalar_dtypes.get(type(x), None)) or
numpy_compat.np_asarray(x).dtype)
def is_sequence(x):
try:
iter(x)
except TypeError:
return False
else:
return True
_default_tolerance = {
onp.dtype(onp.bool_): 0,
onp.dtype(onp.int8): 0,
onp.dtype(onp.int16): 0,
onp.dtype(onp.int32): 0,
onp.dtype(onp.int64): 0,
onp.dtype(onp.uint8): 0,
onp.dtype(onp.uint16): 0,
onp.dtype(onp.uint32): 0,
onp.dtype(onp.uint64): 0,
# TODO(b/154768983): onp.dtype(dtypes.bfloat16): 1e-2,
onp.dtype(onp.float16): 1e-3,
onp.dtype(onp.float32): 1e-6,
onp.dtype(onp.float64): 1e-15,
onp.dtype(onp.complex64): 1e-6,
onp.dtype(onp.complex128): 1e-15,
}
def default_tolerance():
return _default_tolerance
default_gradient_tolerance = {
# TODO(b/154768983): onp.dtype(dtypes.bfloat16): 1e-1,
onp.dtype(onp.float16): 1e-2,
onp.dtype(onp.float32): 2e-3,
onp.dtype(onp.float64): 1e-5,
onp.dtype(onp.complex64): 1e-3,
onp.dtype(onp.complex128): 1e-5,
}
def _assert_numpy_allclose(a, b, atol=None, rtol=None):
# TODO(b/154768983):
# a = a.astype(onp.float32) if a.dtype == dtypes.bfloat16 else a
# b = b.astype(onp.float32) if b.dtype == dtypes.bfloat16 else b
kw = {}
if atol: kw["atol"] = atol
if rtol: kw["rtol"] = rtol
onp.testing.assert_allclose(a, b, **kw)
def tolerance(dtype, tol=None):
tol = {} if tol is None else tol
if not isinstance(tol, dict):
return tol
tol = {onp.dtype(key): value for key, value in tol.items()}
dtype = onp.dtype(dtype)
return tol.get(dtype, default_tolerance()[dtype])
def _normalize_tolerance(tol):
tol = tol or 0
if isinstance(tol, dict):
return {onp.dtype(k): v for k, v in tol.items()}
else:
return {k: tol for k in _default_tolerance}
def join_tolerance(tol1, tol2):
tol1 = _normalize_tolerance(tol1)
tol2 = _normalize_tolerance(tol2)
out = tol1
for k, v in tol2.items():
out[k] = max(v, tol1.get(k, 0))
return out
def _assert_numpy_close(a, b, atol=None, rtol=None):
assert a.shape == b.shape
atol = max(tolerance(a.dtype, atol), tolerance(b.dtype, atol))
rtol = max(tolerance(a.dtype, rtol), tolerance(b.dtype, rtol))
_assert_numpy_allclose(a, b, atol=atol * a.size, rtol=rtol * b.size)
def check_eq(xs, ys):
tree_all(tree_multimap(_assert_numpy_allclose, xs, ys))
def check_close(xs, ys, atol=None, rtol=None):
assert_close = partial(_assert_numpy_close, atol=atol, rtol=rtol)
tree_all(tree_multimap(assert_close, xs, ys))
def inner_prod(xs, ys):
def contract(x, y):
return onp.real(onp.dot(onp.conj(x).reshape(-1), y.reshape(-1)))
return tree_reduce(onp.add, tree_multimap(contract, xs, ys))
add = partial(tree_multimap, lambda x, y: onp.add(x, y, dtype=_dtype(x)))
sub = partial(tree_multimap, lambda x, y: onp.subtract(x, y, dtype=_dtype(x)))
conj = partial(tree_map, lambda x: onp.conj(x, dtype=_dtype(x)))
def scalar_mul(xs, a):
return tree_map(lambda x: onp.multiply(x, a, dtype=_dtype(x)), xs)
def rand_like(rng, x):
shape = onp.shape(x)
dtype = _dtype(x)
randn = lambda: numpy_compat.np_asarray(rng.randn(*shape), dtype=dtype)
if onp.issubdtype(dtype, onp.complexfloating):
return randn() + dtype.type(1.0j) * randn()
else:
return randn()
def numerical_jvp(f, primals, tangents, eps=EPS):
delta = scalar_mul(tangents, eps)
f_pos = f(*add(primals, delta))
f_neg = f(*sub(primals, delta))
return scalar_mul(sub(f_pos, f_neg), 0.5 / eps)
def _merge_tolerance(tol, default):
if tol is None:
return default
if not isinstance(tol, dict):
return tol
out = default.copy()
for k, v in tol.items():
out[onp.dtype(k)] = v
return out
def check_jvp(f, f_jvp, args, atol=None, rtol=None, eps=EPS):
atol = _merge_tolerance(atol, default_gradient_tolerance)
rtol = _merge_tolerance(rtol, default_gradient_tolerance)
rng = onp.random.RandomState(0)
tangent = tree_map(partial(rand_like, rng), args)
v_out, t_out = f_jvp(args, tangent)
v_out_expected = f(*args)
t_out_expected = numerical_jvp(f, args, tangent, eps=eps)
# In principle we should expect exact equality of v_out and v_out_expected,
# but due to nondeterminism especially on GPU (e.g., due to convolution
# autotuning) we only require "close".
check_close(v_out, v_out_expected, atol=atol, rtol=rtol)
check_close(t_out, t_out_expected, atol=atol, rtol=rtol)
def check_vjp(f, f_vjp, args, atol=None, rtol=None, eps=EPS):
atol = _merge_tolerance(atol, default_gradient_tolerance)
rtol = _merge_tolerance(rtol, default_gradient_tolerance)
_rand_like = partial(rand_like, onp.random.RandomState(0))
v_out, vjpfun = f_vjp(*args)
v_out_expected = f(*args)
check_close(v_out, v_out_expected, atol=atol, rtol=rtol)
tangent = tree_map(_rand_like, args)
tangent_out = numerical_jvp(f, args, tangent, eps=eps)
cotangent = tree_map(_rand_like, v_out)
cotangent_out = conj(vjpfun(conj(cotangent)))
ip = inner_prod(tangent, cotangent_out)
ip_expected = inner_prod(tangent_out, cotangent)
check_close(ip, ip_expected, atol=atol, rtol=rtol)
def device_under_test():
return FLAGS.test_dut
def if_device_under_test(device_type: Union[str, Sequence[str]],
if_true, if_false):
"""Chooses `if_true` of `if_false` based on device_under_test."""
if device_under_test() in ([device_type] if isinstance(device_type, str)
else device_type):
return if_true
else:
return if_false
def supported_dtypes():
if device_under_test() == "tpu":
return {onp.bool_, onp.int32, onp.uint32, dtypes.bfloat16, onp.float32,
onp.complex64}
else:
return {onp.bool_, onp.int8, onp.int16, onp.int32, onp.int64,
onp.uint8, onp.uint16, onp.uint32, onp.uint64,
dtypes.bfloat16, onp.float16, onp.float32, onp.float64,
onp.complex64, onp.complex128}
def skip_if_unsupported_type(dtype):
if dtype not in supported_dtypes():
raise unittest.SkipTest(
f"Type {dtype} not supported on {device_under_test()}")
def skip_on_devices(*disabled_devices):
"""A decorator for test methods to skip the test on certain devices."""
def skip(test_method):
@functools.wraps(test_method)
def test_method_wrapper(self, *args, **kwargs):
device = device_under_test()
if device in disabled_devices:
test_name = getattr(test_method, '__name__', '[unknown test]')
raise unittest.SkipTest(
f"{test_name} not supported on {device.upper()}.")
return test_method(self, *args, **kwargs)
return test_method_wrapper
return skip
def skip_on_flag(flag_name, skip_value):
"""A decorator for test methods to skip the test when flags are set."""
def skip(test_method): # pylint: disable=missing-docstring
@functools.wraps(test_method)
def test_method_wrapper(self, *args, **kwargs):
flag_value = getattr(FLAGS, flag_name)
if flag_value == skip_value:
test_name = getattr(test_method, '__name__', '[unknown test]')
raise unittest.SkipTest(
f"{test_name} not supported when FLAGS.{flag_name} is {flag_value}")
return test_method(self, *args, **kwargs)
return test_method_wrapper
return skip
def format_test_name_suffix(opname, shapes, dt):
arg_descriptions = (format_shape_dtype_string(shape, dtype)
for shape, dtype in zip(shapes, dt))
return '{}_{}'.format(opname.capitalize(), '_'.join(arg_descriptions))
# We use special symbols, represented as singleton objects, to distinguish
# between NumPy scalars, Python scalars, and 0-D arrays.
class ScalarShape:
def __len__(self): return 0
def __getitem__(self, i):
raise IndexError(f'index {i} out of range.')
class _NumpyScalar(ScalarShape): pass
class _PythonScalar(ScalarShape): pass
NUMPY_SCALAR_SHAPE = _NumpyScalar()
PYTHON_SCALAR_SHAPE = _PythonScalar()
def _dims_of_shape(shape):
"""Converts `shape` to a tuple of dimensions."""
if type(shape) in (list, tuple):
return shape
elif isinstance(shape, ScalarShape):
return ()
else:
raise TypeError(type(shape))
def _cast_to_shape(value, shape, dtype):
"""Casts `value` to the correct Python type for `shape` and `dtype`."""
if shape is NUMPY_SCALAR_SHAPE:
# explicitly cast to NumPy scalar in case `value` is a Python scalar.
return onp.dtype(dtype).type(value)
elif shape is PYTHON_SCALAR_SHAPE:
# explicitly cast to Python scalar via https://stackoverflow.com/a/11389998
return numpy_compat.np_asarray(value).item()
elif type(shape) in (list, tuple):
assert onp.shape(value) == tuple(shape)
return value
else:
raise TypeError(type(shape))
def dtype_str(dtype):
return onp.dtype(dtype).name
def format_shape_dtype_string(shape, dtype):
if shape is NUMPY_SCALAR_SHAPE:
return dtype_str(dtype)
elif shape is PYTHON_SCALAR_SHAPE:
return 'py' + dtype_str(dtype)
elif type(shape) in (list, tuple):
shapestr = ','.join(str(dim) for dim in shape)
return '{}[{}]'.format(dtype_str(dtype), shapestr)
elif type(shape) is int:
return '{}[{},]'.format(dtype_str(dtype), shape)
elif isinstance(shape, onp.ndarray):
return '{}[{}]'.format(dtype_str(dtype), shape)
else:
raise TypeError(type(shape))
def _rand_dtype(rand, shape, dtype, scale=1., post=lambda x: x):
"""Produce random values given shape, dtype, scale, and post-processor.
Args:
rand: a function for producing random values of a given shape, e.g. a
bound version of either onp.RandomState.randn or onp.RandomState.rand.
shape: a shape value as a tuple of positive integers.
dtype: a numpy dtype.
scale: optional, a multiplicative scale for the random values (default 1).
post: optional, a callable for post-processing the random values (default
identity).
Returns:
An ndarray of the given shape and dtype using random values based on a call
to rand but scaled, converted to the appropriate dtype, and post-processed.
"""
r = lambda: numpy_compat.np_asarray(scale * rand(*_dims_of_shape(shape)),
dtype)
if onp.issubdtype(dtype, onp.complexfloating):
vals = r() + 1.0j * r()
else:
vals = r()
return _cast_to_shape(numpy_compat.np_asarray(post(vals), dtype), shape,
dtype)
def rand_default(scale=3):
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=scale)
def rand_nonzero():
post = lambda x: onp.where(x == 0, onp.array(1, dtype=x.dtype), x)
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=3, post=post)
def rand_positive():
post = lambda x: x + 1
rand = npr.RandomState(0).rand
return partial(_rand_dtype, rand, scale=2, post=post)
def rand_small():
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=1e-3)
def rand_not_small(offset=10.):
post = lambda x: x + onp.where(x > 0, offset, -offset)
randn = npr.RandomState(0).randn
return partial(_rand_dtype, randn, scale=3., post=post)
def rand_small_positive():
rand = npr.RandomState(0).rand
return partial(_rand_dtype, rand, scale=2e-5)
def rand_uniform(low=0.0, high=1.0):
assert low < high
rand = npr.RandomState(0).rand
post = lambda x: x * (high - low) + low
return partial(_rand_dtype, rand, post=post)
def rand_some_equal():
randn = npr.RandomState(0).randn
rng = npr.RandomState(0)
def post(x):
x_ravel = x.ravel()
if len(x_ravel) == 0:
return x
flips = rng.rand(*onp.shape(x)) < 0.5
return onp.where(flips, x_ravel[0], x)
return partial(_rand_dtype, randn, scale=100., post=post)
def rand_some_inf():
"""Return a random sampler that produces infinities in floating types."""
rng = npr.RandomState(1)
base_rand = rand_default()
"""
TODO: Complex numbers are not correctly tested
If blocks should be switched in order, and relevant tests should be fixed
"""
def rand(shape, dtype):
"""The random sampler function."""
if not onp.issubdtype(dtype, onp.floating):
# only float types have inf
return base_rand(shape, dtype)
if onp.issubdtype(dtype, onp.complexfloating):
base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
out = (rand(shape, base_dtype) +
onp.array(1j, dtype) * rand(shape, base_dtype))
return _cast_to_shape(out, shape, dtype)
dims = _dims_of_shape(shape)
posinf_flips = rng.rand(*dims) < 0.1
neginf_flips = rng.rand(*dims) < 0.1
vals = base_rand(shape, dtype)
vals = onp.where(posinf_flips, onp.array(onp.inf, dtype=dtype), vals)
vals = onp.where(neginf_flips, onp.array(-onp.inf, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
def rand_some_nan():
"""Return a random sampler that produces nans in floating types."""
rng = npr.RandomState(1)
base_rand = rand_default()
def rand(shape, dtype):
"""The random sampler function."""
if onp.issubdtype(dtype, onp.complexfloating):
base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
out = (rand(shape, base_dtype) +
onp.array(1j, dtype) * rand(shape, base_dtype))
return _cast_to_shape(out, shape, dtype)
if not onp.issubdtype(dtype, onp.floating):
# only float types have inf
return base_rand(shape, dtype)
dims = _dims_of_shape(shape)
nan_flips = rng.rand(*dims) < 0.1
vals = base_rand(shape, dtype)
vals = onp.where(nan_flips, onp.array(onp.nan, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
def rand_some_inf_and_nan():
"""Return a random sampler that produces infinities in floating types."""
rng = npr.RandomState(1)
base_rand = rand_default()
"""
TODO: Complex numbers are not correctly tested
If blocks should be switched in order, and relevant tests should be fixed
"""
def rand(shape, dtype):
"""The random sampler function."""
if not onp.issubdtype(dtype, onp.floating):
# only float types have inf
return base_rand(shape, dtype)
if onp.issubdtype(dtype, onp.complexfloating):
base_dtype = onp.real(onp.array(0, dtype=dtype)).dtype
out = (rand(shape, base_dtype) +
onp.array(1j, dtype) * rand(shape, base_dtype))
return _cast_to_shape(out, shape, dtype)
dims = _dims_of_shape(shape)
posinf_flips = rng.rand(*dims) < 0.1
neginf_flips = rng.rand(*dims) < 0.1
nan_flips = rng.rand(*dims) < 0.1
vals = base_rand(shape, dtype)
vals = onp.where(posinf_flips, onp.array(onp.inf, dtype=dtype), vals)
vals = onp.where(neginf_flips, onp.array(-onp.inf, dtype=dtype), vals)
vals = onp.where(nan_flips, onp.array(onp.nan, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
# TODO(mattjj): doesn't handle complex types
def rand_some_zero():
"""Return a random sampler that produces some zeros."""
rng = npr.RandomState(1)
base_rand = rand_default()
def rand(shape, dtype):
"""The random sampler function."""
dims = _dims_of_shape(shape)
zeros = rng.rand(*dims) < 0.5
vals = base_rand(shape, dtype)
vals = onp.where(zeros, onp.array(0, dtype=dtype), vals)
return _cast_to_shape(numpy_compat.np_asarray(vals, dtype=dtype), shape,
dtype)
return rand
def rand_int(low, high=None):
randint = npr.RandomState(0).randint
def fn(shape, dtype):
return randint(low, high=high, size=shape, dtype=dtype)
return fn
def rand_unique_int():
randchoice = npr.RandomState(0).choice
def fn(shape, dtype):
return randchoice(onp.arange(onp.prod(shape), dtype=dtype),
size=shape, replace=False)
return fn
def rand_bool():
rng = npr.RandomState(0)
def generator(shape, dtype):
return _cast_to_shape(rng.rand(*_dims_of_shape(shape)) < 0.5, shape, dtype)
return generator
def check_raises(thunk, err_type, msg):
try:
thunk()
assert False
except err_type as e:
assert str(e).startswith(msg), "\n{}\n\n{}\n".format(e, msg)
def check_raises_regexp(thunk, err_type, pattern):
try:
thunk()
assert False
except err_type as e:
assert re.match(pattern, str(e)), "{}\n\n{}\n".format(e, pattern)
def _iter_eqns(jaxpr):
# TODO(necula): why doesn't this search in params?
for eqn in jaxpr.eqns:
yield eqn
for subjaxpr in core.subjaxprs(jaxpr):
yield from _iter_eqns(subjaxpr)
def assert_dot_precision(expected_precision, fun, *args):
jaxpr = api.make_jaxpr(fun)(*args)
precisions = [eqn.params['precision'] for eqn in _iter_eqns(jaxpr.jaxpr)
if eqn.primitive == lax.dot_general_p]
for precision in precisions:
msg = "Unexpected precision: {} != {}".format(expected_precision, precision)
assert precision == expected_precision, msg
_CACHED_INDICES: Dict[int, Sequence[int]] = {}
def cases_from_list(xs):
xs = list(xs)
n = len(xs)
k = min(n, FLAGS.num_generated_cases)
# Random sampling for every parameterized test is expensive. Do it once and
# cache the result.
indices = _CACHED_INDICES.get(n)
if indices is None:
rng = npr.RandomState(42)
_CACHED_INDICES[n] = indices = rng.permutation(n)
return [xs[i] for i in indices[:k]]
def cases_from_gens(*gens):
sizes = [1, 3, 10]
cases_per_size = int(FLAGS.num_generated_cases / len(sizes)) + 1
for size in sizes:
for i in range(cases_per_size):
yield ('_{}_{}'.format(size, i),) + tuple(gen(size) for gen in gens)
def to_np(a):
return nest.map_structure(np_array_ops.asarray, a)
def to_tf_fn(f):
return lambda *args: f(*to_np(args))
class TestCase(parameterized.TestCase):
"""Base class for tests including numerical checks and boilerplate."""
# copied from jax.test_util
def setUp(self):
super().setUp()
self._rng = npr.RandomState(zlib.adler32(self._testMethodName.encode()))
# copied from jax.test_util
def rng(self):
return self._rng
# TODO(mattjj): this obscures the error messages from failures, figure out how
# to re-enable it
# def tearDown(self) -> None:
# assert core.reset_trace_state()
def assertArraysAllClose(self, x, y, check_dtypes, atol=None, rtol=None):
"""Assert that x and y are close (up to numerical tolerances)."""
self.assertEqual(x.shape, y.shape)
atol = max(tolerance(_dtype(x), atol), tolerance(_dtype(y), atol))
rtol = max(tolerance(_dtype(x), rtol), tolerance(_dtype(y), rtol))
_assert_numpy_allclose(x, y, atol=atol, rtol=rtol)
if check_dtypes:
self.assertDtypesMatch(x, y)
def assertDtypesMatch(self, x, y):
if FLAGS.enable_x64:
self.assertEqual(_dtype(x), _dtype(y))
def assertAllClose(self, x, y, check_dtypes, atol=None, rtol=None):
"""Assert that x and y, either arrays or nested tuples/lists, are close."""
if isinstance(x, dict):
self.assertIsInstance(y, dict)
self.assertEqual(set(x.keys()), set(y.keys()))
for k in x:
self.assertAllClose(x[k], y[k], check_dtypes, atol=atol, rtol=rtol)
elif is_sequence(x) and not hasattr(x, '__array__'):
self.assertTrue(is_sequence(y) and not hasattr(y, '__array__'))
self.assertEqual(len(x), len(y))
for x_elt, y_elt in zip(x, y):
self.assertAllClose(x_elt, y_elt, check_dtypes, atol=atol, rtol=rtol)
elif hasattr(x, '__array__') or onp.isscalar(x):
self.assertTrue(hasattr(y, '__array__') or onp.isscalar(y))
if check_dtypes:
self.assertDtypesMatch(x, y)
x = numpy_compat.np_asarray(x)
y = numpy_compat.np_asarray(y)
self.assertArraysAllClose(x, y, check_dtypes=False, atol=atol, rtol=rtol)
elif x == y:
return
else:
raise TypeError((type(x), type(y)))
def assertMultiLineStrippedEqual(self, expected, what):
"""Asserts two strings are equal, after stripping each line."""
ignore_space_re = re.compile(r'\s*\n\s*')
expected_clean = re.sub(ignore_space_re, '\n', expected.strip())
what_clean = re.sub(ignore_space_re, '\n', what.strip())
self.assertMultiLineEqual(expected_clean, what_clean,
msg="Found\n{}\nExpecting\n{}".format(what, expected))
def _CheckAgainstNumpy(self, numpy_reference_op, lax_op, args_maker,
check_dtypes=True, tol=None):
args = args_maker()
lax_ans = lax_op(*args)
numpy_ans = numpy_reference_op(*args)
self.assertAllClose(numpy_ans, lax_ans, check_dtypes=check_dtypes,
atol=tol, rtol=tol)
def _CompileAndCheck(self,
fun,
args_maker,
check_dtypes=True,
rtol=None,
atol=None,
check_eval_on_shapes=True,
check_incomplete_shape=True,
check_unknown_rank=True,
static_argnums=(),
check_experimental_compile=True,
check_xla_forced_compile=True):
"""Compiles the function and checks the results.
Args:
fun: the function to be checked.
args_maker: a callable that returns a tuple which will be used as the
positional arguments.
check_dtypes: whether to check that the result dtypes from non-compiled
and compiled runs agree.
rtol: relative tolerance for allclose assertions.
atol: absolute tolerance for allclose assertions.
check_eval_on_shapes: whether to run `eval_on_shapes` on the function and
check that the result shapes and dtypes are correct.
check_incomplete_shape: whether to check that the function can handle
incomplete shapes (including those with and without a known rank).
check_unknown_rank: (only has effect when check_incomplete_shape is True)
whether to check that the function can handle unknown ranks.
static_argnums: indices of arguments to be treated as static arguments for
`jit` and `eval_on_shapes`.
check_experimental_compile: whether to check compilation with
experimental_compile=True (in addition to compilation without the flag).
check_xla_forced_compile: whether to check compilation with
forced_compile=True (in addition to compilation without the flag). This
flag is different from experimental_compile because it enforces
whole-function compilation while the latter doesn't. TPU requires
whole-function compilation.
"""
args = args_maker()
for x in args:
if not hasattr(x, 'dtype'):
# If there is a input that doesn't have dtype info, jit and
# eval_on_shapes may pick a different dtype for it than numpy, so we
# skip the dtype check.
check_dtypes = False
python_ans = fun(*args)
python_shapes = nest.map_structure(onp.shape, python_ans)
onp_shapes = nest.map_structure(
lambda x: onp.shape(numpy_compat.np_asarray(x)), python_ans
)
self.assertEqual(python_shapes, onp_shapes)
def check_compile(**kwargs):
# `wrapped_fun` and `python_should_be_executing` are used to check that
# when the jitted function is called the second time, the original Python
# function won't be executed.
def wrapped_fun(*args):
self.assertTrue(python_should_be_executing)
return fun(*args)
cfun = nje.jit(wrapped_fun, static_argnums=static_argnums, **kwargs)
python_should_be_executing = True
monitored_ans = cfun(*args)
python_should_be_executing = False
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, monitored_ans, check_dtypes, atol, rtol)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
# Run `cfun` with a different set of arguments to check that changing
# arguments won't cause recompilation.
new_args = args_maker()
skip_retracing_test = False
for old, new in zip(nest.flatten(args), nest.flatten(new_args)):
if nje.most_precise_int_dtype(old) != nje.most_precise_int_dtype(new):
# If the old and new arguments result in different dtypes (because
# they fall into different value ranges), tf-numpy will retrace, so we
# skip the no-retrace test.
skip_retracing_test = True
if not skip_retracing_test:
python_should_be_executing = True
new_python_ans = fun(*new_args)
python_should_be_executing = False
compiled_ans = cfun(*new_args)
self.assertAllClose(new_python_ans, compiled_ans, check_dtypes, atol,
rtol)
check_compile()
if check_experimental_compile:
check_compile(experimental_compile=True)
if check_xla_forced_compile:
check_compile(xla_forced_compile=True)
if check_eval_on_shapes:
# Check that nje.eval_on_shapes can get complete output shapes given
# complete input shapes.
cfun = nje.eval_on_shapes(fun, static_argnums=static_argnums)
compiled_ans = cfun(*args)
flat_python_ans = nest.flatten(python_ans)
flat_compiled_ans = nest.flatten(compiled_ans)
self.assertEqual(len(flat_python_ans), len(flat_compiled_ans))
for a, b in zip(flat_python_ans, flat_compiled_ans):
if hasattr(a, 'shape'):
self.assertEqual(a.shape, b.shape)
if check_dtypes and hasattr(a, 'dtype'):
self.assertEqual(dtypes.as_dtype(a.dtype), b.dtype)
# If some argument doesn't have a `dtype` attr (e.g. a Python scalar), we
# skip incomplete-shape checks, since shape specs need dtype. It's OK to
# skip since the same incomplete-shape checks will run for []-shaped arrays.
if check_incomplete_shape and all(hasattr(x, 'dtype') for x in args):
# Check partial shapes with known ranks.
# Numpy scalars (created by e.g. np.int32(5)) have `dtype` but not
# `shape`.
if all(hasattr(x, 'shape') for x in args):
specs = [
tensor.TensorSpec([None] * len(x.shape), x.dtype) for x in args
]
cfun = nje.jit(
fun, static_argnums=static_argnums, input_signature=specs
)
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
if check_unknown_rank:
# Check unknown ranks.
specs = [tensor.TensorSpec(None, x.dtype) for x in args]
cfun = nje.jit(
fun, static_argnums=static_argnums, input_signature=specs)
compiled_ans = cfun(*args)
self.assertAllClose(python_ans, compiled_ans, check_dtypes, atol, rtol)
def check_grads(self, f, args, atol=None, rtol=None, delta=None):
"""Check gradients against finite differences.
Args:
f: function to check at ``f(*args)``.
args: a list or tuple of argument values.
atol: absolute tolerance for gradient equality.
rtol: relative tolerance for gradient equality.
delta: step size used for finite differences.
"""
if delta is None:
# Optimal stepsize for central difference is O(epsilon^{1/3}).
dtype = np_utils.result_type(*args)
epsilon = onp.finfo(dtype).eps
delta = epsilon ** (1.0 / 3.0)
theoretical, numerical = gradient_checker_v2.compute_gradient(
to_tf_fn(f), args, delta=delta)
self.assertAllClose(theoretical, numerical, check_dtypes=False, atol=atol,
rtol=rtol)
@contextmanager
def ignore_warning(**kw):
with warnings.catch_warnings():
warnings.filterwarnings("ignore", **kw)
yield
def disable(_):
def wrapper(self, *args, **kwargs):
self.skipTest('Test is disabled')
return wrapper