pywrap_dtensor_device.cc

/* Copyright 2022 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 <string>
#include <vector>

#include "pybind11/pybind11.h"  // from @pybind11
#include "pybind11/stl.h"  // from @pybind11
#include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil
#include "pybind11_protobuf/native_proto_caster.h"  // from @pybind11_protobuf
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/dtensor/cc/dtensor_device.h"
#include "tensorflow/dtensor/cc/tensor_layout.h"
#include "tensorflow/python/eager/pywrap_tensor.h"
#include "tensorflow/python/eager/pywrap_tfe.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
#include "tensorflow/python/util/util.h"

namespace py = ::pybind11;
using tensorflow::dtensor::AddMesh;
using tensorflow::dtensor::AllocateDTensorDevice;
using tensorflow::dtensor::ClearTPUCoreIDs;
using tensorflow::dtensor::ExperimentalClearDefaultLayout;
using tensorflow::dtensor::ExperimentalClearDefaultMesh;
using tensorflow::dtensor::ExperimentalSetDefaultLayout;
using tensorflow::dtensor::ExperimentalSetDefaultMesh;
using tensorflow::dtensor::FetchLayout;
using tensorflow::dtensor::GetStats;
using tensorflow::dtensor::IsDTensor;
using tensorflow::dtensor::IsSparseDTensor;
using tensorflow::dtensor::Layout;
using tensorflow::dtensor::Mesh;
using tensorflow::dtensor::Pack;
using tensorflow::dtensor::SetIteratorElementLayouts;
using tensorflow::dtensor::SetTPUCoreIDs;
using tensorflow::dtensor::SparsePack;
using tensorflow::dtensor::TPUCoreIDsToLocations;
using tensorflow::dtensor::TPUCoreLocationsToIDs;
using tensorflow::dtensor::Unpack;

void PyXDecref(PyObject* obj) { Py_XDECREF(obj); }

void CallDelete_Device(PyObject* capsule) {
  delete reinterpret_cast<TFE_CustomDevice*>(
      PyCapsule_GetPointer(capsule, "TFE_CustomDevice"));
}

void CallDelete_DeviceInfo(PyObject* capsule) {
  void (*destructor)(void*) =
      reinterpret_cast<void (*)(void*)>(PyCapsule_GetContext(capsule));
  destructor(PyCapsule_GetPointer(capsule, "TFE_CustomDevice_DeviceInfo"));
}

bool CheckResourceVariable(PyObject* item) {
  if (tensorflow::swig::IsResourceVariable(item)) {
    tensorflow::Safe_PyObjectPtr handle(
        PyObject_GetAttrString(item, "_handle"));
    return EagerTensor_CheckExact(handle.get());
  }

  return false;
}
// Supports 2 cases:
//  i) input is an EagerTensor.
//  ii) input is an arbitrary python list/tuple.
void ConvertToTensor(TFE_Context* ctx, PyObject* input,
                     tensorflow::Safe_PyObjectPtr* output_handle,
                     TF_Status* status) {
  if (EagerTensor_CheckExact(input)) {
    // Input is already a EagerTensor so increment the reference, since the
    // caller will use it through output_handle.
    Py_INCREF(input);
    output_handle->reset(input);
    return;
  }
  if (CheckResourceVariable(input)) {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "Variable input is not supported.");
    return;
  }
  TFE_TensorHandle* handle =
      tensorflow::ConvertToEagerTensor(ctx, input, tensorflow::DT_INVALID);
  if (handle == nullptr) {
    TF_SetStatus(status, TF_INTERNAL, "Failure converting to eager tensor.");
    return;
  }
  output_handle->reset(EagerTensorFromHandle(handle));
}

PYBIND11_MODULE(_pywrap_dtensor_device, m) {
  pybind11_protobuf::ImportNativeProtoCasters();
  m.def("Allocate", [](const std::string& name, bool is_async,
                       int in_flight_nodes_limit) {
    TFE_CustomDevice* device = new TFE_CustomDevice;
    std::unique_ptr<PyObject, decltype(&PyXDecref)> device_capsule(
        PyCapsule_New(device, "TFE_CustomDevice", &CallDelete_Device),
        PyXDecref);
    void* device_info = nullptr;
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    AllocateDTensorDevice(name, device, &device_info, is_async,
                          in_flight_nodes_limit, status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
    std::unique_ptr<PyObject, decltype(&PyXDecref)> device_info_capsule(
        PyCapsule_New(device_info, "TFE_CustomDevice_DeviceInfo",
                      &CallDelete_DeviceInfo),
        PyXDecref);
    // The PyCapsule destructor needs a pointer to the destructor for
    // DeviceInfo.
    PyCapsule_SetContext(device_info_capsule.get(),
                         reinterpret_cast<void*>(device->delete_device));
    if (PyErr_Occurred()) throw py::error_already_set();
    return pybind11::reinterpret_steal<pybind11::object>(
        PyTuple_Pack(2, device_capsule.get(), device_info_capsule.get()));
  });
  m.def("AddMesh", [](const py::capsule& device_info,
                      const std::string& serialized_mesh, bool is_host_mesh) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    AddMesh(
        serialized_mesh,
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
        is_host_mesh, status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
  });
  m.def(
      "ExperimentalSetDefaultLayout",
      [](const py::capsule& device_info, const std::string& serialized_layout) {
        std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
            TF_NewStatus(), TF_DeleteStatus);
        ExperimentalSetDefaultLayout(
            serialized_layout,
            PyCapsule_GetPointer(device_info.ptr(),
                                 "TFE_CustomDevice_DeviceInfo"),
            status.get());
        if (TF_GetCode(status.get()) != TF_OK) {
          PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
          throw py::error_already_set();
        }
      });
  m.def("ExperimentalClearDefaultLayout", [](const py::capsule& device_info) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    ExperimentalClearDefaultLayout(
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
        status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
  });
  m.def("ExperimentalSetDefaultMesh", [](const py::capsule& device_info,
                                         const std::string& serialized_mesh) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    ExperimentalSetDefaultMesh(
        serialized_mesh,
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
        status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
  });
  m.def("ExperimentalClearDefaultMesh", [](const py::capsule& device_info) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    ExperimentalClearDefaultMesh(
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
        status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
  });
  m.def("SetTPUCoreIDs", [](const py::capsule& device_info,
                            const std::string& mesh_name,
                            const std::vector<int>& tpu_core_ids) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    SetTPUCoreIDs(
        mesh_name, tpu_core_ids,
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"),
        status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
  });
  m.def("ClearTPUCoreIDs", [](const py::capsule& device_info) {
    ClearTPUCoreIDs(
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"));
  });
  m.def("TPUCoreIDsToLocations", [](const py::handle& context,
                                    const py::capsule& device_info,
                                    const std::vector<int>& tpu_core_ids) {
    return TPUCoreIDsToLocations(
        static_cast<TFE_Context*>(PyCapsule_GetPointer(context.ptr(), nullptr)),
        tpu_core_ids,
        PyCapsule_GetPointer(device_info.ptr(), "TFE_CustomDevice_DeviceInfo"));
  });
  m.def("TPUCoreLocationsToIDs",
        [](const py::handle& context, const py::capsule& device_info,
           const std::vector<std::vector<int>>& tpu_core_locations) {
          return TPUCoreLocationsToIDs(
              static_cast<TFE_Context*>(
                  PyCapsule_GetPointer(context.ptr(), nullptr)),
              tpu_core_locations,
              PyCapsule_GetPointer(device_info.ptr(),
                                   "TFE_CustomDevice_DeviceInfo"));
        });
  m.def("Pack", [](const py::handle& context, const py::handle& input_tensors,
                   const std::string& string_layout,
                   const py::capsule& device_info, const bool is_sparse) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    TFE_Context* ctx =
        static_cast<TFE_Context*>(PyCapsule_GetPointer(context.ptr(), nullptr));
    // Convert each python object to safe py eagertensors.
    std::vector<tensorflow::Safe_PyObjectPtr> py_eager_tensor_handles;
    Py_ssize_t len = PyList_Size(input_tensors.ptr());
    py_eager_tensor_handles.resize(len);

    for (Py_ssize_t i = 0; i < len; ++i) {
      PyObject* elem = PyList_GetItem(input_tensors.ptr(), i);
      ConvertToTensor(ctx, elem, &py_eager_tensor_handles[i], status.get());

      if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
        return tensorflow::PyoOrThrow(nullptr);
    }
    std::vector<TFE_TensorHandle*> input_vector;
    input_vector.resize(len);
    for (int i = 0; i < len; ++i)
      input_vector[i] = EagerTensor_Handle(py_eager_tensor_handles[i].get());
    TFE_TensorHandle* packed_tensor;
    if (is_sparse) {
      auto size = input_vector.size() / 3;
      packed_tensor = SparsePack(
          ctx,
          /*num_inputs=*/input_vector.size() / 3,
          /*indices=*/
          std::vector<TFE_TensorHandle*>(input_vector.begin(),
                                         input_vector.begin() + size)
              .data(),
          /*values=*/
          std::vector<TFE_TensorHandle*>(input_vector.begin() + size,
                                         input_vector.begin() + 2 * size)
              .data(),
          /*shapes=*/
          std::vector<TFE_TensorHandle*>(input_vector.begin() + 2 * size,
                                         input_vector.end())
              .data(),
          string_layout, device_info, status.get());
    } else {
      packed_tensor = Pack(ctx, input_vector.size(), input_vector.data(),
                           string_layout, device_info, status.get());
    }
    if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
      return tensorflow::PyoOrThrow(nullptr);
    // Convert c++ packed tensor handle into a python eager tensor object.
    tensorflow::Safe_PyObjectPtr flat_result(PyList_New(1));
    PyList_SET_ITEM(flat_result.get(), 0, EagerTensorFromHandle(packed_tensor));
    auto* result = PyList_GET_ITEM(flat_result.get(), 0);
    Py_INCREF(result);
    return tensorflow::PyoOrThrow(result);
  });
  m.def("Unpack", [](const py::handle& context,
                     const py::handle& dtensor_handle,
                     const py::capsule& device_info) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);

    if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
      throw py::type_error(absl::StrFormat("Expecting a Tensor, got %s",
                                           py::str(dtensor_handle.get_type())));
    }
    TFE_TensorHandle* tensor_handle = EagerTensor_Handle(dtensor_handle.ptr());
    std::vector<TFE_TensorHandle*> unpacked_handles = Unpack(
        static_cast<TFE_Context*>(PyCapsule_GetPointer(context.ptr(), nullptr)),
        tensor_handle, device_info, status.get());

    if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
      return tensorflow::PyoOrThrow(nullptr);
    // Convert all TFE_TensorHandles to py EagerTensor and
    // return a python list of them.
    int num_outputs = unpacked_handles.size();
    PyObject* result(PyList_New(num_outputs));
    for (int i = 0; i < num_outputs; ++i) {
      PyList_SET_ITEM(result, i, EagerTensorFromHandle(unpacked_handles[i]));
    }
    return tensorflow::PyoOrThrow(result);
  });
  m.def(
      "FetchLayout",
      [](const py::handle& context, const py::handle& dtensor_handle,
         const py::capsule& device_info) -> py::object {
        std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
            TF_NewStatus(), TF_DeleteStatus);
        if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
          return py::none();
        }
        TFE_TensorHandle* tensor_handle =
            EagerTensor_Handle(dtensor_handle.ptr());
        std::string layout_string =
            FetchLayout(static_cast<TFE_Context*>(
                            PyCapsule_GetPointer(context.ptr(), nullptr)),
                        tensor_handle, device_info, status.get());
        if (tensorflow::MaybeRaiseExceptionFromTFStatus(status.get(), nullptr))
          return tensorflow::PyoOrThrow(nullptr);
        return tensorflow::PyoOrThrow(
            PyUnicode_FromString(layout_string.c_str()));
      });
  m.def("IsDTensor", [](const py::handle& context,
                        const py::handle& dtensor_handle,
                        const py::capsule& device_info) {
    if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
      return false;
    }
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    TFE_TensorHandle* tensor_handle = EagerTensor_Handle(dtensor_handle.ptr());
    bool is_dtensor = IsDTensor(
        static_cast<TFE_Context*>(PyCapsule_GetPointer(context.ptr(), nullptr)),
        tensor_handle, device_info, status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
    return is_dtensor;
  });
  m.def("IsSparseDTensor", [](const py::handle& context,
                              const py::handle& dtensor_handle,
                              const py::capsule& device_info) {
    if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
      return false;
    }
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);

    TFE_TensorHandle* tensor_handle = EagerTensor_Handle(dtensor_handle.ptr());
    bool is_sparse = IsSparseDTensor(
        static_cast<TFE_Context*>(PyCapsule_GetPointer(context.ptr(), nullptr)),
        tensor_handle, device_info, status.get());

    if (TF_GetCode(status.get()) != TF_OK) {
      PyErr_SetString(PyExc_ValueError, TF_Message(status.get()));
      throw py::error_already_set();
    }
    return is_sparse;
  });
  m.def("GetStats", [](const py::handle& context,
                       const py::capsule& device_info) {
    std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
        TF_NewStatus(), TF_DeleteStatus);
    return GetStats(
        static_cast<TFE_Context*>(PyCapsule_GetPointer(context.ptr(), nullptr)),
        device_info, status.get());
  });
  m.def("SetIteratorElementLayouts",
        [](const py::handle& context, const py::handle& dtensor_handle,
           const std::vector<std::string>& element_layouts,
           const py::capsule& device_info) {
          if (!EagerTensor_CheckExact(dtensor_handle.ptr())) {
            throw py::type_error(
                absl::StrFormat("Expecting a Tensor, got %s",
                                py::str(dtensor_handle.get_type())));
          }
          std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
              TF_NewStatus(), TF_DeleteStatus);
          TFE_TensorHandle* tensor_handle =
              EagerTensor_Handle(dtensor_handle.ptr());
          SetIteratorElementLayouts(
              static_cast<TFE_Context*>(
                  PyCapsule_GetPointer(context.ptr(), nullptr)),
              tensor_handle, element_layouts, device_info, status.get());
        });
  py::class_<Mesh>(m, "Mesh")
      .def(py::init([](Mesh& mesh) { return mesh; }), py::arg("mesh"),
           "Create a copy of a mesh.")
      .def(py::init(&Mesh::CreateMesh))
      .def(py::init([](absl::string_view single_device) {
             auto mesh = Mesh::GetSingleDeviceMesh(single_device);
             if (!mesh.ok()) {
               throw py::value_error(std::string(mesh.status().message()));
             }
             return *mesh;
           }),
           py::arg("single_device"), "Creates a single device mesh.")
      .def(py::init([](const tensorflow::dtensor::MeshProto& proto) {
             auto mesh = Mesh::ParseFromProto(proto);
             if (!mesh.ok()) {
               throw py::value_error(std::string(mesh.status().message()));
             }
             return *mesh;
           }),
           py::arg("mesh_proto"), "Returns a Mesh from a MeshProto.")
      .def(py::init([](std::string_view mesh_str) {
             auto mesh = Mesh::FromString(mesh_str);
             if (!mesh.ok()) {
               throw py::value_error(std::string(mesh.status().message()));
             }
             return *mesh;
           }),
           py::arg("mesh_str"), "Returns a Mesh from a string.")
      .def_property_readonly("name", &Mesh::name)
      .def_property_readonly("dim_names", &Mesh::MeshDimNames)
      .def_property_readonly("size", &Mesh::num_devices)
      .def_property_readonly("single_device", &Mesh::single_device)
      .def("__contains__", &Mesh::IsMeshDim, py::arg("dim_name"))
      .def("__eq__", &Mesh::operator==)
      .def("to_string", &Mesh::ToString,
           "Returns string representation of Mesh.")
      .def("is_single_device", &Mesh::IsSingleDevice,
           "Returns True if the mesh represents a non-distributed device.")
      .def("contains_dim", &Mesh::IsMeshDim, py::arg("dim_name"),
           "Returns True if a Mesh contains the given dimension name.")
      .def(
          "dim_size",
          [](const Mesh& mesh, std::string_view name) {
            auto dim_size = mesh.dim_size(name);
            if (!dim_size.ok()) {
              throw py::value_error(std::string(dim_size.status().message()));
            }
            return *dim_size;
          },
          py::arg("dim_name"), "Returns the size of mesh dimension.")
      .def("device_type", &Mesh::device_type,
           "Returns the device_type of a Mesh.")
      .def("host_mesh", &Mesh::host_mesh,
           "Returns a host mesh corresponding to this mesh.")
      .def("num_local_devices", &Mesh::num_local_devices,
           "Returns the number of local devices.")
      .def("min_global_device_id", &Mesh::min_global_device_id,
           "Returns the minimum global device ID.")
      .def("is_remote", &Mesh::is_remote,
           "Returns True if a Mesh contains only remote devices.")
      .def("local_device_ids", &Mesh::local_device_ids,
           "Returns a list of local device IDs.")
      .def("local_devices", &Mesh::local_devices,
           "Returns a list of local device specs "
           "represented as strings.")
      .def("global_device_ids", &Mesh::global_device_ids,
           "Returns a list of global device IDs.")
      .def("global_devices", &Mesh::global_devices,
           "Returns a list of global device specs "
           "represented as strings.")
      .def("shape", &Mesh::dim_sizes, "Returns the shape of the mesh.")
      .def("use_xla_spmd", &Mesh::use_xla_spmd,
           "Returns True if Mesh will use XLA for SPMD "
           "instead of DTensor SPMD.")
      .def(
          "as_proto",
          [](const Mesh& mesh) {
            auto mesh_proto = mesh.ToProto();
            if (!mesh_proto.ok()) {
              throw py::value_error(std::string(mesh_proto.status().message()));
            }
            return *mesh_proto;
          },
          "Returns the MeshProto protobuf message.")
      .def("device_location", [](const Mesh& mesh, int device_id) {
        auto location = mesh.device_location(device_id);
        if (!location.ok()) {
          throw py::value_error(std::string(location.status().message()));
        }
        return std::vector<int64_t>(location->begin(), location->end());
      });

  py::enum_<Layout::LayoutType>(m, "LayoutType")
      .value("STATIC", Layout::LayoutType::kStatic)
      .value("PARTED", Layout::LayoutType::kParted)
      .value("SINGLE_DEVICE", Layout::LayoutType::kSingleDevice);

  py::class_<Layout>(m, "Layout")
      .def(py::init([](Layout& layout) { return layout; }), py::arg("layout"),
           "Create a copy of a layout.")
      .def(py::init([](Layout::LayoutType type,
                       const std::vector<std::string>& sharding_specs,
                       const Mesh& mesh) {
             auto layout = Layout::GetLayout(type, sharding_specs, mesh);
             if (!layout.ok()) {
               throw py::value_error(std::string(layout.status().message()));
             }
             return *layout;
           }),
           py::arg("type"), py::arg("sharding_specs"), py::arg("mesh"))
      .def(py::init([](const tensorflow::dtensor::LayoutProto& proto) {
             auto layout = Layout::FromProto(proto);
             if (!layout.ok()) {
               throw py::value_error(std::string(layout.status().message()));
             }
             return *layout;
           }),
           py::arg("layout_proto"), "Returns a Layout from a LayoutProto.")
      .def(py::init([](std::string_view layout_str) {
             auto layout = Layout::FromString(layout_str);
             if (!layout.ok()) {
               throw py::value_error(std::string(layout.status().message()));
             }
             return *layout;
           }),
           py::arg("layout_str"), "Returns a Layout from a string.")
      .def(py::init(&Layout::ReplicatedOnMesh), py::arg("mesh"),
           py::arg("rank"), "Returns a replicated layout.")
      .def(py::init(&Layout::BatchShardedOnMesh), py::arg("mesh"),
           py::arg("rank"), py::arg("batch_dim"), py::arg("axis"),
           "Returns a batch sharded layout.")
      .def(py::init([](const Mesh& mesh) {
             auto layout = Layout::GetSingleDeviceLayout(mesh);
             if (!layout.ok()) {
               throw py::value_error(std::string(layout.status().message()));
             }
             return *layout;
           }),
           py::arg("mesh"), "Returns a single device layout.")
      .def("__eq__", &Layout::operator==)
      .def(
          "as_proto",
          [](const Layout& layout) {
            auto layout_proto = layout.ToProto();
            if (!layout_proto.ok()) {
              throw py::value_error(
                  std::string(layout_proto.status().message()));
            }
            return *layout_proto;
          },
          "Returns the LayoutProto protobuf message.")
      .def("to_string", &Layout::ToString)
      .def("to_parted", &Layout::ToParted)
      .def_property_readonly("type", &Layout::type)
      .def_property_readonly("sharding_specs", &Layout::sharding_spec_strs)
      .def_property_readonly("rank", &Layout::rank)
      .def_property_readonly("mesh", &Layout::mesh)
      .def("is_fully_replicated", &Layout::IsFullyReplicated,
           "Returns True if all tensor axes are replicated.")
      .def("is_batch_parallel",
           [](const Layout& layout) { return layout.IsBatchParallel(); })
      .def("is_single_device", &Layout::IsSingleDevice,
           "Returns True if the Layout represents a non-distributed device.")
      .def(
          "num_shards",
          [](const Layout& layout, int dim) {
            return layout.num_shards_for_dim(dim);
          },
          py::arg("idx"),
          "Returns the number of shards for tensor dimension `idx`.")
      .def(
          "global_shape_from_local_shape",
          [](const Layout& layout, std::vector<int64_t> local_shape) {
            return py::tuple(
                py::cast(layout.GlobalShapeFromLocalShape(local_shape)));
          },
          py::arg("local_shape"),
          "Returns the global shape computed from this local shape.")
      .def(
          "local_shape_from_global_shape",
          [](const Layout& layout, std::vector<int64_t> global_shape) {
            return py::tuple(
                py::cast(layout.LocalShapeFromGlobalShape(global_shape)));
          },
          py::arg("global_shape"),
          "Returns the local shape computed from this global shape.");
}