/* Copyright 2022 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/pjit.h"

#include <Python.h>
#include <structmember.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <exception>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <thread>  // NOLINT
#include <tuple>
#include <unordered_map>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/cleanup/cleanup.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/inlined_vector.h"
#include "absl/hash/hash.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/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/synchronization/mutex.h"
#include "absl/synchronization/notification.h"
#include "absl/types/span.h"
#include "nanobind/nanobind.h"
#include "nanobind/stl/optional.h"  // IWYU pragma: keep
#include "nanobind/stl/shared_ptr.h"  // IWYU pragma: keep
#include "nanobind/stl/string.h"  // IWYU pragma: keep
#include "nanobind/stl/string_view.h"  // IWYU pragma: keep
#include "nanobind/stl/vector.h"  // IWYU pragma: keep
#include "jaxlib/call_location.h"
#include "jaxlib/config.h"
#include "jaxlib/free_threading.h"
#include "jaxlib/guard_lib.h"
#include "jaxlib/jax_jit.h"
#include "jaxlib/nb_class_ptr.h"
#include "jaxlib/py_array.h"
#include "jaxlib/py_client.h"
#include "jaxlib/py_executable.h"
#include "jaxlib/py_user_context.h"
#include "jaxlib/py_values.h"
#include "jaxlib/python_ref_manager.h"
#include "jaxlib/pytree.h"
#include "jaxlib/reentrant_hash_map.h"
#include "jaxlib/sharding.h"
#include "jaxlib/to_ifrt_sharding.h"
#include "xla/layout.h"
#include "xla/pjrt/exceptions.h"
#include "xla/pjrt/lru_cache.h"
#include "xla/python/ifrt/array.h"
#include "xla/python/ifrt/device.h"
#include "xla/python/ifrt/device_list.h"
#include "xla/python/ifrt/executable.h"
#include "xla/python/ifrt/index.h"
#include "xla/python/ifrt/index_domain.h"
#include "xla/python/ifrt/memory.h"
#include "xla/python/ifrt/shape.h"
#include "xla/python/ifrt/sharding.h"
#include "xla/python/ifrt/user_context.h"
#include "xla/python/nb_helpers.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/tsl/concurrency/ref_count.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"
#include "tsl/profiler/lib/traceme.h"

namespace jax {
namespace {

namespace nb = nanobind;

struct PjitCacheEntry {
  explicit PjitCacheEntry(PyTreeRegistry* registry)
      : out_pytree_def(registry) {}
  std::shared_ptr<PyLoadedExecutable> executable;
  std::vector<nb::object> in_shardings;
  std::vector<nb::object> out_avals;
  std::vector<xla::nb_dtype> out_dtypes;
  std::vector<std::vector<int64_t>> out_shapes;
  std::vector<bool> out_weak_types;
  std::vector<nb::object> out_shardings;
  std::vector<bool> out_committed;
  PyTreeDef out_pytree_def;
  // Bitvector of kept arguments from Jaxpr DCE pass. Used to drop some `args`
  // in PjitFunction::Call before calling into compiled computation.
  std::vector<bool> kept_var_bitvec;
  std::vector<nb::object> in_device_local_layouts;
  std::vector<nb::object> const_args;

  // Ensures a single thread performs the compilation for a given executable.
  //
  // The first thread (holding the GIL) will create the CacheEntry associated to
  // a signature and if the object has been inserted already, other threads
  // will wait for the notification.
  absl::Notification compilation_complete;

  std::thread::id thread_id = std::this_thread::get_id();

  bool fall_back_to_python = false;
};

// A PjitFunctionCache represents a cache of compiled functions that can be
// shared between one or more PjitFunction objects. It serves two goals:
// - reduce the number of lru caches (hash map) across multiple JITs.
// - make the cache global to increase cache hits (e.g. calling jit(f)(3) twice)
//   keeping entries alive as long as the underlying function f is alive.
// Assume the cache is protected by the GIL.
class PjitFunctionCache {
 public:
  static constexpr int kDefaultCapacity = 4096;
  explicit PjitFunctionCache(int capacity);

  // Cache entries are shared_ptr<>s because it's possible the cache entry
  // might be evicted before we finish tracing/compiling.
  typedef xla::LRUCache<CallSignature, std::shared_ptr<PjitCacheEntry>,
                        CallSignature::Hash>
      Cache;

  // We include as part of the cache key `global_cache_key` (and any other
  // fields that aren't subsumed by the CallSignature we compute for each call).
  static std::shared_ptr<Cache> Lookup(nb_class_ptr<PjitFunctionCache> self,
                                       nb::handle function,
                                       nb::object global_cache_key);
  std::shared_ptr<Cache> DefaultCache();

  // These methods require the GIL or the object's lock in no-GIL mode.
  int Size() const { return lru_list_.Size(); }
  int Capacity() const { return lru_list_.Capacity(); }
  void Clear() {
    lru_list_.Clear();
    functions_.clear();
  }

 private:
  struct Key {
    nb::handle function;  // Does not hold a reference.

    // Other fields that are part of the arguments to `jit`, but are not
    // otherwise part of CallSignature.
    nb::object global_cache_key;

    size_t cached_hash;
  };

  template <typename H>
  friend H AbslHashValue(H h, const Key& key) {
    h = H::combine(std::move(h), key.function.ptr());
    Py_hash_t hash;
    try {
      hash = nb::hash(key.global_cache_key);
    } catch (const nanobind::python_error& e) {
      if (!e.matches(PyExc_TypeError)) throw;
      throw std::invalid_argument(absl::StrCat(
          "Hashing global cache key lead to an exception. The error was:\n",
          e.what(), "\n"));
    }
    h = H::combine(std::move(h), hash);
    return h;
  }

  // Helper for fast, Python-free lookup during deletion.
  struct PointerKey {
    nb::handle function;
    nb::handle global_cache_key;
    size_t cached_hash;
  };

  struct KeyEq {
    // Normal lookup: uses Python equal() but short-circuits on pointer
    // mismatch.
    bool operator()(const Key& a, const Key& b) const {
      if (a.cached_hash != b.cached_hash) return false;
      if (a.function.ptr() != b.function.ptr()) return false;
      try {
        return a.global_cache_key.equal(b.global_cache_key);
      } catch (const nanobind::python_error& e) {
        throw std::invalid_argument(
            absl::StrCat("Equality of global cache key led to an exception. "
                         "The error was:\n",
                         e.what(), "\n"));
      }
    }

    // Fast lookup: pure pointer equality, no Python calls.
    // Explicitly compare .ptr() to ensure we do not trigger Python-level
    // equality.
    bool operator()(const Key& a, const PointerKey& b) const {
      if (a.cached_hash != b.cached_hash) return false;
      return a.function.ptr() == b.function.ptr() &&
             a.global_cache_key.ptr() == b.global_cache_key.ptr();
    }
  };

  struct Hash {
    size_t operator()(const Key& key) const { return key.cached_hash; }
    size_t operator()(const PointerKey& pkey) const { return pkey.cached_hash; }
  };

  struct Value {
    explicit Value(std::shared_ptr<Cache> cache) : cache(std::move(cache)) {}
    std::shared_ptr<Cache> cache;

    // A weak reference to the key function. We use the weak reference to
    // register a callback that is triggered when the key function is destroyed.
    // We use a weak pointer because we want to allow caching across multiple
    // calls to `pjit(f)` if `f` remains alive, but we do not want the cache
    // to keep `f` alive if all other references are dropped.
    std::optional<nb::weakref> weakref;
  };

  // lru_list_ and functions_ are protected by the GIL in GIL mode, and by the
  // self object lock in freethreading mode.
  Cache::LRUList lru_list_;
  // We use ReentrantHashMap to be safe against reentrant mutations.
  jax::ReentrantHashMap<Key, std::unique_ptr<Value>, Hash, KeyEq> functions_;
  // mu_ prevents concurrent insertions into functions_ if the gil or critical
  // section lock is released during insertion.
  absl::Mutex mu_;
};

PjitFunctionCache::PjitFunctionCache(int capacity) : lru_list_(capacity) {}

std::shared_ptr<PjitFunctionCache::Cache> PjitFunctionCache::DefaultCache() {
  return std::make_shared<Cache>(&lru_list_);
}

/*static*/ std::shared_ptr<PjitFunctionCache::Cache> PjitFunctionCache::Lookup(
    nb_class_ptr<PjitFunctionCache> self, nb::handle function,
    nb::object global_cache_key) ABSL_NO_THREAD_SAFETY_ANALYSIS {
  // In no-GIL mode, a critical section on self plays the same role that
  // the GIL plays in GIL mode.
  nb::ft_object_guard lock(self);
  {
    // Because the gil (or the critical section lock) can be released during
    // cache insertion, this forces the lock order to be mu_ then gil so we
    // must release the gil first.
    nb::gil_scoped_release release;
    // Acquire a mutex to avoid problems where the gil is released during
    // cache insertion and then a second thread invalidates the cache order.
    self->mu_.lock();
  }
  absl::Cleanup unlock = [&self]() ABSL_UNLOCK_FUNCTION(self->mu_) {
    self->mu_.unlock();
  };
  Key key;
  key.function = function;
  key.global_cache_key = global_cache_key;
  key.cached_hash = absl::HashOf(key);
  auto insert = self->functions_.try_emplace(key, nullptr);
  if (!insert.second) {
    return insert.first->second->cache;
  }
  std::shared_ptr<Cache> cache = std::make_shared<Cache>(&self->lru_list_);
  auto callback =
      nb::cpp_function([self, key{std::move(key)}](nb::handle weakref) {
        nb::ft_object_guard lock(self);
        PointerKey pkey{key.function, key.global_cache_key, key.cached_hash};
        auto it = self->functions_.find(pkey);
        if (it == self->functions_.end()) {
          return;
        }
        // Remove the value from the map before destroying it. Destroying
        // the value may release `lock` since it may call arbitrary Python
        // code.
        std::unique_ptr<Value> value = std::move(it->second);
        self->functions_.erase(it);
        value.reset();
      });
  PyObject* weakref = PyWeakref_NewRef(function.ptr(), callback.ptr());
  if (weakref) {
    std::unique_ptr<Value>& entry = insert.first->second;
    entry = std::make_unique<Value>(cache);
    entry->weakref = nb::steal<nb::weakref>(weakref);
  } else {
    PyErr_Clear();
    // `function` is not weak-referenceable. Don't bother adding it to the
    // shared cache in that case; the `jit` object will hold the only shared
    // reference to the cache entry.
    PointerKey pkey{key.function, key.global_cache_key, key.cached_hash};
    auto it = self->functions_.find(pkey);
    if (it != self->functions_.end()) {
      self->functions_.erase(it);
    }
  }
  return cache;
}

class PjitFunction {
 public:
  PjitFunction(std::string function_name, std::optional<nb::callable> fun,
               nb::callable cache_miss, std::vector<int> static_argnums,
               std::vector<nb::str> static_argnames,
               nb::object global_cache_key,
               nb_class_ptr<PyTreeRegistry> pytree_registry,
               nb::callable shard_arg_fallback,
               nb_class_ptr<PjitFunctionCache> cache);
  ~PjitFunction();

  PjitFunction(const PjitFunction&) = delete;
  PjitFunction& operator=(const PjitFunction&) = delete;
  PjitFunction(PjitFunction&&) = default;
  PjitFunction& operator=(PjitFunction&&) = default;

  // nb::object typed subclass for PjitFunction objects.
  class pyobject : public nb::object {
   public:
    NB_OBJECT(pyobject, nb::object, "PjitFunction",
              PjitFunction::IsPjitFunction);
    pyobject() = default;
    PjitFunction* func() const {
      return PjitFunction::AsPjitFunctionUnchecked(*this);
    }
  };
  // Alias as ::object; outside the scope above we won't confuse nanobind's
  // macros.
  using object = pyobject;

  // Returns true if `h` is a PjitFunction.
  static bool IsPjitFunction(nb::handle handle);
  // Converts `handle` to a PjitFunction*. Does not do any checking.
  static PjitFunction* AsPjitFunctionUnchecked(nb::handle handle);

  absl::StatusOr<nb::object> Call(nb::handle callable, PyObject* const* args,
                                  size_t nargs, PyObject* kwnames);

  void InitExecutables();

  void ClearPythonReferences();

  const std::string& function_name() const { return function_name_; }
  const std::optional<nb::callable>& fun() const { return fun_; }
  const nb::callable& cache_miss() const { return cache_miss_; }
  const nb_class_ptr<PyTreeRegistry>& pytree_registry() const {
    return pytree_registry_;
  }
  const nb::callable& shard_arg_fallback() const { return shard_arg_fallback_; }

  const std::vector<int>& static_argnums() const { return static_argnums_; }
  const std::vector<nb::str>& static_argnames() const {
    return static_argnames_;
  }
  const nb::object& global_cache_key() const { return global_cache_key_; }
  const nb_class_ptr<PjitFunctionCache>& cache() const { return cache_; }

  int cache_capacity() const {
    nb::ft_object_guard lock(cache_);
    return executables_->Size();
  }

  void ClearCache() {
    nb::ft_object_guard lock(cache_);
    executables_->Clear();
  }

  std::shared_ptr<PjitFunctionCache::Cache> executables() {
    nb::ft_object_guard lock(cache_);
    return executables_;
  }

  nb::object PythonSignature() {
    if (!fun_.has_value()) {
      throw nb::value_error(
          absl::StrFormat(
              "Calling __signature__ on PjitFunction(%s) not supported.",
              function_name_)
              .c_str());
    }
    static xla::SafeStatic<nb::object> signature_fn_init;
    const nb::object& signature_fn = signature_fn_init.Get([]() {
      return std::make_unique<nb::object>(
          nb::module_::import_("inspect").attr("signature"));
    });
    return signature_fn(*fun_);
  }

 private:
  absl::Status ComputeCallSignature(
      absl::Span<nb::object const> flat_dynamic_args, bool enable_x64,
      CallSignature& call_signature);

  void PopulateCacheEntry(PjitCacheEntry& cache_entry,
                          const nb::tuple& out_and_fastpath_data);

  std::string function_name_;
  std::optional<nb::callable> fun_;
  nb::callable cache_miss_;
  std::vector<int> static_argnums_;
  std::vector<nb::str> static_argnames_;
  nb::object global_cache_key_;

  nb_class_ptr<PyTreeRegistry> pytree_registry_;
  nb::callable shard_arg_fallback_;
  nb_class_ptr<PjitFunctionCache> cache_;

  // In no-GIL mode executables_ is protected by the object lock on cache_,
  // because it shared an LRU list with cache_.
  std::shared_ptr<PjitFunctionCache::Cache> executables_;
};

PjitFunction::PjitFunction(
    std::string function_name, std::optional<nb::callable> fun,
    nb::callable cache_miss, std::vector<int> static_argnums,
    std::vector<nb::str> static_argnames, nb::object global_cache_key,
    nb_class_ptr<PyTreeRegistry> pytree_registry,
    nb::callable shard_arg_fallback, nb_class_ptr<PjitFunctionCache> cache)
    : function_name_(std::move(function_name)),
      fun_(std::move(fun)),
      cache_miss_(std::move(cache_miss)),
      static_argnums_(std::move(static_argnums)),
      global_cache_key_(std::move(global_cache_key)),
      pytree_registry_(std::move(pytree_registry)),
      shard_arg_fallback_(std::move(shard_arg_fallback)),
      cache_(std::move(cache)) {
  std::sort(static_argnums_.begin(), static_argnums_.end());
  static_argnames_.reserve(static_argnames.size());
  for (nb::str& name : static_argnames) {
    PyObject* s = name.inc_ref().ptr();
    PyUnicode_InternInPlace(&s);
    static_argnames_.push_back(nb::steal<nb::str>(s));
  }
}

void PjitFunction::InitExecutables() {
  // Construction of the object hasn't completed yet, so we don't need to hold
  // the cache lock to mutate executables_.
  if (!fun_.has_value()) {
    executables_ = cache_->DefaultCache();
  } else {
    executables_ = cache_->Lookup(cache_, fun_.value(), global_cache_key_);
  }
}

PjitFunction::~PjitFunction() {
  nb::ft_object_guard lock(cache_);
  executables_ = nullptr;
}

void CallShardArgFallback(nb::handle arg, nb::handle sharding,
                          nb::handle layout, const nb::callable& fallback,
                          std::vector<xla::ifrt::ArrayRef>& num_args_arrays,
                          std::vector<nb::object>& keep_alive_objects) {
  tsl::profiler::TraceMe traceme("cpp_pjit_shard_arg_fallback");
  auto py_array_or_bufs = fallback(arg, sharding, layout);
  auto py_array = nb::cast<PyArray>(py_array_or_bufs);
  num_args_arrays.push_back(py_array.ifrt_array_ref());
  keep_alive_objects.push_back(std::move(py_array_or_bufs));
}

// Prepares the input PjRtBuffers from the python arguments. This is equivalent
// to shard_args() in pxla.py but for only a few supported cases.
absl::StatusOr<std::vector<xla::ifrt::ArrayRef>> PrepareIfrtInputs(
    const PyLoadedExecutable& executable,
    absl::Span<nb::object const> flat_dynamic_args,
    absl::Span<PyArgSignature const> flat_dynamic_arg_signatures,
    bool enable_x64, const std::vector<bool>& kept_args,
    const std::vector<nb::object>& in_shardings,
    const std::vector<nb::object>& in_device_local_layouts,
    const nb::callable& shard_arg_fallback,
    std::vector<nb::object>& keep_alive_objects) {
#if JAX_IFRT_VERSION_NUMBER >= 73
  std::optional<xla::ifrt::DeviceListRef> devices =
      executable.ifrt_loaded_executable()->devices();
  const absl::Span<xla::ifrt::Device* const> addressable_devices =
      devices.has_value() ? (*devices)->AddressableDeviceList()->devices()
                          : absl::Span<xla::ifrt::Device* const>();
#else
  const auto& addressable_devices =
      executable.ifrt_loaded_executable()->addressable_devices();
#endif
#if JAX_IFRT_VERSION_NUMBER >= 74
  const int num_global_devices = devices.has_value() ? (*devices)->size() : 1;
#else
  const auto& num_global_devices =
      executable.ifrt_loaded_executable()->num_devices();
#endif
  int num_args = flat_dynamic_args.size();

  std::vector<xla::ifrt::ArrayRef> num_args_arrays;
  num_args_arrays.reserve(num_args);

  struct CopyGroup {
    std::vector<int> indices;
    std::vector<xla::ifrt::ArrayRef> arrays;
  };
  absl::flat_hash_map<std::tuple<xla::ifrt::Device*, xla::ifrt::MemoryKind,
                                 xla::ifrt::MemoryKind>,
                      CopyGroup>
      copy_groups;

  DevicePutOptions options;
  options.squash_64bit_types = !enable_x64;
  options.allow_zero_copy = true;
  xla::ifrt::Device* data_device = nullptr;
#if JAX_IFRT_VERSION_NUMBER >= 74
  if (num_global_devices == 1 && !addressable_devices.empty()) {
    data_device = addressable_devices[0];
  }
#else
  if (executable.ifrt_loaded_executable()->num_devices() == 1 &&
      !addressable_devices.empty()) {
    data_device = addressable_devices[0];
  }
#endif
  int dce_i = 0;
  for (int i = 0; i < num_args; ++i) {
    if (!kept_args[i]) {
      continue;
    }
    int dce_index = dce_i;
    ++dce_i;

    const nb::object& arg = flat_dynamic_args[i];
    const nb::object& in_device_local_layout =
        in_device_local_layouts[dce_index];

    auto transfer_guard_formatter = [] { return std::string(""); };

    if (arg.type().ptr() != PyArray::type().ptr()) {
      if (data_device != nullptr && in_device_local_layout.is_none()) {
        ABSL_RETURN_IF_ERROR(
            ApplyTransferGuardToHostToDevice(transfer_guard_formatter));
        ABSL_ASSIGN_OR_RETURN(
            auto device_put_result,
            DevicePutWithDevice(arg,
                                executable.ifrt_loaded_executable()->client(),
                                data_device, xla::ifrt::MemoryKind(), options));
        num_args_arrays.push_back(std::move(device_put_result.ifrt_array));
        continue;
      } else {
        CallShardArgFallback(arg, in_shardings[dce_index],
                             in_device_local_layout, shard_arg_fallback,
                             num_args_arrays, keep_alive_objects);
        continue;
      }
    }

    PyArray py_array = nb::borrow<PyArray>(arg);
    const auto& sharding = py_array.sharding();
    int sharding_num_devices =
        nb::cast<const Sharding*>(sharding)->num_devices();

    // Currently only committed PyArray inputs or uncommitted PyArray on a
    // single device inputs are allowed. This is checked previously in the entry
    // point of PjitFunction::Call().
    DCHECK(py_array.committed() ||
           (!py_array.committed() && sharding_num_devices == 1));

    if (!in_device_local_layout.is_none()) {
      xla::ifrt::ArrayRef ifrt_array = py_array.ifrt_array_ref();
      ABSL_ASSIGN_OR_RETURN(auto arr_layout, ifrt_array->pjrt_layout());
      if (arr_layout == nullptr) {
        ABSL_ASSIGN_OR_RETURN(
            xla::ifrt::Shape shard_shape,
            ifrt_array->sharding().GetShardShape(ifrt_array->shape()));
        ABSL_ASSIGN_OR_RETURN(
            arr_layout,
            executable.ifrt_loaded_executable()->client()->GetDefaultPjRtLayout(
                ifrt_array->dtype(), shard_shape.dims(),
                ifrt_array->sharding().devices()->devices().front(),
                ifrt_array->sharding().memory_kind()));
      }
      xla::Layout in_xc_layout = nb::cast<xla::Layout>(
          in_device_local_layout.attr("_to_xla_layout")(py_array.dtype()));
      if (in_xc_layout != arr_layout->xla_layout()) {
        CallShardArgFallback(arg, in_shardings[dce_index],
                             in_device_local_layout, shard_arg_fallback,
                             num_args_arrays, keep_alive_objects);
        continue;
      }
    }

    if (sharding_num_devices != num_global_devices) {
      CallShardArgFallback(arg, in_shardings[dce_index], in_device_local_layout,
                           shard_arg_fallback, num_args_arrays,
                           keep_alive_objects);
      continue;
    }

    xla::ifrt::ArrayRef ifrt_array = py_array.ifrt_array_ref();
    // PyArray inputs should have already been checked in
    // `PyArgSignatureOfValue()` called by
    // `PjitFunction::ComputeCallSignature()`.
    DCHECK(ifrt_array != nullptr) << "PyArray has been unexpectedly deleted.";

    const auto& ifrt_sharding = ifrt_array->sharding();
    if (sharding_num_devices == 1 && !addressable_devices.empty() &&
        ifrt_sharding.devices()->devices().front() != addressable_devices[0]) {
      auto& copy_group = copy_groups[std::make_tuple(
          ifrt_sharding.devices()->devices().front(),
          ifrt_sharding.memory_kind(), GetMemoryKind(in_shardings[dce_index]))];
      copy_group.indices.push_back(num_args_arrays.size());
      copy_group.arrays.push_back(ifrt_array);
      num_args_arrays.push_back({});
    } else {
      num_args_arrays.push_back(ifrt_array);
    }

    keep_alive_objects.push_back(arg);
  }

  if (!copy_groups.empty() && !addressable_devices.empty()) {
    xla::ifrt::Client* const ifrt_client =
        executable.ifrt_loaded_executable()->client();
    ABSL_ASSIGN_OR_RETURN(
        xla::ifrt::DeviceListRef ifrt_devices,
        ifrt_client->MakeDeviceList({addressable_devices[0]}));
    for (auto& [key, group] : copy_groups) {
      ABSL_ASSIGN_OR_RETURN(
          auto copied_ifrt_arrays,
          ifrt_client->CopyArrays(absl::MakeSpan(group.arrays), ifrt_devices,
                                  std::get<2>(key),
                                  xla::ifrt::ArrayCopySemantics::kReuseInput));
      for (int i = 0; i < copied_ifrt_arrays.size(); ++i) {
        num_args_arrays[group.indices[i]] = std::move(copied_ifrt_arrays[i]);
      }
    }
  }

  return num_args_arrays;
}

absl::StatusOr<nb::object> PjitFunction::Call(nb::handle callable,
                                              PyObject* const* args,
                                              size_t nargs, PyObject* kwnames) {
  tsl::profiler::TraceMe traceme(
      [&] { return absl::StrCat("PjitFunction(", function_name_, ")"); });

  // Make sure we trigger a garbage collection on JIT function calls. Otherwise
  // code like
  // f = jit(...)
  // while True:
  //   f(x)
  // may never free temporary buffers for copies of arguments.
  GlobalPyRefManager()->MaybeCollectGarbage();

  if (GetDisableJit()) {
    if (!fun_.has_value()) {
      throw nb::value_error(
          absl::StrFormat("Disable jit is not supported in the AOT path since "
                          "the function is not available for (%s)",
                          function_name_)
              .c_str());
    }
    return nb::steal<nb::object>(
        PyObject_Vectorcall(fun_.value().ptr(), args, nargs, kwnames));
  }

  // Calls the cache_miss_ function. This just calls the Python function; it may
  // return nullptr value if a Python exception is thrown.
  auto cache_miss = [&]() -> nb::tuple {
    return nb::steal<nb::tuple>(
        PyObject_Vectorcall(cache_miss_.ptr(), args, nargs, kwnames));
  };

  // Call the cache_miss() function, extracting the output data and ignoring
  // the fastpath data. If the cache miss returns a Python error, returns
  // nullptr and leaves the Python error set.
  auto fallback_to_cache_miss = [&]() {
    nb::tuple cache_miss_output = cache_miss();
    if (!cache_miss_output.ptr()) {
      return nb::object();
    }
    return nb::object(cache_miss_output[0]);
  };

  size_t num_positional_args = PyVectorcall_NARGS(nargs);
  size_t num_keyword_args = kwnames ? PyTuple_GET_SIZE(kwnames) : 0;
  absl::Span<PyObject* const> positional_args(args, num_positional_args);
  absl::Span<PyObject* const> keyword_args(args + num_positional_args,
                                           num_keyword_args);

  CallSignature call_signature;
  std::vector<nb::object> keep_alive_objects;
  absl::InlinedVector<nb::object, 2> flat_dynamic_args;
  auto status = ParseArguments(
      positional_args, keyword_args, kwnames, static_argnums_, static_argnames_,
      pytree_registry_.get(), call_signature.arg_signature, flat_dynamic_args);
  if (!status.ok()) {
    VLOG(2) << "ParseArguments failed: " << status;
    return fallback_to_cache_miss();
  }

  // Perform a few checks for the arguments. Currently we are only allowing
  // committed PyArray inputs. For other cases, e.g. Tracers or ShapedArray, it
  // will fallback to python. For jit, numpy arrays and scalars are also
  // allowed, which we will check later.
  for (const auto& arg : flat_dynamic_args) {
    if (arg.type().ptr() != PyArray::type().ptr()) {
      continue;
    }

    PyArray py_array = nb::borrow<PyArray>(arg);

    // Only allow committed PyArray in cpp pjit for now as the logic on handling
    // sharding for uncommitted PyArray is complicated and still under
    // development.
    //
    // TODO(chky): Consider support uncommitted PyArray in cpp when the python
    // side stabilizes.
    int sharding_num_devices =
        nb::cast<const Sharding*>(py_array.sharding())->num_devices();
    if (!py_array.committed() && sharding_num_devices > 1) {
      VLOG(2) << "PyArray argument is not committed and number of global "
                 "devices is more than 1; fallback to python.";
      return fallback_to_cache_miss();
    }
  }

  bool enable_x64 = GetEnableX64();
  status = ComputeCallSignature(flat_dynamic_args, enable_x64, call_signature);
  if (!status.ok()) {
    VLOG(2) << "ComputeCallSignature failed: " << status;
    return fallback_to_cache_miss();
  }

  VLOG(2) << "CallSignature:\n" << call_signature.DebugString();
  bool inserted = false;
  std::shared_ptr<PjitCacheEntry> cache_entry;
  {
    nb::ft_object_guard lock(cache_);
    cache_entry = executables_->GetOrCreateIfAbsent(
        call_signature, [this, &inserted](const CallSignature& unused) {
          inserted = true;
          return std::make_shared<PjitCacheEntry>(pytree_registry_.get());
        });
  }

  if (!cache_entry->compilation_complete.HasBeenNotified()) {
    // In case of several threads attempting to compile the executable, only
    // the one that inserted the item will perform the compilation.
    if (inserted) {
      nb::object out_and_fastpath_data;
      nb::tuple out_tuple;
      VLOG(2) << "Cache miss for " << call_signature.DebugString();
      bool remove_cache = false;
      try {
        // Calls Python and may release the GIL. May also throw if
        // compilation/tracing fails.
        out_and_fastpath_data = cache_miss();
        if (!out_and_fastpath_data.ptr()) {
          throw nb::python_error();
        }
        out_tuple = nb::cast<nb::tuple>(out_and_fastpath_data);

        PopulateCacheEntry(*cache_entry, out_tuple);

        if (out_tuple.size() > 2 && out_tuple[2].is_valid()) {
          remove_cache = nb::cast<bool>(out_tuple[2]);
        }
      } catch (const std::exception& e) {
        VLOG(2) << "cache miss fail: " << e.what();
        cache_entry->fall_back_to_python = true;
        cache_entry->compilation_complete.Notify();
        throw;
      }
      cache_entry->compilation_complete.Notify();

      if (remove_cache) {
        nb::ft_object_guard lock(cache_);
        executables_->Remove(call_signature);
      }

      // We have already computed the result in the miss path so we can return
      // it. We are even *required* to do so if there are donated arguments,
      // because any donated buffers will now be invalid.
      return nb::object(out_tuple[0]);
    } else {
      if (cache_entry->thread_id == std::this_thread::get_id()) {
        auto error_string = absl::StrCat("Recursively calling jit: ",
                                         call_signature.DebugString());
        PyErr_SetString(PyExc_RecursionError, error_string.c_str());
        throw nb::python_error();
      }
      // Release the GIL while we wait, making sure the compile thread can
      // lock it.
      nb::gil_scoped_release release;
      cache_entry->compilation_complete.WaitForNotification();
    }
  }

  if (cache_entry->fall_back_to_python) {
    VLOG(2) << "cpp pjit fallback to python.";
    return fallback_to_cache_miss();
  }

  absl::InlinedVector<PyArgSignature, 2> dynamic_arg_signatures;
  dynamic_arg_signatures.reserve(cache_entry->const_args.size() +
                                 flat_dynamic_args.size());
  if (!cache_entry->const_args.empty()) {
    flat_dynamic_args.reserve(cache_entry->const_args.size() +
                              flat_dynamic_args.size());
    flat_dynamic_args.insert(flat_dynamic_args.begin(),
                             cache_entry->const_args.begin(),
                             cache_entry->const_args.end());

    for (nb::handle const_arg : cache_entry->const_args) {
      ABSL_ASSIGN_OR_RETURN(auto const_arg_signature,
                            PyArgSignatureOfValue(const_arg, enable_x64));
      dynamic_arg_signatures.push_back(std::move(const_arg_signature));
    }
  }
  for (const auto& arg : call_signature.dynamic_arg_signatures) {
    dynamic_arg_signatures.push_back(std::move(arg));
  }

  PyUserContextScope user_context_scope;
  // A vector of [num_inputs].
  auto num_args_arrays = PrepareIfrtInputs(
      *cache_entry->executable, flat_dynamic_args, dynamic_arg_signatures,
      enable_x64, cache_entry->kept_var_bitvec, cache_entry->in_shardings,
      cache_entry->in_device_local_layouts, shard_arg_fallback_,
      keep_alive_objects);

  if (!num_args_arrays.ok()) {
    VLOG(2) << "Failed to prepare IFRT inputs: " << num_args_arrays.status();
    return fallback_to_cache_miss();
  }

  xla::ifrt::ExecuteOptions execute_options =
      cache_entry->executable->options();
  execute_options.launch_id = cache_entry->executable->GetNextLaunchId();
  execute_options.execution_stream_id = GetExecutionStreamId();
  if (execute_options.execution_stream_id == 0) {
    execute_options.execution_stream_id =
        tsl::Env::Default()->GetCurrentThreadId();
  }
  PopulateCallLocation(execute_options,
                       xla::ifrt::UserContextScope::current().get());

  // Check if the thread guard is active and should prevent execution.
  // Skipped for portable executables.
  if (cache_entry->executable->ifrt_executable()->devices().has_value()) {
    ABSL_RETURN_IF_ERROR(CheckThreadGuard(
        *cache_entry->executable->ifrt_executable()->devices()));
  }

  // A vector of [num_outputs].
  std::vector<xla::ifrt::ArrayRef> output_arrays;
  {
    nb::gil_scoped_release gil_release;
    ABSL_ASSIGN_OR_RETURN(auto result,
                          cache_entry->executable->ifrt_executable()->Execute(
                              absl::MakeSpan(*num_args_arrays), execute_options,
                              /*devices=*/std::nullopt));
    output_arrays = std::move(result.outputs);
  }

  // Convert the ifrt::Array objects to PyArray.
  int num_outputs = output_arrays.size();
  absl::InlinedVector<nb::object, 4> outputs;
  outputs.reserve(num_outputs);
  for (int i = 0; i < num_outputs; ++i) {
    // Creating the PyArray result. In addition to the IFRT arrays, the metadata
    // like `aval` and `sharding` are retrieved from the cache for this
    // function, which are produced by the python path in `cache_miss`.
    PyArray py_array(cache_entry->out_avals[i], cache_entry->out_weak_types[i],
                     cache_entry->out_dtypes[i], cache_entry->out_shapes[i],
                     cache_entry->out_shardings[i],
                     cache_entry->executable->client(),
                     std::move(output_arrays[i]),
                     /*committed=*/cache_entry->out_committed.at(i));

    outputs.push_back(std::move(py_array));
  }

  nb::object out = nb::steal<nb::object>(
      cache_entry->out_pytree_def.Unflatten(outputs).release().ptr());

  // If there is a post-hook function, call it with the inputs and the outputs.
  std::optional<nb::object> post_hook = GetPostHook();
  if (post_hook) {
    nb::tuple_builder args_tuple(num_positional_args);
    for (size_t i = 0; i < num_positional_args; ++i) {
      args_tuple.put(nb::handle(args[i]));
    }
    nb::dict kwargs;
    if (kwnames) {
      for (size_t i = 0; i < num_keyword_args; ++i) {
        kwargs[nb::handle(PyTuple_GET_ITEM(kwnames, i))] =
            nb::borrow(args[num_positional_args + i]);
      }
    }
    (*post_hook)(nb::handle(callable.ptr()), args_tuple.commit(), kwargs,
                 nb::handle(out.ptr()));
  }

  return out;
}

absl::Status PjitFunction::ComputeCallSignature(
    absl::Span<nb::object const> flat_dynamic_args, bool enable_x64,
    CallSignature& signature) {
  signature.function_name = function_name_;

  // Get dynamic argument signatures.
  auto& dynamic_arg_signatures = signature.dynamic_arg_signatures;
  dynamic_arg_signatures.reserve(flat_dynamic_args.size());
  auto& dynamic_arg_shardings = signature.dynamic_arg_shardings;
  dynamic_arg_shardings.reserve(flat_dynamic_args.size());
  auto& dynamic_arg_layouts = signature.dynamic_arg_layouts;
  dynamic_arg_layouts.reserve(flat_dynamic_args.size());

  for (nb::handle arg : flat_dynamic_args) {
    ABSL_ASSIGN_OR_RETURN(auto arg_signature,
                          PyArgSignatureOfValue(arg, enable_x64));
    signature.dynamic_arg_signatures.push_back(std::move(arg_signature));

    // It should be already checked previously in the entry point of
    // PjitFunction::Call().
    if (arg.type().ptr() == PyArray::type().ptr()) {
      auto py_array = nb::borrow<PyArray>(arg);
      signature.dynamic_arg_shardings.push_back(py_array.sharding());
      auto layout = py_array.layout();
      if (layout.ok()) {
        signature.dynamic_arg_layouts.push_back(*std::move(layout));
      } else if (absl::IsUnimplemented(layout.status())) {
        signature.dynamic_arg_layouts.push_back(nullptr);
      } else {
        return layout.status();
      }
      signature.committed_args.push_back(py_array.committed());
    } else {
      signature.dynamic_arg_shardings.push_back(nb::none());
      signature.dynamic_arg_layouts.push_back(nullptr);
      signature.committed_args.push_back(false);
    }
  }

  signature.configs = JitConfigs();
  signature.cached_hash = absl::HashOf(signature);

  return absl::OkStatus();
}

void PjitFunction::PopulateCacheEntry(PjitCacheEntry& cache_entry,
                                      const nb::tuple& out_and_fastpath_data) {
  DCHECK_GE(out_and_fastpath_data.size(), 2);

  if (out_and_fastpath_data[1].is_none()) {
    VLOG(2) << "fastpath_data is none";
    cache_entry.fall_back_to_python = true;
    return;
  }

  nb::tuple fastpath_data = nb::cast<nb::tuple>(out_and_fastpath_data[1]);

  cache_entry.executable = nb::cast<std::shared_ptr<PyLoadedExecutable>>(
      fastpath_data.attr("xla_executable"));

  nb::sequence in_shardings = fastpath_data.attr("in_shardings");
  cache_entry.in_shardings.reserve(nb::len(in_shardings));
  for (nb::handle sharding : in_shardings) {
    cache_entry.in_shardings.push_back(nb::borrow(sharding));
  }

  nb::sequence out_shardings = fastpath_data.attr("out_shardings");
  cache_entry.out_shardings.reserve(nb::len(out_shardings));
  for (nb::handle sharding : out_shardings) {
    cache_entry.out_shardings.push_back(nb::borrow(sharding));
  }

  nb::sequence out_committed = fastpath_data.attr("out_committed");
  cache_entry.out_committed.reserve(nb::len(out_committed));
  for (nb::handle c : out_committed) {
    cache_entry.out_committed.push_back(nb::cast<bool>(c));
  }

  nb::sequence out_avals = fastpath_data.attr("out_avals");
  cache_entry.out_avals.reserve(nb::len(out_avals));
  cache_entry.out_dtypes.reserve(nb::len(out_avals));
  cache_entry.out_shapes.reserve(nb::len(out_avals));
  cache_entry.out_weak_types.reserve(nb::len(out_avals));
  for (nb::handle aval : out_avals) {
    cache_entry.out_avals.push_back(nb::borrow(aval));
    cache_entry.out_dtypes.push_back(aval.attr("dtype"));
    cache_entry.out_shapes.push_back(
        nb::cast<std::vector<int64_t>>(aval.attr("shape")));
    cache_entry.out_weak_types.push_back(
        nb::cast<bool>(aval.attr("weak_type")));
  }

  cache_entry.out_pytree_def = nb::cast<PyTreeDef>(
      nb::handle(fastpath_data.attr("out_pytree_def").ptr()));

  nb::sequence kept_var_bitvec = fastpath_data.attr("kept_var_bitvec");
  cache_entry.kept_var_bitvec.reserve(nb::len(kept_var_bitvec));
  for (nb::handle k : kept_var_bitvec) {
    cache_entry.kept_var_bitvec.push_back(nb::cast<bool>(k));
  }

  nb::sequence in_device_local_layouts =
      fastpath_data.attr("in_device_local_layouts");
  cache_entry.in_device_local_layouts.reserve(nb::len(in_device_local_layouts));
  for (nb::handle dll : in_device_local_layouts) {
    cache_entry.in_device_local_layouts.push_back(nb::borrow(dll));
  }

  nb::sequence const_args = fastpath_data.attr("const_args");
  cache_entry.const_args.reserve(nb::len(const_args));
  for (nb::handle ca : const_args) {
    cache_entry.const_args.push_back(nb::borrow(ca));
  }
}

// Helper function used by the tp_clear GC method.
void PjitFunction::ClearPythonReferences() {
  nb::callable cache_miss;
  std::optional<nb::callable> fun;
  nb::callable shard_arg_fallback;
  nb::object global_cache_key;
  nb_class_ptr<PyTreeRegistry> pytree_registry;
  // Swap values for nulls before they are destroyed. See the Python
  // Py_CLEAR() documentation for a discussion of this topic.
  std::swap(cache_miss_, cache_miss);
  std::swap(fun_, fun);
  std::swap(shard_arg_fallback_, shard_arg_fallback);
  std::swap(global_cache_key_, global_cache_key);
  std::swap(pytree_registry_, pytree_registry);
}

struct PjitFunctionObject {
  PyObject_HEAD;
  vectorcallfunc vectorcall;
  PjitFunction fun;

  // Doubly-linked list of PjitFunctionObjects, protected by
  // PjitFunctionStore::mu_ or the GIL in GIL mode.
  PjitFunctionObject* next;
  PjitFunctionObject* prev;
};

// Contains a list of all PjitFunctionObjects.
// Thread-safe.
class PjitFunctionStore {
 public:
  void Insert(PjitFunctionObject* o) {
    ft_lock_guard lock(mu_);
    o->next = compiled_functions_;
    o->prev = nullptr;
    if (o->next) {
      o->next->prev = o;
    }
    compiled_functions_ = o;
  }

  void Remove(PjitFunctionObject* o) {
    ft_lock_guard lock(mu_);
    if (o->next) {
      o->next->prev = o->prev;
    }
    if (o->prev) {
      o->prev->next = o->next;
    } else {
      compiled_functions_ = o->next;
    }
  }

  void ClearCaches() {
    std::vector<
        std::pair<nb::object, std::shared_ptr<PjitFunctionCache::Cache>>>
        caches;
    {
      ft_lock_guard lock(mu_);
      for (PjitFunctionObject* fn = compiled_functions_; fn != nullptr;
           fn = fn->next) {
        caches.emplace_back(fn->fun.cache(), fn->fun.executables());
      }
    }
    for (auto& [cache, executables] : caches) {
      nb::ft_object_guard lock(cache);
      executables->Clear();
    }
  };

 private:
  ft_mutex mu_;
  PjitFunctionObject* compiled_functions_ ABSL_GUARDED_BY(mu_);
};

PjitFunctionStore pjit_function_store;

PyObject* PjitFunction_Type = nullptr;

bool PjitFunction::IsPjitFunction(nb::handle handle) {
  return handle.type().ptr() == PjitFunction_Type;
}

PjitFunction* PjitFunction::AsPjitFunctionUnchecked(nb::handle handle) {
  return &(reinterpret_cast<PjitFunctionObject*>(handle.ptr())->fun);
}

PjitFunction* AsPjitFunction(nb::handle handle) {
  if (!PjitFunction::IsPjitFunction(handle)) {
    throw xla::XlaRuntimeError(xla::InvalidArgument("Expected a PjitFunction"));
  }
  return PjitFunction::AsPjitFunctionUnchecked(handle);
}

extern "C" {

PyObject* PjitFunction_tp_vectorcall(PyObject* callable, PyObject* const* args,
                                     size_t nargs, PyObject* kwnames) {
  PjitFunctionObject* o = reinterpret_cast<PjitFunctionObject*>(callable);
  tsl::profiler::TraceMe traceme([&] {
    return absl::StrCat("PjitFunction(", o->fun.function_name(), ")");
  });
  try {
    absl::StatusOr<nb::object> out =
        o->fun.Call(callable, args, nargs, kwnames);
    if (!out.ok()) {
      PyErr_SetString(PyExc_ValueError, out.status().ToString().c_str());
      return nullptr;
    }
    return out.value().release().ptr();
  } catch (nb::python_error& e) {
    e.restore();
    return nullptr;
  } catch (nb::cast_error& e) {
    PyErr_SetString(PyExc_ValueError, e.what());
    return nullptr;
  } catch (std::invalid_argument& e) {
    PyErr_SetString(PyExc_ValueError, e.what());
    return nullptr;
  } catch (std::runtime_error& e) {
    PyErr_SetString(PyExc_ValueError, e.what());
    return nullptr;
  }
}

PyObject* PjitFunction_tp_new(PyTypeObject* subtype, PyObject* args,
                              PyObject* kwds) {
  PjitFunctionObject* self =
      reinterpret_cast<PjitFunctionObject*>(subtype->tp_alloc(subtype, 0));
  if (!self) return nullptr;

  self->vectorcall = PjitFunction_tp_vectorcall;
  return reinterpret_cast<PyObject*>(self);
}

void PjitFunction_tp_dealloc(PyObject* self) {
  PyObject_GC_UnTrack(self);
  PyTypeObject* tp = Py_TYPE(self);
  PjitFunctionObject* o = reinterpret_cast<PjitFunctionObject*>(self);
  pjit_function_store.Remove(o);
  PyObject_ClearWeakRefs(self);
#if PY_VERSION_HEX < 0x030D0000
  _PyObject_ClearManagedDict(self);
#else
  PyObject_ClearManagedDict(self);
#endif
  o->fun.~PjitFunction();
  tp->tp_free(self);
  Py_DECREF(tp);
}

int PjitFunction_tp_traverse(PyObject* self, visitproc visit, void* arg) {
  PjitFunctionObject* o = reinterpret_cast<PjitFunctionObject*>(self);
  // https://docs.python.org/3/c-api/typeobj.html#c.PyTypeObject.tp_traverse
  Py_VISIT(Py_TYPE(self));
#if PY_VERSION_HEX < 0x030D0000
  _PyObject_VisitManagedDict(self, visit, arg);
#else
  PyObject_VisitManagedDict(self, visit, arg);
#endif
  Py_VISIT(o->fun.cache_miss().ptr());
  Py_VISIT(o->fun.shard_arg_fallback().ptr());
  Py_VISIT(o->fun.global_cache_key().ptr());
  Py_VISIT(o->fun.pytree_registry().ptr());
  if (o->fun.fun()) {
    Py_VISIT(o->fun.fun()->ptr());
  }
  return 0;
}

int PjitFunction_tp_clear(PyObject* self) {
  PjitFunctionObject* o = reinterpret_cast<PjitFunctionObject*>(self);
#if PY_VERSION_HEX < 0x030D0000
  _PyObject_ClearManagedDict(self);
#else
  PyObject_ClearManagedDict(self);
#endif
  o->fun.ClearPythonReferences();
  return 0;
}

// Implements the Python descriptor protocol so JIT-compiled functions can be
// used as bound methods. See:
// https://docs.python.org/3/howto/descriptor.html#functions-and-methods
PyObject* PjitFunction_tp_descr_get(PyObject* self, PyObject* obj,
                                    PyObject* type) {
  if (obj == nullptr || obj == Py_None) {
    Py_INCREF(self);
    return self;
  }
  return PyMethod_New(self, obj);
}

static PyGetSetDef PjitFunction_tp_getset[] = {
    // Having a __dict__ seems necessary to allow !functool.wraps to override
    // __doc__.
    {const_cast<char*>("__dict__"), PyObject_GenericGetDict,
     PyObject_GenericSetDict, nullptr, nullptr},
    {nullptr, nullptr, nullptr, nullptr, nullptr}};

PyObject* PjitFunction_tp_repr(PyObject* self) {
  try {
    const std::string& repr = absl::StrFormat(
        "<PjitFunction of %s>",
        nb::cast<std::string_view>(nb::repr(nb::getattr(self, "__wrapped__"))));
    return PyUnicode_FromString(repr.c_str());
  } catch (...) {
    // Ignore all errors when accessing a repr.
    return PyUnicode_FromString("<PjitFunction>");
  }
}

}  // extern "C"

void InitializePjitFunction(
    PjitFunctionObject* fn_obj, std::string function_name,
    std::optional<nb::callable> fun, nb::callable cache_miss,
    std::vector<int> static_argnums, std::vector<nb::str> static_argnames,
    nb::object global_cache_key, nb_class_ptr<PyTreeRegistry> pytree_registry,
    nb::callable shard_arg_fallback, nb_class_ptr<PjitFunctionCache> cache) {
  fn_obj->next = fn_obj->prev = nullptr;
  if (nb::isinstance<nb::list>(global_cache_key)) {
    global_cache_key = nb::tuple(global_cache_key);
  }
  new (&fn_obj->fun) PjitFunction(
      std::move(function_name), std::move(fun), std::move(cache_miss),
      std::move(static_argnums), std::move(static_argnames),
      std::move(global_cache_key), std::move(pytree_registry),
      std::move(shard_arg_fallback), std::move(cache));
  // Handled separately because it is not exception safe to call this
  // in the constructor because it leaves the object improperly constructed.
  fn_obj->fun.InitExecutables();

  // Only add the executable to the store after executables_ has been
  // initialized. We want only fully constructed executables in the store.
  pjit_function_store.Insert(fn_obj);
}

nb::object MakePjitFunction(
    std::string function_name, std::optional<nb::callable> fun,
    nb::callable cache_miss, std::vector<int> static_argnums,
    std::vector<nb::str> static_argnames, nb::object global_cache_key,
    nb_class_ptr<PyTreeRegistry> pytree_registry,
    nb::callable shard_arg_fallback,
    std::optional<nb_class_ptr<PjitFunctionCache>> cache) {
  nb::object obj = nb::steal<nb::object>(PjitFunction_tp_new(
      reinterpret_cast<PyTypeObject*>(PjitFunction_Type), nullptr, nullptr));
  PjitFunctionObject* fn_obj = reinterpret_cast<PjitFunctionObject*>(obj.ptr());
  if (!cache) {
    cache =
        make_nb_class<PjitFunctionCache>(PjitFunctionCache::kDefaultCapacity);
  }
  InitializePjitFunction(
      fn_obj, std::move(function_name), std::move(fun), std::move(cache_miss),
      std::move(static_argnums), std::move(static_argnames),
      std::move(global_cache_key), std::move(pytree_registry),
      std::move(shard_arg_fallback), std::move(*cache));
  return obj;
}

// Version numbers for the pickled representations of
// PjitFunction. Increment these if changing them.
const int kPjitFunctionPickleVersion = 1;

PyMemberDef PjitFunction_members[] = {
    {"__vectorcalloffset__", T_PYSSIZET,
     static_cast<Py_ssize_t>(offsetof(PjitFunctionObject, vectorcall)),
     READONLY, nullptr},
    {nullptr, 0, 0, 0, nullptr},
};

PyType_Slot PjitFunction_slots[] = {
    {Py_tp_new, reinterpret_cast<void*>(PjitFunction_tp_new)},
    {Py_tp_dealloc, reinterpret_cast<void*>(PjitFunction_tp_dealloc)},
    {Py_tp_traverse, reinterpret_cast<void*>(PjitFunction_tp_traverse)},
    {Py_tp_clear, reinterpret_cast<void*>(PjitFunction_tp_clear)},
    {Py_tp_getset, reinterpret_cast<void*>(PjitFunction_tp_getset)},
    {Py_tp_descr_get, reinterpret_cast<void*>(PjitFunction_tp_descr_get)},
    {Py_tp_call, reinterpret_cast<void*>(PyVectorcall_Call)},
    {Py_tp_repr, reinterpret_cast<void*>(PjitFunction_tp_repr)},
    {Py_tp_members, reinterpret_cast<void*>(PjitFunction_members)},
    {0, nullptr},
};

class ShardArgsHandler {
 public:
  virtual ~ShardArgsHandler() = default;

  virtual bool TryAdd(size_t result_idx, nb::handle arg, nb::handle sharding,
                      nb::handle layout, nb::handle copy_semantics) = 0;

  virtual void Run(std::vector<nb::object>& results) = 0;
};

class FallbackHandler : public ShardArgsHandler {
 public:
  explicit FallbackHandler(std::optional<nb::callable> fallback)
      : fallback_(std::move(fallback)) {}

  bool TryAdd(size_t result_idx, nb::handle arg, nb::handle sharding,
              nb::handle layout, nb::handle copy_semantics) override {
    indices_.push_back(result_idx);
    args_.append(arg);
    shardings_.append(sharding);
    layouts_.append(layout);
    copy_semantics_.append(copy_semantics);
    return true;
  }

  void Run(std::vector<nb::object>& results) override {
    if (indices_.empty()) {
      return;
    }
    if (!fallback_.has_value() || fallback_->is_none()) {
      throw nb::value_error("Fallback handler has no valid fallback function.");
    }
    nb::sequence fallback_results = nb::cast<nb::sequence>(
        (*fallback_)(shardings_, layouts_, copy_semantics_, args_));

    for (size_t i = 0; i < indices_.size(); ++i) {
      results[indices_[i]] = fallback_results[i];
    }
  }

 private:
  std::optional<nb::callable> fallback_;
  std::vector<size_t> indices_;
  nb::list args_;
  nb::list shardings_;
  nb::list layouts_;
  nb::list copy_semantics_;
};

class NumpyHandler : public ShardArgsHandler {
 private:
  struct GroupKey {
    nb_class_ptr<PyClient> py_client;
    xla::ifrt::DeviceListRef devices;
    xla::ifrt::MemoryKind memory_kind;

    template <typename H>
    friend H AbslHashValue(H h, const GroupKey& key) {
      return H::combine(std::move(h), key.py_client.ptr(), key.devices,
                        key.memory_kind);
    }

    bool operator==(const GroupKey& other) const {
      return py_client.ptr() == other.py_client.ptr() &&
             devices == other.devices && memory_kind == other.memory_kind;
    }
  };

  struct Group {
    std::vector<xla::ifrt::Client::MakeArraysFromHostBufferShardsSpec> specs;
    std::vector<size_t> indices;
    std::vector<nb::object> shardings;
    std::vector<bool> weak_types;
  };

 public:
  NumpyHandler() = default;

  bool TryAdd(size_t result_idx, nb::handle arg, nb::handle sharding,
              nb::handle layout, nb::handle copy_semantics) override {
    if (!layout.is_none()) {
      return false;
    }
    static nb::object typed_ndarray_type = []() -> nb::object {
      try {
        return nb::module_::import_("jax._src.literals").attr("TypedNdArray");
      } catch (...) {
        return nb::object();
      }
    }();
    PyObject* typ = reinterpret_cast<PyObject*>(arg.type().ptr());
    if (typ != reinterpret_cast<PyObject*>(&PyArray_Type) &&
        (typed_ndarray_type.ptr() == nullptr ||
         typ != typed_ndarray_type.ptr())) {
      return false;
    }

    auto device_list_or = GetPyDeviceList(sharding);
    if (!device_list_or.ok()) {
      VLOG(2) << "GetPyDeviceList failed: " << device_list_or.status()
              << "; fallback to python.";
      return false;
    }
    nb_class_ptr<PyClient> arg_py_client = (*device_list_or)->py_client();
    xla::ifrt::Client* arg_client = arg_py_client->ifrt_client();
    if (arg_client == nullptr) {
      VLOG(2) << "IFRT client is null; fallback to python.";
      return false;
    }

    bool enable_x64 = GetEnableX64();
    auto signature_or = PyArgSignatureOfValue(arg, enable_x64);
    if (!signature_or.ok()) {
      VLOG(2) << "PyArgSignatureOfValue failed: " << signature_or.status()
              << "; fallback to python.";
      return false;
    }
    const PyArgSignature& signature = *signature_or;

    auto ifrt_dtype_or = xla::ifrt::ToDType(signature.dtype);
    if (!ifrt_dtype_or.ok()) {
      VLOG(2) << "ToDType failed: " << ifrt_dtype_or.status()
              << "; fallback to python.";
      return false;
    }
    xla::ifrt::DType ifrt_dtype = *ifrt_dtype_or;
    xla::ifrt::Shape ifrt_shape(signature.shape);

    xla::nb_numpy_ndarray ndarray = nb::borrow<xla::nb_numpy_ndarray>(arg);

    auto actual_type_or = xla::DtypeToPrimitiveType(ndarray.dtype());
    if (!actual_type_or.ok() || *actual_type_or != signature.dtype) {
      VLOG(2) << "Type mismatch or conversion needed: actual="
              << (actual_type_or.ok() ? static_cast<int>(*actual_type_or) : -1)
              << " expected=" << static_cast<int>(signature.dtype)
              << "; fallback to python.";
      return false;
    }

    std::shared_ptr<PythonRefManager::ManagedPyObjects> py_buffer_ref =
        GlobalPyRefManager()->ManageReference(nb::borrow<nb::object>(arg));

    auto ifrt_sharding_or = GetIfrtHloSharding(sharding, ifrt_shape);
    if (!ifrt_sharding_or.ok()) {
      VLOG(2) << "GetIfrtHloSharding failed: " << ifrt_sharding_or.status()
              << "; fallback to python.";
      return false;
    }
    xla::ifrt::ShardingRef ifrt_sharding = std::move(*ifrt_sharding_or);
    xla::ifrt::DeviceListRef devices = ifrt_sharding->devices();
    xla::ifrt::MemoryKind memory_kind = ifrt_sharding->memory_kind();

    auto index_domains_or = ifrt_sharding->IndexDomains(
        ifrt_shape, xla::ifrt::SingleDeviceShardSemantics::kAddressableShards);
    if (!index_domains_or.ok()) {
      VLOG(2) << "IndexDomains failed: " << index_domains_or.status()
              << "; fallback to python.";
      return false;
    }
    auto index_domains = std::move(*index_domains_or);

    xla::ifrt::Client::MakeArraysFromHostBufferShardsSpec spec = {
        .array_spec = {
            .dtype = ifrt_dtype,
            .shape = ifrt_shape,
            .sharding = std::move(ifrt_sharding),
            .layout = nullptr,
        }};

    absl::flat_hash_map<xla::ifrt::IndexDomain, size_t> shards;
    for (int i = 0; i < index_domains.size(); ++i) {
      const auto& index_domain = index_domains[i];
      auto [it, inserted] =
          shards.try_emplace(index_domain, spec.buffers.size());
      if (!inserted) {
        spec.buffers[it->second].first.push_back(i);
        continue;
      }
      const char* base_ptr = reinterpret_cast<const char*>(ndarray.data());
      int64_t offset = 0;
      for (int d = 0; d < ndarray.ndim(); ++d) {
        offset += index_domain.origin().elements()[d] * ndarray.strides(d);
      }
      const char* shard_ptr = base_ptr + offset;

      xla::ifrt::Shape shard_shape = index_domain.shape();

      std::vector<int64_t> byte_strides(ndarray.ndim());
      for (int d = 0; d < ndarray.ndim(); ++d) {
        byte_strides[d] = ndarray.strides(d);
      }

      std::function<void()> on_done = [py_buffer_ref]() {
        // keeps py_buffer_ref alive
      };

      xla::ifrt::Client::HostBuffer host_buffer = {
          .data = shard_ptr,
          .dtype = ifrt_dtype,
          .shape = shard_shape,
          .byte_strides = std::move(byte_strides),
          .on_done = std::move(on_done),
      };
      spec.buffers.push_back({{i}, std::move(host_buffer)});
    }

    auto transfer_guard_status = ApplyTransferGuardToHostToDevice(
        [&arg]() { return nb::cast<std::string>(nb::repr(arg)); });
    if (!transfer_guard_status.ok()) {
      throw nb::value_error(
          std::string(transfer_guard_status.message()).c_str());
    }

    GroupKey key{std::move(arg_py_client), std::move(devices),
                 std::move(memory_kind)};
    Group& group = groups_[key];
    group.specs.push_back(std::move(spec));
    group.indices.push_back(result_idx);
    group.shardings.push_back(nb::borrow<nb::object>(sharding));
    group.weak_types.push_back(signature.weak_type);
    return true;
  }

  void Run(std::vector<nb::object>& results) override {
    auto status = RunInternal(results);
    if (!status.ok()) {
      throw nb::value_error(
          absl::StrCat("NumpyHandler failed: ", status.message()).c_str());
    }
  }

 private:
  absl::Status RunInternal(std::vector<nb::object>& results) {
    if (groups_.empty()) {
      return absl::OkStatus();
    }
    PyUserContextScope user_context_scope;
    for (auto& [key, group] : groups_) {
      xla::ifrt::Client* client = key.py_client->ifrt_client();
      if (client == nullptr) {
        return xla::InvalidArgument("Numpy handler has no valid IFRT client.");
      }
      std::vector<xla::ifrt::ArrayRef> arrays;
      {
        nb::gil_scoped_release gil_release;
        ABSL_ASSIGN_OR_RETURN(
            arrays,
            client->MakeArraysFromHostBufferShards(
                absl::MakeSpan(group.specs),
                xla::ifrt::Client::HostBufferSemantics::kImmutableZeroCopy));
      }

      for (size_t i = 0; i < arrays.size(); ++i) {
        size_t orig_idx = group.indices[i];
        PyArray py_array = PyArray::MakeFromIfrtArrayAndSharding(
            key.py_client, std::move(arrays[i]), group.shardings[i],
            /*weak_type=*/group.weak_types[i], /*committed=*/true,
            /*skip_checks=*/false);
        results[orig_idx] = py_array;
      }
    }

    return absl::OkStatus();
  }

  absl::flat_hash_map<GroupKey, Group> groups_;
};

}  // namespace

void BuildPjitSubmodule(nb::module_& m) {
  m.attr("_PyTreeRegistry") = m.attr("pytree").attr("PyTreeRegistry");

  nb::class_<PjitFunctionCache> cache(m, "PjitFunctionCache");
  cache.def(nb::init<int>(),
            nb::arg("capacity") = PjitFunctionCache::kDefaultCapacity);
  cache.def("size", &PjitFunctionCache::Size, nb::lock_self());
  cache.def("capacity", &PjitFunctionCache::Capacity, nb::lock_self());
  cache.def("clear", &PjitFunctionCache::Clear, nb::lock_self());
  cache.def_static("clear_all", []() { pjit_function_store.ClearCaches(); });
  cache.def(
      "__getstate__",
      // Pickles as an empty cache; the client can repopulate as needed.
      [](const PjitFunctionCache& cache) {
        nb::dict pickle;
        pickle["version"] = kPjitFunctionPickleVersion;
        pickle["capacity"] = cache.Capacity();
        return pickle;
      },
      nb::lock_self());
  cache.def("__setstate__",
            [](PjitFunctionCache* cache, const nb::dict& pickle) {
              int version = nb::cast<int>(pickle["version"]);
              if (version != kPjitFunctionPickleVersion) {
                throw std::invalid_argument(absl::StrFormat(
                    "Invalid PjitFunction pickle version, got %d, expected %d",
                    version, kPjitFunctionPickleVersion));
              }
              int capacity = nb::cast<int>(pickle["capacity"]);
              new (cache) PjitFunctionCache(capacity);
            });

  // We need to use heap-allocated type objects because we want to add
  // additional methods dynamically.
  std::string name =
      absl::StrCat(nb::cast<std::string>(m.attr("__name__")), ".PjitFunction");
  PyType_Spec PjitFunction_spec = {
      /*.name=*/name.c_str(),
      /*.basicsize=*/static_cast<int>(sizeof(PjitFunctionObject)),
      /*.itemsize=*/0,
      /*.flags=*/Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
          Py_TPFLAGS_HAVE_VECTORCALL | Py_TPFLAGS_MANAGED_DICT |
          Py_TPFLAGS_MANAGED_WEAKREF,
      /*.slots=*/PjitFunction_slots,
  };
  PjitFunction_Type = PyType_FromSpec(&PjitFunction_spec);
  if (!PjitFunction_Type) {
    throw nb::python_error();
  }
  nb::object cfun = nb::borrow<nb::object>(PjitFunction_Type);

  // Add PjitFunction to the _jax module so it can be pickled.
  m.attr("PjitFunction") = cfun;
  cfun.attr("__getstate__") = nb::cpp_function(
      [](const PjitFunction::object& self) {
        PjitFunction* fn = self.func();
        nb::dict pickle;
        pickle["version"] = kPjitFunctionPickleVersion;
        pickle["function_name"] = fn->function_name();
        if (fn->fun().has_value()) {
          pickle["fun"] = *fn->fun();
        }
        pickle["cache_miss"] = fn->cache_miss();
        pickle["static_argnums"] = fn->static_argnums();
        pickle["static_argnames"] = nb::cast(fn->static_argnames());
        pickle["global_cache_key"] = fn->global_cache_key();
        pickle["pytree_registry"] = nb::cast(fn->pytree_registry());
        pickle["shard_arg_fallback"] = fn->shard_arg_fallback();
        pickle["cache"] = fn->cache();
        return pickle;
      },
      nb::is_method());
  cfun.attr("__setstate__") = nb::cpp_function(
      [](nb::object& self, const nb::dict& pickle) {
        int version = nb::cast<int>(pickle["version"]);
        if (version != kPjitFunctionPickleVersion) {
          throw std::invalid_argument(absl::StrFormat(
              "Invalid PjitFunction pickle version, got %d, expected %d. "
              "Pickling/Unpickling jitted functions using different JAX "
              "versions is not supported.",
              version, kPjitFunctionPickleVersion));
        }
        std::string function_name =
            nb::cast<std::string>(pickle["function_name"]);
        std::optional<nb::callable> fun;
        if (pickle.contains("fun")) {
          fun = nb::cast<nb::callable>(pickle["fun"]);
        }
        nb::callable cache_miss = nb::cast<nb::callable>(pickle["cache_miss"]);
        std::vector<int> static_argnums =
            nb::cast<std::vector<int>>(pickle["static_argnums"]);
        std::vector<nb::str> static_argnames =
            nb::cast<std::vector<nb::str>>(pickle["static_argnames"]);
        nb::object global_cache_key = pickle["global_cache_key"];
        nb_class_ptr<PyTreeRegistry> pytree_registry =
            nb::cast<nb_class_ptr<PyTreeRegistry>>(
                nb::handle(pickle["pytree_registry"].ptr()));
        nb::callable shard_arg_fallback =
            nb::cast<nb::callable>(pickle["shard_arg_fallback"]);
        nb_class_ptr<PjitFunctionCache> cache =
            nb::cast<nb_class_ptr<PjitFunctionCache>>(pickle["cache"]);
        InitializePjitFunction(
            reinterpret_cast<PjitFunctionObject*>(self.ptr()),
            std::move(function_name), std::move(fun), std::move(cache_miss),
            std::move(static_argnums), std::move(static_argnames),
            std::move(global_cache_key), std::move(pytree_registry),
            std::move(shard_arg_fallback), std::move(cache));
      },
      nb::is_method());
  cfun.attr("__signature__") = xla::nb_property_readonly(
      [](nb::handle self) { return AsPjitFunction(self)->PythonSignature(); },
      nb::sig("def __signature__(self) -> inspect.Signature"));
  cfun.attr("_cache_miss") = xla::nb_property_readonly(
      [](nb::handle self) { return AsPjitFunction(self)->cache_miss(); });
  // All private members are only for testing/debugging purposes
  cfun.attr("_cache_size") = nb::cpp_function(
      [](nb::handle self) -> int {
        return AsPjitFunction(self)->cache_capacity();
      },
      nb::is_method());
  cfun.attr("_clear_cache") = nb::cpp_function(
      [](nb::handle self) { AsPjitFunction(self)->ClearCache(); },
      nb::is_method());

  m.def(
      "shard_args",
      [](nb::sequence shardings, nb::sequence layouts,
         nb::sequence copy_semantics, nb::sequence args,
         std::optional<nb::callable> fallback) -> std::vector<nb::object> {
        size_t num_args = nb::len(args);
        if (nb::len(shardings) != num_args || nb::len(layouts) != num_args ||
            nb::len(copy_semantics) != num_args) {
          throw nb::value_error(
              "shard_args arguments must have the same length.");
        }

        GlobalPyRefManager()->CollectGarbage();

        PyUserContextScope user_context_scope;

        FallbackHandler fallback_handler(fallback);
        NumpyHandler numpy_handler;

        absl::InlinedVector<ShardArgsHandler*, 2> handlers = {
            &numpy_handler, &fallback_handler};

        for (size_t i = 0; i < num_args; ++i) {
          nb::object arg = args[i];
          nb::object sharding = shardings[i];
          nb::object layout = layouts[i];
          nb::object cs = copy_semantics[i];

          for (auto* handler : handlers) {
            if (handler->TryAdd(i, arg, sharding, layout, cs)) {
              break;
            }
          }
        }

        std::vector<nb::object> results_list(num_args);

        for (auto* handler : handlers) {
          handler->Run(results_list);
        }

        return results_list;
      },
      nb::arg("shardings"), nb::arg("layouts"), nb::arg("copy_semantics"),
      nb::arg("args"), nb::arg("fallback").none() = nb::none());

  m.def(
      "pjit",
      [](std::string function_name, std::optional<nb::callable> fun,
         nb::callable cache_miss, std::vector<int> static_argnums,
         std::vector<nb::str> static_argnames, nb::object global_cache_key,
         nb::object pytree_registry, nb::callable shard_arg_fallback,
         std::optional<nb_class_ptr<PjitFunctionCache>> cache) {
        nb_class_ptr<PyTreeRegistry> registry =
            nb::cast<nb_class_ptr<PyTreeRegistry>>(
                nb::handle(pytree_registry.ptr()));
        return MakePjitFunction(
            std::move(function_name), std::move(fun), std::move(cache_miss),
            std::move(static_argnums), std::move(static_argnames),
            std::move(global_cache_key), std::move(registry),
            std::move(shard_arg_fallback), std::move(cache));
      },
      nb::arg("function_name"), nb::arg("fun").none(), nb::arg("cache_miss"),
      nb::arg("static_argnums"), nb::arg("static_argnames"),
      nb::arg("global_cache_key"), nb::arg("pytree_registry"),
      nb::arg("shard_arg_fallback"), nb::arg("cache").none() = nb::none(),
      nb::sig(
          // clang-format off
      "def pjit("
      "function_name: str, "
      "fun: Callable[..., Any] | None, "
      "cache_miss: Callable[..., Any], "
      "static_argnums: Sequence[int], "
      "static_argnames: Sequence[str], "
      "global_cache_key: Any, "
      "pytree_registry: _PyTreeRegistry, "
      "shard_arg_fallback: Callable[..., Any], "
      "cache: PjitFunctionCache | None = ..."
      ") -> PjitFunction"
          // clang-format on
          ));
}

}  // namespace jax