Repository navigation
feat(numpy_ops): implement unravel_index and ravel_multi_index in tf.experimental.numpy - #130340
Yogyabhilash wants to merge 1 commit into
Conversation
…experimental.numpy
![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 introduces the unravel_index and ravel_multi_index functions to TensorFlow's NumPy-compatible API, along with corresponding unit tests. The feedback highlights critical compatibility issues with tf.function and graph mode due to Python conditional statements depending on runtime Tensor values; using context.executing_eagerly() and control_flow_ops.Assert is recommended instead. Additionally, the reviewer suggests adding validation checks for empty dimensions to match NumPy's behavior and adding test cases that execute these functions within a tf.function context.
| if math_ops.reduce_any(indices < 0) or math_ops.reduce_any( | ||
| indices >= total_size | ||
| ): | ||
| raise ValueError('index is out of bounds') |
There was a problem hiding this comment.
Using Python if statements with conditions that depend on the runtime values of TensorFlow Tensors (such as math_ops.reduce_any(indices < 0)) will raise an OperatorNotAllowedInGraphError when executed inside a tf.function or in graph mode.
To support both eager and graph execution modes, we should check if we are executing eagerly. If so, we can raise a ValueError to match NumPy's behavior. Otherwise, we should use control_flow_ops.Assert to perform the check at runtime in the graph, and use ops.control_dependencies to ensure the assertion is executed.
if context.executing_eagerly():
if math_ops.reduce_any(indices < 0) or math_ops.reduce_any(
indices >= total_size
):
raise ValueError('index is out of bounds')
else:
assert_op = control_flow_ops.Assert(
math_ops.reduce_all(indices >= 0) & math_ops.reduce_all(indices < total_size),
['index is out of bounds']
)
with ops.control_dependencies([assert_op]):
indices = array_ops.identity(indices)References
- Ensure compatibility with tf.function and avoid retracing issues. (link)
| if m == 'raise': | ||
| if math_ops.reduce_any(arr < 0) or math_ops.reduce_any(arr >= d): | ||
| raise ValueError('invalid entry in coordinates array') | ||
| processed.append(arr) |
There was a problem hiding this comment.
Similar to unravel_index, using Python if with Tensor-dependent conditions will fail in graph mode or inside tf.function with OperatorNotAllowedInGraphError. We should use context.executing_eagerly() to guard the Python exception and use control_flow_ops.Assert for graph mode.
| if m == 'raise': | |
| if math_ops.reduce_any(arr < 0) or math_ops.reduce_any(arr >= d): | |
| raise ValueError('invalid entry in coordinates array') | |
| processed.append(arr) | |
| if m == 'raise': | |
| if context.executing_eagerly(): | |
| if math_ops.reduce_any(arr < 0) or math_ops.reduce_any(arr >= d): | |
| raise ValueError('invalid entry in coordinates array') | |
| else: | |
| assert_op = control_flow_ops.Assert( | |
| math_ops.reduce_all(arr >= 0) & math_ops.reduce_all(arr < d), | |
| ['invalid entry in coordinates array'] | |
| ) | |
| with ops.control_dependencies([assert_op]): | |
| arr = array_ops.identity(arr) | |
| processed.append(arr) |
References
- Ensure compatibility with tf.function and avoid retracing issues. (link)
| def unravel_index(indices, shape, order='C'): # pylint: disable=missing-docstring | ||
| if order not in ('C', 'F'): | ||
| raise ValueError(f"order must be 'C' or 'F', got {order}") | ||
| indices = asarray(indices) |
There was a problem hiding this comment.
NumPy's unravel_index raises a ValueError if the shape argument is empty (i.e., 0-dimensional). We should add a check to ensure shape has at least one dimension.
def unravel_index(indices, shape, order='C'): # pylint: disable=missing-docstring
if order not in ('C', 'F'):
raise ValueError(f"order must be 'C' or 'F', got {order}")
if not shape:
raise ValueError('number of dimensions must be positive')
indices = asarray(indices)| if len(multi_index) != len(dims): | ||
| raise ValueError('multi_index and dims must have the same length') |
There was a problem hiding this comment.
NumPy's ravel_multi_index raises a ValueError if dims is empty. We should add a check to ensure dims is not empty.
| if len(multi_index) != len(dims): | |
| raise ValueError('multi_index and dims must have the same length') | |
| if len(multi_index) != len(dims): | |
| raise ValueError('multi_index and dims must have the same length') | |
| if not dims: | |
| raise ValueError('number of dimensions must be positive') |
| with self.assertRaises(ValueError): | ||
| np_array_ops.unravel_index(9, (3, 3)) | ||
| with self.assertRaises(ValueError): | ||
| np_array_ops.unravel_index(0, (3, 3), order='X') |
There was a problem hiding this comment.
To ensure full compatibility with tf.function (graph mode), we should add test cases that wrap unravel_index and ravel_multi_index in a tf.function and execute them, verifying that they do not raise OperatorNotAllowedInGraphError.
References
- Verify that unit tests are included for any new logic, feature, or bug fix. (link)
Problem & Motivation
tf.experimental.numpyprovides drop-in NumPy parity for TensorFlow users. However, two widely used coordinate and indexing functions—unravel_indexandravel_multi_index—were completely absent fromtf.experimental.numpy.tf.unravel_indexop, it returns a single 2D tensor rather than a tuple of coordinate arrays, does not accept arbitrary-dimensional indices, and lacks Fortran ('F') order support.ravel_multi_index(converting multi-dimensional coordinates to flat indices with boundary wrapping and clipping modes) had no equivalent intf.experimental.numpy.Substantive Changes
tensorflow/python/ops/numpy_ops/np_array_ops.py:unravel_index:@tf_export.tf_export('experimental.numpy.unravel_index').indices.'C', default) and Fortran-style column-major ('F') unravelling.ValueError.ravel_multi_index:@tf_export.tf_export('experimental.numpy.ravel_multi_index').'raise'(bounds check),'wrap'(modulo wrap), and'clip'(clamped into valid range[0, dim - 1]).'C'and'F'unravelling orders.tensorflow/python/ops/numpy_ops/np_array_ops_test.py:testUnravelIndexandtestRavelMultiIndex.Verification
All functionality and unit tests were validated locally against NumPy:
unravel_indexandravel_multi_indexmatch NumPy exactly across scalars, vectors, matrices, and Fortran layout.