<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/python/py_buffer.h"

#include <functional>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include "absl/base/casts.h"
#include "pybind11/pybind11.h"
#include "pybind11/pytypes.h"
#include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
#include "tensorflow/compiler/xla/python/py_client.h"
#include "tensorflow/compiler/xla/python/python_ref_manager.h"
#include "tensorflow/compiler/xla/python/python_utils.h"
#include "tensorflow/compiler/xla/python/transfer_guard_lib.h"
#include "tensorflow/compiler/xla/python/types.h"
#include "tensorflow/compiler/xla/python/util.h"
#include "tensorflow/compiler/xla/util.h"

namespace xla {

namespace py = pybind11;

namespace {

// Representation of a DeviceArrayBase as a Python object. Since
// a DeviceArrayBase has no fields, this is just a PyObject.
struct PyBufferBasePyObject {
  PyObject_HEAD;
};
static_assert(std::is_standard_layout<PyBufferBasePyObject>::value,
              "PyBufferBasePyObject must be standard layout");

// Representation of a DeviceArray as a Python object.
struct PyBufferPyObject {
  PyBufferBasePyObject base;
  PyBuffer buffer;
  // Used by the Python interpreter to maintain a list of weak references to
  // this object.
  PyObject* weakrefs;
};
static_assert(std::is_standard_layout<PyBufferPyObject>::value,
              "PyBufferPyObject must be standard layout");

PyObject* PyBuffer_tp_new(PyTypeObject* subtype, PyObject* args,
                          PyObject* kwds) {
  PyBufferPyObject* self =
      reinterpret_cast<PyBufferPyObject*>(subtype->tp_alloc(subtype, 0));
  if (!self) return nullptr;
  self->weakrefs = nullptr;
  return reinterpret_cast<PyObject*>(self);
}

void PyBuffer_tp_dealloc(PyObject* self) {
  PyTypeObject* tp = Py_TYPE(self);
  PyBufferPyObject* o = reinterpret_cast<PyBufferPyObject*>(self);
  if (o->weakrefs) {
    PyObject_ClearWeakRefs(self);
  }
  o->buffer.~PyBuffer();
  tp->tp_free(self);
  Py_DECREF(tp);
}

}  // namespace

/*static*/ PyBuffer::object PyBuffer::Make(
    std::shared_ptr<PyClient> client, std::shared_ptr<PjRtBuffer> buffer,
    std::shared_ptr<Traceback> traceback) {
  py::object obj = py::reinterpret_steal<py::object>(PyBuffer_tp_new(
      reinterpret_cast<PyTypeObject*>(type_), nullptr, nullptr));
  PyBufferPyObject* buf = reinterpret_cast<PyBufferPyObject*>(obj.ptr());
  new (&buf->buffer)
      PyBuffer(std::move(client), std::move(buffer), std::move(traceback));
  return py::reinterpret_borrow<PyBuffer::object>(obj);
}

bool PyBuffer::IsPyBuffer(py::handle handle) {
  return handle.get_type() == PyBuffer::type();
}

/*static*/ PyBuffer* PyBuffer::AsPyBufferUnchecked(pybind11::handle handle) {
  return &(reinterpret_cast<PyBufferPyObject*>(handle.ptr())->buffer);
}

/*static*/ StatusOr<PyBuffer*> PyBuffer::AsPyBuffer(pybind11::handle handle) {
  if (!IsPyBuffer(handle)) {
    return InvalidArgument("Expected a DeviceArray");
  }
  return AsPyBufferUnchecked(handle);
}

py::handle PyBuffer::AsHandle() {
  return reinterpret_cast<PyObject*>(reinterpret_cast<char*>(this) -
                                     offsetof(PyBufferPyObject, buffer));
}

PyBuffer::PyBuffer(std::shared_ptr<PyClient> client,
                   std::shared_ptr<PjRtBuffer> buffer,
                   std::shared_ptr<Traceback> traceback)
    : client_(std::move(client)),
      buffer_(std::move(buffer)),
      traceback_(std::move(traceback)) {
  CHECK(PyGILState_Check());
  next_ = client_->buffers_[buffer_->device()->id()];<commit_msg>Also displays the type of the object when the conversion to a `PyBuffer` fails.<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/xla/python/py_buffer.h"

#include <functional>
#include <memory>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include "absl/base/casts.h"
#include "pybind11/pybind11.h"
#include "pybind11/pytypes.h"
#include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
#include "tensorflow/compiler/xla/python/py_client.h"
#include "tensorflow/compiler/xla/python/python_ref_manager.h"
#include "tensorflow/compiler/xla/python/python_utils.h"
#include "tensorflow/compiler/xla/python/transfer_guard_lib.h"
#include "tensorflow/compiler/xla/python/types.h"
#include "tensorflow/compiler/xla/python/util.h"
#include "tensorflow/compiler/xla/util.h"

namespace xla {

namespace py = pybind11;

namespace {

// Representation of a DeviceArrayBase as a Python object. Since
// a DeviceArrayBase has no fields, this is just a PyObject.
struct PyBufferBasePyObject {
  PyObject_HEAD;
};
static_assert(std::is_standard_layout<PyBufferBasePyObject>::value,
              "PyBufferBasePyObject must be standard layout");

// Representation of a DeviceArray as a Python object.
struct PyBufferPyObject {
  PyBufferBasePyObject base;
  PyBuffer buffer;
  // Used by the Python interpreter to maintain a list of weak references to
  // this object.
  PyObject* weakrefs;
};
static_assert(std::is_standard_layout<PyBufferPyObject>::value,
              "PyBufferPyObject must be standard layout");

PyObject* PyBuffer_tp_new(PyTypeObject* subtype, PyObject* args,
                          PyObject* kwds) {
  PyBufferPyObject* self =
      reinterpret_cast<PyBufferPyObject*>(subtype->tp_alloc(subtype, 0));
  if (!self) return nullptr;
  self->weakrefs = nullptr;
  return reinterpret_cast<PyObject*>(self);
}

void PyBuffer_tp_dealloc(PyObject* self) {
  PyTypeObject* tp = Py_TYPE(self);
  PyBufferPyObject* o = reinterpret_cast<PyBufferPyObject*>(self);
  if (o->weakrefs) {
    PyObject_ClearWeakRefs(self);
  }
  o->buffer.~PyBuffer();
  tp->tp_free(self);
  Py_DECREF(tp);
}

}  // namespace

/*static*/ PyBuffer::object PyBuffer::Make(
    std::shared_ptr<PyClient> client, std::shared_ptr<PjRtBuffer> buffer,
    std::shared_ptr<Traceback> traceback) {
  py::object obj = py::reinterpret_steal<py::object>(PyBuffer_tp_new(
      reinterpret_cast<PyTypeObject*>(type_), nullptr, nullptr));
  PyBufferPyObject* buf = reinterpret_cast<PyBufferPyObject*>(obj.ptr());
  new (&buf->buffer)
      PyBuffer(std::move(client), std::move(buffer), std::move(traceback));
  return py::reinterpret_borrow<PyBuffer::object>(obj);
}

bool PyBuffer::IsPyBuffer(py::handle handle) {
  return handle.get_type() == PyBuffer::type();
}

/*static*/ PyBuffer* PyBuffer::AsPyBufferUnchecked(pybind11::handle handle) {
  return &(reinterpret_cast<PyBufferPyObject*>(handle.ptr())->buffer);
}

/*static*/ StatusOr<PyBuffer*> PyBuffer::AsPyBuffer(pybind11::handle handle) {
  if (!IsPyBuffer(handle)) {
    return InvalidArgument("Expected a DeviceArray, got object of type %s",
                           py::cast<std::string>(py::str(handle.get_type())));
  }
  return AsPyBufferUnchecked(handle);
}

py::handle PyBuffer::AsHandle() {
  return reinterpret_cast<PyObject*>(reinterpret_cast<char*>(this) -
                                     offsetof(PyBufferPyObject, buffer));
}

PyBuffer::PyBuffer(std::shared_ptr<PyClient> client,
                   std::shared_ptr<PjRtBuffer> buffer,
                   std::shared_ptr<Traceback> traceback)
    : client_(std::move(client)),
      buffer_(std::move(buffer)),
      traceback_(std::move(traceback)) {
  CHECK(PyGILState_Check());
  next_ = client_->buffers_[buffer_->device()->id()];
{'fraction_non_alphanumeric': 0.14023333718378192, 'fraction_numerical': 0.004851565207346656, 'mean_word_length': 4.942347288949897, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 73, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4544667', 'n_tokens_mistral': 2659, 'n_tokens_neox': 2649, 'n_words': 608}