Repository navigation
Conversation
`np_asarray`, `np_where` and `np_reshape` parsed `np.__version__` with `np.lib.NumpyVersion` on every call. `np_asarray` runs for every value that `tf.numpy_function`/`tf.py_func` returns and for every component that `tf.data.Dataset.from_generator` yields, so the version parsing showed up as the largest cost in those per-element paths. The installed NumPy version doesn't change at runtime, so compute the flags at import. `tf.experimental.numpy.sign` and `linspace` did the same per call; they get a module-level flag too.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes NumPy version checks by caching the results in module-level constants instead of parsing the version string on every function call. The review feedback suggests avoiding code duplication by importing the centralized constants from numpy_compat into the other modules, and replacing the deprecated np.lib.NumpyVersion with packaging.version.Version to ensure future compatibility.
| import numpy as np | ||
|
|
||
| # The installed NumPy version does not change while the process runs, so | ||
| # check it once here instead of parsing `np.__version__` on every call. | ||
| _NP_2_0_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0' | ||
| _NP_2_1_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.1.0.rc0' |
There was a problem hiding this comment.
np.lib.NumpyVersion is deprecated in NumPy 2.0 and is scheduled for removal in a future release. To avoid deprecation warnings and ensure future compatibility, use packaging.version.Version instead.
| import numpy as np | |
| # The installed NumPy version does not change while the process runs, so | |
| # check it once here instead of parsing `np.__version__` on every call. | |
| _NP_2_0_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0' | |
| _NP_2_1_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.1.0.rc0' | |
| import numpy as np | |
| from packaging.version import Version | |
| # The installed NumPy version does not change while the process runs, so | |
| # check it once here instead of parsing `np.__version__` on every call. | |
| _NP_2_0_OR_LATER = Version(np.__version__) >= Version('2.0.0.dev0') | |
| _NP_2_1_OR_LATER = Version(np.__version__) >= Version('2.1.0.rc0') |
| from tensorflow.python.util import nest | ||
| from tensorflow.python.util import tf_export | ||
|
|
||
| # The installed NumPy version does not change while the process runs. | ||
| _NP_2_0_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0' |
There was a problem hiding this comment.
Instead of duplicating the NumPy version check and parsing logic across multiple files, import the centralized _NP_2_0_OR_LATER constant from tensorflow.python.util.numpy_compat.
| from tensorflow.python.util import nest | |
| from tensorflow.python.util import tf_export | |
| # The installed NumPy version does not change while the process runs. | |
| _NP_2_0_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0' | |
| from tensorflow.python.util import nest | |
| from tensorflow.python.util import tf_export | |
| from tensorflow.python.util.numpy_compat import _NP_2_0_OR_LATER |
| from tensorflow.python.ops.numpy_ops import np_utils | ||
| from tensorflow.python.util import tf_export | ||
|
|
||
| # The installed NumPy version does not change while the process runs. | ||
| _NP_2_0_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0' |
There was a problem hiding this comment.
Instead of duplicating the NumPy version check and parsing logic across multiple files, import the centralized _NP_2_0_OR_LATER constant from tensorflow.python.util.numpy_compat.
| from tensorflow.python.ops.numpy_ops import np_utils | |
| from tensorflow.python.util import tf_export | |
| # The installed NumPy version does not change while the process runs. | |
| _NP_2_0_OR_LATER = np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0' | |
| from tensorflow.python.ops.numpy_ops import np_utils | |
| from tensorflow.python.util import tf_export | |
| from tensorflow.python.util.numpy_compat import _NP_2_0_OR_LATER |
numpy_compat.np_asarray,np_whereandnp_reshapebuild twonp.lib.NumpyVersionobjects (a regex parse ofnp.__version__and of the bound) on every call.np_asarrayruns for every value returned fromtf.numpy_function/tf.compat.v1.py_func, andDataset.from_generatorcalls it for every component of every element, so this check was the largest single cost in those per-element Python paths. The installed NumPy version can't change while the process runs, so this computes the flags once at import.tf.experimental.numpy.signandlinspacehad the same per-call check and get a module-level flag too. Behaviour is unchanged.Measured on Windows 11, i7-9700, Python 3.11.9, tf-nightly 2.23.0.dev20261112, NumPy 2.4.6, CPU only. Base and patched
numpy_compat.pywere swapped in alternately, 4 runs each, one process per run; median [min-max]:batch(64))Dataset.from_generatoryielding(float32[32], int64)withoutput_signatureDataset.from_tensor_slices(...).map(lambda a: tf.numpy_function(f, [a], tf.float32))np_asarray(np.float32 array of 32, dtype=np.float32)alone goes from 7.0 to 0.54 us per call.numpy_compat_test.py,script_ops_test.py,from_generator_test.pyandtensor_util_test.pypass with the patched module, andtnp.sign/tnp.linspacegive the same results as NumPy 2.4 on tf-nightly.py_func_test.pyhas the same result before and after: everything passes excepttestRenamedDeviceInTestClusterCorrectlyIdentifiedAsLocalhost, which fails on this Windows machine either way.