/* Copyright 2020 The JAX Authors

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 "jaxlib/py_values.h"

#include <Python.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <exception>
#include <functional>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>

#include "absl/base/no_destructor.h"
#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/inlined_vector.h"
#include "absl/functional/any_invocable.h"
#include "absl/hash/hash.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/types/span.h"
#include "nanobind/nanobind.h"
#include "nanobind/stl/complex.h"  // IWYU pragma: keep
#include "nanobind/stl/string_view.h"  // IWYU pragma: keep
#include "jaxlib/config.h"
#include "jaxlib/free_threading.h"
#include "jaxlib/nb_class_ptr.h"
#include "jaxlib/numpy.h"
#include "jaxlib/py_array.h"
#include "jaxlib/python_ref_manager.h"
#include "jaxlib/sharding.h"
#include "jaxlib/to_ifrt_sharding.h"
#include "jaxlib/weak_key_weak_value_cache.h"
#include "xla/primitive_util.h"
#include "xla/python/ifrt/array.h"
#include "xla/python/ifrt/array_spec.h"
#include "xla/python/ifrt/device.h"
#include "xla/python/ifrt/device_list.h"
#include "xla/python/ifrt/dtype.h"
#include "xla/python/ifrt/memory.h"
#include "xla/python/ifrt/shape.h"
#include "xla/python/ifrt/sharding.h"
#include "xla/python/nb_numpy.h"
#include "xla/python/pjrt_ifrt/pjrt_dtype.h"
#include "xla/python/safe_static_init.h"
#include "xla/python/types.h"
#include "xla/python/version.h"
#include "xla/shape.h"
#include "xla/tsl/concurrency/ref_count.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/python/lib/core/numpy.h"
#include "xla/types.h"
#include "xla/util.h"
#include "xla/xla_data.pb.h"
#include "tsl/platform/ml_dtypes.h"
#include "tsl/profiler/lib/traceme.h"

namespace nb = nanobind;
namespace ifrt = xla::ifrt;

namespace jax {

namespace {

nb_class_ptr<Config>& enable_x64_state = *new nb_class_ptr<Config>();

// The TypedInt, TypedFloat, TypedComplex, and TypedNdArray types.
nb::object& typed_int_type = *new nb::object();
nb::object& typed_float_type = *new nb::object();
nb::object& typed_complex_type = *new nb::object();
nb::object& typed_ndarray_type = *new nb::object();
nb::object& invalid_input_exception = *new nb::object();
nb::object& valid_dtypes = *new nb::object();

using CanonicalizeValueHandler = std::function<nb::object(nb::handle)>;

static ft_mutex canonicalize_value_handlers_mutex;

static absl::NoDestructor<
    absl::flat_hash_map<PyObject*, CanonicalizeValueHandler>>
    canonicalize_value_handlers
        ABSL_GUARDED_BY(canonicalize_value_handlers_mutex);

// For xla::S64/U64/F64/C128 types, returns the largest 32-bit equivalent.
xla::PrimitiveType Squash64BitType(xla::PrimitiveType type) {
  switch (type) {
    case xla::S64:
      return xla::S32;
    case xla::U64:
      return xla::U32;
    case xla::F64:
      return xla::F32;
    case xla::C128:
      return xla::C64;
    default:
      return type;
  }
}

// Gets the thread-local instance.
static DevicePutInfo& GetDevicePutInfo() {
  thread_local DevicePutInfo device_put_info;
  return device_put_info;
}

// Prepared data for creating a single shard of an array. Holds a single-device
// IFRT array or a host buffer.
struct Shard {
  explicit Shard(ifrt::ArrayRef ifrt_array, bool weak_type)
      : ifrt_array_or_host_buffer(std::move(ifrt_array)),
        weak_type(weak_type),
        // host_buffer_semantics is not meaningful when
        // `ifrt_array_or_host_buffer` is an IFRT Array.
        host_buffer_semantics(
            ifrt::Client::HostBufferSemantics::kImmutableOnlyDuringCall) {}

  Shard(ifrt::Client::HostBuffer ifrt_host_buffer, bool weak_type,
        ifrt::Client::HostBufferSemantics host_buffer_semantics)
      : ifrt_array_or_host_buffer(std::move(ifrt_host_buffer)),
        weak_type(weak_type),
        host_buffer_semantics(host_buffer_semantics) {}

  Shard(const Shard&) = delete;
  Shard& operator=(const Shard&) = delete;
  Shard(Shard&&) noexcept = default;
  Shard& operator=(Shard&&) noexcept = default;

  bool is_ifrt_array() const {
    return std::holds_alternative<ifrt::ArrayRef>(ifrt_array_or_host_buffer);
  }
  ifrt::DType ifrt_dtype() const;
  const ifrt::Shape& ifrt_shape() const;

  // Points to the on-device array or on-host buffer.
  std::variant<ifrt::ArrayRef, ifrt::Client::HostBuffer>
      ifrt_array_or_host_buffer;
  bool weak_type;
  ifrt::Client::HostBufferSemantics host_buffer_semantics;
};

// A function that creates a `Shard` from a Python object when called.
using ShardFn = absl::AnyInvocable<absl::StatusOr<Shard>() &&>;

absl::StatusOr<std::vector<absl::Cord>> StringDTypeArrayToCords(
    PyArrayObject* py_array_obj) {
  if (PyArray_SIZE(py_array_obj) == 0) {
    return absl::InvalidArgumentError("empty numpy array");
  }

  std::vector<absl::Cord> cords;
  cords.reserve(PyArray_SIZE(py_array_obj));

  auto iter =
      nb::steal(PyArray_IterNew(reinterpret_cast<PyObject*>(py_array_obj)));
  while (PyArray_ITER_NOTDONE(iter.ptr())) {
    auto* iter_data = PyArray_ITER_DATA(iter.ptr());
    auto* item = PyArray_GETITEM(py_array_obj, static_cast<char*>(iter_data));
    if (!item) {
      return absl::InternalError(
          "Failed to get elements out of the ndarray iter.");
    }
    Py_ssize_t len;
    auto str = PyUnicode_AsUTF8AndSize(item, &len);
    cords.push_back(absl::Cord(std::string_view(str, len)));
    PyArray_ITER_NEXT(iter.ptr());
  }
  return cords;
}

// Handler that creates a `Shard` from a Python object.
using DevicePutHandler = std::function<absl::StatusOr<ShardFn>(
    nb::handle obj, ifrt::Client* client, ifrt::Device* to_device,
    ifrt::MemoryKind to_memory_kind, const DevicePutOptions& options)>;

// Shared logic that makes an IFRT array (either single-device or multi-device)
// from a fully-replicated `shard` that is created from a host buffer (not from
// an existing IFRT array). `shard` will be consumed.
//
// Expected to be called without holding GIL.
absl::StatusOr<tsl::RCReference<ifrt::Array>>
MakeIfrtArrayFromFullyReplicatedShard(ifrt::Client* ifrt_client,
                                      ifrt::ShardingRef ifrt_sharding,
                                      Shard& shard) {
  auto host_buffer_shard = std::get<ifrt::Client::HostBuffer>(
      std::move(shard.ifrt_array_or_host_buffer));
  return ifrt_client->MakeArrayFromHostBuffer(
      host_buffer_shard.data, host_buffer_shard.dtype,
      std::move(host_buffer_shard.shape),
      std::move(host_buffer_shard.byte_strides), std::move(ifrt_sharding),
      /*layout=*/nullptr, shard.host_buffer_semantics,
      std::move(host_buffer_shard.on_done));
}

// Shared logic that makes a single-device IFRT array from a `shard`. `shard`
// will be consumed.
//
// Expected to be called without holding GIL.
absl::StatusOr<ifrt::ArrayRef> MakeSingleDeviceIfrtArrayFromShard(
    xla::ifrt::Client* ifrt_client, xla::ifrt::Device* ifrt_device,
    xla::ifrt::MemoryKind ifrt_memory_kind, Shard& shard) {
  if (auto* ifrt_array =
          std::get_if<ifrt::ArrayRef>(&shard.ifrt_array_or_host_buffer)) {
    return std::move(*ifrt_array);
  }
  ifrt::ShardingRef ifrt_sharding =
      ifrt::SingleDeviceSharding::Create(ifrt_device, ifrt_memory_kind);
  return MakeIfrtArrayFromFullyReplicatedShard(ifrt_client,
                                               std::move(ifrt_sharding), shard);
}

using HostBuffer = ifrt::Client::HostBuffer;

struct HostBufferHash {
  // We hash on everything except for on_done.
  size_t operator()(const HostBuffer* h) const {
    return absl::Hash<
        std::tuple<const void*, xla::ifrt::DType, const xla::ifrt::Shape&,
                   const std::optional<HostBuffer::ByteStrides>&>>()(
        {h->data, h->dtype, h->shape, h->byte_strides});
  }
};

struct HostBufferEq {
  bool operator()(const HostBuffer* lhs, const HostBuffer* rhs) const {
    return lhs->data == rhs->data && lhs->dtype == rhs->dtype &&
           lhs->shape == rhs->shape && lhs->byte_strides == rhs->byte_strides;
  }
};

// Makes an IFRT Array from `shards` using a batched array creation API (fast
// path). `shards` will be consumed.
//
// Expected to be called without holding GIL.
absl::StatusOr<ifrt::ArrayRef> MakeIfrtArrayFromShardsInBatch(
    ifrt::Client* ifrt_client, ifrt::DType ifrt_dtype, ifrt::Shape ifrt_shape,
    ifrt::ShardingRef ifrt_sharding, absl::Span<Shard> shards) {
  absl::InlinedVector<
      std::pair<absl::InlinedVector<int64_t, 1>, ifrt::Client::HostBuffer>, 1>
      host_buffers;
  // Note: Dedup map relies on this reserve to give pointer stability.
  host_buffers.reserve(shards.size());
  ifrt::Client::HostBufferSemantics safe_host_semantics =
      ifrt::Client::HostBufferSemantics::kImmutableZeroCopy;

  // TODO(hyeontaek): Consider performing this deduplication earlier to avoid
  // even constructing the HostBuffers.
  absl::flat_hash_map<const HostBuffer*, size_t, HostBufferHash, HostBufferEq>
      host_buffer_dedup_map;
  host_buffer_dedup_map.reserve(shards.size());

  for (int64_t i = 0; i < shards.size(); ++i) {
    auto insert_idx = host_buffers.size();
    host_buffers.push_back(
        {{},
         std::move(std::get<ifrt::Client::HostBuffer>(
             std::move(shards[i].ifrt_array_or_host_buffer)))});
    auto [it, insert_happened] = host_buffer_dedup_map.emplace(
        std::make_pair(&host_buffers.back().second, insert_idx));
    if (!insert_happened) {
      std::move(host_buffers.back().second.on_done)();
      host_buffers.pop_back();
    }
    host_buffers[it->second].first.push_back(i);
    // The minimum host buffer semantics is a safe semantics that can be used
    // for all shards when they are created in a single batch.
    safe_host_semantics =
        std::min(safe_host_semantics, shards[i].host_buffer_semantics);
  }

  std::vector<ifrt::Client::MakeArraysFromHostBufferShardsSpec> specs;
  specs.push_back(ifrt::Client::MakeArraysFromHostBufferShardsSpec{
      std::move(host_buffers),
      ifrt::ArraySpec{/*dtype=*/ifrt_dtype,
                      /*shape=*/std::move(ifrt_shape),
                      /*sharding=*/std::move(ifrt_sharding),
                      /*layout=*/nullptr}});
  ABSL_ASSIGN_OR_RETURN(auto arrays,
                        ifrt_client->MakeArraysFromHostBufferShards(
                            absl::MakeSpan(specs), safe_host_semantics));
  return std::move(arrays.front());
}

// Makes an IFRT Array from `shards` using an array assembly API (slow path).
// `shards` will be consumed.
//
// Expected to be called without holding GIL.
absl::StatusOr<ifrt::ArrayRef> MakeIfrtArrayFromShardsWithAssembly(
    ifrt::Client* ifrt_client, ifrt::DType ifrt_dtype, ifrt::Shape ifrt_shape,
    ifrt::ShardingRef ifrt_sharding,
    ifrt::DeviceList* ifrt_addressable_device_list,
    ifrt::MemoryKind ifrt_memory_kind, absl::Span<Shard> shards) {
  absl::Span<ifrt::Device* const> ifrt_addressable_devices =
      ifrt_addressable_device_list->devices();
  std::vector<ifrt::ArrayRef> ifrt_array_shards;
  ifrt_array_shards.reserve(shards.size());
  for (int64_t i = 0; i < shards.size(); ++i) {
    ABSL_ASSIGN_OR_RETURN(ifrt::ArrayRef ifrt_array_shard,
                          MakeSingleDeviceIfrtArrayFromShard(
                              ifrt_client, ifrt_addressable_devices[i],
                              ifrt_memory_kind, shards[i]));
    ifrt_array_shards.push_back(std::move(ifrt_array_shard));
  }
  return ifrt_client->AssembleArrayFromSingleDeviceArrays(
      ifrt_dtype, std::move(ifrt_shape), std::move(ifrt_sharding),
      absl::MakeSpan(ifrt_array_shards), ifrt::ArrayCopySemantics::kReuseInput,
      ifrt::SingleDeviceShardSemantics::kAddressableShards);
}

template <typename T, typename SquashedT>
absl::StatusOr<ShardFn> HandlePythonScalar(nb::handle obj, ifrt::Client* client,
                                           ifrt::Device* to_device,
                                           ifrt::MemoryKind to_memory_kind,
                                           const DevicePutOptions& options) {
  T value;
  try {
    value = nb::cast<T>(obj);
  } catch (const std::exception& e) {
    return xla::InvalidArgument(
        "Unable to convert Python scalar to %s. This most likely means the "
        "value (%s) overflows the range of the type.",
        xla::PrimitiveType_Name(
            xla::primitive_util::NativeToPrimitiveType<T>()),
        nb::cast<std::string_view>(nb::repr(obj)));
  }

  std::variant<T, SquashedT> data;
  xla::Shape shape;
  xla::PrimitiveType type;
  if (std::is_same<T, SquashedT>() || !options.squash_64bit_types) {
    data.template emplace<0>(value);
    type = xla::primitive_util::NativeToPrimitiveType<T>();
  } else {
    // TODO(phawkins): we should check for overflow here, e.g., because of bugs
    // like https://github.com/google/jax/issues/2006
    data.template emplace<1>(static_cast<SquashedT>(value));
    type = xla::primitive_util::NativeToPrimitiveType<SquashedT>();
  }
  ABSL_ASSIGN_OR_RETURN(ifrt::DType ifrt_dtype, ifrt::ToDType(type));

  return [data, ifrt_dtype]() -> absl::StatusOr<Shard> {
    const void* ptr = std::visit(
        [](const auto& v) { return static_cast<const void*>(&v); }, data);
    ifrt::Client::HostBuffer ifrt_host_buffer{
        ptr, ifrt_dtype, ifrt::Shape({}),
        /*byte_strides=*/std::nullopt,
        /*on_done_with_host_buffer=*/nullptr};
    return Shard(std::move(ifrt_host_buffer), /*weak_type=*/true,
                 ifrt::Client::HostBufferSemantics::kImmutableOnlyDuringCall);
  };
}

absl::StatusOr<ShardFn> HandlePythonInt(nb::handle obj, ifrt::Client* client,
                                        ifrt::Device* to_device,
                                        ifrt::MemoryKind to_memory_kind,
                                        const DevicePutOptions& options) {
  xla::PrimitiveType type;
  std::variant<int64_t, int32_t> data;

  if (options.squash_64bit_types) {
    try {
      data.emplace<1>(nb::cast<int32_t>(obj));
    } catch (const std::exception& e) {
      return xla::InvalidArgument(
          "Unable to convert Python scalar to %s. This most likely means the "
          "value (%s) overflows the range of the type.",
          PrimitiveType_Name(
              xla::primitive_util::NativeToPrimitiveType<int32_t>()),
          nb::cast<std::string_view>(nb::repr(obj)));
    }
    type = xla::S32;
  } else {
    try {
      data.emplace<0>(nb::cast<int64_t>(obj));
    } catch (const std::exception& e) {
      return xla::InvalidArgument(
          "Unable to convert Python scalar to %s. This most likely means the "
          "value (%s) overflows the range of the type.",
          xla::PrimitiveType_Name(
              xla::primitive_util::NativeToPrimitiveType<int64_t>()),
          nb::cast<std::string_view>(nb::repr(obj)));
    }
    type = xla::S64;
  }
  ABSL_ASSIGN_OR_RETURN(ifrt::DType ifrt_dtype, ifrt::ToDType(type));
  return [data, ifrt_dtype]() -> absl::StatusOr<Shard> {
    const void* ptr = std::visit(
        [](const auto& v) { return static_cast<const void*>(&v); }, data);
    ifrt::Client::HostBuffer ifrt_host_buffer{
        ptr, ifrt_dtype, ifrt::Shape({}),
        /*byte_strides=*/std::nullopt,
        /*on_done_with_host_buffer=*/nullptr};
    return Shard(std::move(ifrt_host_buffer), /*weak_type=*/true,
                 ifrt::Client::HostBufferSemantics::kImmutableOnlyDuringCall);
  };
}

template <typename T, typename SquashedT = T>
absl::StatusOr<ShardFn> HandleNumpyScalar(nb::handle h, ifrt::Client* client,
                                          ifrt::Device* to_device,
                                          ifrt::MemoryKind to_memory_kind,
                                          const DevicePutOptions& options) {
  std::variant<T, SquashedT, void*> data;
  xla::PrimitiveType type;
  // For extension types, ScalarAsCtype returns a pointer to the data.
  if (std::is_same<T, xla::s1>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::S1;
  } else if (std::is_same<T, xla::s2>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::S2;
  } else if (std::is_same<T, xla::s4>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::S4;
  } else if (std::is_same<T, xla::u1>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::U1;
  } else if (std::is_same<T, xla::u2>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::U2;
  } else if (std::is_same<T, xla::u4>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::U4;
  } else if (std::is_same<T, xla::bfloat16>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::BF16;
  } else if (std::is_same<T, tsl::float4_e2m1fn>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F4E2M1FN;
  } else if (std::is_same<T, tsl::float6_e2m3fn>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F6E2M3FN;
  } else if (std::is_same<T, tsl::float6_e3m2fn>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F6E3M2FN;
  } else if (std::is_same<T, tsl::float8_e3m4>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E3M4;
  } else if (std::is_same<T, tsl::float8_e4m3>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E4M3;
  } else if (std::is_same<T, tsl::float8_e4m3fn>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E4M3FN;
  } else if (std::is_same<T, tsl::float8_e4m3b11fnuz>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E4M3B11FNUZ;
  } else if (std::is_same<T, tsl::float8_e5m2>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E5M2;
  } else if (std::is_same<T, tsl::float8_e4m3fnuz>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E4M3FNUZ;
  } else if (std::is_same<T, tsl::float8_e5m2fnuz>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E5M2FNUZ;
  } else if (std::is_same<T, tsl::float8_e8m0fnu>()) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<2>());
    type = xla::F8E8M0FNU;
  } else if (std::is_same<T, SquashedT>() || !options.squash_64bit_types) {
    PyArray_ScalarAsCtype(h.ptr(), &data.template emplace<0>());
    type = xla::primitive_util::NativeToPrimitiveType<T>();
  } else {
    T value;
    PyArray_ScalarAsCtype(h.ptr(), &value);
    data.template emplace<1>(static_cast<SquashedT>(value));
    type = xla::primitive_util::NativeToPrimitiveType<SquashedT>();
  }
  std::shared_ptr<PythonRefManager::ManagedPyObjects> py_buffer_ref;
  if (data.index() == 2) {
    py_buffer_ref =
        GlobalPyRefManager()->ManageReference(nb::cast<nb::object>(h));
  }
  ABSL_ASSIGN_OR_RETURN(ifrt::DType ifrt_dtype, ifrt::ToDType(type));
  return [data, py_buffer_ref = std::move(py_buffer_ref),
          ifrt_dtype]() mutable -> absl::StatusOr<Shard> {
    const void* ptr = std::visit(
        [](const auto& v) -> const void* {
          if constexpr (std::is_same_v<std::decay_t<decltype(v)>, void*>) {
            return v;
          } else {
            return static_cast<const void*>(&v);
          }
        },
        data);
    ifrt::Client::HostBuffer ifrt_host_buffer{
        ptr, ifrt_dtype, ifrt::Shape({}),
        /*byte_strides=*/std::nullopt,
        /*on_done_with_host_buffer=*/
        [py_buffer_ref =
             std::move(py_buffer_ref)]() { /* keeps py_buffer_ref alive */ }};
    return Shard(std::move(ifrt_host_buffer), /*weak_type=*/false,
                 ifrt::Client::HostBufferSemantics::kImmutableOnlyDuringCall);
  };
}

absl::StatusOr<ShardFn> HandleStringNumpyArray(
    nb::handle h, ifrt::Client* client, ifrt::Device* to_device,
    ifrt::MemoryKind to_memory_kind, const DevicePutOptions& options) {
  xla::nb_numpy_ndarray array = nb::cast<xla::nb_numpy_ndarray>(h);
  auto py_array_obj = reinterpret_cast<PyArrayObject*>(array.ptr());
  ABSL_ASSIGN_OR_RETURN(auto cords, StringDTypeArrayToCords(py_array_obj));

  // Assemble all the parameters of MakeArrayFromHostBuffer
  const void* data = cords.data();

  // Make an explicit copy of the shape elements so we won't run into complex
  // endianness and precision issues that might arise if we reinterpret-casted
  // from npy_intp, that can be just 32 bits-wide in some environments
  // such as macos_arm64 to const int64_t* that must be 64 bits-wide.
  ifrt::Shape::Dimensions dims;
  dims.reserve(array.ndim());
  for (int i = 0; i < array.ndim(); ++i) {
    dims.push_back(array.shape(i));
  }
  ifrt::Shape shape(std::move(dims));

  auto on_done_with_host_buffer = [cords = std::move(cords)] {};

  return [data, shape = std::move(shape),
          on_done_with_host_buffer = std::move(
              on_done_with_host_buffer)]() mutable -> absl::StatusOr<Shard> {
    ifrt::Client::HostBuffer ifrt_host_buffer{
        data, ifrt::DType(ifrt::DType::kString), std::move(shape),
        /*byte_strides=*/std::nullopt, std::move(on_done_with_host_buffer)};
    return Shard(
        std::move(ifrt_host_buffer), /*weak_type=*/false,
        ifrt::Client::HostBufferSemantics::kImmutableUntilTransferCompletes);
  };
}

absl::StatusOr<ShardFn> HandleNumpyArray(nb::handle h, ifrt::Client* client,
                                         ifrt::Device* to_device,
                                         ifrt::MemoryKind to_memory_kind,
                                         const DevicePutOptions& options) {
  xla::nb_numpy_ndarray array = nb::cast<xla::nb_numpy_ndarray>(h);

  // String numpy arrays require substantially different processing.
  if (array.dtype().char_() == (int)'T' || array.dtype().kind() == 'T') {
    return HandleStringNumpyArray(h, client, to_device, to_memory_kind,
                                  options);
  }

  ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType type,
                        DtypeToPrimitiveType(array.dtype()));

  xla::PrimitiveType squashed_type;
  if (options.squash_64bit_types) {
    squashed_type = Squash64BitType(type);
    if (squashed_type != type) {
      ABSL_ASSIGN_OR_RETURN(xla::nb_dtype squashed_dtype,
                            PrimitiveTypeToNbDtype(squashed_type));
      array = nb::steal<xla::nb_numpy_ndarray>(PyArray_CastToType(
          reinterpret_cast<PyArrayObject*>(array.ptr()),
          reinterpret_cast<PyArray_Descr*>(squashed_dtype.release().ptr()),
          /*fortran=*/0));
    }
  } else {
    squashed_type = type;
  }

  absl::InlinedVector<int64_t, 4> dims(array.ndim());
  ifrt::Client::HostBuffer::ByteStrides byte_strides(array.ndim());
  for (int i = 0; i < array.ndim(); ++i) {
    dims[i] = array.shape(i);
    byte_strides[i] = array.strides(i);
  }
  const void* data = array.data();
  std::shared_ptr<PythonRefManager::ManagedPyObjects> py_buffer_ref =
      GlobalPyRefManager()->ManageReference(std::move(array));
  ABSL_ASSIGN_OR_RETURN(ifrt::DType ifrt_dtype, ifrt::ToDType(squashed_type));
  return [data, ifrt_dtype, dims = std::move(dims),
          byte_strides = std::move(byte_strides),
          py_buffer_ref = std::move(py_buffer_ref),
          allow_zero_copy =
              options.allow_zero_copy]() mutable -> absl::StatusOr<Shard> {
    ifrt::Client::HostBufferSemantics host_buffer_semantics =
        ifrt::Client::HostBufferSemantics::kImmutableOnlyDuringCall;
    std::function<void()> on_done_with_host_buffer;
    if (allow_zero_copy) {
      on_done_with_host_buffer =
          [py_buffer_ref{
              std::move(py_buffer_ref)}]() { /* keeps py_buffer_ref alive */ };
      host_buffer_semantics =
          ifrt::Client::HostBufferSemantics::kImmutableZeroCopy;
    }

    ifrt::Client::HostBuffer ifrt_host_buffer{
        data, ifrt_dtype, ifrt::Shape(dims), std::move(byte_strides),
        std::move(on_done_with_host_buffer)};
    return Shard(std::move(ifrt_host_buffer), /*weak_type=*/false,
                 host_buffer_semantics);
  };
}

absl::StatusOr<ShardFn> HandleTypedInt(nb::handle h, ifrt::Client* client,
                                       ifrt::Device* to_device,
                                       ifrt::MemoryKind to_memory_kind,
                                       const DevicePutOptions& options) {
  xla::nb_dtype dtype = nb::cast<xla::nb_dtype>(h.attr("dtype"));
  ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType type,
                        xla::DtypeToPrimitiveType(dtype));
  switch (type) {
    case xla::S64:
      return HandlePythonScalar<int64_t, int64_t>(h, client, to_device,
                                                  to_memory_kind, options);
    case xla::S32:
      return HandlePythonScalar<int32_t, int32_t>(h, client, to_device,
                                                  to_memory_kind, options);
    default:
      return xla::InvalidArgument("Unsupported type: %s",
                                  xla::PrimitiveType_Name(type));
  }
}

absl::StatusOr<ShardFn> HandleTypedFloat(nb::handle h, ifrt::Client* client,
                                         ifrt::Device* to_device,
                                         ifrt::MemoryKind to_memory_kind,
                                         const DevicePutOptions& options) {
  xla::nb_dtype dtype = nb::cast<xla::nb_dtype>(h.attr("dtype"));
  ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType type,
                        xla::DtypeToPrimitiveType(dtype));
  switch (type) {
    case xla::F64:
      return HandlePythonScalar<double, double>(h, client, to_device,
                                                to_memory_kind, options);
    case xla::F32:
      return HandlePythonScalar<float, float>(h, client, to_device,
                                              to_memory_kind, options);
    default:
      return xla::InvalidArgument("Unsupported type: %s",
                                  xla::PrimitiveType_Name(type));
  }
}

absl::StatusOr<ShardFn> HandleTypedComplex(nb::handle h, ifrt::Client* client,
                                           ifrt::Device* to_device,
                                           ifrt::MemoryKind to_memory_kind,
                                           const DevicePutOptions& options) {
  xla::nb_dtype dtype = nb::cast<xla::nb_dtype>(h.attr("dtype"));
  ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType type,
                        xla::DtypeToPrimitiveType(dtype));
  switch (type) {
    case xla::C128:
      return HandlePythonScalar<xla::complex128, xla::complex128>(
          h, client, to_device, to_memory_kind, options);
    case xla::C64:
      return HandlePythonScalar<xla::complex64, xla::complex64>(
          h, client, to_device, to_memory_kind, options);
    default:
      return xla::InvalidArgument("Unsupported type: %s",
                                  xla::PrimitiveType_Name(type));
  }
}

absl::StatusOr<ShardFn> HandleTypedNdArray(nb::handle h, ifrt::Client* client,
                                           ifrt::Device* to_device,
                                           ifrt::MemoryKind to_memory_kind,
                                           const DevicePutOptions& options) {
  DevicePutOptions o = options;
  o.squash_64bit_types = false;
  return HandleNumpyArray(h.attr("val"), client, to_device, to_memory_kind, o);
}

absl::StatusOr<ShardFn> HandlePyArray(nb::handle obj, ifrt::Client* client,
                                      ifrt::Device* to_device,
                                      ifrt::MemoryKind to_memory_kind,
                                      const DevicePutOptions& options) {
  auto py_array = nb::borrow<PyArray>(obj);

  // We only allow single device case for PyArray in device put.
  if (py_array.num_shards() != 1) {
    return xla::InvalidArgument(
        "device_put expects an array with exactly one shard, got an array with "
        "with %d shards.",
        py_array.num_shards());
  }

  xla::ifrt::ArrayRef ifrt_array = py_array.ifrt_array_ref();
  if (ifrt_array == nullptr) {
    return xla::InvalidArgument("Array has been deleted.");
  }

  if (ifrt_array->sharding().devices()->devices().front()->client() !=
      to_device->client()) {
    return HandleNumpyArray(obj.attr("_value"), client, to_device,
                            to_memory_kind, options);
  }

#if JAX_IFRT_VERSION_NUMBER >= 64
  if (ifrt_array->sharding().devices()->devices().front() == to_device &&
      options.allow_zero_copy &&
      ifrt_array->sharding().memory_kind() == to_memory_kind) {
#else
  if (ifrt_array->sharding().devices()->devices().front() == to_device &&
      options.allow_zero_copy &&
      (!to_memory_kind.memory_kind().has_value() ||
       !ifrt_array->sharding().memory_kind().memory_kind().has_value() ||
       ifrt_array->sharding().memory_kind() == to_memory_kind)) {
#endif
    Shard result(ifrt_array, py_array.weak_type());
    return [result = std::move(result)]() mutable -> absl::StatusOr<Shard> {
      return std::move(result);
    };
  } else {
    return [ifrt_array, to_device, to_memory_kind,
            weak_type = py_array.weak_type(),
            allow_zero_copy =
                options.allow_zero_copy]() mutable -> absl::StatusOr<Shard> {
      auto* ifrt_client = ifrt_array->client();
      ABSL_ASSIGN_OR_RETURN(ifrt::DeviceListRef device_list,
                            ifrt_client->MakeDeviceList({to_device}));
      ABSL_ASSIGN_OR_RETURN(
          auto copied_ifrt_arrays,
          ifrt_client->CopyArrays(absl::MakeSpan(&ifrt_array, 1),
                                  std::move(device_list), to_memory_kind,
                                  allow_zero_copy
                                      ? ifrt::ArrayCopySemantics::kReuseInput
                                      : ifrt::ArrayCopySemantics::kAlwaysCopy));
      return Shard(std::move(copied_ifrt_arrays.front()), weak_type);
    };
  }
}

ifrt::DType Shard::ifrt_dtype() const {
  if (is_ifrt_array()) {
    return std::get<ifrt::ArrayRef>(ifrt_array_or_host_buffer)->dtype();
  } else {
    return std::get<ifrt::Client::HostBuffer>(ifrt_array_or_host_buffer).dtype;
  }
}

const ifrt::Shape& Shard::ifrt_shape() const {
  if (is_ifrt_array()) {
    return std::get<ifrt::ArrayRef>(ifrt_array_or_host_buffer)->shape();
  } else {
    return std::get<ifrt::Client::HostBuffer>(ifrt_array_or_host_buffer).shape;
  }
}

// Creates a `ShardFn` that copies `arg` to `to_device` and `to_memory_kind`.
//
// Requires GIL. The returned `ShardFn` should be called without GIL held.
absl::StatusOr<ShardFn> MakeShardFn(nb::handle arg, ifrt::Client* client,
                                    ifrt::Device* to_device,
                                    ifrt::MemoryKind to_memory_kind,
                                    const DevicePutOptions& options) {
  using PyObjectDeviceHandlerMap =
      absl::flat_hash_map<PyObject*, DevicePutHandler>;

  auto init_fn = []() {
    std::unique_ptr<PyObjectDeviceHandlerMap> p =
        std::make_unique<PyObjectDeviceHandlerMap>();

    const xla::NumpyScalarTypes& dtypes = xla::GetNumpyScalarTypes();
    // Python scalar types.
    static_assert(sizeof(bool) == 1, "Conversion code assumes bool is 1 byte");
    (*p)[reinterpret_cast<PyObject*>(&PyBool_Type)] =
        HandlePythonScalar<bool, bool>;
    (*p)[reinterpret_cast<PyObject*>(&PyLong_Type)] = HandlePythonInt;
    (*p)[reinterpret_cast<PyObject*>(&PyFloat_Type)] =
        HandlePythonScalar<double, float>;
    (*p)[reinterpret_cast<PyObject*>(&PyComplex_Type)] =
        HandlePythonScalar<xla::complex128, xla::complex64>;

    if (typed_int_type.ptr() != nullptr) {
      (*p)[typed_int_type.ptr()] = HandleTypedInt;
    }
    if (typed_float_type.ptr() != nullptr) {
      (*p)[typed_float_type.ptr()] = HandleTypedFloat;
    }
    if (typed_complex_type.ptr() != nullptr) {
      (*p)[typed_complex_type.ptr()] = HandleTypedComplex;
    }
    (*p)[reinterpret_cast<PyObject*>(&PyArray_Type)] = HandleNumpyArray;

    if (typed_ndarray_type.ptr() != nullptr) {
      (*p)[typed_ndarray_type.ptr()] = HandleTypedNdArray;
    }
    // Numpy scalar types. For some of them, we share the handler with
    // Python types (np_int64, np_float64, np_complex128).
    (*p)[dtypes.np_bool.ptr()] = HandleNumpyScalar<bool>;
    if (dtypes.np_int1.has_value()) {
      (*p)[dtypes.np_int1->ptr()] = HandleNumpyScalar<xla::s1>;
    }
    (*p)[dtypes.np_int2.ptr()] = HandleNumpyScalar<xla::s2>;
    (*p)[dtypes.np_int4.ptr()] = HandleNumpyScalar<xla::s4>;
    (*p)[dtypes.np_int8.ptr()] = HandleNumpyScalar<int8_t>;
    (*p)[dtypes.np_int16.ptr()] = HandleNumpyScalar<int16_t>;
    (*p)[dtypes.np_int32.ptr()] = HandleNumpyScalar<int32_t>;
    (*p)[dtypes.np_int64.ptr()] = HandleNumpyScalar<int64_t, int32_t>;
    if (dtypes.np_uint1.has_value()) {
      (*p)[dtypes.np_uint1->ptr()] = HandleNumpyScalar<xla::u1>;
    }
    (*p)[dtypes.np_uint2.ptr()] = HandleNumpyScalar<xla::u2>;
    (*p)[dtypes.np_uint4.ptr()] = HandleNumpyScalar<xla::u4>;
    (*p)[dtypes.np_uint8.ptr()] = HandleNumpyScalar<uint8_t>;
    (*p)[dtypes.np_uint16.ptr()] = HandleNumpyScalar<uint16_t>;
    (*p)[dtypes.np_uint32.ptr()] = HandleNumpyScalar<uint32_t>;
    (*p)[dtypes.np_uint64.ptr()] = HandleNumpyScalar<uint64_t, uint32_t>;
    (*p)[dtypes.np_float4_e2m1fn.ptr()] = HandleNumpyScalar<tsl::float4_e2m1fn>;
    (*p)[dtypes.np_float6_e2m3fn.ptr()] = HandleNumpyScalar<tsl::float6_e2m3fn>;
    (*p)[dtypes.np_float6_e3m2fn.ptr()] = HandleNumpyScalar<tsl::float6_e3m2fn>;
    (*p)[dtypes.np_float8_e3m4.ptr()] = HandleNumpyScalar<tsl::float8_e3m4>;
    (*p)[dtypes.np_float8_e4m3.ptr()] = HandleNumpyScalar<tsl::float8_e4m3>;
    (*p)[dtypes.np_float8_e4m3fn.ptr()] = HandleNumpyScalar<tsl::float8_e4m3fn>;
    (*p)[dtypes.np_float8_e4m3b11fnuz.ptr()] =
        HandleNumpyScalar<tsl::float8_e4m3b11fnuz>;
    (*p)[dtypes.np_float8_e5m2.ptr()] = HandleNumpyScalar<tsl::float8_e5m2>;
    (*p)[dtypes.np_float8_e4m3fnuz.ptr()] =
        HandleNumpyScalar<tsl::float8_e4m3fnuz>;
    (*p)[dtypes.np_float8_e5m2fnuz.ptr()] =
        HandleNumpyScalar<tsl::float8_e5m2fnuz>;
    (*p)[dtypes.np_float8_e8m0fnu.ptr()] =
        HandleNumpyScalar<tsl::float8_e8m0fnu>;
    (*p)[dtypes.np_bfloat16.ptr()] = HandleNumpyScalar<xla::bfloat16>;
    (*p)[dtypes.np_float16.ptr()] = HandleNumpyScalar<xla::half>;
    (*p)[dtypes.np_float32.ptr()] = HandleNumpyScalar<float>;
    (*p)[dtypes.np_float64.ptr()] = HandleNumpyScalar<double, float>;
    (*p)[dtypes.np_complex64.ptr()] = HandleNumpyScalar<xla::complex64>;
    (*p)[dtypes.np_complex128.ptr()] =
        HandleNumpyScalar<xla::complex128, xla::complex64>;
    static_assert(sizeof(long long) == sizeof(int64_t),  // NOLINT
                  "long long must be the same size as int64_t");
    (*p)[dtypes.np_longlong.ptr()] = HandleNumpyScalar<int64_t, int32_t>;
    static_assert(sizeof(int) == sizeof(int32_t),
                  "int must be the same size as int32_t");
    (*p)[dtypes.np_intc.ptr()] = HandleNumpyScalar<int32_t>;
    return p;
  };
  static xla::SafeStatic<PyObjectDeviceHandlerMap> handlers_init;
  const PyObjectDeviceHandlerMap& handlers = handlers_init.Get(init_fn);

  if (arg.type().ptr() == PyArray::type().ptr()) {
    auto array = nb::borrow<PyArray>(arg);
    return HandlePyArray(arg, client, to_device, to_memory_kind, options);
  }

  auto res = handlers.find(arg.type().ptr());
  if (res == handlers.end()) {
    PyTypeObject* type = reinterpret_cast<PyTypeObject*>(arg.type().ptr());
    PyObject* mro = type->tp_mro;
    if (mro != nullptr) {
      Py_ssize_t n = PyTuple_GET_SIZE(mro);
      for (Py_ssize_t i = 0; i < n; ++i) {
        PyObject* base = PyTuple_GET_ITEM(mro, i);
        res = handlers.find(base);
        if (res != handlers.end()) {
          return res->second(arg, client, to_device, to_memory_kind, options);
        }
      }
    }
    return xla::InvalidArgument(
        "%s", absl::StrCat(
                  "Not supported: The C++ jax jit execution path, only accepts "
                  "DeviceArray, Numpy arrays scalars of supported types "
                  "(see implementation), or Python scalars. Got type ",
                  nb::cast<std::string_view>(nb::str(arg.type()))));
  }
  return res->second(arg, client, to_device, to_memory_kind, options);
}

nb::object IdentityHandler(nb::handle x) {
  return nb::borrow<nb::object>(x.ptr());
}

xla::nb_dtype CanonicalizeDtype(xla::nb_dtype dtype) {
  if (GetEnableX64()) {
    return dtype;
  }
  char kind = dtype.kind();
  ssize_t itemsize = dtype.itemsize();
  if (kind == 'i' && itemsize == 8) {
    return NumpyTypes::Get().int32_dtype;
  } else if (kind == 'u' && itemsize == 8) {
    return NumpyTypes::Get().uint32_dtype;
  } else if (kind == 'f' && itemsize == 8) {
    return NumpyTypes::Get().float32_dtype;
  } else if (kind == 'c' && itemsize == 16) {
    return NumpyTypes::Get().complex64_dtype;
  }
  return dtype;
}

nb::object CanonicalizeInt(nb::handle x) {
  if (GetEnableX64()) {
    int64_t val;
    if (!nb::try_cast<int64_t>(x, val)) {
      throw std::overflow_error(absl::StrCat("Python int ",
                                             nb::cast<std::string>(nb::str(x)),
                                             " too large to convert to int64")
                                    .c_str());
    }
    return typed_int_type(x, NumpyTypes::Get().int64_dtype);
  } else {
    int32_t val;
    if (!nb::try_cast<int32_t>(x, val)) {
      throw std::overflow_error(absl::StrCat("Python int ",
                                             nb::cast<std::string>(nb::str(x)),
                                             " too large to convert to int32")
                                    .c_str());
    }
    return typed_int_type(x, NumpyTypes::Get().int32_dtype);
  }
}

nb::object CanonicalizeFloat(nb::handle x) {
  const NumpyTypes& dtypes = NumpyTypes::Get();
  return typed_float_type(
      x, GetEnableX64() ? dtypes.float64_dtype : dtypes.float32_dtype);
}

nb::object CanonicalizeComplex(nb::handle x) {
  const NumpyTypes& dtypes = NumpyTypes::Get();
  return typed_complex_type(
      x, GetEnableX64() ? dtypes.complex128_dtype : dtypes.complex64_dtype);
}

void CheckValidDtype(nb::handle dtype) {
  if (valid_dtypes.ptr() == nullptr) {
    throw std::runtime_error("valid_dtypes is not set");
  }
  int res = PySequence_Contains(valid_dtypes.ptr(), dtype.ptr());
  if (res == -1) {
    throw nb::python_error();
  }
  if (res == 0) {
    std::string msg = absl::StrCat(
        "Dtype ", nb::cast<std::string>(nb::str(dtype)),
        " is not a valid JAX array type. Only arrays of numeric types are "
        "supported by JAX.");
    PyErr_SetString(PyExc_TypeError, msg.c_str());
    throw nb::python_error();
  }
}

nb::object CanonicalizeNumpyArray(nb::handle x) {
  auto array = nb::borrow<xla::nb_numpy_ndarray>(x);
  xla::nb_dtype raw_dtype = array.dtype();
  CheckValidDtype(raw_dtype);
  xla::nb_dtype canonical_dtype = CanonicalizeDtype(raw_dtype);
  if (raw_dtype.ptr() == canonical_dtype.ptr()) {
    return typed_ndarray_type(x);
  }
  Py_INCREF(canonical_dtype.ptr());
  PyObject* casted = PyArray_CastToType(
      reinterpret_cast<PyArrayObject*>(x.ptr()),
      reinterpret_cast<PyArray_Descr*>(canonical_dtype.ptr()),
      /*fortran=*/0);
  if (!casted) {
    throw nb::python_error();
  }
  auto coerced = nb::steal<xla::nb_numpy_ndarray>(casted);
  return typed_ndarray_type(coerced);
}

nb::object CanonicalizeNumpyScalar(nb::handle x) {
  PyArray_Descr* descr = PyArray_DescrFromObject(x.ptr(), nullptr);
  if (!descr) {
    throw nb::python_error();
  }
  xla::nb_dtype raw_dtype =
      nb::steal<xla::nb_dtype>(reinterpret_cast<PyObject*>(descr));
  CheckValidDtype(raw_dtype);
  xla::nb_dtype canonical_dtype = CanonicalizeDtype(raw_dtype);
  if (raw_dtype.ptr() == canonical_dtype.ptr()) {
    return typed_ndarray_type(x);
  }
  Py_INCREF(canonical_dtype.ptr());
  PyObject* coerced_val = PyArray_FromScalar(
      x.ptr(), reinterpret_cast<PyArray_Descr*>(canonical_dtype.ptr()));
  if (!coerced_val) {
    throw nb::python_error();
  }
  auto coerced = nb::steal<xla::nb_numpy_ndarray>(coerced_val);
  return typed_ndarray_type(coerced);
}

}  // namespace

bool GetEnableX64() {
  if (!enable_x64_state.ptr()) {
    throw std::runtime_error("enable_x64_state is not set");
  }
  bool out = nb::cast<bool>(enable_x64_state->Get());
  return out;
}

void SetEnableX64State(nb_class_ptr<Config> config) {
  enable_x64_state = config;
}

void SetTypedIntType(nb::object t) {
  typed_int_type = t;
  RegisterCanonicalizeValueHandler(reinterpret_cast<PyObject*>(&PyLong_Type),
                                   CanonicalizeInt);
  RegisterCanonicalizeValueHandler(t.ptr(), IdentityHandler);
}

void SetTypedFloatType(nb::object t) {
  typed_float_type = t;
  RegisterCanonicalizeValueHandler(reinterpret_cast<PyObject*>(&PyFloat_Type),
                                   CanonicalizeFloat);
  RegisterCanonicalizeValueHandler(t.ptr(), IdentityHandler);
}

void SetTypedComplexType(nb::object t) {
  typed_complex_type = t;
  RegisterCanonicalizeValueHandler(reinterpret_cast<PyObject*>(&PyComplex_Type),
                                   CanonicalizeComplex);
  RegisterCanonicalizeValueHandler(t.ptr(), IdentityHandler);
}

void SetTypedNdArrayType(nb::object t) {
  typed_ndarray_type = t;
  auto cache = WeakKeyWeakValueCache::Create(CanonicalizeNumpyArray);
  RegisterCanonicalizeValueHandler(reinterpret_cast<PyObject*>(&PyArray_Type),
                                   [cache](nb::handle x) -> nb::object {
                                     return cache->Call(cache.ptr(), x);
                                   });
  RegisterCanonicalizeValueHandler(t.ptr(), IdentityHandler);
  RegisterCanonicalizeValueHandler(NumpyTypes::Get().numpy_generic.ptr(),
                                   CanonicalizeNumpyScalar);
}

void InitCanonicalizeValueHandlers() {
  RegisterCanonicalizeValueHandler(PyArray::type().ptr(), IdentityHandler);
  RegisterCanonicalizeValueHandler(reinterpret_cast<PyObject*>(&PyBool_Type),
                                   IdentityHandler);
}

std::string PyArgSignature::DebugString() const {
  std::string result = "";
  if (weak_type) {
    absl::StrAppend(&result, "weak_");
  }
  absl::StrAppend(&result, xla::PrimitiveType_Name(dtype));
  absl::StrAppend(&result, "[", absl::StrJoin(shape, ","), "]");
  return result;
}

using ToPyArgSignatureHandler =
    std::function<absl::StatusOr<PyArgSignature>(nb::handle, bool)>;

absl::StatusOr<PyArgSignature> PyArgSignatureOfValue(nb::handle arg,
                                                     bool jax_enable_x64) {
  static xla::SafeStatic<
      absl::flat_hash_map<PyObject*, ToPyArgSignatureHandler>>
      handlers_init;
  const absl::flat_hash_map<PyObject*, ToPyArgSignatureHandler>& handlers =
      handlers_init.Get([] {
        auto p = std::make_unique<
            absl::flat_hash_map<PyObject*, ToPyArgSignatureHandler>>();

        const xla::NumpyScalarTypes& dtypes = xla::GetNumpyScalarTypes();

        // The 4 Python native types.
        ToPyArgSignatureHandler bool_handler =
            [](nb::handle, bool) -> absl::StatusOr<PyArgSignature> {
          return PyArgSignature(xla::PrimitiveType::PRED, {}, true);
        };
        ToPyArgSignatureHandler int_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          // TODO(phawkins): we should consider checking for integer
          // overflow.
          if (jax_enable_x64) {
            return PyArgSignature(xla::PrimitiveType::S64, {}, true);
          } else {
            return PyArgSignature(xla::PrimitiveType::S32, {}, true);
          }
        };
        ToPyArgSignatureHandler float_handler =
            [&dtypes](nb::handle h,
                      bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          // Only Python native types has a True weak_type.
          bool weak_type = !nb::isinstance(h, dtypes.np_float64);
          if (jax_enable_x64) {
            return PyArgSignature(xla::PrimitiveType::F64, {}, weak_type);
          } else {
            return PyArgSignature(xla::PrimitiveType::F32, {}, weak_type);
          }
        };
        ToPyArgSignatureHandler complex_handler =
            [&dtypes](nb::handle h,
                      bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          // Note that this branch is also taken  for np.complex128:
          // isinstance(np.complex128(3), complex) returns True
          // isinstance(np.complex64(3), complex) returns False
          bool weak_type = !nb::isinstance(h, dtypes.np_complex128);
          if (jax_enable_x64) {
            return PyArgSignature(xla::PrimitiveType::C128, {}, weak_type);
          } else {
            return PyArgSignature(xla::PrimitiveType::C64, {}, weak_type);
          }
        };

        (*p)[reinterpret_cast<PyObject*>(&PyBool_Type)] = bool_handler;
        (*p)[reinterpret_cast<PyObject*>(&PyLong_Type)] = int_handler;
        (*p)[reinterpret_cast<PyObject*>(&PyFloat_Type)] = float_handler;
        (*p)[reinterpret_cast<PyObject*>(&PyComplex_Type)] = complex_handler;

        ToPyArgSignatureHandler typed_scalar_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          ABSL_ASSIGN_OR_RETURN(
              xla::PrimitiveType dtype,
              DtypeToPrimitiveType(nb::cast<xla::nb_dtype>(h.attr("dtype"))));
          return PyArgSignature(dtype, {}, true);
        };
        if (typed_int_type.ptr() != nullptr) {
          (*p)[typed_int_type.ptr()] = typed_scalar_handler;
        }
        if (typed_float_type.ptr() != nullptr) {
          (*p)[typed_float_type.ptr()] = typed_scalar_handler;
        }
        if (typed_complex_type.ptr() != nullptr) {
          (*p)[typed_complex_type.ptr()] = typed_scalar_handler;
        }

        ToPyArgSignatureHandler numpy_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          xla::nb_numpy_ndarray numpy_array =
              nb::cast<xla::nb_numpy_ndarray>(h);
          ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType dtype,
                                DtypeToPrimitiveType(numpy_array.dtype()));
          if (!jax_enable_x64) {
            dtype = Squash64BitType(dtype);
          }
          // We use reinterpret_cast<> to defend against environments where
          // ssize_t may not be precisely the same type as int64_t, even if
          // it is the same size (long vs long long).
          static_assert(sizeof(int64_t) == sizeof(ssize_t),
                        "Code assumes ssize_t is the same as int64_t");
          return PyArgSignature(
              dtype,
              absl::MakeConstSpan(
                  reinterpret_cast<const int64_t*>(numpy_array.shape()),
                  numpy_array.ndim()),
              /*weak_type=*/false);
        };
        (*p)[reinterpret_cast<PyObject*>(&PyArray_Type)] = numpy_handler;

        ToPyArgSignatureHandler typed_ndarray_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          bool weak_type = nb::cast<bool>(h.attr("weak_type"));
          xla::nb_numpy_ndarray numpy_array =
              nb::cast<xla::nb_numpy_ndarray>(h.attr("val"));
          ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType dtype,
                                DtypeToPrimitiveType(numpy_array.dtype()));
          static_assert(sizeof(int64_t) == sizeof(ssize_t),
                        "Code assumes ssize_t is the same as int64_t");
          return PyArgSignature(
              dtype,
              absl::MakeConstSpan(
                  reinterpret_cast<const int64_t*>(numpy_array.shape()),
                  numpy_array.ndim()),
              weak_type);
        };
        if (typed_ndarray_type.ptr() != nullptr) {
          (*p)[typed_ndarray_type.ptr()] = typed_ndarray_handler;
        }

        ToPyArgSignatureHandler np_uint64_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          if (jax_enable_x64) {
            return PyArgSignature(xla::PrimitiveType::U64, {},
                                  /*weak_type=*/false);
          } else {
            return PyArgSignature(xla::PrimitiveType::U32, {},
                                  /*weak_type=*/false);
          }
        };
        ToPyArgSignatureHandler np_int_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          if (jax_enable_x64) {
            return PyArgSignature(xla::PrimitiveType::S64, {},
                                  /*weak_type=*/false);
          } else {
            return PyArgSignature(xla::PrimitiveType::S32, {},
                                  /*weak_type=*/false);
          }
        };
        ToPyArgSignatureHandler numpy_array_handler =
            [](nb::handle h,
               bool jax_enable_x64) -> absl::StatusOr<PyArgSignature> {
          // This block deals with all numpy scalar types, except for
          // int64_dt, float64_dt and complex128_dt which are taken care of
          // in previous if blocks.
          ABSL_ASSIGN_OR_RETURN(auto dtype,
                                xla::DtypeToPrimitiveType(h.attr("dtype")));
          return PyArgSignature(dtype, {}, /*weak_type=*/false);
        };

        // This block deals with all numpy scalar types, except for
        // int64_dt, float64_dt and complex128_dt which are taken care of in
        // previous if blocks.
        (*p)[dtypes.np_bool.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_int4.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_int8.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_int16.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_int32.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_int64.ptr()] = np_int_handler;
        (*p)[dtypes.np_uint4.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_uint8.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_uint16.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_uint32.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_uint64.ptr()] = np_uint64_handler;
        (*p)[dtypes.np_float4_e2m1fn.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float6_e2m3fn.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float6_e3m2fn.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e3m4.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e4m3.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e4m3fn.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e4m3b11fnuz.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e4m3fnuz.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e5m2.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e5m2fnuz.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float8_e8m0fnu.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float16.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_bfloat16.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float32.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_float64.ptr()] = float_handler;
        (*p)[dtypes.np_complex64.ptr()] = numpy_array_handler;
        (*p)[dtypes.np_complex128.ptr()] = complex_handler;
        (*p)[dtypes.np_longlong.ptr()] = np_int_handler;
        (*p)[dtypes.np_intc.ptr()] = numpy_array_handler;

        return p;
      });

  if (arg.type().ptr() == PyArray::type().ptr()) {
    auto array = nb::borrow<PyArray>(arg);
    xla::ifrt::ArrayRef ifrt_array = array.ifrt_array_ref();
    if (ifrt_array == nullptr) {
      return xla::InvalidArgument("Array has been deleted.");
    }
    ABSL_ASSIGN_OR_RETURN(auto primitive_type,
                          ifrt::ToPrimitiveType(ifrt_array->dtype()));
    return PyArgSignature(primitive_type, array.shape(), array.weak_type());
  }

  auto res = handlers.find(arg.type().ptr());
  if (res == handlers.end()) {
    // We attempt to look at the MRO classes
    PyTypeObject* type = reinterpret_cast<PyTypeObject*>(arg.type().ptr());
    PyObject* mro = type->tp_mro;
    if (mro != nullptr) {
      Py_ssize_t n = PyTuple_GET_SIZE(mro);
      for (Py_ssize_t i = 0; i < n; ++i) {
        PyObject* base = PyTuple_GET_ITEM(mro, i);
        res = handlers.find(base);
        if (res != handlers.end()) {
          return res->second(arg, jax_enable_x64);
        }
      }
    }
    return xla::InvalidArgument(
        "%s",
        absl::StrCat("Not supported: The C++ ToPyArgSignature only accepts "
                     "JAX Arrays, Numpy arrays and scalars of supported types "
                     "(see implementation), or Python scalars. Got type ",
                     nb::cast<std::string_view>(nb::str(arg.type()))));
  }
  return res->second(arg, jax_enable_x64);
}

absl::StatusOr<DevicePutResult> DevicePutWithDevice(
    nanobind::handle addressable_shard, ifrt::Client* ifrt_client,
    ifrt::Device* ifrt_device, ifrt::MemoryKind ifrt_memory_kind,
    const DevicePutOptions& options) {
  tsl::profiler::TraceMe traceme("DevicePut");
  ++GetDevicePutInfo().device_put_with_device;

  if (!ifrt_device->IsAddressable()) {
    return xla::InvalidArgument(
        "Cannot copy array to non-addressable device: %s",
        ifrt_device->DebugString());
  }

  ABSL_ASSIGN_OR_RETURN(ShardFn shard_fn,
                        MakeShardFn(addressable_shard, ifrt_client, ifrt_device,
                                    ifrt_memory_kind, options));

  nb::gil_scoped_release gil_release;

  ABSL_ASSIGN_OR_RETURN(Shard shard, std::move(shard_fn)());
  ABSL_ASSIGN_OR_RETURN(ifrt::ArrayRef ifrt_array,
                        MakeSingleDeviceIfrtArrayFromShard(
                            ifrt_client, ifrt_device, ifrt_memory_kind, shard));
  return DevicePutResult(std::move(ifrt_array), shard.weak_type);
}

absl::StatusOr<DevicePutResult> DevicePutWithSharding(
    absl::Span<const nanobind::handle> src_shards,
    xla::ifrt::DeviceListRef dst_devices, ifrt::Client* ifrt_client,
    const xla::nb_dtype& dtype, absl::Span<const int64_t> shape,
    nanobind::handle sharding, const DevicePutOptions& options) {
  tsl::profiler::TraceMe traceme("DevicePutWithSharding");
  ++GetDevicePutInfo().device_put_with_sharding;

  ABSL_ASSIGN_OR_RETURN(ifrt::DeviceListRef ifrt_device_list,
                        GetIfrtDeviceList(sharding));
  ifrt::DeviceList* ifrt_addressable_device_list =
      ifrt_device_list->AddressableDeviceList();

  if (src_shards.size() != dst_devices->size()) {
    return xla::InvalidArgument(
        "Number of input shards does not match the number of destination "
        "devices: %d vs. %d",
        src_shards.size(), dst_devices->size());
  }

  ABSL_ASSIGN_OR_RETURN(ifrt::DType ifrt_dtype, DtypeToIfRtDType(dtype));
  ifrt::Shape ifrt_shape(shape);
  ifrt::MemoryKind ifrt_memory_kind = GetMemoryKind(sharding);

  std::vector<ShardFn> shard_fns;
  shard_fns.reserve(src_shards.size());
  for (int i = 0; i < src_shards.size(); ++i) {
    ABSL_ASSIGN_OR_RETURN(
        ShardFn shard,
        MakeShardFn(src_shards[i], ifrt_client, dst_devices->devices()[i],
                    ifrt_memory_kind, options));
    shard_fns.push_back(std::move(shard));
  }

  ifrt::ShardingRef ifrt_sharding;
  ABSL_ASSIGN_OR_RETURN(ifrt_sharding,
                        GetIfrtHloSharding(sharding, ifrt_shape));
  // Fully-replicated shardings enable additional optimizations of using a
  // single host buffer.
  // TODO(hyeontaek): Enable a similar optimization for partially replicated
  // cases to reduce the number of host buffers to obtain.
  bool is_fully_replicated = ifrt_sharding->IsFullyReplicated();

  nb::gil_scoped_release gil_release;

  // Whether to build an IFRT array from host buffers as a single batch. We do
  // not batch any shard is already an IFRT array.
  bool should_batch = true;

  std::vector<Shard> shards;
  shards.reserve(shard_fns.size());
  for (int64_t i = 0; i < shard_fns.size(); ++i) {
    ABSL_ASSIGN_OR_RETURN(Shard shard, std::move(shard_fns[i])());
    if (shard.is_ifrt_array()) {
      // If any shard is an IFRT array, we should assemble shards.
      should_batch = false;
    }
    if (dst_devices->devices()[i]->IsAddressable()) {
      shards.push_back(std::move(shard));
      if (should_batch && is_fully_replicated) {
        // We need only one host buffer for a fully-replicated array.
        break;
      }
    }
  }
  // While we have finished calling `shard_fns`, we cannot destroy them until we
  // make a call to IFRT array creation. Destroying `shard_fns` would release
  // host buffers prematurely and can cause the array creation API to see
  // garbage data.

  // TODO(emilyaf): Remove the following and just use ifrt_dtype when tokens are
  // supported.
  if (!shards.empty()) {
    ifrt_dtype = shards.front().ifrt_dtype();
  }

  ifrt::ArrayRef ifrt_array;
  if (should_batch) {
    if (is_fully_replicated && shards.size() == 1) {
      ++GetDevicePutInfo().device_put_fully_replicated;
      ABSL_ASSIGN_OR_RETURN(
          ifrt_array,
          MakeIfrtArrayFromFullyReplicatedShard(
              ifrt_client, std::move(ifrt_sharding), shards.front()));
    } else {
      ++GetDevicePutInfo().device_put_batched;
      ABSL_ASSIGN_OR_RETURN(
          ifrt_array, MakeIfrtArrayFromShardsInBatch(
                          ifrt_client, ifrt_dtype, std::move(ifrt_shape),
                          std::move(ifrt_sharding), absl::MakeSpan(shards)));
    }
  } else {
    ++GetDevicePutInfo().device_put_assembled;
    ABSL_ASSIGN_OR_RETURN(
        ifrt_array, MakeIfrtArrayFromShardsWithAssembly(
                        ifrt_client, ifrt_dtype, std::move(ifrt_shape),
                        std::move(ifrt_sharding), ifrt_addressable_device_list,
                        ifrt_memory_kind, absl::MakeSpan(shards)));
  }
  const bool weak_type = shards.empty() ? false : shards.front().weak_type;
  return DevicePutResult(std::move(ifrt_array), weak_type);
}

void RegisterCanonicalizeValueHandler(
    PyObject* type, std::function<nb::object(nb::handle)> handler) {
  ft_lock_guard lock(canonicalize_value_handlers_mutex);
  auto& handlers = *canonicalize_value_handlers;
  handlers[type] = std::move(handler);
}

void SetInvalidInputException(nb::object exc) { invalid_input_exception = exc; }

void SetValidDtypes(nb::object dtypes) { valid_dtypes = dtypes; }

nb::object CanonicalizeValue(nb::handle x) {
  CanonicalizeValueHandler handler;
  PyObject* typ = reinterpret_cast<PyObject*>(x.type().ptr());
  {
    ft_lock_guard lock(canonicalize_value_handlers_mutex);
    auto& handlers = *canonicalize_value_handlers;

    // 1. Direct lookup
    auto it = handlers.find(typ);
    if (it != handlers.end()) {
      handler = it->second;
    } else {
      // 2. MRO walk
      PyTypeObject* type = reinterpret_cast<PyTypeObject*>(typ);
      PyObject* mro = type->tp_mro;
      if (mro != nullptr) {
        Py_ssize_t n = PyTuple_GET_SIZE(mro);
        for (Py_ssize_t i = 0; i < n; ++i) {
          PyObject* base = PyTuple_GET_ITEM(mro, i);
          it = handlers.find(base);
          if (it != handlers.end()) {
            handler = it->second;
            handlers[typ] = handler;
            break;
          }
        }
      }
    }
  }

  if (handler) {
    return handler(x);
  }

  // 3. Check __jax_array__
  if (nb::hasattr(x, "__jax_array__")) {
    throw nb::value_error(
        "Triggering __jax_array__() during abstractification is no longer"
        " supported. To avoid this error, either explicitly convert your object"
        " using jax.numpy.array(), or register your object as a pytree.");
  }

  // 4. Error
  if (invalid_input_exception.ptr() == nullptr) {
    throw std::runtime_error("invalid_input_exception not initialized");
  }
  std::string msg = absl::StrCat(
      "Argument '", nb::cast<std::string>(nb::str(x)), "' of type ",
      nb::cast<std::string>(nb::str(x.type())), " is not a valid JAX type.");
  PyErr_SetString(invalid_input_exception.ptr(), msg.c_str());
  throw nb::python_error();
}

absl::flat_hash_map<std::string, int64_t> DevicePutInfo::GetInfo() {
  const DevicePutInfo& info = GetDevicePutInfo();
  return absl::flat_hash_map<std::string, int64_t>({
      {"device_put_with_device", info.device_put_with_device},
      {"device_put_with_sharding", info.device_put_with_sharding},
      {"device_put_fully_replicated", info.device_put_fully_replicated},
      {"device_put_batched", info.device_put_batched},
      {"device_put_assembled", info.device_put_assembled},
  });
}

}  // namespace jax