Function LiteralToPython

tensorflow/compiler/xla/python/types.cc:197–241

}

StatusOr<py::object> LiteralToPython(std::shared_ptr<xla::Literal> literal) {
  xla::Literal& m = *literal;
  if (m.shape().IsTuple()) {
    std::vector<Literal> elems = m.DecomposeTuple();
    std::vector<py::object> arrays(elems.size());
    for (int i = 0; i < elems.size(); ++i) {
      TF_ASSIGN_OR_RETURN(
          arrays[i],
          LiteralToPython(absl::make_unique<Literal>(std::move(elems[i]))));
    }
    py::tuple result(elems.size());
    for (int i = 0; i < elems.size(); ++i) {
      PyTuple_SET_ITEM(result.ptr(), i, arrays[i].release().ptr());
    }
    return result;
  }
  TF_RET_CHECK(m.shape().IsArray());

  py::object literal_object = py::cast(literal);
  TF_ASSIGN_OR_RETURN(std::string format, FormatDescriptorForPrimitiveType(
                                              m.shape().element_type()));
  py::buffer_info info(
      m.untyped_data(),  // Pointer to buffer
      xla::ShapeUtil::ByteSizeOfPrimitiveType(
          m.shape().element_type()),  // Size of one scalar
      format,                         // Python struct-style format descriptor
      m.shape().dimensions_size(),    // Number of dimensions
      m.shape().dimensions(),         // Buffer dimensions
      ByteStridesForShape(m.shape())  // Strides (in bytes) for each index
  );

  py::array array(pybind11::dtype(info), info.shape, info.strides, info.ptr,
                  literal_object);
  if (m.shape().element_type() == xla::BF16) {
    // We requested an array of uint16 since NumPy doesn't know how
    // to produce our custom bfloat16 type. Reinterpret the array as bfloat16
    // before handing it back to the caller.
    TF_ASSIGN_OR_RETURN(py::object bfloat16, Bfloat16Dtype());
    array = py::reinterpret_steal<py::array>(
        PyArray_View(reinterpret_cast<PyArrayObject*>(array.ptr()),
                     reinterpret_cast<PyArray_Descr*>(bfloat16.release().ptr()),
                     static_cast<PyTypeObject*>(nullptr)));
  }
  return array;
}

StatusOr<PythonBufferTree> GetPythonBufferTree(const py::object& argument) {
  PythonBufferTree tree;