Repository navigation
feat(numpy_ops): implement cov and corrcoef in tf.experimental.numpy - #130339
Yogyabhilash wants to merge 1 commit into
Conversation
![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 implements the cov (covariance) and corrcoef (correlation coefficient) functions in experimental.numpy along with their corresponding unit tests. Feedback on the implementation highlights several areas for improvement: clipping both real and imaginary parts of complex correlation coefficients in corrcoef to match NumPy behavior, using c.ndim instead of c.shape.rank to avoid potential AttributeErrors, removing redundant dimension checks for m and y when they are already guaranteed to be 2D, and adding static shape validation for fweights and aweights to prevent silent broadcasting errors.
| if c.dtype in (dtypes.complex64, dtypes.complex128): | ||
| real_dtype = math_ops.real(c).dtype | ||
| real_part = clip_ops.clip_by_value( | ||
| math_ops.real(c), | ||
| math_ops.cast(-1.0, real_dtype), | ||
| math_ops.cast(1.0, real_dtype), | ||
| ) | ||
| c = math_ops.complex(real_part, math_ops.imag(c)) |
There was a problem hiding this comment.
In NumPy, corrcoef clips both the real and imaginary parts of complex correlation coefficients to the range [-1.0, 1.0]. Currently, only the real part is being clipped, which can lead to values outside the valid range for the imaginary part of complex inputs. We should clip both parts to ensure exact parity with NumPy.
| if c.dtype in (dtypes.complex64, dtypes.complex128): | |
| real_dtype = math_ops.real(c).dtype | |
| real_part = clip_ops.clip_by_value( | |
| math_ops.real(c), | |
| math_ops.cast(-1.0, real_dtype), | |
| math_ops.cast(1.0, real_dtype), | |
| ) | |
| c = math_ops.complex(real_part, math_ops.imag(c)) | |
| if c.dtype in (dtypes.complex64, dtypes.complex128): | |
| real_dtype = math_ops.real(c).dtype | |
| real_part = clip_ops.clip_by_value( | |
| math_ops.real(c), | |
| math_ops.cast(-1.0, real_dtype), | |
| math_ops.cast(1.0, real_dtype), | |
| ) | |
| imag_part = clip_ops.clip_by_value( | |
| math_ops.imag(c), | |
| math_ops.cast(-1.0, real_dtype), | |
| math_ops.cast(1.0, real_dtype), | |
| ) | |
| c = math_ops.complex(real_part, imag_part) |
References
- Challenge the implementation by identifying edge cases, failure scenarios, and incorrect assumptions. (link)
| if c.shape.rank == 0 or c.shape == (1, 1): | ||
| return math_ops.cast(1.0, c.dtype) |
There was a problem hiding this comment.
Using c.shape.rank can raise an AttributeError if c.shape is a standard Python tuple (which is the case for tf.experimental.numpy.ndarray to maintain NumPy compatibility). It is safer and more idiomatic to use c.ndim to check the rank of the array, which is consistent with how rank is checked elsewhere in this file.
| if c.shape.rank == 0 or c.shape == (1, 1): | |
| return math_ops.cast(1.0, c.dtype) | |
| if c.ndim == 0 or c.shape == (1, 1): | |
| return math_ops.cast(1.0, c.dtype) |
References
- Detect Low-Quality or Non-Idiomatic Code: Flag patterns that do not align with established TensorFlow practices. (link)
| if not rowvar and m.ndim == 2: | ||
| m = np_array_ops.transpose(m) |
There was a problem hiding this comment.
Since m is guaranteed to be 2D after the preceding if m.ndim == 1: block, checking m.ndim == 2 is redundant. We can simplify this condition to just if not rowvar:.
| if not rowvar and m.ndim == 2: | |
| m = np_array_ops.transpose(m) | |
| if not rowvar: | |
| m = np_array_ops.transpose(m) |
References
- Detect Low-Quality or Non-Idiomatic Code: Flag patterns that do not align with established TensorFlow practices. (link)
| if not rowvar and y.ndim == 2: | ||
| y = np_array_ops.transpose(y) |
There was a problem hiding this comment.
Similarly, since y is guaranteed to be 2D after the preceding if y.ndim == 1: block, checking y.ndim == 2 is redundant. We can simplify this condition to just if not rowvar:.
| if not rowvar and y.ndim == 2: | |
| y = np_array_ops.transpose(y) | |
| if not rowvar: | |
| y = np_array_ops.transpose(y) |
References
- Detect Low-Quality or Non-Idiomatic Code: Flag patterns that do not align with established TensorFlow practices. (link)
| if fweights is not None: | ||
| fweights = np_array_ops.array(fweights) | ||
| if fweights.ndim != 1: | ||
| raise ValueError('fweights must be 1D') | ||
| fweights = math_ops.cast(fweights, real_dtype) | ||
| w = fweights | ||
| if aweights is not None: | ||
| aweights = np_array_ops.array(aweights) | ||
| if aweights.ndim != 1: | ||
| raise ValueError('aweights must be 1D') | ||
| aweights = math_ops.cast(aweights, real_dtype) | ||
| w = aweights if w is None else w * aweights |
There was a problem hiding this comment.
To ensure exact parity with NumPy and prevent silent broadcasting errors or incorrect results, we should validate that the lengths of fweights and aweights match the number of observations (m.shape[1]) when known statically.
if fweights is not None:
fweights = np_array_ops.array(fweights)
if fweights.ndim != 1:
raise ValueError('fweights must be 1D')
if fweights.shape[0] is not None and m.shape[1] is not None:
if fweights.shape[0] != m.shape[1]:
raise ValueError('fweights and m have incompatible shapes')
fweights = math_ops.cast(fweights, real_dtype)
w = fweights
if aweights is not None:
aweights = np_array_ops.array(aweights)
if aweights.ndim != 1:
raise ValueError('aweights must be 1D')
if aweights.shape[0] is not None and m.shape[1] is not None:
if aweights.shape[0] != m.shape[1]:
raise ValueError('aweights and m have incompatible shapes')
aweights = math_ops.cast(aweights, real_dtype)
w = aweights if w is None else w * aweightsReferences
- Validate tensor shapes, dtypes, and execution paths. Ensure proper error handling and defensive checks. (link)
Problem & Motivation
tf.experimental.numpyprovides drop-in NumPy-compatible routines for TensorFlow pipelines. However, two foundational statistics operations from NumPy—np.cov(covariance matrix) andnp.corrcoef(Pearson correlation coefficients)—were missing fromtf.experimental.numpy.Users analyzing covariance/correlation structures or computing empirical statistics in machine learning pipelines had to manually implement these operations or convert to host NumPy arrays.
Substantive Changes
tensorflow/python/ops/numpy_ops/np_math_ops.py:cov:@tf_export.tf_export('experimental.numpy.cov').mand optionaly.rowvar=True(variables are rows, observations are columns) androwvar=False.fweights) and observation weights (aweights), including combined frequency-observation weighting.bias=True(population normalization byddofdegree-of-freedom normalization.real_dtypeand returnsNaNwhen degrees of freedomm.shape == (1, 1)).corrcoef:@tf_export.tf_export('experimental.numpy.corrcoef').cov(x, y, rowvar=rowvar).clip_ops.clip_by_valueto eliminate floating-point precision overshoots.tensorflow/python/ops/numpy_ops/np_math_ops_test.py:testCovandtestCorrcoef.rowvar=False, two-input combinations(x, y), integer, float32, float64, and complex128 inputs.bias=True,ddof=2,fweights,aweights, and combined weights.Verification
All functionality and unit tests were validated locally against NumPy:
np.covandnp.corrcoefwithin standard floating-point tolerances.