MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / PyUnpackTensorIndex

Function PyUnpackTensorIndex

oneflow/api/python/functional/common.cpp:314–400  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

312 || numpy::PyArrayCheckLongScalar(obj);
313}
314TensorIndex PyUnpackTensorIndex(PyObject* obj) {
315 TensorIndex tensor_index;
316 // Obvious single-entry cases.
317 if (PySlice_Check(obj) // NOLINT
318 || PyLong_Check(obj) // NOLINT
319 || obj == Py_Ellipsis // NOLINT
320 || obj == Py_None // NOLINT
321 || PyTensor_Check(obj) // NOLINT
322 || !PySequence_Check(obj) // NOLINT
323 || numpy::PyArrayCheckLongScalar(obj) // NOLINT
324 || PyUnicode_Check(obj)) {
325 tensor_index.emplace_back(detail::UnpackIndexItem(obj));
326 return tensor_index;
327 }
328 PyObject* tup = NULL;
329 Py_ssize_t n = 0;
330 if (PyTuple_Check(obj)) {
331 tup = PySequence_Tuple(obj);
332 n = PySequence_Size(tup);
333 } else {
334 // The follow comments are from numpy:
335 // https://github.com/numpy/numpy/blob/main/numpy/core/src/multiarray/mapping.c#L266
336 /*
337 * At this point, we're left with a non-tuple, non-array, sequence:
338 * typically, a list. We use some somewhat-arbitrary heuristics from here
339 * onwards to decided whether to treat that list as a single index, or a
340 * list of indices.
341 */
342 n = PySequence_Size(obj);
343 // Negative size indicates a Python error in the PySequence_Size call.
344 if (n < 0) {
345 PyErr_Clear();
346 tensor_index.emplace_back(detail::UnpackIndexItem(obj));
347 return tensor_index;
348 }
349 // The follow comments are from numpy:
350 // https://github.com/numpy/numpy/blob/main/numpy/core/src/multiarray/mapping.c#L280
351 /*
352 * Backwards compatibility only takes effect for short sequences - otherwise
353 * we treat it like any other scalar.
354 *
355 * Sequences < NPY_MAXDIMS with any slice objects
356 * or newaxis, Ellipsis or other arrays or sequences
357 * embedded, are considered equivalent to an indexing
358 * tuple. (`a[[[1,2], [3,4]]] == a[[1,2], [3,4]]`)
359 */
360 if (n >= /*NPY_MAXDIMS=*/32) {
361 tensor_index.emplace_back(detail::UnpackIndexItem(obj));
362 return tensor_index;
363 }
364 // Check whether we should unpack the index like a tuple.
365 bool commit_to_unpack = false;
366 for (Py_ssize_t i = 0; i < n; ++i) {
367 PyObject* item = PySequence_GetItem(obj, i);
368 if (commit_to_unpack) {
369 CHECK_OR_THROW(item) << "Sequence index is required.";
370 } else {
371 if (!item) {

Callers 2

PyTensorObject_setitemFunction · 0.85
ObjectAs<TensorIndex>Method · 0.85

Calls 5

PyTensor_CheckFunction · 0.85
PyArrayCheckLongScalarFunction · 0.85
UnpackIndexItemFunction · 0.85
emplace_backMethod · 0.45
resizeMethod · 0.45

Tested by

no test coverage detected